server: Disallow ServeConn usage after Shutdown

This reverts commit 48c67bcead.

client.go

@@ -14,10 +14,10 @@ import (
 
 // A Client is a Gemini client. Its zero value is a usable client.
 type Client struct {
-	// TrustCertificate is called to determine whether the client
-	// should trust the certificate provided by the server.
-	// If TrustCertificate is nil, the client will accept any certificate.
-	// If the returned error is not nil, the certificate will not be trusted
+	// TrustCertificate is called to determine whether the client should
+	// trust the certificate provided by the server.
+	// If TrustCertificate is nil or returns nil, the client will accept
+	// any certificate. Otherwise, the certificate will not be trusted
 	// and the request will be aborted.
 	//
 	// See the tofu submodule for an implementation of trust on first use.
@@ -29,14 +29,14 @@ type Client struct {
 }
 
 // Get sends a Gemini request for the given URL.
-// If the provided context is canceled or times out, the request
-// will be aborted and the context's error will be returned.
+// The context controls the entire lifetime of a request and its response:
+// obtaining a connection, sending the request, and reading the response
+// header and body.
 //
 // An error is returned if there was a Gemini protocol error.
 // A non-2x status code doesn't cause an error.
 //
-// If the returned error is nil, the Response will contain a non-nil Body
-// which the user is expected to close.
+// If the returned error is nil, the user is expected to close the Response.
 //
 // For more control over requests, use NewRequest and Client.Do.
 func (c *Client) Get(ctx context.Context, url string) (*Response, error) {
@@ -48,16 +48,14 @@ func (c *Client) Get(ctx context.Context, url string) (*Response, error) {
 }
 
 // Do sends a Gemini request and returns a Gemini response.
-// If the provided context is canceled or times out, the request
-// will be aborted and the context's error will be returned.
+// The context controls the entire lifetime of a request and its response:
+// obtaining a connection, sending the request, and reading the response
+// header and body.
 //
 // An error is returned if there was a Gemini protocol error.
 // A non-2x status code doesn't cause an error.
 //
-// If the returned error is nil, the Response will contain a non-nil Body
-// which the user is expected to close.
-//
-// Generally Get will be used instead of Do.
+// If the returned error is nil, the user is expected to close the Response.
 func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
 	if ctx == nil {
 		panic("nil context")
@@ -79,9 +77,10 @@ func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
 
 		// Use the new URL in the request so that the server gets
 		// the punycoded hostname
-		req = &Request{
-			URL: u,
-		}
+		r := new(Request)
+		*r = *req
+		r.URL = u
+		req = r
 	}
 
 	// Use request host if provided
@@ -124,7 +123,7 @@ func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
 
 	res := make(chan result, 1)
 	go func() {
-		resp, err := c.do(conn, req)
+		resp, err := c.do(ctx, conn, req)
 		res <- result{resp, err}
 	}()
 
@@ -137,23 +136,33 @@ func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
 	}
 }
 
-func (c *Client) do(conn net.Conn, req *Request) (*Response, error) {
+func (c *Client) do(ctx context.Context, conn net.Conn, req *Request) (*Response, error) {
+	ctx, cancel := context.WithCancel(ctx)
+	done := ctx.Done()
+	w := &contextWriter{
+		ctx:    ctx,
+		done:   done,
+		cancel: cancel,
+		wc:     conn,
+	}
+	rc := &contextReader{
+		ctx:    ctx,
+		done:   done,
+		cancel: cancel,
+		rc:     conn,
+	}
+
 	// Write the request
-	if err := req.Write(conn); err != nil {
+	if err := req.Write(w); err != nil {
 		return nil, err
 	}
 
 	// Read the response
-	resp, err := ReadResponse(conn)
+	resp, err := ReadResponse(rc)
 	if err != nil {
 		return nil, err
 	}
-
-	// Store TLS connection state
-	if tlsConn, ok := conn.(*tls.Conn); ok {
-		state := tlsConn.ConnectionState()
-		resp.TLS = &state
-	}
+	resp.conn = conn
 
 	return resp, nil
 }

doc.go

@@ -4,7 +4,8 @@ Package gemini provides Gemini client and server implementations.
 Client is a Gemini client.
 
 	client := &gemini.Client{}
-	resp, err := client.Get("gemini://example.com")
+	ctx := context.Background()
+	resp, err := client.Get(ctx, "gemini://example.com")
 	if err != nil {
 		// handle error
 	}
@@ -21,11 +22,12 @@ Server is a Gemini server.
 Servers should be configured with certificates:
 
 	certificates := &certificate.Store{}
+	certificates.Register("localhost")
 	err := certificates.Load("/var/lib/gemini/certs")
 	if err != nil {
 		// handle error
 	}
-	server.GetCertificate = certificates.GetCertificate
+	server.GetCertificate = certificates.Get
 
 ServeMux is a Gemini request multiplexer.
 ServeMux can handle requests for multiple hosts and schemes.
@@ -44,7 +46,8 @@ ServeMux can handle requests for multiple hosts and schemes.
 
 To start the server, call ListenAndServe:
 
-	err := server.ListenAndServe(context.Background())
+	ctx := context.Background()
+	err := server.ListenAndServe(ctx)
 	if err != nil {
 		// handle error
 	}

examples/auth.go

@@ -52,11 +52,12 @@ func fingerprint(cert *x509.Certificate) string {
 }
 
 func profile(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
-	if len(r.TLS.PeerCertificates) == 0 {
+	tls := r.TLS()
+	if len(tls.PeerCertificates) == 0 {
 		w.WriteHeader(gemini.StatusCertificateRequired, "Certificate required")
 		return
 	}
-	fingerprint := fingerprint(r.TLS.PeerCertificates[0])
+	fingerprint := fingerprint(tls.PeerCertificates[0])
 	user, ok := users[fingerprint]
 	if !ok {
 		user = &User{}
@@ -67,7 +68,8 @@ func profile(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
 }
 
 func changeUsername(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
-	if len(r.TLS.PeerCertificates) == 0 {
+	tls := r.TLS()
+	if len(tls.PeerCertificates) == 0 {
 		w.WriteHeader(gemini.StatusCertificateRequired, "Certificate required")
 		return
 	}
@@ -77,7 +79,7 @@ func changeUsername(ctx context.Context, w gemini.ResponseWriter, r *gemini.Requ
 		w.WriteHeader(gemini.StatusInput, "Username")
 		return
 	}
-	fingerprint := fingerprint(r.TLS.PeerCertificates[0])
+	fingerprint := fingerprint(tls.PeerCertificates[0])
 	user, ok := users[fingerprint]
 	if !ok {
 		user = &User{}

examples/client.go

@@ -97,7 +97,8 @@ func do(req *gemini.Request, via []*gemini.Request) (*gemini.Response, error) {
 	client := gemini.Client{
 		TrustCertificate: trustCertificate,
 	}
-	resp, err := client.Do(context.Background(), req)
+	ctx := context.Background()
+	resp, err := client.Do(ctx, req)
 	if err != nil {
 		return resp, err
 	}
@@ -156,11 +157,10 @@ func main() {
 
 	// Handle response
 	if resp.Status.Class() == gemini.StatusSuccess {
-		body, err := io.ReadAll(resp.Body)
+		_, err := io.Copy(os.Stdout, resp.Body)
 		if err != nil {
 			log.Fatal(err)
 		}
-		fmt.Print(string(body))
 	} else {
 		fmt.Printf("%d %s\n", resp.Status, resp.Meta)
 		os.Exit(1)

examples/server.go

@@ -26,7 +26,7 @@ func main() {
 	mux.Handle("/", gemini.FileServer(os.DirFS("/var/www")))
 
 	server := &gemini.Server{
-		Handler:        mux,
+		Handler:        logMiddleware(mux),
 		ReadTimeout:    30 * time.Second,
 		WriteTimeout:   1 * time.Minute,
 		GetCertificate: certificates.Get,
@@ -48,10 +48,57 @@ func main() {
 	case <-c:
 		// Shutdown the server
 		log.Println("Shutting down...")
-		ctx, _ := context.WithTimeout(context.Background(), 30*time.Second)
+		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+		defer cancel()
 		err := server.Shutdown(ctx)
 		if err != nil {
 			log.Fatal(err)
 		}
 	}
 }
+
+func logMiddleware(h gemini.Handler) gemini.Handler {
+	return gemini.HandlerFunc(func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
+		lw := &logResponseWriter{rw: w}
+		h.ServeGemini(ctx, lw, r)
+		host := r.TLS().ServerName
+		log.Printf("gemini: %s %q %d %d", host, r.URL, lw.status, lw.wrote)
+	})
+}
+
+type logResponseWriter struct {
+	rw          gemini.ResponseWriter
+	status      gemini.Status
+	meta        string
+	mediatype   string
+	wroteHeader bool
+	wrote       int
+}
+
+func (w *logResponseWriter) SetMediaType(mediatype string) {
+	w.mediatype = mediatype
+}
+
+func (w *logResponseWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		w.WriteHeader(gemini.StatusSuccess, w.mediatype)
+	}
+	n, err := w.rw.Write(b)
+	w.wrote += n
+	return n, err
+}
+
+func (w *logResponseWriter) WriteHeader(status gemini.Status, meta string) {
+	if w.wroteHeader {
+		return
+	}
+	w.status = status
+	w.meta = meta
+	w.wroteHeader = true
+	w.rw.WriteHeader(status, meta)
+	w.wrote += len(meta) + 5
+}
+
+func (w *logResponseWriter) Flush() error {
+	return nil
+}

examples/stream.go

@@ -39,33 +39,16 @@ func main() {
 
 // stream writes an infinite stream to w.
 func stream(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
-	ch := make(chan string)
-	ctx, cancel := context.WithCancel(ctx)
-
-	go func(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		default:
-				ch <- fmt.Sprint(time.Now().UTC())
+		}
+		fmt.Fprintln(w, time.Now().UTC())
+		if err := w.Flush(); err != nil {
+			return
 		}
 		time.Sleep(time.Second)
 	}
-		// Close channel when finished.
-		// In this example this will never be reached.
-		close(ch)
-	}(ctx)
-
-	for {
-		s, ok := <-ch
-		if !ok {
-			break
-		}
-		fmt.Fprintln(w, s)
-		if err := w.Flush(); err != nil {
-			cancel()
-			return
-		}
-	}
 }

fs.go

@@ -1,3 +1,5 @@
+// +build go1.16
+
 package gemini
 
 import (
@@ -13,12 +15,6 @@ import (
 	"strings"
 )
 
-func init() {
-	// Add Gemini mime types
-	mime.AddExtensionType(".gmi", "text/gemini")
-	mime.AddExtensionType(".gemini", "text/gemini")
-}
-
 // FileServer returns a handler that serves Gemini requests with the contents
 // of the provided file system.
 //
@@ -65,8 +61,8 @@ func serveContent(w ResponseWriter, name string, content io.Reader) {
 //
 // As a precaution, ServeFile will reject requests where r.URL.Path contains a
 // ".." path element; this protects against callers who might unsafely use
-// filepath.Join on r.URL.Path without sanitizing it and then use that
-// filepath.Join result as the name argument.
+// path.Join on r.URL.Path without sanitizing it and then use that
+// path.Join result as the name argument.
 //
 // As another special case, ServeFile redirects any request where r.URL.Path
 // ends in "/index.gmi" to the same path, without the final "index.gmi". To
@@ -81,7 +77,7 @@ func ServeFile(w ResponseWriter, r *Request, fsys fs.FS, name string) {
 		// serveFile. Reject the request under the assumption that happened
 		// here and ".." may not be wanted.
 		// Note that name might not contain "..", for example if code (still
-		// incorrectly) used filepath.Join(myDir, r.URL.Path).
+		// incorrectly) used path.Join(myDir, r.URL.Path).
 		w.WriteHeader(StatusBadRequest, "invalid URL path")
 		return
 	}

gemini.go

@@ -2,8 +2,15 @@ package gemini
 
 import (
 	"errors"
+	"mime"
 )
 
+func init() {
+	// Add Gemini mime types
+	mime.AddExtensionType(".gmi", "text/gemini")
+	mime.AddExtensionType(".gemini", "text/gemini")
+}
+
 var crlf = []byte("\r\n")
 
 // Errors.
@@ -16,8 +23,4 @@ var (
 	// ErrBodyNotAllowed is returned by ResponseWriter.Write calls
 	// when the response status code does not permit a body.
 	ErrBodyNotAllowed = errors.New("gemini: response status code does not allow body")
-
-	// ErrHandlerTimeout is returned on ResponseWriter Write calls
-	// in handlers which have timed out.
-	ErrHandlerTimeout = errors.New("gemini: Handler timeout")
 )

go.mod

@@ -1,5 +1,5 @@
 module git.sr.ht/~adnano/go-gemini
 
-go 1.16
+go 1.15
 
 require golang.org/x/net v0.0.0-20210119194325-5f4716e94777

handler.go

@@ -1,9 +1,12 @@
 package gemini
 
 import (
+	"bytes"
 	"context"
+	"io"
 	"net/url"
 	"strings"
+	"time"
 )
 
 // A Handler responds to a Gemini request.
@@ -13,6 +16,9 @@ import (
 // valid to use the ResponseWriter after or concurrently with the completion
 // of the ServeGemini call.
 //
+// The provided context is canceled when the client's connection is closed
+// or the ServeGemini method returns.
+//
 // Handlers should not modify the provided Request.
 type Handler interface {
 	ServeGemini(context.Context, ResponseWriter, *Request)
@@ -31,16 +37,9 @@ func (f HandlerFunc) ServeGemini(ctx context.Context, w ResponseWriter, r *Reque
 // StatusHandler returns a request handler that responds to each request
 // with the provided status code and meta.
 func StatusHandler(status Status, meta string) Handler {
-	return &statusHandler{status, meta}
-}
-
-type statusHandler struct {
-	status Status
-	meta   string
-}
-
-func (h *statusHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
-	w.WriteHeader(h.status, h.meta)
+	return HandlerFunc(func(ctx context.Context, w ResponseWriter, r *Request) {
+		w.WriteHeader(status, meta)
+	})
 }
 
 // NotFoundHandler returns a simple request handler that replies to each
@@ -75,3 +74,82 @@ func StripPrefix(prefix string, h Handler) Handler {
 		}
 	})
 }
+
+// TimeoutHandler returns a Handler that runs h with the given time limit.
+//
+// The new Handler calls h.ServeGemini to handle each request, but
+// if a call runs for longer than its time limit, the handler responds with a
+// 40 Temporary Failure error.  After such a timeout, writes by h to
+// its ResponseWriter will return context.DeadlineExceeded.
+func TimeoutHandler(h Handler, dt time.Duration) Handler {
+	return &timeoutHandler{
+		h:  h,
+		dt: dt,
+	}
+}
+
+type timeoutHandler struct {
+	h  Handler
+	dt time.Duration
+}
+
+func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
+	ctx, cancel := context.WithTimeout(ctx, t.dt)
+	defer cancel()
+
+	buf := &bytes.Buffer{}
+	tw := &timeoutWriter{
+		wr: &contextWriter{
+			ctx:    ctx,
+			cancel: cancel,
+			done:   ctx.Done(),
+			wc:     nopCloser{buf},
+		},
+	}
+
+	done := make(chan struct{})
+	go func() {
+		t.h.ServeGemini(ctx, tw, r)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		w.WriteHeader(tw.status, tw.meta)
+		w.Write(buf.Bytes())
+	case <-ctx.Done():
+		w.WriteHeader(StatusTemporaryFailure, "Timeout")
+	}
+}
+
+type timeoutWriter struct {
+	wr          io.Writer
+	status      Status
+	meta        string
+	mediatype   string
+	wroteHeader bool
+}
+
+func (w *timeoutWriter) SetMediaType(mediatype string) {
+	w.mediatype = mediatype
+}
+
+func (w *timeoutWriter) Write(b []byte) (int, error) {
+	if !w.wroteHeader {
+		w.WriteHeader(StatusSuccess, w.mediatype)
+	}
+	return w.wr.Write(b)
+}
+
+func (w *timeoutWriter) WriteHeader(status Status, meta string) {
+	if w.wroteHeader {
+		return
+	}
+	w.status = status
+	w.meta = meta
+	w.wroteHeader = true
+}
+
+func (w *timeoutWriter) Flush() error {
+	return nil
+}

io.go

@@ -0,0 +1,104 @@
+package gemini
+
+import (
+	"bufio"
+	"context"
+	"io"
+)
+
+type contextReader struct {
+	ctx    context.Context
+	done   <-chan struct{}
+	cancel func()
+	rc     io.ReadCloser
+}
+
+func (r *contextReader) Read(p []byte) (int, error) {
+	select {
+	case <-r.done:
+		r.rc.Close()
+		return 0, r.ctx.Err()
+	default:
+	}
+	n, err := r.rc.Read(p)
+	if err != nil {
+		r.cancel()
+	}
+	return n, err
+}
+
+func (r *contextReader) Close() error {
+	r.cancel()
+	return r.rc.Close()
+}
+
+type contextWriter struct {
+	ctx    context.Context
+	done   <-chan struct{}
+	cancel func()
+	wc     io.WriteCloser
+}
+
+func (w *contextWriter) Write(b []byte) (int, error) {
+	select {
+	case <-w.done:
+		w.wc.Close()
+		return 0, w.ctx.Err()
+	default:
+	}
+	n, err := w.wc.Write(b)
+	if err != nil {
+		w.cancel()
+	}
+	return n, err
+}
+
+func (w *contextWriter) Close() error {
+	w.cancel()
+	return w.wc.Close()
+}
+
+type nopCloser struct {
+	io.Writer
+}
+
+func (nopCloser) Close() error {
+	return nil
+}
+
+type nopReadCloser struct{}
+
+func (nopReadCloser) Read(p []byte) (int, error) {
+	return 0, io.EOF
+}
+
+func (nopReadCloser) Close() error {
+	return nil
+}
+
+type bufReadCloser struct {
+	br *bufio.Reader // used until empty
+	io.ReadCloser
+}
+
+func newBufReadCloser(br *bufio.Reader, rc io.ReadCloser) io.ReadCloser {
+	body := &bufReadCloser{ReadCloser: rc}
+	if br.Buffered() != 0 {
+		body.br = br
+	}
+	return body
+}
+
+func (b *bufReadCloser) Read(p []byte) (n int, err error) {
+	if b.br != nil {
+		if n := b.br.Buffered(); len(p) > n {
+			p = p[:n]
+		}
+		n, err = b.br.Read(p)
+		if b.br.Buffered() == 0 {
+			b.br = nil
+		}
+		return n, err
+	}
+	return b.ReadCloser.Read(p)
+}

query.go

@@ -12,7 +12,8 @@ func QueryEscape(query string) string {
 	return strings.ReplaceAll(url.PathEscape(query), "+", "%2B")
 }
 
-// QueryUnescape is identical to url.PathUnescape.
+// QueryUnescape unescapes a Gemini URL query.
+// It is identical to url.PathUnescape.
 func QueryUnescape(query string) (string, error) {
 	return url.PathUnescape(query)
 }

request.go

@@ -10,11 +10,8 @@ import (
 
 // A Request represents a Gemini request received by a server or to be sent
 // by a client.
-//
-// The field semantics differ slightly between client and server usage.
 type Request struct {
-	// URL specifies the URL being requested (for server
-	// requests) or the URL to access (for client requests).
+	// URL specifies the URL being requested.
 	URL *url.URL
 
 	// For client requests, Host optionally specifies the server to
@@ -23,9 +20,7 @@ type Request struct {
 	// For international domain names, Host may be in Punycode or
 	// Unicode form. Use golang.org/x/net/idna to convert it to
 	// either format if needed.
-	//
-	// For server requests, Host specifies the host on which the URL
-	// is sought.
+	// This field is ignored by the Gemini server.
 	Host string
 
 	// For client requests, Certificate optionally specifies the
@@ -33,24 +28,10 @@ type Request struct {
 	// This field is ignored by the Gemini server.
 	Certificate *tls.Certificate
 
-	// RemoteAddr allows Gemini servers and other software to record
-	// the network address that sent the request, usually for
-	// logging. This field is not filled in by ReadRequest.
-	// This field is ignored by the Gemini client.
-	RemoteAddr net.Addr
-
-	// TLS allows Gemini servers and other software to record
-	// information about the TLS connection on which the request
-	// was received. This field is not filled in by ReadRequest.
-	// The Gemini server in this package sets the field for
-	// TLS-enabled connections before invoking a handler;
-	// otherwise it leaves the field nil.
-	// This field is ignored by the Gemini client.
-	TLS *tls.ConnectionState
+	conn net.Conn
 }
 
 // NewRequest returns a new request.
-//
 // The returned Request is suitable for use with Client.Do.
 //
 // Callers should be careful that the URL query is properly escaped.
@@ -111,3 +92,18 @@ func (r *Request) Write(w io.Writer) error {
 	}
 	return bw.Flush()
 }
+
+// Conn returns the network connection on which the request was received.
+func (r *Request) Conn() net.Conn {
+	return r.conn
+}
+
+// TLS returns information about the TLS connection on which the
+// request was received.
+func (r *Request) TLS() *tls.ConnectionState {
+	if tlsConn, ok := r.conn.(*tls.Conn); ok {
+		state := tlsConn.ConnectionState()
+		return &state
+	}
+	return nil
+}

response.go

@@ -3,8 +3,8 @@ package gemini
 import (
 	"bufio"
 	"crypto/tls"
-	"fmt"
 	"io"
+	"net"
 	"strconv"
 )
 
@@ -17,13 +17,12 @@ const defaultMediaType = "text/gemini; charset=utf-8"
 // header has been received. The response body is streamed on demand
 // as the Body field is read.
 type Response struct {
-	// Status contains the response status code.
+	// Status is the response status code.
 	Status Status
 
-	// Meta contains more information related to the response status.
-	// For successful responses, Meta should contain the media type of the response.
-	// For failure responses, Meta should contain a short description of the failure.
-	// Meta should not be longer than 1024 bytes.
+	// Meta returns the response meta.
+	// For successful responses, the meta should contain the media type of the response.
+	// For failure responses, the meta should contain a short description of the failure.
 	Meta string
 
 	// Body represents the response body.
@@ -38,15 +37,13 @@ type Response struct {
 	// close Body.
 	Body io.ReadCloser
 
-	// TLS contains information about the TLS connection on which the
-	// response was received. It is nil for unencrypted responses.
-	TLS *tls.ConnectionState
+	conn net.Conn
 }
 
 // ReadResponse reads a Gemini response from the provided io.ReadCloser.
-func ReadResponse(rc io.ReadCloser) (*Response, error) {
+func ReadResponse(r io.ReadCloser) (*Response, error) {
 	resp := &Response{}
-	br := bufio.NewReader(rc)
+	br := bufio.NewReader(r)
 
 	// Read the status
 	statusB := make([]byte, 2)
@@ -91,65 +88,25 @@ func ReadResponse(rc io.ReadCloser) (*Response, error) {
 	}
 
 	if resp.Status.Class() == StatusSuccess {
-		resp.Body = newReadCloserBody(br, rc)
+		resp.Body = newBufReadCloser(br, r)
 	} else {
 		resp.Body = nopReadCloser{}
-		rc.Close()
+		r.Close()
 	}
 	return resp, nil
 }
 
-type nopReadCloser struct{}
-
-func (nopReadCloser) Read(p []byte) (int, error) {
-	return 0, io.EOF
+// Conn returns the network connection on which the response was received.
+func (r *Response) Conn() net.Conn {
+	return r.conn
 }
 
-func (nopReadCloser) Close() error {
-	return nil
-}
-
-type readCloserBody struct {
-	br *bufio.Reader // used until empty
-	io.ReadCloser
-}
-
-func newReadCloserBody(br *bufio.Reader, rc io.ReadCloser) io.ReadCloser {
-	body := &readCloserBody{ReadCloser: rc}
-	if br.Buffered() != 0 {
-		body.br = br
-	}
-	return body
-}
-
-func (b *readCloserBody) Read(p []byte) (n int, err error) {
-	if b.br != nil {
-		if n := b.br.Buffered(); len(p) > n {
-			p = p[:n]
-		}
-		n, err = b.br.Read(p)
-		if b.br.Buffered() == 0 {
-			b.br = nil
-		}
-		return n, err
-	}
-	return b.ReadCloser.Read(p)
-}
-
-// Write writes r to w in the Gemini response format, including the
-// header and body.
-//
-// This method consults the Status, Meta, and Body fields of the response.
-// The Response Body is closed after it is sent.
-func (r *Response) Write(w io.Writer) error {
-	if _, err := fmt.Fprintf(w, "%02d %s\r\n", r.Status, r.Meta); err != nil {
-		return err
-	}
-	if r.Body != nil {
-		defer r.Body.Close()
-		if _, err := io.Copy(w, r.Body); err != nil {
-			return err
-		}
+// TLS returns information about the TLS connection on which the
+// response was received.
+func (r *Response) TLS() *tls.ConnectionState {
+	if tlsConn, ok := r.conn.(*tls.Conn); ok {
+		state := tlsConn.ConnectionState()
+		return &state
 	}
 	return nil
 }
@@ -166,7 +123,7 @@ type ResponseWriter interface {
 	//
 	// Setting the media type after a call to Write or WriteHeader has
 	// no effect.
-	SetMediaType(string)
+	SetMediaType(mediatype string)
 
 	// Write writes the data to the connection as part of a Gemini response.
 	//
@@ -193,20 +150,16 @@ type ResponseWriter interface {
 }
 
 type responseWriter struct {
-	b           *bufio.Writer
+	bw          *bufio.Writer
+	cl          io.Closer
 	mediatype   string
 	wroteHeader bool
 	bodyAllowed bool
 }
 
-// NewResponseWriter returns a ResponseWriter that uses the provided io.Writer.
-func NewResponseWriter(w io.Writer) ResponseWriter {
-	return newResponseWriter(w)
-}
-
 func newResponseWriter(w io.Writer) *responseWriter {
 	return &responseWriter{
-		b: bufio.NewWriter(w),
+		bw: bufio.NewWriter(w),
 	}
 }
 
@@ -226,7 +179,7 @@ func (w *responseWriter) Write(b []byte) (int, error) {
 	if !w.bodyAllowed {
 		return 0, ErrBodyNotAllowed
 	}
-	return w.b.Write(b)
+	return w.bw.Write(b)
 }
 
 func (w *responseWriter) WriteHeader(status Status, meta string) {
@@ -238,10 +191,10 @@ func (w *responseWriter) WriteHeader(status Status, meta string) {
 		w.bodyAllowed = true
 	}
 
-	w.b.WriteString(strconv.Itoa(int(status)))
-	w.b.WriteByte(' ')
-	w.b.WriteString(meta)
-	w.b.Write(crlf)
+	w.bw.WriteString(strconv.Itoa(int(status)))
+	w.bw.WriteByte(' ')
+	w.bw.WriteString(meta)
+	w.bw.Write(crlf)
 	w.wroteHeader = true
 }
 
@@ -249,6 +202,6 @@ func (w *responseWriter) Flush() error {
 	if !w.wroteHeader {
 		w.WriteHeader(StatusTemporaryFailure, "Temporary failure")
 	}
-	// Write errors from writeHeader will be returned here.
-	return w.b.Flush()
+	// Write errors from WriteHeader will be returned here.
+	return w.bw.Flush()
 }

response_test.go

@@ -2,6 +2,7 @@ package gemini
 
 import (
 	"io"
+	"io/ioutil"
 	"strings"
 	"testing"
 )
@@ -82,7 +83,7 @@ func TestReadWriteResponse(t *testing.T) {
 
 	for _, test := range tests {
 		t.Logf("%#v", test.Raw)
-		resp, err := ReadResponse(io.NopCloser(strings.NewReader(test.Raw)))
+		resp, err := ReadResponse(ioutil.NopCloser(strings.NewReader(test.Raw)))
 		if err != test.Err {
 			t.Errorf("expected err = %v, got %v", test.Err, err)
 		}
@@ -96,7 +97,7 @@ func TestReadWriteResponse(t *testing.T) {
 		if resp.Meta != test.Meta {
 			t.Errorf("expected meta = %s, got %s", test.Meta, resp.Meta)
 		}
-		b, _ := io.ReadAll(resp.Body)
+		b, _ := ioutil.ReadAll(resp.Body)
 		body := string(b)
 		if body != test.Body {
 			t.Errorf("expected body = %#v, got %#v", test.Body, body)
@@ -107,14 +108,12 @@ func TestReadWriteResponse(t *testing.T) {
 		if test.Err != nil || test.SkipWrite {
 			continue
 		}
-		resp := &Response{
-			Status: test.Status,
-			Meta:   test.Meta,
-			Body:   io.NopCloser(strings.NewReader(test.Body)),
-		}
 
 		var b strings.Builder
-		if err := resp.Write(&b); err != nil {
+		w := newResponseWriter(nopCloser{&b})
+		w.WriteHeader(test.Status, test.Meta)
+		io.Copy(w, strings.NewReader(test.Body))
+		if err := w.Flush(); err != nil {
 			t.Error(err)
 			continue
 		}

server.go

@@ -52,7 +52,7 @@ type Server struct {
 
 	listeners map[*net.Listener]context.CancelFunc
 	conns     map[*net.Conn]context.CancelFunc
-	closed    bool // true if Closed or Shutdown called
+	closed    bool // true if Close or Shutdown called
 	shutdown  bool // true if Shutdown called
 	doneChan  chan struct{}
 	mu        sync.Mutex
@@ -229,8 +229,8 @@ func (srv *Server) deleteListener(l *net.Listener) {
 }
 
 // Serve accepts incoming connections on the Listener l, creating a new
-// service goroutine for each. The service goroutines reads the request and
-// then calls the appropriate Handler to reply to them. If the provided
+// service goroutine for each. The service goroutines read the requests and
+// then call the appropriate Handler to reply to them. If the provided
 // context expires, Serve closes l and returns the context's error.
 //
 // Serve always closes l and returns a non-nil error.
@@ -282,14 +282,17 @@ func (srv *Server) serve(ctx context.Context, l net.Listener) error {
 			return err
 		}
 		tempDelay = 0
-		go srv.ServeConn(ctx, rw)
+		go srv.serveConn(ctx, rw, false)
 	}
 }
 
-func (srv *Server) trackConn(conn *net.Conn, cancel context.CancelFunc) bool {
+func (srv *Server) trackConn(conn *net.Conn, cancel context.CancelFunc, external bool) bool {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
-	if srv.closed && !srv.shutdown {
+	// Reject the connection under the following conditions:
+	// - Shutdown or Close has been called and conn is external (from ServeConn)
+	// - Close (not Shutdown) has been called and conn is internal (from Serve)
+	if srv.closed && (external || !srv.shutdown) {
 		return false
 	}
 	if srv.conns == nil {
@@ -309,15 +312,17 @@ func (srv *Server) deleteConn(conn *net.Conn) {
 // It closes the connection when the response has been completed.
 // If the provided context expires before the response has completed,
 // ServeConn closes the connection and returns the context's error.
-//
-// Note that ServeConn can be used during a Shutdown.
 func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
+	return srv.serveConn(ctx, conn, true)
+}
+
+func (srv *Server) serveConn(ctx context.Context, conn net.Conn, external bool) error {
 	defer conn.Close()
 
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	if !srv.trackConn(&conn, cancel) {
+	if !srv.trackConn(&conn, cancel, external) {
 		return context.Canceled
 	}
 	defer srv.tryCloseDone()
@@ -332,7 +337,7 @@ func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
 
 	errch := make(chan error, 1)
 	go func() {
-		errch <- srv.serveConn(ctx, conn)
+		errch <- srv.goServeConn(ctx, conn)
 	}()
 
 	select {
@@ -343,24 +348,30 @@ func (srv *Server) ServeConn(ctx context.Context, conn net.Conn) error {
 	}
 }
 
-func (srv *Server) serveConn(ctx context.Context, conn net.Conn) error {
-	w := newResponseWriter(conn)
+func (srv *Server) goServeConn(ctx context.Context, conn net.Conn) error {
+	ctx, cancel := context.WithCancel(ctx)
+	done := ctx.Done()
+	cw := &contextWriter{
+		ctx:    ctx,
+		done:   done,
+		cancel: cancel,
+		wc:     conn,
+	}
+	r := &contextReader{
+		ctx:    ctx,
+		done:   done,
+		cancel: cancel,
+		rc:     conn,
+	}
 
-	req, err := ReadRequest(conn)
+	w := newResponseWriter(cw)
+
+	req, err := ReadRequest(r)
 	if err != nil {
 		w.WriteHeader(StatusBadRequest, "Bad request")
 		return w.Flush()
 	}
-
-	// Store the TLS connection state
-	if tlsConn, ok := conn.(*tls.Conn); ok {
-		state := tlsConn.ConnectionState()
-		req.TLS = &state
-		req.Host = state.ServerName
-	}
-
-	// Store remote address
-	req.RemoteAddr = conn.RemoteAddr()
+	req.conn = conn
 
 	h := srv.Handler
 	if h == nil {

timeout.go

@@ -1,112 +0,0 @@
-package gemini
-
-import (
-	"bytes"
-	"context"
-	"sync"
-	"time"
-)
-
-// TimeoutHandler returns a Handler that runs h with the given time limit.
-//
-// The new Handler calls h.ServeGemini to handle each request, but
-// if a call runs for longer than its time limit, the handler responds with a
-// 40 Temporary Failure error. After such a timeout, writes by h to its
-// ResponseWriter will return ErrHandlerTimeout.
-func TimeoutHandler(h Handler, dt time.Duration) Handler {
-	return &timeoutHandler{
-		h:  h,
-		dt: dt,
-	}
-}
-
-type timeoutHandler struct {
-	h  Handler
-	dt time.Duration
-}
-
-func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
-	ctx, cancel := context.WithTimeout(ctx, t.dt)
-	defer cancel()
-
-	done := make(chan struct{})
-	tw := &timeoutWriter{}
-	panicChan := make(chan interface{}, 1)
-	go func() {
-		defer func() {
-			if p := recover(); p != nil {
-				panicChan <- p
-			}
-		}()
-		t.h.ServeGemini(ctx, tw, r)
-		close(done)
-	}()
-
-	select {
-	case p := <-panicChan:
-		panic(p)
-	case <-done:
-		tw.mu.Lock()
-		defer tw.mu.Unlock()
-		if !tw.wroteHeader {
-			tw.status = StatusSuccess
-		}
-		w.WriteHeader(tw.status, tw.meta)
-		w.Write(tw.b.Bytes())
-	case <-ctx.Done():
-		tw.mu.Lock()
-		defer tw.mu.Unlock()
-		w.WriteHeader(StatusTemporaryFailure, "Timeout")
-		tw.timedOut = true
-	}
-}
-
-type timeoutWriter struct {
-	mu          sync.Mutex
-	b           bytes.Buffer
-	status      Status
-	meta        string
-	mediatype   string
-	wroteHeader bool
-	timedOut    bool
-}
-
-func (w *timeoutWriter) SetMediaType(mediatype string) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	w.mediatype = mediatype
-}
-
-func (w *timeoutWriter) Write(b []byte) (int, error) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	if w.timedOut {
-		return 0, ErrHandlerTimeout
-	}
-	if !w.wroteHeader {
-		w.writeHeaderLocked(StatusSuccess, w.mediatype)
-	}
-	return w.b.Write(b)
-}
-
-func (w *timeoutWriter) WriteHeader(status Status, meta string) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-	if w.timedOut {
-		return
-	}
-	w.writeHeaderLocked(status, meta)
-}
-
-func (w *timeoutWriter) writeHeaderLocked(status Status, meta string) {
-	if w.wroteHeader {
-		return
-	}
-	w.status = status
-	w.meta = meta
-	w.wroteHeader = true
-}
-
-func (w *timeoutWriter) Flush() error {
-	return nil
-}