Make Response implement io.WriterTo

certificate/store.go

@@ -6,53 +6,58 @@ import (
 	"crypto/x509/pkix"
 	"errors"
 	"fmt"
+	"io/fs"
+	"os"
+	"path"
 	"path/filepath"
 	"strings"
 	"sync"
 	"time"
 )
 
-// A Store represents a certificate store.
-// It generates certificates as needed and automatically rotates expired certificates.
+// A Store represents a TLS certificate store.
 // The zero value for Store is an empty store ready to use.
 //
-// Certificate scopes must be registered with Register before calling Get or Load.
-// This prevents the Store from creating or loading unnecessary certificates.
+// Store can be used to store server certificates.
+// Servers should provide a hostname or wildcard pattern as a certificate scope.
+// Servers will most likely use the methods Register, Load and Get.
+//
+// Store can also be used to store client certificates.
+// Clients should provide the hostname and path of a URL as a certificate scope
+// (without a trailing slash).
+// Clients will most likely use the methods Add, Load, and Lookup.
 //
 // 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, if not nil, is called by Get to create a new
+	// certificate to replace a missing or expired certificate.
 	// The provided scope is suitable for use in a certificate's DNSNames.
 	CreateCertificate func(scope string) (tls.Certificate, error)
 
-	certs map[string]tls.Certificate
-	path  string
-	mu    sync.RWMutex
+	scopes map[string]struct{}
+	certs  map[string]tls.Certificate
+	path   string
+	mu     sync.RWMutex
 }
 
 // Register registers the provided scope with the certificate store.
 // The scope can either be a hostname or a wildcard pattern (e.g. "*.example.com").
 // To accept all hostnames, use the special pattern "*".
+//
+// Calls to Get will only succeed for registered scopes.
+// Other methods are not affected.
 func (s *Store) Register(scope string) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	if s.certs == nil {
-		s.certs = make(map[string]tls.Certificate)
+	if s.scopes == nil {
+		s.scopes = make(map[string]struct{})
 	}
-	s.certs[scope] = tls.Certificate{}
+	s.scopes[scope] = struct{}{}
 }
 
-// Add adds a certificate with the given scope to the certificate store.
-// If a certificate for the given scope already exists, Add will overwrite it.
+// Add registers the certificate for the given scope.
+// If a certificate already exists for scope, Add will overwrite it.
 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])
@@ -62,43 +67,65 @@ func (s *Store) Add(scope string, cert tls.Certificate) error {
 		cert.Leaf = parsed
 	}
 
+	if err := s.write(scope, cert); err != nil {
+		return err
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.certs == nil {
+		s.certs = make(map[string]tls.Certificate)
+	}
+	s.certs[scope] = cert
+	return nil
+}
+
+func (s *Store) write(scope string, cert tls.Certificate) error {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
 	if s.path != "" {
 		certPath := filepath.Join(s.path, scope+".crt")
 		keyPath := filepath.Join(s.path, scope+".key")
+
+		dir := filepath.Dir(certPath)
+		os.MkdirAll(dir, 0755)
+
 		if err := Write(cert, certPath, keyPath); err != nil {
 			return err
 		}
 	}
-
-	s.certs[scope] = cert
 	return nil
 }
 
 // Get retrieves a certificate for the given hostname.
 // If no matching scope has been registered, Get returns an error.
 // Get generates new certificates as needed and rotates expired certificates.
