Update documentation

Use a separate context to cancel listeners so that cancelling the
listener does not cancel it's connections.

.gitignore

@@ -1,2 +0,0 @@
-*.crt
-*.key

certificate/create.go

@@ -14,103 +14,9 @@ import (
 	"math/big"
 	"net"
 	"os"
-	"path/filepath"
-	"strings"
-	"sync"
 	"time"
 )
 
-// Dir represents a directory of certificates.
-// The zero value for Dir is an empty directory ready to use.
-//
-// Dir is safe for concurrent use by multiple goroutines.
-type Dir struct {
-	certs map[string]tls.Certificate
-	path  *string
-	mu    sync.RWMutex
-}
-
-// Add adds a certificate for the given scope to the directory.
-// It tries to parse the certificate if it is not already parsed.
-func (d *Dir) Add(scope string, cert tls.Certificate) error {
-	d.mu.Lock()
-	defer d.mu.Unlock()
-	if d.certs == nil {
-		d.certs = map[string]tls.Certificate{}
-	}
-	// Parse certificate if not already parsed
-	if cert.Leaf == nil {
-		parsed, err := x509.ParseCertificate(cert.Certificate[0])
-		if err == nil {
-			cert.Leaf = parsed
-		}
-	}
-
-	if d.path != nil {
-		// Escape slash character
-		scope = strings.ReplaceAll(scope, "/", ":")
-		certPath := filepath.Join(*d.path, scope+".crt")
-		keyPath := filepath.Join(*d.path, scope+".key")
-		if err := Write(cert, certPath, keyPath); err != nil {
-			return err
-		}
-	}
-
-	d.certs[scope] = cert
-	return nil
-}
-
-// Lookup returns the certificate for the provided scope.
-func (d *Dir) Lookup(scope string) (tls.Certificate, bool) {
-	d.mu.RLock()
-	defer d.mu.RUnlock()
-	cert, ok := d.certs[scope]
-	return cert, ok
-}
-
-// Entries returns a map of hostnames to certificates.
-func (d *Dir) Entries() map[string]tls.Certificate {
-	certs := map[string]tls.Certificate{}
-	for key := range d.certs {
-		certs[key] = d.certs[key]
-	}
-	return certs
-}
-
-// Load loads certificates from the provided path.
-// Add will write certificates to this path.
-//
-// The directory should contain certificates and private keys
-// named scope.crt and scope.key respectively, where scope is
-// the scope of the certificate.
-func (d *Dir) Load(path string) error {
-	matches, err := filepath.Glob(filepath.Join(path, "*.crt"))
-	if err != nil {
-		return err
-	}
-	for _, crtPath := range matches {
-		keyPath := strings.TrimSuffix(crtPath, ".crt") + ".key"
-		cert, err := tls.LoadX509KeyPair(crtPath, keyPath)
-		if err != nil {
-			continue
-		}
-		scope := strings.TrimSuffix(filepath.Base(crtPath), ".crt")
-		// Unescape slash character
-		scope = strings.ReplaceAll(scope, ":", "/")
-		d.Add(scope, cert)
-	}
-	d.SetPath(path)
-	return nil
-}
-
-// SetPath sets the directory path.
-// Add will write certificates to this path.
-func (d *Dir) SetPath(path string) {
-	d.mu.Lock()
-	defer d.mu.Unlock()
-	d.path = &path
-}
-
 // CreateOptions configures the creation of a TLS certificate.
 type CreateOptions struct {
 	// Subject Alternate Name values.

certificate/store.go

@@ -0,0 +1,170 @@
+package certificate
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"errors"
+	"fmt"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+)
+
+// Store represents a certificate store.
+// The zero value for Store is an empty store ready to use.
+//
+// Store is safe for concurrent use by multiple goroutines.
+type Store struct {
+	// CreateCertificate, if not nil, is called to create a new certificate
+	// to replace a missing or expired certificate. If CreateCertificate
+	// is nil, a certificate with a duration of 1 year will be created.
+	CreateCertificate func(scope string) (tls.Certificate, error)
+
+	certs map[string]tls.Certificate
+	path  string
+	mu    sync.RWMutex
+}
+
+// Register registers the provided scope in the certificate store.
+// The certificate will be created upon calling GetCertificate.
+func (s *Store) Register(scope string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.certs == nil {
+		s.certs = make(map[string]tls.Certificate)
+	}
+	s.certs[scope] = tls.Certificate{}
+}
+
+// Add adds a certificate for the given scope to the certificate store.
+func (s *Store) Add(scope string, cert tls.Certificate) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.certs == nil {
+		s.certs = make(map[string]tls.Certificate)
+	}
+
+	// Parse certificate if not already parsed
+	if cert.Leaf == nil {
+		parsed, err := x509.ParseCertificate(cert.Certificate[0])
+		if err != nil {
+			return err
+		}
+		cert.Leaf = parsed
+	}
+
+	if s.path != "" {
+		// Escape slash character
+		path := strings.ReplaceAll(scope, "/", ":")
+		certPath := filepath.Join(s.path, path+".crt")
+		keyPath := filepath.Join(s.path, path+".key")
+		if err := Write(cert, certPath, keyPath); err != nil {
+			return err
+		}
+	}
+
+	s.certs[scope] = cert
+	return nil
+}
+
+// Lookup returns the certificate for the provided scope.
+func (s *Store) Lookup(scope string) (tls.Certificate, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	cert, ok := s.certs[scope]
+	return cert, ok
+}
+
+// GetCertificate retrieves the certificate for the given scope.
+// If the retrieved certificate is expired or the scope is registered but
+// has no certificate, it calls CreateCertificate to create a new certificate.
+func (s *Store) GetCertificate(scope string) (*tls.Certificate, error) {
+	cert, ok := s.Lookup(scope)
+	if !ok {
+		// Try wildcard
+		wildcard := strings.SplitN(scope, ".", 2)
+		if len(wildcard) == 2 {
+			cert, ok = s.Lookup("*." + wildcard[1])
+		}
+	}
+	if !ok {
+		// Try "*"
+		_, ok = s.Lookup("*")
+	}
+	if !ok {
+		return nil, errors.New("unrecognized scope")
+	}
+
+	// If the certificate is empty or expired, generate a new one.
+	if cert.Leaf == nil || cert.Leaf.NotAfter.Before(time.Now()) {
+		var err error
+		cert, err = s.createCertificate(scope)
+		if err != nil {
+			return nil, err
+		}
+		if err := s.Add(scope, cert); err != nil {
+			return nil, fmt.Errorf("failed to add certificate for %s: %w", scope, err)
+		}
+	}
+
+	return &cert, nil
+}
+
+func (s *Store) createCertificate(scope string) (tls.Certificate, error) {
+	if s.CreateCertificate != nil {
+		return s.CreateCertificate(scope)
+	}
+	return Create(CreateOptions{
+		DNSNames: []string{scope},
+		Subject: pkix.Name{
+			CommonName: scope,
+		},
+		Duration: 365 * 24 * time.Hour,
+	})
+}
+
+// Load loads certificates from the provided path.
+// New certificates will be written to this path.
+//
+// The path should lead to a directory containing certificates
+// and private keys named "scope.crt" and "scope.key" respectively,
+// where "scope" is the scope of the certificate.
+func (s *Store) Load(path string) error {
+	matches, err := filepath.Glob(filepath.Join(path, "*.crt"))
+	if err != nil {
+		return err
+	}
+	for _, crtPath := range matches {
+		keyPath := strings.TrimSuffix(crtPath, ".crt") + ".key"
+		cert, err := tls.LoadX509KeyPair(crtPath, keyPath)
+		if err != nil {
+			continue
+		}
+		scope := strings.TrimSuffix(filepath.Base(crtPath), ".crt")
+		// Unescape slash character
+		scope = strings.ReplaceAll(scope, ":", "/")
+		s.Add(scope, cert)
+	}
+	s.SetPath(path)
+	return nil
+}
+
+// Entries returns a map of scopes to certificates.
+func (s *Store) Entries() map[string]tls.Certificate {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	certs := make(map[string]tls.Certificate)
+	for key := range s.certs {
+		certs[key] = s.certs[key]
+	}
+	return certs
+}
+
+// SetPath sets the path that new certificates will be written to.
+func (s *Store) SetPath(path string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.path = path
+}

client.go

@@ -1,13 +1,12 @@
 package gemini
 
 import (
-	"bufio"
 	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"errors"
-	"fmt"
 	"net"
+	"net/url"
 	"time"
 )
 
@@ -29,6 +28,10 @@ type Client struct {
 	//
 	// A Timeout of zero means no timeout.
 	Timeout time.Duration
+
+	// DialContext specifies the dial function for creating TCP connections.
+	// If DialContext is nil, the client dials using package net.
+	DialContext func(ctx context.Context, network, addr string) (net.Conn, error)
 }
 
 // Get sends a Gemini request for the given URL.