+// It calls CreateCertificate to create a new certificate if it is not nil,
+// otherwise it creates certificates with a duration of 1 year.
 //
 // Get is suitable for use in a gemini.Server's GetCertificate field.
 func (s *Store) Get(hostname string) (*tls.Certificate, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
-	cert, ok := s.certs[hostname]
+	_, ok := s.certs[hostname]
 	if !ok {
 		// Try wildcard
 		wildcard := strings.SplitN(hostname, ".", 2)
 		if len(wildcard) == 2 {
 			hostname = "*." + wildcard[1]
-			cert, ok = s.certs[hostname]
+			_, ok = s.scopes[hostname]
 		}
 	}
 	if !ok {
 		// Try "*"
-		cert, ok = s.certs["*"]
+		_, ok = s.scopes["*"]
 	}
 	if !ok {
 		return nil, errors.New("unrecognized scope")
 	}
 
+	cert := s.certs[hostname]
+
 	// If the certificate is empty or expired, generate a new one.
 	if cert.Leaf == nil || cert.Leaf.NotAfter.Before(time.Now()) {
 		var err error
@@ -114,6 +141,29 @@ func (s *Store) Get(hostname string) (*tls.Certificate, error) {
 	return &cert, nil
 }
 
+// Lookup returns the certificate for the provided scope.
+// Lookup also checks for certificates in parent scopes.
+// For example, given the scope "example.com/a/b/c", Lookup will first check
+// "example.com/a/b/c", then "example.com/a/b", then "example.com/a", and
+// finally "example.com" for a certificate. As a result, a certificate with
+// scope "example.com" will match all scopes beginning with "example.com".
+func (s *Store) Lookup(scope string) (tls.Certificate, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	cert, ok := s.certs[scope]
+	if !ok {
+		scope = path.Dir(scope)
+		for scope != "." {
+			cert, ok = s.certs[scope]
+			if ok {
+				break
+			}
+			scope = path.Dir(scope)
+		}
+	}
+	return cert, ok
+}
+
 func (s *Store) createCertificate(scope string) (tls.Certificate, error) {
 	if s.CreateCertificate != nil {
 		return s.CreateCertificate(scope)
@@ -132,29 +182,34 @@ func (s *Store) createCertificate(scope string) (tls.Certificate, error) {
 // 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.
-// Certificates with scopes that have not been registered will be ignored.
 func (s *Store) Load(path string) error {
-	matches, err := filepath.Glob(filepath.Join(path, "*.crt"))
-	if err != nil {
-		return err
-	}
+	matches := findCertificates(path)
 	for _, crtPath := range matches {
-		scope := strings.TrimSuffix(filepath.Base(crtPath), ".crt")
-		if _, ok := s.certs[scope]; !ok {
-			continue
-		}
-
 		keyPath := strings.TrimSuffix(crtPath, ".crt") + ".key"
 		cert, err := tls.LoadX509KeyPair(crtPath, keyPath)
 		if err != nil {
 			continue
 		}
+
+		scope := strings.TrimPrefix(crtPath, path)
+		scope = strings.TrimPrefix(scope, "/")
+		scope = strings.TrimSuffix(scope, ".crt")
 		s.Add(scope, cert)
 	}
 	s.SetPath(path)
 	return nil
 }
 
+func findCertificates(path string) (matches []string) {
+	filepath.Walk(path, func(path string, _ fs.FileInfo, err error) error {
+		if filepath.Ext(path) == ".crt" {
+			matches = append(matches, path)
+		}
+		return nil
+	})
+	return
+}
+
 // Entries returns a map of scopes to certificates.
 func (s *Store) Entries() map[string]tls.Certificate {
 	s.mu.RLock()

client.go

@@ -153,7 +153,7 @@ func (c *Client) do(ctx context.Context, conn net.Conn, req *Request) (*Response
 	}
 
 	// Write the request
-	if err := req.Write(w); err != nil {
+	if _, err := req.WriteTo(w); err != nil {
 		return nil, err
 	}
 

examples/server.go

@@ -26,7 +26,7 @@ func main() {
 	mux.Handle("/", gemini.FileServer(os.DirFS("/var/www")))
 
 	server := &gemini.Server{
-		Handler:        logMiddleware(mux),
+		Handler:        LoggingMiddleware(mux),
 		ReadTimeout:    30 * time.Second,
 		WriteTimeout:   1 * time.Minute,
 		GetCertificate: certificates.Get,
@@ -57,22 +57,21 @@ func main() {
 	}
 }
 
-func logMiddleware(h gemini.Handler) gemini.Handler {
+func LoggingMiddleware(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)
+		log.Printf("gemini: %s %q %d %d", host, r.URL, lw.Status, lw.Wrote)
 	})
 }
 
 type logResponseWriter struct {
+	Status      gemini.Status
+	Wrote       int
 	rw          gemini.ResponseWriter
-	status      gemini.Status
-	meta        string
 	mediatype   string
 	wroteHeader bool
-	wrote       int
 }
 
 func (w *logResponseWriter) SetMediaType(mediatype string) {
@@ -84,7 +83,7 @@ func (w *logResponseWriter) Write(b []byte) (int, error) {
 		w.WriteHeader(gemini.StatusSuccess, w.mediatype)
 	}
 	n, err := w.rw.Write(b)
-	w.wrote += n
+	w.Wrote += n
 	return n, err
 }
 
@@ -92,11 +91,10 @@ func (w *logResponseWriter) WriteHeader(status gemini.Status, meta string) {
 	if w.wroteHeader {
 		return
 	}
-	w.status = status
-	w.meta = meta
 	w.wroteHeader = true
+	w.Status = status
+	w.Wrote += len(meta) + 5
 	w.rw.WriteHeader(status, meta)
-	w.wrote += len(meta) + 5
 }
 
 func (w *logResponseWriter) Flush() error {

fs.go

@@ -1,5 +1,3 @@
-// +build go1.16
-
 package gemini
 
 import (

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

@@ -79,18 +79,21 @@ func StripPrefix(prefix string, h Handler) Handler {
 //
 // 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 {
+// 40 Temporary Failure status code and the given message in its response meta.
+// After such a timeout, writes by h to its ResponseWriter will return
+// context.DeadlineExceeded.
+func TimeoutHandler(h Handler, dt time.Duration, message string) Handler {
 	return &timeoutHandler{
-		h:  h,
-		dt: dt,
+		h:   h,
+		dt:  dt,
+		msg: message,
 	}
 }
 
 type timeoutHandler struct {
-	h  Handler
-	dt time.Duration
+	h   Handler
+	dt  time.Duration
+	msg string
 }
 
 func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *Request) {
@@ -118,7 +121,7 @@ func (t *timeoutHandler) ServeGemini(ctx context.Context, w ResponseWriter, r *R
 		w.WriteHeader(tw.status, tw.meta)
 		w.Write(buf.Bytes())
 	case <-ctx.Done():
-		w.WriteHeader(StatusTemporaryFailure, "Timeout")
+		w.WriteHeader(StatusTemporaryFailure, t.msg)
 	}
 }
 

request.go

@@ -29,6 +29,7 @@ type Request struct {
 	Certificate *tls.Certificate
 
 	conn net.Conn
+	tls  *tls.ConnectionState
 }
 
 // NewRequest returns a new request.
@@ -76,34 +77,53 @@ func ReadRequest(r io.Reader) (*Request, error) {
 	return &Request{URL: u}, nil
 }
 
-// Write writes a Gemini request in wire format.
+// WriteTo writes r to w in the Gemini request format.
 // This method consults the request URL only.
-func (r *Request) Write(w io.Writer) error {
+func (r *Request) WriteTo(w io.Writer) (int64, error) {
 	bw := bufio.NewWriterSize(w, 1026)
 	url := r.URL.String()
 	if len(url) > 1024 {
-		return ErrInvalidRequest
+		return 0, ErrInvalidRequest
 	}
-	if _, err := bw.WriteString(url); err != nil {
-		return err
+	var wrote int64
+	n, err := bw.WriteString(url)
+	wrote += int64(n)
+	if err != nil {
+		return wrote, err
 	}
-	if _, err := bw.Write(crlf); err != nil {
-		return err
+	n, err = bw.Write(crlf)
+	wrote += int64(n)
+	if err != nil {
+		return wrote, err
 	}
-	return bw.Flush()
+	return wrote, bw.Flush()
 }
 
 // Conn returns the network connection on which the request was received.
+// Conn returns nil for client requests.
 func (r *Request) Conn() net.Conn {
 	return r.conn
 }
 
 // TLS returns information about the TLS connection on which the
 // request was received.
+// TLS returns nil for client requests.
 func (r *Request) TLS() *tls.ConnectionState {
-	if tlsConn, ok := r.conn.(*tls.Conn); ok {
-		state := tlsConn.ConnectionState()
-		return &state
+	if r.tls == nil {
+		if tlsConn, ok := r.conn.(*tls.Conn); ok {
+			state := tlsConn.ConnectionState()
+			r.tls = &state
+		}
 	}
-	return nil
+	return r.tls
+}
+
+// ServerName returns the value of the TLS Server Name Indication extension
+// sent by the client.
+// ServerName returns an empty string for client requests.
+func (r *Request) ServerName() string {
+	if tls := r.TLS(); tls != nil {
+		return tls.ServerName
+	}
+	return ""
 }

request_test.go

@@ -119,7 +119,7 @@ func TestWriteRequest(t *testing.T) {
 		t.Logf("%s", test.Req.URL)
 		var b strings.Builder
 		bw := bufio.NewWriter(&b)
-		err := test.Req.Write(bw)
+		_, err := test.Req.WriteTo(bw)
 		if err != test.Err {
 			t.Errorf("expected err = %v, got %v", test.Err, err)
 		}

response.go

@@ -3,6 +3,7 @@ package gemini
 import (
 	"bufio"
 	"crypto/tls"
+	"fmt"
 	"io"
 	"net"
 	"strconv"
@@ -111,6 +112,29 @@ func (r *Response) TLS() *tls.ConnectionState {
 	return nil
 }
 
+// WriteTo 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) WriteTo(w io.Writer) (int64, error) {
+	var wrote int64
+	n, err := fmt.Fprintf(w, "%02d %s\r\n", r.Status, r.Meta)
+	wrote += int64(n)
+	if err != nil {
+		return wrote, err
+	}
+	if r.Body != nil {
+		defer r.Body.Close()
+		n, err := io.Copy(w, r.Body)
+		wrote += n
+		if err != nil {
+			return wrote, err
+		}
+	}
+	return wrote, nil
+}
+
 // A ResponseWriter interface is used by a Gemini handler to construct
 // a Gemini response.
 //

response_test.go

@@ -2,7 +2,6 @@ package gemini
 
 import (
 	"io"
-	"io/ioutil"
 	"strings"
 	"testing"
 )
@@ -38,6 +37,15 @@ func TestReadWriteResponse(t *testing.T) {
 			Meta:      "/redirect",
 			SkipWrite: true, // skip write test since result won't match Raw
 		},
+		{
+			Raw:    "32 " + maxURL + "\r\n",
+			Status: 32,
+			Meta:   maxURL,
+		},
+		{
+			Raw: "33 " + maxURL + "xxxx" + "\r\n",
+			Err: ErrInvalidResponse,
+		},
 		{
 			Raw:    "99 Unknown status code\r\n",
 			Status: 99,
@@ -83,11 +91,11 @@ func TestReadWriteResponse(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)
 		}
-		if test.Err != nil {
+		if err != nil {
 			// No response
 			continue
 		}
@@ -97,7 +105,7 @@ func TestReadWriteResponse(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)