@@ -40,12 +43,12 @@ type Client struct {
 // which the user is expected to close.
 //
 // For more control over requests, use NewRequest and Client.Do.
-func (c *Client) Get(url string) (*Response, error) {
+func (c *Client) Get(ctx context.Context, url string) (*Response, error) {
 	req, err := NewRequest(url)
 	if err != nil {
 		return nil, err
 	}
-	return c.Do(req)
+	return c.Do(ctx, req)
 }
 
 // Do sends a Gemini request and returns a Gemini response, following
@@ -58,47 +61,56 @@ func (c *Client) Get(url string) (*Response, error) {
 // which the user is expected to close.
 //
 // Generally Get will be used instead of Do.
-func (c *Client) Do(req *Request) (*Response, error) {
-	// Punycode request URL host
-	hostname, port, err := net.SplitHostPort(req.URL.Host)
-	if err != nil {
-		// Likely no port
-		hostname = req.URL.Host
-		port = "1965"
+func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
+	if ctx == nil {
+		panic("nil context")
 	}
-	punycode, err := punycodeHostname(hostname)
+
+	// Punycode request URL host
+	host, port := splitHostPort(req.URL.Host)
+	punycode, err := punycodeHostname(host)
 	if err != nil {
 		return nil, err
 	}
-	if hostname != punycode {
-		hostname = punycode
+	if host != punycode {
+		host = punycode
 
 		// Make a copy of the request
-		_req := *req
-		req = &_req
-		_url := *req.URL
-		req.URL = &_url
+		r2 := new(Request)
+		*r2 = *req
+		r2.URL = new(url.URL)
+		*r2.URL = *req.URL
+		req = r2
 
 		// Set the host
-		req.URL.Host = net.JoinHostPort(hostname, port)
+		req.URL.Host = net.JoinHostPort(host, port)
 	}
 
 	// Use request host if provided
 	if req.Host != "" {
-		hostname, port, err = net.SplitHostPort(req.Host)
-		if err != nil {
-			// Port is required
-			return nil, err
-		}
-		// Punycode hostname
-		hostname, err = punycodeHostname(hostname)
+		host, port = splitHostPort(req.Host)
+		host, err = punycodeHostname(host)
 		if err != nil {
 			return nil, err
 		}
 	}
 
+	addr := net.JoinHostPort(host, port)
+
 	// Connect to the host
-	config := &tls.Config{
+	start := time.Now()
+	conn, err := c.dialContext(ctx, "tcp", addr)
+	if err != nil {
+		return nil, err
+	}
+
+	// Set the connection deadline
+	if c.Timeout != 0 {
+		conn.SetDeadline(start.Add(c.Timeout))
+	}
+
+	// Setup TLS
+	conn = tls.Client(conn, &tls.Config{
 		InsecureSkipVerify: true,
 		MinVersion:         tls.VersionTLS12,
 		GetClientCertificate: func(_ *tls.CertificateRequestInfo) (*tls.Certificate, error) {
@@ -108,63 +120,34 @@ func (c *Client) Do(req *Request) (*Response, error) {
 			return &tls.Certificate{}, nil
 		},
 		VerifyConnection: func(cs tls.ConnectionState) error {
-			return c.verifyConnection(hostname, punycode, cs)
+			return c.verifyConnection(cs, host)
 		},
-		ServerName: hostname,
+		ServerName: host,
+	})
+
+	type result struct {
+		resp *Response
+		err  error
 	}
 
-	ctx := req.Context
-	if ctx == nil {
-		ctx = context.Background()
+	res := make(chan result, 1)
+	go func() {
+		resp, err := c.do(conn, req)
+		res <- result{resp, err}
+	}()
+
+	select {
+	case <-ctx.Done():
+		conn.Close()
+		return nil, ctx.Err()
+	case r := <-res:
+		return r.resp, r.err
 	}
-
-	start := time.Now()
-	dialer := net.Dialer{
-		Timeout: c.Timeout,
-	}
-
-	address := net.JoinHostPort(hostname, port)
-	netConn, err := dialer.DialContext(ctx, "tcp", address)
-	if err != nil {
-		return nil, err
-	}
-
-	conn := tls.Client(netConn, config)
-
-	// Set connection deadline
-	if c.Timeout != 0 {
-		err := conn.SetDeadline(start.Add(c.Timeout))
-		if err != nil {
-			return nil, fmt.Errorf("failed to set connection deadline: %w", err)
-		}
-	}
-
-	resp, err := c.do(conn, req)
-	if err != nil {
-		// If we fail to perform the request/response we have
-		// to take responsibility for closing the connection.
-		_ = conn.Close()
-
-		return nil, err
-	}
-
-	// Store connection state
-	state := conn.ConnectionState()
-	resp.TLS = &state
-
-	return resp, nil
 }
 
-func (c *Client) do(conn *tls.Conn, req *Request) (*Response, error) {
+func (c *Client) do(conn net.Conn, req *Request) (*Response, error) {
 	// Write the request
-	w := bufio.NewWriter(conn)
-
-	err := req.Write(w)
-	if err != nil {
-		return nil, fmt.Errorf("failed to write request: %w", err)
-	}
-
-	if err := w.Flush(); err != nil {
+	if err := req.Write(conn); err != nil {
 		return nil, err
 	}
 
@@ -174,23 +157,48 @@ func (c *Client) do(conn *tls.Conn, req *Request) (*Response, error) {
 		return nil, err
 	}
 
+	// Store TLS connection state
+	if tlsConn, ok := conn.(*tls.Conn); ok {
+		state := tlsConn.ConnectionState()
+		resp.TLS = &state
+	}
+
 	return resp, nil
 }
 
-func (c *Client) verifyConnection(hostname, punycode string, cs tls.ConnectionState) error {
+func (c *Client) dialContext(ctx context.Context, network, addr string) (net.Conn, error) {
+	if c.DialContext != nil {
+		return c.DialContext(ctx, network, addr)
+	}
+	return (&net.Dialer{
+		Timeout: c.Timeout,
+	}).DialContext(ctx, network, addr)
+}
+
+func (c *Client) verifyConnection(cs tls.ConnectionState, hostname string) error {
 	cert := cs.PeerCertificates[0]
-	// Verify punycoded hostname
-	if err := verifyHostname(cert, punycode); err != nil {
+	// Verify hostname
+	if err := verifyHostname(cert, hostname); err != nil {
 		return err
 	}
 	// Check expiration date
 	if !time.Now().Before(cert.NotAfter) {
 		return errors.New("gemini: certificate expired")
 	}
-
 	// See if the client trusts the certificate
 	if c.TrustCertificate != nil {
 		return c.TrustCertificate(hostname, cert)
 	}
 	return nil
 }
+
+func splitHostPort(hostport string) (host, port string) {
+	var err error
+	host, port, err = net.SplitHostPort(hostport)
+	if err != nil {
+		// Likely no port
+		host = hostport
+		port = "1965"
+	}
+	return
+}

doc.go

@@ -1,5 +1,5 @@
 /*
-Package gemini implements the Gemini protocol.
+Package gemini provides Gemini client and server implementations.
 
 Client is a Gemini client.
 
@@ -20,26 +20,31 @@ Server is a Gemini server.
 
 Servers should be configured with certificates:
 
-	err := server.Certificates.Load("/var/lib/gemini/certs")
+	certificates := &certificate.Store{}
+	err := certificates.Load("/var/lib/gemini/certs")
 	if err != nil {
 		// handle error
 	}
+	server.GetCertificate = certificates.GetCertificate
 
-Servers can accept requests for multiple hosts and schemes:
+ServeMux is a Gemini request multiplexer.
+ServeMux can handle requests for multiple hosts and schemes.
 
-	server.RegisterFunc("example.com", func(w gemini.ResponseWriter, r *gemini.Request) {
+	mux := &gemini.ServeMux{}
+	mux.HandleFunc("example.com", func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
 		fmt.Fprint(w, "Welcome to example.com")
 	})
-	server.RegisterFunc("example.org", func(w gemini.ResponseWriter, r *gemini.Request) {
-		fmt.Fprint(w, "Welcome to example.org")
+	mux.HandleFunc("example.org/about.gmi", func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
+		fmt.Fprint(w, "About example.org")
 	})
-	server.RegisterFunc("http://example.net", func(w gemini.ResponseWriter, r *gemini.Request) {
+	mux.HandleFunc("http://example.net", func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
 		fmt.Fprint(w, "Proxied content from http://example.net")
 	})
+	server.Handler = mux
 
 To start the server, call ListenAndServe:
 
-	err := server.ListenAndServe()
+	err := server.ListenAndServe(context.Background())
 	if err != nil {
 		// handle error
 	}

examples/auth.go

@@ -3,10 +3,9 @@
 package main
 
 import (
+	"context"
 	"crypto/sha512"
-	"crypto/tls"
 	"crypto/x509"
-	"crypto/x509/pkix"
 	"fmt"
 	"log"
 	"time"
@@ -25,26 +24,24 @@ var (
 )
 
 func main() {
-	var mux gemini.ServeMux
+	certificates := &certificate.Store{}
+	certificates.Register("localhost")
+	if err := certificates.Load("/var/lib/gemini/certs"); err != nil {
+		log.Fatal(err)
+	}
+
+	mux := &gemini.ServeMux{}
 	mux.HandleFunc("/", profile)
 	mux.HandleFunc("/username", changeUsername)
 
-	var server gemini.Server
-	if err := server.Certificates.Load("/var/lib/gemini/certs"); err != nil {
-		log.Fatal(err)
+	server := &gemini.Server{
+		Handler:        mux,
+		ReadTimeout:    30 * time.Second,
+		WriteTimeout:   1 * time.Minute,
+		GetCertificate: certificates.GetCertificate,
 	}
-	server.GetCertificate = func(hostname string) (tls.Certificate, error) {
-		return certificate.Create(certificate.CreateOptions{
-			Subject: pkix.Name{
-				CommonName: hostname,
-			},
-			DNSNames: []string{hostname},
-			Duration: time.Hour,
-		})
-	}
-	server.Handle("localhost", &mux)
 
-	if err := server.ListenAndServe(); err != nil {
+	if err := server.ListenAndServe(context.Background()); err != nil {
 		log.Fatal(err)
 	}
 }
@@ -54,9 +51,9 @@ func fingerprint(cert *x509.Certificate) string {
 	return string(b[:])
 }
 
-func profile(w gemini.ResponseWriter, r *gemini.Request) {
+func profile(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
 	if len(r.TLS.PeerCertificates) == 0 {
-		w.Status(gemini.StatusCertificateRequired)
+		w.WriteHeader(gemini.StatusCertificateRequired, "Certificate required")
 		return
 	}
 	fingerprint := fingerprint(r.TLS.PeerCertificates[0])
@@ -69,15 +66,15 @@ func profile(w gemini.ResponseWriter, r *gemini.Request) {
 	fmt.Fprintln(w, "=> /username Change username")
 }
 
-func changeUsername(w gemini.ResponseWriter, r *gemini.Request) {
+func changeUsername(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
 	if len(r.TLS.PeerCertificates) == 0 {
-		w.Status(gemini.StatusCertificateRequired)
+		w.WriteHeader(gemini.StatusCertificateRequired, "Certificate required")
 		return
 	}
 
 	username, err := gemini.QueryUnescape(r.URL.RawQuery)
 	if err != nil || username == "" {
-		w.Header(gemini.StatusInput, "Username")
+		w.WriteHeader(gemini.StatusInput, "Username")
 		return
 	}
 	fingerprint := fingerprint(r.TLS.PeerCertificates[0])
@@ -87,5 +84,5 @@ func changeUsername(w gemini.ResponseWriter, r *gemini.Request) {
 		users[fingerprint] = user
 	}
 	user.Name = username
-	w.Header(gemini.StatusRedirect, "/")
+	w.WriteHeader(gemini.StatusRedirect, "/")
 }

examples/client.go

@@ -7,10 +7,11 @@ package main
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"crypto/x509"
 	"errors"
 	"fmt"
-	"io/ioutil"
+	"io"
 	"log"
 	"net/url"
 	"os"
@@ -19,7 +20,6 @@ import (
 
 	"git.sr.ht/~adnano/go-gemini"
 	"git.sr.ht/~adnano/go-gemini/tofu"
-	"git.sr.ht/~adnano/go-xdg"
 )
 
 var (
@@ -28,9 +28,16 @@ var (
 	scanner   *bufio.Scanner
 )
 
+func xdgDataHome() string {
+	if s, ok := os.LookupEnv("XDG_DATA_HOME"); ok {
+		return s
+	}
+	return filepath.Join(os.Getenv("HOME"), ".local", "share")
+}
+
 func init() {
 	// Load known hosts file
-	path := filepath.Join(xdg.DataHome(), "gemini", "known_hosts")
+	path := filepath.Join(xdgDataHome(), "gemini", "known_hosts")
 	err := hosts.Load(path)
 	if err != nil {
 		log.Fatal(err)
@@ -90,12 +97,12 @@ func do(req *gemini.Request, via []*gemini.Request) (*gemini.Response, error) {
 	client := gemini.Client{
 		TrustCertificate: trustCertificate,
 	}
-	resp, err := client.Do(req)
+	resp, err := client.Do(context.Background(), req)
 	if err != nil {
 		return resp, err
 	}
 
-	switch gemini.StatusClass(resp.Status) {
+	switch resp.Status.Class() {
 	case gemini.StatusInput:
 		input, ok := getInput(resp.Meta, resp.Status == gemini.StatusSensitiveInput)
 		if !ok {
@@ -148,8 +155,8 @@ func main() {
 	defer resp.Body.Close()
 
 	// Handle response
-	if gemini.StatusClass(resp.Status) == gemini.StatusSuccess {
-		body, err := ioutil.ReadAll(resp.Body)
+	if resp.Status.Class() == gemini.StatusSuccess {
+		body, err := io.ReadAll(resp.Body)
 		if err != nil {
 			log.Fatal(err)
 		}

examples/server.go

@@ -5,9 +5,9 @@
 package main
 
 import (
-	"crypto/tls"
-	"crypto/x509/pkix"
+	"context"
 	"log"
+	"os"
 	"time"
 
 	"git.sr.ht/~adnano/go-gemini"
@@ -15,27 +15,23 @@ import (
 )
 
 func main() {
-	var server gemini.Server
-	server.ReadTimeout = 30 * time.Second
-	server.WriteTimeout = 1 * time.Minute
-	if err := server.Certificates.Load("/var/lib/gemini/certs"); err != nil {
+	certificates := &certificate.Store{}
+	certificates.Register("localhost")
+	if err := certificates.Load("/var/lib/gemini/certs"); err != nil {
 		log.Fatal(err)
 	}
-	server.GetCertificate = func(hostname string) (tls.Certificate, error) {
-		return certificate.Create(certificate.CreateOptions{
-			Subject: pkix.Name{
-				CommonName: hostname,
-			},
-			DNSNames: []string{hostname},
-			Duration: 365 * 24 * time.Hour,
-		})
+
+	mux := &gemini.ServeMux{}
+	mux.Handle("/", gemini.FileServer(os.DirFS("/var/www")))
+
+	server := &gemini.Server{
+		Handler:        mux,
+		ReadTimeout:    30 * time.Second,
+		WriteTimeout:   1 * time.Minute,
+		GetCertificate: certificates.GetCertificate,
 	}
 
-	var mux gemini.ServeMux
-	mux.Handle("/", gemini.FileServer(gemini.Dir("/var/www")))
-
-	server.Handle("localhost", &mux)
-	if err := server.ListenAndServe(); err != nil {
+	if err := server.ListenAndServe(context.Background()); err != nil {
 		log.Fatal(err)
 	}
 }

examples/stream.go

@@ -6,10 +6,9 @@ package main
 
 import (
 	"context"
-	"crypto/tls"
-	"crypto/x509/pkix"
 	"fmt"
 	"log"
+	"sync"
 	"time"
 
 	"git.sr.ht/~adnano/go-gemini"
@@ -17,30 +16,52 @@ import (
 )
 
 func main() {
-	var server gemini.Server
-	if err := server.Certificates.Load("/var/lib/gemini/certs"); err != nil {
+	certificates := &certificate.Store{}
+	certificates.Register("localhost")
+	if err := certificates.Load("/var/lib/gemini/certs"); err != nil {
 		log.Fatal(err)
 	}
-	server.GetCertificate = func(hostname string) (tls.Certificate, error) {
-		return certificate.Create(certificate.CreateOptions{
-			Subject: pkix.Name{
-				CommonName: hostname,
-			},
-			DNSNames: []string{hostname},
-			Duration: 365 * 24 * time.Hour,
-		})
-	}
 
-	server.HandleFunc("localhost", stream)
-	if err := server.ListenAndServe(); err != nil {
-		log.Fatal(err)
+	mux := &gemini.ServeMux{}
+	mux.HandleFunc("/", stream)
+
+	server := &gemini.Server{
+		Handler:        mux,
+		ReadTimeout:    30 * time.Second,
+		WriteTimeout:   1 * time.Minute,
+		GetCertificate: certificates.GetCertificate,
+	}
+
+	var shutdownOnce sync.Once
+	var wg sync.WaitGroup
+	wg.Add(1)
+	defer wg.Wait()
+	mux.HandleFunc("/shutdown", func(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
+		fmt.Fprintln(w, "Shutting down...")
+		if flusher, ok := w.(gemini.Flusher); ok {
+			flusher.Flush()
+		}
+		go shutdownOnce.Do(func() {
+			server.Shutdown(context.Background())
+			wg.Done()
+		})
+	})
+
+	if err := server.ListenAndServe(context.Background()); err != nil {
+		log.Println(err)
 	}
 }
 
 // stream writes an infinite stream to w.
-func stream(w gemini.ResponseWriter, r *gemini.Request) {
+func stream(ctx context.Context, w gemini.ResponseWriter, r *gemini.Request) {
+	flusher, ok := w.(gemini.Flusher)
+	if !ok {
+		w.WriteHeader(gemini.StatusTemporaryFailure, "Internal error")
+		return
+	}
+
 	ch := make(chan string)
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 
 	go func(ctx context.Context) {
 		for {
@@ -63,7 +84,7 @@ func stream(w gemini.ResponseWriter, r *gemini.Request) {
 			break
 		}
 		fmt.Fprintln(w, s)
-		if err := w.Flush(); err != nil {
+		if err := flusher.Flush(); err != nil {
 			cancel()
 			return
 		}

fs.go

@@ -1,10 +1,15 @@
 package gemini
 
 import (
+	"context"
+	"fmt"
 	"io"
+	"io/fs"
 	"mime"
-	"os"
+	"net/url"
 	"path"
+	"sort"
+	"strings"
 )
 
 func init() {
@@ -13,64 +18,40 @@ func init() {
 	mime.AddExtensionType(".gemini", "text/gemini")
 }
 
-// A FileSystem implements access to a collection of named files. The elements
-// in a file path are separated by slash ('/', U+002F) characters, regardless
-// of host operating system convention.
-type FileSystem interface {
-	Open(name string) (File, error)
-}
-
-// A File is returned by a FileSystem's Open method and can be served by the
-// FileServer implementation.
-//
-// The methods should behave the same as those on an *os.File.
-type File interface {
-	Stat() (os.FileInfo, error)
-	Read([]byte) (int, error)
-	Close() error
-}
-
-// A Dir implements FileSystem using the native file system restricted
-// to a specific directory tree.
-//
-// While the FileSystem.Open method takes '/'-separated paths, a Dir's string
-// value is a filename on the native file system, not a URL, so it is separated
-// by filepath.Separator, which isn't necessarily '/'.
-//
-// Note that Dir could expose sensitive files and directories. Dir will follow
-// symlinks pointing out of the directory tree, which can be especially
-// dangerous if serving from a directory in which users are able to create
-// arbitrary symlinks. Dir will also allow access to files and directories
-// starting with a period, which could expose sensitive directories like .git
-// or sensitive files like .htpasswd. To exclude files with a leading period,
-// remove the files/directories from the server or create a custom FileSystem
-// implementation.
-//
-// An empty Dir is treated as ".".
-type Dir string
-
-// Open implements FileSystem using os.Open, opening files for reading
-// rooted and relative to the directory d.
-func (d Dir) Open(name string) (File, error) {
-	return os.Open(path.Join(string(d), name))
-}
-
 // FileServer returns a handler that serves Gemini requests with the contents
 // of the provided file system.
 //
-// To use the operating system's file system implementation, use gemini.Dir:
+// To use the operating system's file system implementation, use os.DirFS:
 //
-//     gemini.FileServer(gemini.Dir("/tmp"))
-func FileServer(fsys FileSystem) Handler {
+//     gemini.FileServer(os.DirFS("/tmp"))
+func FileServer(fsys fs.FS) Handler {
 	return fileServer{fsys}
 }
 
 type fileServer struct {
-	FileSystem
+	fs.FS
 }
 
-func (fs fileServer) ServeGemini(w ResponseWriter, r *Request) {
-	ServeFile(w, fs, r.URL.Path)
+func (fs fileServer) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
+	serveFile(ctx, w, r, fs, path.Clean(r.URL.Path), true)
+}
+
+// ServeContent replies to the request using the content in the
+// provided Reader. The main benefit of ServeContent over io.Copy
+// is that it sets the MIME type of the response.
+//
+// ServeContent tries to deduce the type from name's file extension.
+// The name is otherwise unused; it is never sent in the response.
+func ServeContent(ctx context.Context, w ResponseWriter, r *Request, name string, content io.Reader) {
+	serveContent(ctx, w, name, content)
+}
+
+func serveContent(ctx context.Context, w ResponseWriter, name string, content io.Reader) {
+	// Detect mimetype from file extension
+	ext := path.Ext(name)
+	mimetype := mime.TypeByExtension(ext)
+	w.MediaType(mimetype)
+	io.Copy(w, content)
 }
 
 // ServeFile responds to the request with the contents of the named file
@@ -80,48 +61,146 @@ func (fs fileServer) ServeGemini(w ResponseWriter, r *Request) {
 // relative to the current directory and may ascend to parent directories. If
 // the provided name is constructed from user input, it should be sanitized
 // before calling ServeFile.
-func ServeFile(w ResponseWriter, fsys FileSystem, name string) {
-	f, err := openFile(fsys, name)
-	if err != nil {
-		w.Status(StatusNotFound)
+//
+// 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.
+//
+// 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
+// avoid such redirects either modify the path or use ServeContent.
+//
+// Outside of those two special cases, ServeFile does not use r.URL.Path for
+// selecting the file or directory to serve; only the file or directory
+// provided in the name argument is used.
+func ServeFile(ctx context.Context, w ResponseWriter, r *Request, fsys fs.FS, name string) {
+	if containsDotDot(r.URL.Path) {
+		// Too many programs use r.URL.Path to construct the argument to
+		// 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).
+		w.WriteHeader(StatusBadRequest, "invalid URL path")
 		return
 	}
-	// Detect mimetype
-	ext := path.Ext(name)
-	mimetype := mime.TypeByExtension(ext)
-	w.Meta(mimetype)
-	// Copy file to response writer
-	_, _ = io.Copy(w, f)
+	serveFile(ctx, w, r, fsys, name, false)
 }
 
-func openFile(fsys FileSystem, name string) (File, error) {
+func containsDotDot(v string) bool {
+	if !strings.Contains(v, "..") {
+		return false
+	}
+	for _, ent := range strings.FieldsFunc(v, isSlashRune) {
+		if ent == ".." {
+			return true
+		}
+	}
+	return false
+}
+
+func isSlashRune(r rune) bool { return r == '/' || r == '\\' }
+
+func serveFile(ctx context.Context, w ResponseWriter, r *Request, fsys fs.FS, name string, redirect bool) {
+	const indexPage = "/index.gmi"
+
+	// Redirect .../index.gmi to .../
+	if strings.HasSuffix(r.URL.Path, indexPage) {
+		w.WriteHeader(StatusPermanentRedirect, "./")
+		return
+	}
+
+	if name == "/" {
+		name = "."
+	} else {
+		name = strings.Trim(name, "/")
+	}
+
 	f, err := fsys.Open(name)
 	if err != nil {
-		return nil, err
+		w.WriteHeader(StatusNotFound, "Not found")
+		return
 	}
+	defer f.Close()
 
 	stat, err := f.Stat()
 	if err != nil {
-		return nil, err
+		w.WriteHeader(StatusTemporaryFailure, "Temporary failure")
+		return
 	}
-	if stat.Mode().IsRegular() {
-		return f, nil
+
+	// Redirect to canonical path
+	if redirect {
+		url := r.URL.Path
+		if stat.IsDir() {
+			// Add trailing slash
+			if url[len(url)-1] != '/' {
+				w.WriteHeader(StatusPermanentRedirect, path.Base(url)+"/")
+				return
+			}
+		} else {
+			// Remove trailing slash
+			if url[len(url)-1] == '/' {
+				w.WriteHeader(StatusPermanentRedirect, "../"+path.Base(url))
+				return
+			}
+		}
 	}
 
 	if stat.IsDir() {
-		// Try opening index.gmi
-		f, err := fsys.Open(path.Join(name, "index.gmi"))
-		if err != nil {
-			return nil, err
+		// Redirect if the directory name doesn't end in a slash
+		url := r.URL.Path
+		if url[len(url)-1] != '/' {
+			w.WriteHeader(StatusRedirect, path.Base(url)+"/")
+			return
 		}
-		stat, err := f.Stat()
-		if err != nil {
-			return nil, err
-		}
-		if stat.Mode().IsRegular() {
-			return f, nil
+
+		// Use contents of index.gmi if present
+		index, err := fsys.Open(path.Join(name, indexPage))
+		if err == nil {
+			defer index.Close()
+			istat, err := index.Stat()
+			if err == nil {
+				f = index
+				stat = istat
+			}
 		}
 	}
 
-	return nil, os.ErrNotExist
+	if stat.IsDir() {
+		// Failed to find index file
+		dirList(w, f)
+		return
+	}
+
+	serveContent(ctx, w, name, f)
+}
+
+func dirList(w ResponseWriter, f fs.File) {
+	var entries []fs.DirEntry
+	var err error
+	d, ok := f.(fs.ReadDirFile)
+	if ok {
+		entries, err = d.ReadDir(-1)
+	}
+	if !ok || err != nil {
+		w.WriteHeader(StatusTemporaryFailure, "Error reading directory")
+		return
+	}
+
+	sort.Slice(entries, func(i, j int) bool {
+		return entries[i].Name() < entries[j].Name()
+	})
+
+	for _, entry := range entries {
+		name := entry.Name()
+		if entry.IsDir() {
+			name += "/"
+		}
+		link := LineLink{
+			Name: name,
+			URL:  (&url.URL{Path: name}).EscapedPath(),
+		}
+		fmt.Fprintln(w, link.String())
+	}
 }

gemini.go

@@ -24,5 +24,9 @@ var (
 	// While any panic from ServeGemini aborts the response to the client,
 	// panicking with ErrAbortHandler also suppresses logging of a stack
 	// trace to the server's error log.
-	ErrAbortHandler = errors.New("net/http: abort Handler")
+	ErrAbortHandler = errors.New("gemini: abort Handler")
+
+	// 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.15
+go 1.16
 
 require golang.org/x/net v0.0.0-20210119194325-5f4716e94777

handler.go

@@ -0,0 +1,86 @@
+package gemini
+
+import (
+	"context"
+	"net/url"
+	"strings"
+)
+
+// A Handler responds to a Gemini request.
+//
+// ServeGemini should write the response header and data to the ResponseWriter
+// and then return. Returning signals that the request is finished; it is not
+// valid to use the ResponseWriter after or concurrently with the completion
+// of the ServeGemini call.
+//
+// Handlers should not modify the provided Request.
+//
+// If ServeGemini panics, the server (the caller of ServeGemini) assumes that
+// the effect of the panic was isolated to the active request. It recovers
+// the panic, logs a stack trace to the server error log, and closes the
+// network connection. To abort a handler so the client sees an interrupted
+// response but the server doesn't log an error, panic with the value
+// ErrAbortHandler.
+type Handler interface {
+	ServeGemini(context.Context, ResponseWriter, *Request)
+}
+
+// The HandlerFunc type is an adapter to allow the use of ordinary functions
+// as Gemini handlers. If f is a function with the appropriate signature,
+// HandlerFunc(f) is a Handler that calls f.
+type HandlerFunc func(context.Context, ResponseWriter, *Request)
+
+// ServeGemini calls f(ctx, w, r).
+func (f HandlerFunc) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
+	f(ctx, w, r)
+}
+
+// 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)
+}
+
+// NotFoundHandler returns a simple request handler that replies to each
+// request with a “51 Not found” reply.
+func NotFoundHandler() Handler {
+	return HandlerFunc(func(ctx context.Context, w ResponseWriter, r *Request) {
+		w.WriteHeader(StatusNotFound, "Not found")
+	})
+}
+
+// StripPrefix returns a handler that serves Gemini requests by removing the
+// given prefix from the request URL's Path (and RawPath if set) and invoking
+// the handler h. StripPrefix handles a request for a path that doesn't begin
+// with prefix by replying with a Gemini 51 not found error. The prefix must
+// match exactly: if the prefix in the request contains escaped characters the
+// reply is also a Gemini 51 not found error.
+func StripPrefix(prefix string, h Handler) Handler {
+	if prefix == "" {
+		return h
+	}
+	return HandlerFunc(func(ctx context.Context, w ResponseWriter, r *Request) {
+		p := strings.TrimPrefix(r.URL.Path, prefix)
+		rp := strings.TrimPrefix(r.URL.RawPath, prefix)
+		if len(p) < len(r.URL.Path) && (r.URL.RawPath == "" || len(rp) < len(r.URL.RawPath)) {
+			r2 := new(Request)
+			*r2 = *r
+			r2.URL = new(url.URL)
+			*r2.URL = *r.URL
+			r2.URL.Path = p
+			r2.URL.RawPath = rp
+			h.ServeGemini(ctx, w, r2)
+		} else {
+			w.WriteHeader(StatusNotFound, "Not found")
+		}
+	})
+}

mux.go

@@ -1,6 +1,8 @@
 package gemini
 
 import (
+	"context"
+	"net"
 	"net/url"
 	"path"
 	"sort"
@@ -8,12 +10,6 @@ import (
 	"sync"
 )
 
-// The following code is modified from the net/http package.
-
-// Copyright 2009 The Go Authors. All rights reserved.
-// Use of this source code is governed by a BSD-style
-// license that can be found in the LICENSE file.
-
 // ServeMux is a Gemini request multiplexer.
 // It matches the URL of each incoming request against a list of registered
 // patterns and calls the handler for the pattern that
@@ -32,6 +28,32 @@ import (
 // the pattern "/" matches all paths not matched by other registered
 // patterns, not just the URL with Path == "/".
 //
+// Patterns may also contain schemes and hostnames.
+// Wildcard patterns can be used to match multiple hostnames (e.g. "*.example.com").
+//
+// The following are examples of valid patterns, along with the scheme,
+// hostname, and path that they match.
+//
+//     Pattern                      │ Scheme │ Hostname │ Path
+//     ─────────────────────────────┼────────┼──────────┼─────────────
+//     /file                        │ gemini │ *        │ /file
+//     /directory/                  │ gemini │ *        │ /directory/*
+//     hostname/file                │ gemini │ hostname │ /file
+//     hostname/directory/          │ gemini │ hostname │ /directory/*
+//     scheme://hostname/file       │ scheme │ hostname │ /file
+//     scheme://hostname/directory/ │ scheme │ hostname │ /directory/*
+//     //hostname/file              │ *      │ hostname │ /file
+//     //hostname/directory/        │ *      │ hostname │ /directory/*
+//     scheme:///file               │ scheme │ *        │ /file
+//     scheme:///directory/         │ scheme │ *        │ /directory/*
+//     ///file                      │ *      │ *        │ /file
+//     ///directory/                │ *      │ *        │ /directory/*
+//
+// A pattern without a hostname will match any hostname.
+// If a pattern begins with "//", it will match any scheme.
+// Otherwise, a pattern with no scheme is treated as though it has a
+// scheme of "gemini".
+//
 // If a subtree has been registered and a request is received naming the
 // subtree root without its trailing slash, ServeMux redirects that
 // request to the subtree root (adding the trailing slash). This behavior can
@@ -45,13 +67,19 @@ import (
 // to an equivalent, cleaner URL.
 type ServeMux struct {
 	mu sync.RWMutex
-	m  map[string]muxEntry
-	es []muxEntry // slice of entries sorted from longest to shortest.
+	m  map[muxKey]Handler
+	es []muxEntry // slice of entries sorted from longest to shortest
+}
+
+type muxKey struct {
+	scheme string
+	host   string
+	path   string
 }
 
 type muxEntry struct {
-	r       Handler
-	pattern string
+	handler Handler
+	key     muxKey
 }
 
 // cleanPath returns the canonical path for p, eliminating . and .. elements.
@@ -78,18 +106,25 @@ func cleanPath(p string) string {
 
 // Find a handler on a handler map given a path string.
 // Most-specific (longest) pattern wins.
-func (mux *ServeMux) match(path string) Handler {
+func (mux *ServeMux) match(key muxKey) Handler {
 	// Check for exact match first.
-	v, ok := mux.m[path]
-	if ok {
-		return v.r
+	if r, ok := mux.m[key]; ok {
+		return r
+	} else if r, ok := mux.m[muxKey{"", key.host, key.path}]; ok {
+		return r
+	} else if r, ok := mux.m[muxKey{key.scheme, "", key.path}]; ok {
+		return r
+	} else if r, ok := mux.m[muxKey{"", "", key.path}]; ok {
+		return r
 	}
 
 	// Check for longest valid match.  mux.es contains all patterns
 	// that end in / sorted from longest to shortest.
 	for _, e := range mux.es {
-		if strings.HasPrefix(path, e.pattern) {
-			return e.r
+		if (e.key.scheme == "" || key.scheme == e.key.scheme) &&
+			(e.key.host == "" || key.host == e.key.host) &&
+			strings.HasPrefix(key.path, e.key.path) {
+			return e.handler
 		}
 	}
 	return nil
@@ -99,89 +134,144 @@ func (mux *ServeMux) match(path string) Handler {
 // This occurs when a handler for path + "/" was already registered, but
 // not for path itself. If the path needs appending to, it creates a new
 // URL, setting the path to u.Path + "/" and returning true to indicate so.
-func (mux *ServeMux) redirectToPathSlash(path string, u *url.URL) (*url.URL, bool) {
+func (mux *ServeMux) redirectToPathSlash(key muxKey, u *url.URL) (*url.URL, bool) {
 	mux.mu.RLock()
-	shouldRedirect := mux.shouldRedirectRLocked(path)
+	shouldRedirect := mux.shouldRedirectRLocked(key)
 	mux.mu.RUnlock()
 	if !shouldRedirect {
 		return u, false
 	}
-	path = path + "/"
-	u = &url.URL{Path: path, RawQuery: u.RawQuery}
-	return u, true
+	return u.ResolveReference(&url.URL{Path: key.path + "/"}), true
 }
 
 // shouldRedirectRLocked reports whether the given path and host should be redirected to
 // path+"/". This should happen if a handler is registered for path+"/" but
 // not path -- see comments at ServeMux.
-func (mux *ServeMux) shouldRedirectRLocked(path string) bool {
-	if _, exist := mux.m[path]; exist {
+func (mux *ServeMux) shouldRedirectRLocked(key muxKey) bool {
+	if _, exist := mux.m[key]; exist {
 		return false
 	}
 
-	n := len(path)
+	n := len(key.path)
 	if n == 0 {
 		return false
 	}
-	if _, exist := mux.m[path+"/"]; exist {
-		return path[n-1] != '/'
+	if _, exist := mux.m[muxKey{key.scheme, key.host, key.path + "/"}]; exist {
+		return key.path[n-1] != '/'
 	}
-
 	return false
 }
 
-// ServeGemini dispatches the request to the handler whose
-// pattern most closely matches the request URL.
-func (mux *ServeMux) ServeGemini(w ResponseWriter, r *Request) {
+func getWildcard(hostname string) (string, bool) {
+	if net.ParseIP(hostname) == nil {
+		split := strings.SplitN(hostname, ".", 2)
+		if len(split) == 2 {
+			return "*." + split[1], true
+		}
+	}
+	return "", false
+}
+
+// Handler returns the handler to use for the given request, consulting
+// r.URL.Scheme, r.URL.Host, and r.URL.Path. It always returns a non-nil handler. If
+// the path is not in its canonical form, the handler will be an
+// internally-generated handler that redirects to the canonical path. If the
+// host contains a port, it is ignored when matching handlers.
+func (mux *ServeMux) Handler(r *Request) Handler {
+	scheme := r.URL.Scheme
+	host := r.URL.Hostname()
 	path := cleanPath(r.URL.Path)
 
 	// If the given path is /tree and its handler is not registered,
 	// redirect for /tree/.
-	if u, ok := mux.redirectToPathSlash(path, r.URL); ok {
-		w.Header(StatusRedirect, u.String())
-		return
+	if u, ok := mux.redirectToPathSlash(muxKey{scheme, host, path}, r.URL); ok {
+		return StatusHandler(StatusPermanentRedirect, u.String())
 	}
 
 	if path != r.URL.Path {
 		u := *r.URL
 		u.Path = path
-		w.Header(StatusRedirect, u.String())
-		return
+		return StatusHandler(StatusPermanentRedirect, u.String())
 	}
 
 	mux.mu.RLock()
 	defer mux.mu.RUnlock()
 
-	resp := mux.match(path)
-	if resp == nil {
-		w.Status(StatusNotFound)
-		return
+	h := mux.match(muxKey{scheme, host, path})
+	if h == nil {
+		// Try wildcard
+		if wildcard, ok := getWildcard(host); ok {
+			h = mux.match(muxKey{scheme, wildcard, path})
+		}
 	}
-	resp.ServeGemini(w, r)
+	if h == nil {
+		h = NotFoundHandler()
+	}
+	return h
+}
+
+// ServeGemini dispatches the request to the handler whose
+// pattern most closely matches the request URL.
+func (mux *ServeMux) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
+	h := mux.Handler(r)
+	h.ServeGemini(ctx, w, r)
 }
 
 // Handle registers the handler for the given pattern.
 // If a handler already exists for pattern, Handle panics.
 func (mux *ServeMux) Handle(pattern string, handler Handler) {
-	mux.mu.Lock()
-	defer mux.mu.Unlock()
-
 	if pattern == "" {
 		panic("gemini: invalid pattern")
 	}
 	if handler == nil {
 		panic("gemini: nil handler")
 	}
-	if _, exist := mux.m[pattern]; exist {
+
+	mux.mu.Lock()
+	defer mux.mu.Unlock()
+
+	var key muxKey
+	if strings.HasPrefix(pattern, "//") {
+		// match any scheme
+		key.scheme = ""
+		pattern = pattern[2:]
+	} else {
+		// extract scheme
+		cut := strings.Index(pattern, "://")
+		if cut == -1 {
+			// default scheme of gemini
+			key.scheme = "gemini"
+		} else {
+			key.scheme = pattern[:cut]
+			pattern = pattern[cut+3:]
+		}
+	}
+
+	// extract hostname and path
+	cut := strings.Index(pattern, "/")
+	if cut == -1 {
+		key.host = pattern
+		key.path = "/"
+	} else {
+		key.host = pattern[:cut]
+		key.path = pattern[cut:]
+	}
+
+	// strip port from hostname
+	if hostname, _, err := net.SplitHostPort(key.host); err == nil {
+		key.host = hostname
+	}
+
+	if _, exist := mux.m[key]; exist {
 		panic("gemini: multiple registrations for " + pattern)
 	}
 
 	if mux.m == nil {
-		mux.m = make(map[string]muxEntry)
+		mux.m = make(map[muxKey]Handler)
 	}
-	e := muxEntry{handler, pattern}
-	mux.m[pattern] = e
-	if pattern[len(pattern)-1] == '/' {
+	mux.m[key] = handler
+	e := muxEntry{handler, key}
+	if key.path[len(key.path)-1] == '/' {
 		mux.es = appendSorted(mux.es, e)
 	}
 }
@@ -189,7 +279,9 @@ func (mux *ServeMux) Handle(pattern string, handler Handler) {
 func appendSorted(es []muxEntry, e muxEntry) []muxEntry {
 	n := len(es)
 	i := sort.Search(n, func(i int) bool {
-		return len(es[i].pattern) < len(e.pattern)
+		return len(es[i].key.scheme) < len(e.key.scheme) ||
+			len(es[i].key.host) < len(es[i].key.host) ||
+			len(es[i].key.path) < len(e.key.path)
 	})
 	if i == n {
 		return append(es, e)
@@ -202,7 +294,7 @@ func appendSorted(es []muxEntry, e muxEntry) []muxEntry {
 }
 
 // HandleFunc registers the handler function for the given pattern.
-func (mux *ServeMux) HandleFunc(pattern string, handler func(ResponseWriter, *Request)) {
+func (mux *ServeMux) HandleFunc(pattern string, handler func(context.Context, ResponseWriter, *Request)) {
 	if handler == nil {
 		panic("gemini: nil handler")
 	}

mux_test.go

@@ -0,0 +1,315 @@
+package gemini
+
+import (
+	"context"
+	"net/url"
+	"testing"
+)
+
+type nopHandler struct{}
+
+func (*nopHandler) ServeGemini(context.Context, ResponseWriter, *Request) {}
+
+func TestServeMuxMatch(t *testing.T) {
+	type Match struct {
+		URL string
+		Ok  bool
+	}
+
+	tests := []struct {
+		Pattern string
+		Matches []Match
+	}{
+		{
+			// scheme: gemini, hostname: *, path: /*
+			Pattern: "/",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", true},
+				{"gemini://example.com/path.gmi", true},
+				{"gemini://example.com/path/", true},
+				{"gemini://example.org/path", true},
+				{"http://example.com/path", false},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: gemini, hostname: *, path: /path
+			Pattern: "/path",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", false},
+				{"gemini://example.com/path.gmi", false},
+				{"gemini://example.com/path/", false},
+				{"gemini://example.org/path", true},
+				{"http://example.com/path", false},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: gemini, hostname: *, path: /subtree/*
+			Pattern: "/subtree/",
+			Matches: []Match{
+				{"gemini://example.com/subtree/", true},
+				{"gemini://example.com/subtree/nested/", true},
+				{"gemini://example.com/subtree/nested/file", true},
+				{"gemini://example.org/subtree/", true},
+				{"gemini://example.org/subtree/nested/", true},
+				{"gemini://example.org/subtree/nested/file", true},
+				{"gemini://example.com/subtree", false},
+				{"gemini://www.example.com/subtree/", true},
+				{"http://example.com/subtree/", false},
+			},
+		},
+		{
+			// scheme: gemini, hostname: example.com, path: /*
+			Pattern: "example.com",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", true},
+				{"gemini://example.com/path.gmi", true},
+				{"gemini://example.com/path/", true},
+				{"gemini://example.org/path", false},
+				{"http://example.com/path", false},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: gemini, hostname: example.com, path: /path
+			Pattern: "example.com/path",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", false},
+				{"gemini://example.com/path.gmi", false},
+				{"gemini://example.com/path/", false},
+				{"gemini://example.org/path", false},
+				{"http://example.com/path", false},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: gemini, hostname: example.com, path: /subtree/*
+			Pattern: "example.com/subtree/",
+			Matches: []Match{
+				{"gemini://example.com/subtree/", true},
+				{"gemini://example.com/subtree/nested/", true},
+				{"gemini://example.com/subtree/nested/file", true},
+				{"gemini://example.org/subtree/", false},
+				{"gemini://example.org/subtree/nested/", false},
+				{"gemini://example.org/subtree/nested/file", false},
+				{"gemini://example.com/subtree", false},
+				{"gemini://www.example.com/subtree/", false},
+				{"http://example.com/subtree/", false},
+			},
+		},
+		{
+			// scheme: http, hostname: example.com, path: /*
+			Pattern: "http://example.com",
+			Matches: []Match{
+				{"http://example.com/path", true},
+				{"http://example.com/", true},
+				{"http://example.com/path.gmi", true},
+				{"http://example.com/path/", true},
+				{"http://example.org/path", false},
+				{"gemini://example.com/path", false},
+				{"gemini://example.org/path", false},
+			},
+		},
+		{
+			// scheme: http, hostname: example.com, path: /path
+			Pattern: "http://example.com/path",
+			Matches: []Match{
+				{"http://example.com/path", true},
+				{"http://example.com/", false},
+				{"http://example.com/path.gmi", false},
+				{"http://example.com/path/", false},
+				{"http://example.org/path", false},
+				{"gemini://example.com/path", false},
+				{"gemini://example.org/path", false},
+			},
+		},
+		{
+			// scheme: http, hostname: example.com, path: /subtree/*
+			Pattern: "http://example.com/subtree/",
+			Matches: []Match{
+				{"http://example.com/subtree/", true},
+				{"http://example.com/subtree/nested/", true},
+				{"http://example.com/subtree/nested/file", true},
+				{"http://example.org/subtree/", false},
+				{"http://example.org/subtree/nested/", false},
+				{"http://example.org/subtree/nested/file", false},
+				{"http://example.com/subtree", false},
+				{"http://www.example.com/subtree/", false},
+				{"gemini://example.com/subtree/", false},
+			},
+		},
+		{
+			// scheme: *, hostname: example.com, path: /*
+			Pattern: "//example.com",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", true},
+				{"gemini://example.com/path.gmi", true},
+				{"gemini://example.com/path/", true},
+				{"gemini://example.org/path", false},
+				{"http://example.com/path", true},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: *, hostname: example.com, path: /path
+			Pattern: "//example.com/path",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", false},
+				{"gemini://example.com/path.gmi", false},
+				{"gemini://example.com/path/", false},
+				{"gemini://example.org/path", false},
+				{"http://example.com/path", true},
+				{"http://example.org/path", false},
+			},
+		},
+		{
+			// scheme: *, hostname: example.com, path: /subtree/*
+			Pattern: "//example.com/subtree/",
+			Matches: []Match{
+				{"gemini://example.com/subtree/", true},
+				{"gemini://example.com/subtree/nested/", true},
+				{"gemini://example.com/subtree/nested/file", true},
+				{"gemini://example.org/subtree/", false},
+				{"gemini://example.org/subtree/nested/", false},
+				{"gemini://example.org/subtree/nested/file", false},
+				{"gemini://example.com/subtree", false},
+				{"gemini://www.example.com/subtree/", false},
+				{"http://example.com/subtree/", true},
+			},
+		},
+		{
+			// scheme: http, hostname: *, path: /*
+			Pattern: "http://",
+			Matches: []Match{
+				{"http://example.com/path", true},
+				{"http://example.com/", true},
+				{"http://example.com/path.gmi", true},
+				{"http://example.com/path/", true},
+				{"http://example.org/path", true},
+				{"gemini://example.com/path", false},
+				{"gemini://example.org/path", false},
+			},
+		},
+		{
+			// scheme: http, hostname: *, path: /path
+			Pattern: "http:///path",
+			Matches: []Match{
+				{"http://example.com/path", true},
+				{"http://example.com/", false},
+				{"http://example.com/path.gmi", false},
+				{"http://example.com/path/", false},
+				{"http://example.org/path", true},
+				{"gemini://example.com/path", false},
+				{"gemini://example.org/path", false},
+			},
+		},
+		{
+			// scheme: http, hostname: *, path: /subtree/*
+			Pattern: "http:///subtree/",
+			Matches: []Match{
+				{"http://example.com/subtree/", true},
+				{"http://example.com/subtree/nested/", true},
+				{"http://example.com/subtree/nested/file", true},
+				{"http://example.org/subtree/", true},
+				{"http://example.org/subtree/nested/", true},
+				{"http://example.org/subtree/nested/file", true},
+				{"http://example.com/subtree", false},
+				{"http://www.example.com/subtree/", true},
+				{"gemini://example.com/subtree/", false},
+			},
+		},
+		{
+			// scheme: *, hostname: *, path: /*
+			Pattern: "//",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", true},
+				{"gemini://example.com/path.gmi", true},
+				{"gemini://example.com/path/", true},
+				{"gemini://example.org/path", true},
+				{"http://example.com/path", true},
+				{"http://example.org/path", true},
+			},
+		},
+		{
+			// scheme: *, hostname: *, path: /path
+			Pattern: "///path",
+			Matches: []Match{
+				{"gemini://example.com/path", true},
+				{"gemini://example.com/", false},
+				{"gemini://example.com/path.gmi", false},
+				{"gemini://example.com/path/", false},
+				{"gemini://example.org/path", true},
+				{"http://example.com/path", true},
+				{"http://example.org/path", true},
+			},
+		},
+		{
+			// scheme: *, hostname: *, path: /subtree/*
+			Pattern: "///subtree/",
+			Matches: []Match{
+				{"gemini://example.com/subtree/", true},
+				{"gemini://example.com/subtree/nested/", true},
+				{"gemini://example.com/subtree/nested/file", true},
+				{"gemini://example.org/subtree/", true},
+				{"gemini://example.org/subtree/nested/", true},
+				{"gemini://example.org/subtree/nested/file", true},
+				{"gemini://example.com/subtree", false},
+				{"gemini://www.example.com/subtree/", true},
+				{"http://example.com/subtree/", true},
+			},
+		},
+		{
+			// scheme: gemini, hostname: *.example.com, path: /*
+			Pattern: "*.example.com",
+			Matches: []Match{
+				{"gemini://mail.example.com/", true},
+				{"gemini://www.example.com/index.gmi", true},
+				{"gemini://example.com/", false},
+				{"gemini://a.b.example.com/", false},
+				{"http://www.example.com/", false},
+			},
+		},
+		{
+			// scheme: http, hostname: *.example.com, path: /*
+			Pattern: "http://*.example.com",
+			Matches: []Match{
+				{"http://mail.example.com/", true},
+				{"http://www.example.com/index.gmi", true},
+				{"http://example.com/", false},
+				{"http://a.b.example.com/", false},
+				{"gemini://www.example.com/", false},
+			},
+		},
+	}
+
+	for i, test := range tests {
+		h := &nopHandler{}
+		var mux ServeMux
+		mux.Handle(test.Pattern, h)
+
+		for _, match := range tests[i].Matches {
+			u, err := url.Parse(match.URL)
+			if err != nil {
+				panic(err)
+			}
+			got := mux.Handler(&Request{URL: u})
+			if match.Ok {
+				if h != got {
+					t.Errorf("expected %s to match %s", test.Pattern, match.URL)
+				}
+			} else {
+				if h == got {
+					t.Errorf("expected %s not to match %s", test.Pattern, match.URL)
+				}
+			}
+		}
+	}
+}

request.go

@@ -2,7 +2,6 @@ package gemini
 
 import (
 	"bufio"
-	"context"
 	"crypto/tls"
 	"io"
 	"net"
@@ -19,12 +18,14 @@ type Request struct {
 	URL *url.URL
 
 	// For client requests, Host optionally specifies the server to
-	// connect to. It must be of the form "host:port".
+	// connect to. It may be of the form "host" or "host:port".
 	// If empty, the value of URL.Host is used.
 	// 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.
-	// This field is ignored by the Gemini server.
+	//
+	// For server requests, Host specifies the host on which the URL
+	// is sought.
 	Host string
 
 	// For client requests, Certificate optionally specifies the
@@ -34,10 +35,7 @@ type Request struct {
 
 	// 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 and
-	// has no defined format. The Gemini server in this package
-	// sets RemoteAddr to an "IP:port" address before invoking a
-	// handler.
+	// logging. This field is not filled in by ReadRequest.
 	// This field is ignored by the Gemini client.
 	RemoteAddr net.Addr
 
@@ -49,14 +47,6 @@ type Request struct {
 	// otherwise it leaves the field nil.
 	// This field is ignored by the Gemini client.
 	TLS *tls.ConnectionState
-
-	// Context specifies the context to use for outgoing requests.
-	// 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.
-	// If Context is nil, the background context will be used.
-	// This field is ignored by the Gemini server.
-	Context context.Context
 }
 
 // NewRequest returns a new request.
@@ -107,16 +97,17 @@ func ReadRequest(r io.Reader) (*Request, error) {
 
 // Write writes a Gemini request in wire format.
 // This method consults the request URL only.
-func (r *Request) Write(w *bufio.Writer) error {
+func (r *Request) Write(w io.Writer) error {
+	bw := bufio.NewWriterSize(w, 1026)
 	url := r.URL.String()
 	if len(url) > 1024 {
 		return ErrInvalidRequest
 	}
-	if _, err := w.WriteString(url); err != nil {
+	if _, err := bw.WriteString(url); err != nil {
 		return err
 	}
-	if _, err := w.Write(crlf); err != nil {
+	if _, err := bw.Write(crlf); err != nil {
 		return err
 	}
-	return nil
+	return bw.Flush()
 }

response.go

@@ -3,10 +3,14 @@ package gemini
 import (
 	"bufio"
 	"crypto/tls"
+	"fmt"
 	"io"
 	"strconv"
 )
 
+// The default media type for responses.
+const defaultMediaType = "text/gemini; charset=utf-8"
+
 // Response represents the response from a Gemini request.
 //
 // The Client returns Responses from servers once the response
@@ -14,7 +18,7 @@ import (
 // as the Body field is read.
 type Response struct {
 	// Status contains the response status code.
-	Status int
+	Status Status
 
 	// Meta contains more information related to the response status.
 	// For successful responses, Meta should contain the media type of the response.
@@ -53,7 +57,7 @@ func ReadResponse(rc io.ReadCloser) (*Response, error) {
 	if err != nil {
 		return nil, ErrInvalidResponse
 	}
-	resp.Status = status
+	resp.Status = Status(status)
 
 	// Read one space
 	if b, err := br.ReadByte(); err != nil {
@@ -73,9 +77,9 @@ func ReadResponse(rc io.ReadCloser) (*Response, error) {
 	if len(meta) > 1024 {
 		return nil, ErrInvalidResponse
 	}
-	// Default mime type of text/gemini; charset=utf-8
-	if StatusClass(status) == StatusSuccess && meta == "" {
-		meta = "text/gemini; charset=utf-8"
+	if resp.Status.Class() == StatusSuccess && meta == "" {
+		// Use default media type
+		meta = defaultMediaType
 	}
 	resp.Meta = meta
 
@@ -86,7 +90,7 @@ func ReadResponse(rc io.ReadCloser) (*Response, error) {
 		return nil, ErrInvalidResponse
 	}
 
-	if StatusClass(status) == StatusSuccess {
+	if resp.Status.Class() == StatusSuccess {
 		resp.Body = newReadCloserBody(br, rc)
 	} else {
 		resp.Body = nopReadCloser{}
@@ -132,70 +136,100 @@ func (b *readCloserBody) Read(p []byte) (n int, err error) {
 	return b.ReadCloser.Read(p)
 }
 
-// A ResponseWriter interface is used by a Gemini handler
-// to construct a Gemini response.
+// 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
+		}
+	}
+	return nil
+}
+
+// A ResponseWriter interface is used by a Gemini handler to construct
+// a Gemini response.
+//
+// A ResponseWriter may not be used after the Handler.ServeGemini method
+// has returned.
 type ResponseWriter interface {
-	// Header sets the response header.
-	Header(status int, meta string)
-
-	// Status sets the response status code.
-	// It also sets the response meta to Meta(status).
-	Status(status int)
-
-	// Meta sets the response meta.
+	// MediaType sets the media type that will be sent by Write for a
+	// successful response. If no media type is set, a default of
+	// "text/gemini; charset=utf-8" will be used.
 	//
-	// For successful responses, meta should contain the media type of the response.
-	// For failure responses, meta should contain a short description of the failure.
-	// The response meta should not be greater than 1024 bytes.
-	Meta(meta string)
+	// Setting the media type after a call to Write or WriteHeader has
+	// no effect.
+	MediaType(string)
 
-	// Write writes data to the connection as part of the response body.
-	// If the response status does not allow for a response body, Write returns
-	// ErrBodyNotAllowed.
+	// Write writes the data to the connection as part of a Gemini response.
 	//
-	// Write writes the response header if it has not already been written.
-	// It writes a successful status code if one is not set.
+	// If WriteHeader has not yet been called, Write calls WriteHeader with
+	// StatusSuccess and the media type set in MediaType before writing the data.
+	// If no media type was set, Write uses a default media type of
+	// "text/gemini; charset=utf-8".
 	Write([]byte) (int, error)
 
-	// Flush writes any buffered data to the underlying io.Writer.
+	// WriteHeader sends a Gemini response header with the provided
+	// status code and meta.
 	//
-	// Flush writes the response header if it has not already been written.
-	// It writes a failure status code if one is not set.
+	// If WriteHeader is not called explicitly, the first call to Write
+	// will trigger an implicit call to WriteHeader with a successful
+	// status code and the media type set in MediaType.
+	//
+	// The provided code must be a valid Gemini status code.
+	// The provided meta must not be longer than 1024 bytes.
+	// Only one header may be written.
+	WriteHeader(status Status, meta string)
+}
+
+// The Flusher interface is implemented by ResponseWriters that allow a
+// Gemini handler to flush buffered data to the client.
+//
+// The default Gemini ResponseWriter implementation supports Flusher,
+// but ResponseWriter wrappers may not. Handlers should always test
+// for this ability at runtime.
+type Flusher interface {
+	// Flush sends any buffered data to the client.
 	Flush() error
 }
 
 type responseWriter struct {
 	b           *bufio.Writer
-	status      int
-	meta        string
+	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),
 	}
 }
 
-func (w *responseWriter) Header(status int, meta string) {
-	w.status = status
-	w.meta = meta
-}
-
-func (w *responseWriter) Status(status int) {
-	w.status = status
-	w.meta = Meta(status)
-}
-
-func (w *responseWriter) Meta(meta string) {
-	w.meta = meta
+func (w *responseWriter) MediaType(mediatype string) {
+	w.mediatype = mediatype
 }
 
 func (w *responseWriter) Write(b []byte) (int, error) {
 	if !w.wroteHeader {
-		w.writeHeader(StatusSuccess)
+		meta := w.mediatype
+		if meta == "" {
+			// Use default media type
+			meta = defaultMediaType
+		}
+		w.WriteHeader(StatusSuccess, meta)
 	}
 	if !w.bodyAllowed {
 		return 0, ErrBodyNotAllowed
@@ -203,22 +237,16 @@ func (w *responseWriter) Write(b []byte) (int, error) {
 	return w.b.Write(b)
 }
 
-func (w *responseWriter) writeHeader(defaultStatus int) {
-	status := w.status
-	if status == 0 {
-		status = defaultStatus
+func (w *responseWriter) WriteHeader(status Status, meta string) {
+	if w.wroteHeader {
+		return
 	}
 
-	meta := w.meta
-	if StatusClass(status) == StatusSuccess {
+	if status.Class() == StatusSuccess {
 		w.bodyAllowed = true
-
-		if meta == "" {
-			meta = "text/gemini"
-		}
 	}
 
-	w.b.WriteString(strconv.Itoa(status))
+	w.b.WriteString(strconv.Itoa(int(status)))
 	w.b.WriteByte(' ')
 	w.b.WriteString(meta)
 	w.b.Write(crlf)
@@ -227,7 +255,7 @@ func (w *responseWriter) writeHeader(defaultStatus int) {
 
 func (w *responseWriter) Flush() error {
 	if !w.wroteHeader {
-		w.writeHeader(StatusTemporaryFailure)
+		w.WriteHeader(StatusTemporaryFailure, "Temporary failure")
 	}
 	// Write errors from writeHeader will be returned here.
 	return w.b.Flush()

response_test.go

@@ -2,18 +2,18 @@ package gemini
 
 import (
 	"io"
-	"io/ioutil"
 	"strings"
 	"testing"
 )
 
-func TestReadResponse(t *testing.T) {
+func TestReadWriteResponse(t *testing.T) {
 	tests := []struct {
-		Raw    string
-		Status int
-		Meta   string
-		Body   string
-		Err    error
+		Raw       string
+		Status    Status
+		Meta      string
+		Body      string
+		Err       error
+		SkipWrite bool
 	}{
 		{
 			Raw:    "20 text/gemini\r\nHello, world!\nWelcome to my capsule.",
@@ -32,9 +32,10 @@ func TestReadResponse(t *testing.T) {
 			Meta:   "/redirect",
 		},
 		{
-			Raw:    "31 /redirect\r\nThis body is ignored.",
-			Status: 31,
-			Meta:   "/redirect",
+			Raw:       "31 /redirect\r\nThis body is ignored.",
+			Status:    31,
+			Meta:      "/redirect",
+			SkipWrite: true, // skip write test since result won't match Raw
 		},
 		{
 			Raw:    "99 Unknown status code\r\n",
@@ -81,7 +82,7 @@ func TestReadResponse(t *testing.T) {
 
 	for _, test := range tests {
 		t.Logf("%#v", test.Raw)
-		resp, err := ReadResponse(ioutil.NopCloser(strings.NewReader(test.Raw)))
+		resp, err := ReadResponse(io.NopCloser(strings.NewReader(test.Raw)))
 		if err != test.Err {
 			t.Errorf("expected err = %v, got %v", test.Err, err)
 		}
@@ -95,10 +96,32 @@ func TestReadResponse(t *testing.T) {
 		if resp.Meta != test.Meta {
 			t.Errorf("expected meta = %s, got %s", test.Meta, resp.Meta)
 		}
-		b, _ := ioutil.ReadAll(resp.Body)
+		b, _ := io.ReadAll(resp.Body)
 		body := string(b)
 		if body != test.Body {
 			t.Errorf("expected body = %#v, got %#v", test.Body, body)
 		}
 	}
+
+	for _, test := range tests {
+		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 {
+			t.Error(err)
+			continue
+		}
+
+		got := b.String()
+		if got != test.Raw {
+			t.Errorf("expected %#v, got %#v", test.Raw, got)
+		}
+	}
 }

server.go

@@ -7,12 +7,9 @@ import (
 	"log"
 	"net"
 	"runtime"
-	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
-
-	"git.sr.ht/~adnano/go-gemini/certificate"
 )
 
 // A Server defines parameters for running a Gemini server. The zero value for
@@ -23,6 +20,9 @@ type Server struct {
 	// See net.Dial for details of the address format.
 	Addr string
 
+	// The Handler to invoke.
+	Handler Handler
+
 	// ReadTimeout is the maximum duration for reading the entire
 	// request.
 	//
@@ -35,77 +35,108 @@ type Server struct {
 	// A WriteTimeout of zero means no timeout.
 	WriteTimeout time.Duration
 
-	// Certificates contains one or more certificates to present to the
-	// other side of the connection.
-	Certificates certificate.Dir
-
-	// GetCertificate, if not nil, will be called to retrieve a new certificate
-	// if the current one is expired or missing.
-	GetCertificate func(hostname string) (tls.Certificate, error)
+	// GetCertificate returns a TLS certificate based on the given
+	// hostname.
+	//
+	// If GetCertificate is nil or returns nil, then no certificate
+	// will be used and the connection will be aborted.
+	GetCertificate func(hostname string) (*tls.Certificate, error)
 
 	// ErrorLog specifies an optional logger for errors accepting connections,
 	// unexpected behavior from handlers, and underlying file system errors.
 	// If nil, logging is done via the log package's standard logger.
 	ErrorLog *log.Logger
 
-	// registered handlers
-	handlers map[handlerKey]Handler
-	hosts    map[string]bool
-	hmu      sync.Mutex
-
-	listeners map[*net.Listener]struct{}
-	conns     map[*net.Conn]struct{}
-	done      int32
+	listeners map[*net.Listener]context.CancelFunc
+	conns     map[*net.Conn]context.CancelFunc
+	doneChan  chan struct{}
+	closed    int32
 	mu        sync.Mutex
 }
 
-type handlerKey struct {
-	scheme   string
-	hostname string
+// done returns a channel that's closed when the server has finished closing.
+func (srv *Server) done() chan struct{} {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+	return srv.doneLocked()
 }
 
-// Handle registers the handler for the given pattern.
-// If a handler already exists for pattern, Handle panics.
+func (srv *Server) doneLocked() chan struct{} {
+	if srv.doneChan == nil {
+		srv.doneChan = make(chan struct{})
+	}
+	return srv.doneChan
+}
+
+// tryFinishShutdown closes srv.done() if there are no active listeners or requests.
+func (srv *Server) tryFinishShutdown() {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+	if len(srv.listeners) == 0 && len(srv.conns) == 0 {
+		done := srv.doneLocked()
+		select {
+		case <-done:
+		default:
+			close(done)
+		}
+	}
+}
+
+// Close immediately closes all active net.Listeners and connections.
+// For a graceful shutdown, use Shutdown.
+func (srv *Server) Close() error {
+	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
+		return ErrServerClosed
+	}
+
+	// Close active listeners and connections.
+	srv.mu.Lock()
+	for _, cancel := range srv.listeners {
+		cancel()
+	}
+	for _, cancel := range srv.conns {
+		cancel()
+	}
+	srv.mu.Unlock()
+
+	select {
+	case <-srv.done():
+		return nil
+	}
+}
+
+// Shutdown gracefully shuts down the server without interrupting any
+// active connections. Shutdown works by first closing all open
+// listeners and then waiting indefinitely for connections
+// to close and then shut down.
+// If the provided context expires before the shutdown is complete,
+// Shutdown returns the context's error.
 //
-// The pattern must be in the form of "hostname" or "scheme://hostname".
-// If no scheme is specified, a scheme of "gemini://" is implied.
-// Wildcard patterns are supported (e.g. "*.example.com").
-// To handle any hostname, use the wildcard pattern "*".
-func (srv *Server) Handle(pattern string, handler Handler) {
-	srv.hmu.Lock()
-	defer srv.hmu.Unlock()
-
-	if pattern == "" {
-		panic("gemini: invalid pattern")
-	}
-	if handler == nil {
-		panic("gemini: nil handler")
-	}
-	if srv.handlers == nil {
-		srv.handlers = map[handlerKey]Handler{}
-		srv.hosts = map[string]bool{}
+// When Shutdown is called, Serve and ListenAndServer immediately
+// return ErrServerClosed. Make sure the program doesn't exit and
+// waits instead for Shutdown to return.
+//
+// Once Shutdown has been called on a server, it may not be reused;
+// future calls to methods such as Serve will return ErrServerClosed.
+func (srv *Server) Shutdown(ctx context.Context) error {
+	if !atomic.CompareAndSwapInt32(&srv.closed, 0, 1) {
+		return ErrServerClosed
 	}
 
-	split := strings.SplitN(pattern, "://", 2)
-	var key handlerKey
-	if len(split) == 2 {
-		key.scheme = split[0]
-		key.hostname = split[1]
-	} else {
-		key.scheme = "gemini"
-		key.hostname = split[0]
+	// Close active listeners.
+	srv.mu.Lock()
+	for _, cancel := range srv.listeners {
+		cancel()
 	}
+	srv.mu.Unlock()
 
-	if _, ok := srv.handlers[key]; ok {
-		panic("gemini: multiple registrations for " + pattern)
+	// Wait for active connections to finish.
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-srv.done():
+		return nil
 	}
-	srv.handlers[key] = handler
-	srv.hosts[key.hostname] = true
-}
-
-// HandleFunc registers the handler function for the given pattern.
-func (srv *Server) HandleFunc(pattern string, handler func(ResponseWriter, *Request)) {
-	srv.Handle(pattern, HandlerFunc(handler))
 }
 
 // ListenAndServe listens for requests at the server's configured address.
@@ -116,8 +147,8 @@ func (srv *Server) HandleFunc(pattern string, handler func(ResponseWriter, *Requ
 //
 // ListenAndServe always returns a non-nil error. After Shutdown or Close, the
 // returned error is ErrServerClosed.
-func (srv *Server) ListenAndServe() error {
-	if atomic.LoadInt32(&srv.done) == 1 {
+func (srv *Server) ListenAndServe(ctx context.Context) error {
+	if atomic.LoadInt32(&srv.closed) == 1 {
 		return ErrServerClosed
 	}
 
@@ -126,26 +157,33 @@ func (srv *Server) ListenAndServe() error {
 		addr = ":1965"
 	}
 
-	ln, err := net.Listen("tcp", addr)
+	l, err := net.Listen("tcp", addr)
 	if err != nil {
 		return err
 	}
-	defer ln.Close()
 
-	return srv.Serve(tls.NewListener(ln, &tls.Config{
+	l = tls.NewListener(l, &tls.Config{
 		ClientAuth:     tls.RequestClientCert,
 		MinVersion:     tls.VersionTLS12,
 		GetCertificate: srv.getCertificate,
-	}))
+	})
+	return srv.Serve(ctx, l)
 }
 
-func (srv *Server) trackListener(l *net.Listener) {
+func (srv *Server) getCertificate(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
+	if srv.GetCertificate == nil {
+		return nil, errors.New("gemini: GetCertificate is nil")
+	}
+	return srv.GetCertificate(h.ServerName)
+}
+
+func (srv *Server) trackListener(l *net.Listener, cancel context.CancelFunc) {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
 	if srv.listeners == nil {
-		srv.listeners = make(map[*net.Listener]struct{})
+		srv.listeners = make(map[*net.Listener]context.CancelFunc)
 	}
-	srv.listeners[l] = struct{}{}
+	srv.listeners[l] = cancel
 }
 
 func (srv *Server) deleteListener(l *net.Listener) {
@@ -160,24 +198,52 @@ func (srv *Server) deleteListener(l *net.Listener) {
 //
 // Serve always returns a non-nil error and closes l. After Shutdown or Close,
 // the returned error is ErrServerClosed.
-func (srv *Server) Serve(l net.Listener) error {
+func (srv *Server) Serve(ctx context.Context, l net.Listener) error {
 	defer l.Close()
 
-	srv.trackListener(&l)
-	defer srv.deleteListener(&l)
-
-	if atomic.LoadInt32(&srv.done) == 1 {
+	if atomic.LoadInt32(&srv.closed) == 1 {
 		return ErrServerClosed
 	}
 
-	var tempDelay time.Duration // how long to sleep on accept failure
+	lnctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	srv.trackListener(&l, cancel)
+	defer srv.tryFinishShutdown()
+	defer srv.deleteListener(&l)
+
+	errch := make(chan error, 1)
+	go func() {
+		errch <- srv.serve(ctx, l)
+	}()
+
+	select {
+	case <-lnctx.Done():
+		if atomic.LoadInt32(&srv.closed) == 1 {
+			return ErrServerClosed
+		}
+		return lnctx.Err()
+	case err := <-errch:
+		return err
+	}
+}
+
+func (srv *Server) serve(ctx context.Context, l net.Listener) error {
+	// how long to sleep on accept failure
+	var tempDelay time.Duration
 
 	for {
 		rw, err := l.Accept()
 		if err != nil {
-			if atomic.LoadInt32(&srv.done) == 1 {
-				return ErrServerClosed
+			select {
+			case <-ctx.Done():
+				if atomic.LoadInt32(&srv.closed) == 1 {
+					return ErrServerClosed
+				}
+				return ctx.Err()
+			default:
 			}
+
 			// If this is a temporary error, sleep
 			if ne, ok := err.(net.Error); ok && ne.Temporary() {
 				if tempDelay == 0 {
@@ -193,158 +259,21 @@ func (srv *Server) Serve(l net.Listener) error {
 				continue
 			}
 
-			// Otherwise, return the error
 			return err
 		}
 
 		tempDelay = 0
-		go srv.respond(rw)
+		go srv.serveConn(ctx, rw)
 	}
 }
 
-func (srv *Server) closeListenersLocked() error {
-	var err error
-	for ln := range srv.listeners {
-		if cerr := (*ln).Close(); cerr != nil && err == nil {
-			err = cerr
-		}
-		delete(srv.listeners, ln)
-	}
-	return err
-}
-
-// Close immediately closes all active net.Listeners and connections.
-// For a graceful shutdown, use Shutdown.
-//
-// Close returns any error returned from closing the Server's
-// underlying Listener(s).
-func (srv *Server) Close() error {
-	srv.mu.Lock()
-	defer srv.mu.Unlock()
-	if !atomic.CompareAndSwapInt32(&srv.done, 0, 1) {
-		return ErrServerClosed
-	}
-	err := srv.closeListenersLocked()
-
-	// Close active connections
-	for conn := range srv.conns {
-		(*conn).Close()
-		delete(srv.conns, conn)
-	}
-	return err
-}
-
-func (srv *Server) numConns() int {
-	srv.mu.Lock()
-	defer srv.mu.Unlock()
-	return len(srv.conns)
-}
-
-// shutdownPollInterval is how often we poll for quiescence
-// during Server.Shutdown. This is lower during tests, to
-// speed up tests.
-// Ideally we could find a solution that doesn't involve polling,
-// but which also doesn't have a high runtime cost (and doesn't
-// involve any contentious mutexes), but that is left as an
-// exercise for the reader.
-var shutdownPollInterval = 500 * time.Millisecond
-
-// Shutdown gracefully shuts down the server without interrupting any
-// active connections. Shutdown works by first closing all open
-// listeners and then waiting indefinitely for connections
-// to close and then shut down.
-// If the provided context expires before the shutdown is complete,
-// Shutdown returns the context's error, otherwise it returns any
-// error returned from closing the Server's underlying Listener(s).
-//
-// When Shutdown is called, Serve, ListenAndServe, and
-// ListenAndServeTLS immediately return ErrServerClosed. Make sure the
-// program doesn't exit and waits instead for Shutdown to return.
-//
-// Once Shutdown has been called on a server, it may not be reused;
-// future calls to methods such as Serve will return ErrServerClosed.
-func (srv *Server) Shutdown(ctx context.Context) error {
-	if !atomic.CompareAndSwapInt32(&srv.done, 0, 1) {
-		return ErrServerClosed
-	}
-
-	srv.mu.Lock()
-	err := srv.closeListenersLocked()
-	srv.mu.Unlock()
-
-	// Wait for active connections to close
-	ticker := time.NewTicker(shutdownPollInterval)
-	defer ticker.Stop()
-	for {
-		if srv.numConns() == 0 {
-			return err
-		}
-
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-ticker.C:
-		}
-	}
-}
-
-// getCertificate retrieves a certificate for the given client hello.
-func (srv *Server) getCertificate(h *tls.ClientHelloInfo) (*tls.Certificate, error) {
-	cert, err := srv.lookupCertificate(h.ServerName, h.ServerName)
-	if err != nil {
-		// Try wildcard
-		wildcard := strings.SplitN(h.ServerName, ".", 2)
-		if len(wildcard) == 2 {
-			// Use the wildcard pattern as the hostname.
-			hostname := "*." + wildcard[1]
-			cert, err = srv.lookupCertificate(hostname, hostname)
-		}
-		// Try "*" wildcard
-		if err != nil {
-			// Use the server name as the hostname
-			// since "*" is not a valid hostname.
-			cert, err = srv.lookupCertificate("*", h.ServerName)
-		}
-	}
-	return cert, err
-}
-
-// lookupCertificate retrieves the certificate for the given hostname,
-// if and only if the provided pattern is registered.
-// If no certificate is found in the certificate store or the certificate
-// is expired, it calls GetCertificate to retrieve a new certificate.
-func (srv *Server) lookupCertificate(pattern, hostname string) (*tls.Certificate, error) {
-	srv.hmu.Lock()
-	_, ok := srv.hosts[pattern]
-	srv.hmu.Unlock()
-	if !ok {
-		return nil, errors.New("hostname not registered")
-	}
-
-	cert, ok := srv.Certificates.Lookup(hostname)
-	if !ok || cert.Leaf != nil && cert.Leaf.NotAfter.Before(time.Now()) {
-		if srv.GetCertificate != nil {
-			cert, err := srv.GetCertificate(hostname)
-			if err == nil {
-				if err := srv.Certificates.Add(hostname, cert); err != nil {
-					srv.logf("gemini: Failed to write new certificate for %s: %s", hostname, err)
-				}
-			}
-			return &cert, err
-		}
-		return nil, errors.New("no certificate")
-	}
-
-	return &cert, nil
-}
-
-func (srv *Server) trackConn(conn *net.Conn) {
+func (srv *Server) trackConn(conn *net.Conn, cancel context.CancelFunc) {
 	srv.mu.Lock()
 	defer srv.mu.Unlock()
 	if srv.conns == nil {
-		srv.conns = make(map[*net.Conn]struct{})
+		srv.conns = make(map[*net.Conn]context.CancelFunc)
 	}
-	srv.conns[conn] = struct{}{}
+	srv.conns[conn] = cancel
 }
 
 func (srv *Server) deleteConn(conn *net.Conn) {
@@ -353,10 +282,22 @@ func (srv *Server) deleteConn(conn *net.Conn) {
 	delete(srv.conns, conn)
 }
 
-// respond responds to a connection.
-func (srv *Server) respond(conn net.Conn) {
+// serveConn serves a Gemini response over the provided connection.
+// It closes the connection when the response has been completed.
+func (srv *Server) serveConn(ctx context.Context, conn net.Conn) {
 	defer conn.Close()
 
+	if atomic.LoadInt32(&srv.closed) == 1 {
+		return
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	srv.trackConn(&conn, cancel)
+	defer srv.tryFinishShutdown()
+	defer srv.deleteConn(&conn)
+
 	defer func() {
 		if err := recover(); err != nil && err != ErrAbortHandler {
 			const size = 64 << 10
@@ -366,9 +307,6 @@ func (srv *Server) respond(conn net.Conn) {
 		}
 	}()
 
-	srv.trackConn(&conn)
-	defer srv.deleteConn(&conn)
-
 	if d := srv.ReadTimeout; d != 0 {
 		conn.SetReadDeadline(time.Now().Add(d))
 	}
@@ -376,51 +314,45 @@ func (srv *Server) respond(conn net.Conn) {
 		conn.SetWriteDeadline(time.Now().Add(d))
 	}
 
-	w := NewResponseWriter(conn)
+	done := make(chan struct{})
+	go func() {
+		srv.respond(ctx, conn)
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+	case <-done:
+	}
+}
+
+func (srv *Server) respond(ctx context.Context, conn net.Conn) {
+	w := newResponseWriter(conn)
+	defer w.Flush()
 
 	req, err := ReadRequest(conn)
 	if err != nil {
-		w.Status(StatusBadRequest)
-		w.Flush()
+		w.WriteHeader(StatusBadRequest, "Bad request")
 		return
 	}
 
-	// Store information about the TLS connection
+	// 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()
 
-	h := srv.handler(req)
+	h := srv.Handler
 	if h == nil {
-		w.Status(StatusNotFound)
-		w.Flush()
+		w.WriteHeader(StatusNotFound, "Not found")
 		return
 	}
 
-	h.ServeGemini(w, req)
-	w.Flush()
-}
-
-func (srv *Server) handler(r *Request) Handler {
-	srv.hmu.Lock()
-	defer srv.hmu.Unlock()
-	if h, ok := srv.handlers[handlerKey{r.URL.Scheme, r.URL.Hostname()}]; ok {
-		return h
-	}
-	wildcard := strings.SplitN(r.URL.Hostname(), ".", 2)
-	if len(wildcard) == 2 {
-		if h, ok := srv.handlers[handlerKey{r.URL.Scheme, "*." + wildcard[1]}]; ok {
-			return h
-		}
-	}
-	if h, ok := srv.handlers[handlerKey{r.URL.Scheme, "*"}]; ok {
-		return h
-	}
-	return nil
+	h.ServeGemini(ctx, w, req)
 }
 
 func (srv *Server) logf(format string, args ...interface{}) {
@@ -430,32 +362,3 @@ func (srv *Server) logf(format string, args ...interface{}) {
 		log.Printf(format, args...)
 	}
 }
-
-// A Handler responds to a Gemini request.
-//
-// ServeGemini should write the response header and data to the ResponseWriter
-// and then return. Returning signals that the request is finished; it is not
-// valid to use the ResponseWriter after or concurrently with the completion
-// of the ServeGemini call.
-//
-// Handlers should not modify the provided Request.
-//
-// If ServeGemini panics, the server (the caller of ServeGemini) assumes that
-// the effect of the panic was isolated to the active request. It recovers
-// the panic, logs a stack trace to the server error log, and closes the
-// network connection. To abort a handler so the client sees an interrupted
-// response but the server doesn't log an error, panic with the value
-// ErrAbortHandler.
-type Handler interface {
-	ServeGemini(ResponseWriter, *Request)
-}
-
-// The HandlerFunc type is an adapter to allow the use of ordinary functions
-// as Gemini handlers. If f is a function with the appropriate signature,
-// HandlerFunc(f) is a Handler that calls f.
-type HandlerFunc func(ResponseWriter, *Request)
-
-// ServeGemini calls f(w, r).
-func (f HandlerFunc) ServeGemini(w ResponseWriter, r *Request) {
-	f(w, r)
-}

status.go

@@ -1,39 +1,50 @@
 package gemini
 
+// Status represents a Gemini status code.
+type Status int
+
 // Gemini status codes.
 const (
-	StatusInput                    = 10
-	StatusSensitiveInput           = 11
-	StatusSuccess                  = 20
-	StatusRedirect                 = 30
-	StatusPermanentRedirect        = 31
-	StatusTemporaryFailure         = 40
-	StatusServerUnavailable        = 41
-	StatusCGIError                 = 42
-	StatusProxyError               = 43
-	StatusSlowDown                 = 44
-	StatusPermanentFailure         = 50
-	StatusNotFound                 = 51
-	StatusGone                     = 52
-	StatusProxyRequestRefused      = 53
-	StatusBadRequest               = 59
-	StatusCertificateRequired      = 60
-	StatusCertificateNotAuthorized = 61
-	StatusCertificateNotValid      = 62
+	StatusInput                    Status = 10
+	StatusSensitiveInput           Status = 11
+	StatusSuccess                  Status = 20
+	StatusRedirect                 Status = 30
+	StatusPermanentRedirect        Status = 31
+	StatusTemporaryFailure         Status = 40
+	StatusServerUnavailable        Status = 41
+	StatusCGIError                 Status = 42
+	StatusProxyError               Status = 43
+	StatusSlowDown                 Status = 44
+	StatusPermanentFailure         Status = 50
+	StatusNotFound                 Status = 51
+	StatusGone                     Status = 52
+	StatusProxyRequestRefused      Status = 53
+	StatusBadRequest               Status = 59
+	StatusCertificateRequired      Status = 60
+	StatusCertificateNotAuthorized Status = 61
+	StatusCertificateNotValid      Status = 62
 )
 
-// StatusClass returns the status class for this status code.
+// Class returns the status class for the status code.
 // 1x becomes 10, 2x becomes 20, and so on.
-func StatusClass(status int) int {
-	return (status / 10) * 10
+func (s Status) Class() Status {
+	return (s / 10) * 10
 }
 
-// Meta returns a description of the provided status code appropriate
-// for use in a response.
-//
-// Meta returns an empty string for input, success, and redirect status codes.
-func Meta(status int) string {
-	switch status {
+// String returns a text for the status code.
+// It returns the empty string if the status code is unknown.
+func (s Status) String() string {
+	switch s {
+	case StatusInput:
+		return "Input"
+	case StatusSensitiveInput:
+		return "Sensitive input"
+	case StatusSuccess:
+		return "Success"
+	case StatusRedirect:
+		return "Redirect"
+	case StatusPermanentRedirect:
+		return "Permanent redirect"
 	case StatusTemporaryFailure:
 		return "Temporary failure"
 	case StatusServerUnavailable:

timeout.go

@@ -0,0 +1,110 @@
+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.
+//
+// TimeoutHandler does not support the Hijacker or Flusher interfaces.
+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) MediaType(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
+}