1 files changed, 33 insertions, 37 deletions

diff --git a/gemini/serve.go b/gemini/serve.go
index bc13531..c148558 100644
--- a/gemini/serve.go
+++ b/gemini/serve.go
@@ -6,27 +6,28 @@ import (
 	"io"
 	"net"
 	"sync"
+
+	"tildegit.org/tjp/gus"
 )
 
-// Server listens on a network and serves the gemini protocol.
-type Server struct {
+type server struct {
 	ctx      context.Context
 	network  string
 	address  string
 	cancel   context.CancelFunc
 	wg       *sync.WaitGroup
 	listener net.Listener
-	handler  Handler
+	handler  gus.Handler
 }
 
-// NewServer builds a server.
+// NewServer builds a gemini server.
 func NewServer(
 	ctx context.Context,
 	tlsConfig *tls.Config,
 	network string,
 	address string,
-	handler Handler,
-) (*Server, error) {
+	handler gus.Handler,
+) (gus.Server, error) {
 	listener, err := net.Listen(network, address)
 	if err != nil {
 		return nil, err
@@ -34,7 +35,7 @@ func NewServer(
 
 	addr := listener.Addr()
 
-	s := &Server{
+	s := &server{
 		ctx:      ctx,
 		network:  addr.Network(),
 		address:  addr.String(),
@@ -54,7 +55,7 @@ func NewServer(
 // It will respect cancellation of the context the server was created with,
 // but be aware that Close() must still be called in that case to avoid
 // dangling goroutines.
-func (s *Server) Serve() error {
+func (s *server) Serve() error {
 	s.wg.Add(1)
 	defer s.wg.Done()
 
@@ -66,7 +67,7 @@ func (s *Server) Serve() error {
 	for {
 		conn, err := s.listener.Accept()
 		if err != nil {
-			if s.closed() {
+			if s.Closed() {
 				err = nil
 			}
 			return err
@@ -77,62 +78,57 @@ func (s *Server) Serve() error {
 	}
 }
 
-// Close begins a graceful shutdown of the server.
-//
-// It cancels the server's context which interrupts all concurrently running
-// request handlers, if they support it. It then blocks until all resources
-// have been cleaned up and all request handlers have completed.
-func (s *Server) Close() {
+func (s *server) Close() {
 	s.cancel()
 	s.wg.Wait()
 }
 
-// Network returns the network type on which the server is running.
-func (s *Server) Network() string {
+func (s *server) Network() string {
 	return s.network
 }
 
-// Address returns the address on which the server is listening.
-func (s *Server) Address() string {
+func (s *server) Address() string {
 	return s.address
 }
 
-// Hostname returns just the hostname portion of the listen address.
-func (s *Server) Hostname() string {
+func (s *server) Hostname() string {
 	host, _, _ := net.SplitHostPort(s.address)
 	return host
 }
 
-// Port returns the port on which the server is listening.
-func (s *Server) Port() string {
+func (s *server) Port() string {
 	_, portStr, _ := net.SplitHostPort(s.address)
 	return portStr
 }
 
-func (s *Server) handleConn(conn net.Conn) {
+func (s *server) handleConn(conn net.Conn) {
 	defer s.wg.Done()
 	defer conn.Close()
 
+	var response *gus.Response
 	req, err := ParseRequest(conn)
 	if err != nil {
-		_, _ = io.Copy(conn, BadRequest(err.Error()))
+		response = BadRequest(err.Error())
 		return
-	}
+	} else {
+		req.Server = s
+		req.RemoteAddr = conn.RemoteAddr()
+		if tlsconn, ok := conn.(*tls.Conn); req != nil && ok {
+			state := tlsconn.ConnectionState()
+			req.TLSState = &state
+		}
 
-	req.Server = s
-	req.RemoteAddr = conn.RemoteAddr()
-	if tlsconn, ok := conn.(*tls.Conn); req != nil && ok {
-		state := tlsconn.ConnectionState()
-		req.TLSState = &state
+		response = s.handler(s.ctx, req)
+		if response == nil {
+			response = NotFound("Resource does not exist.")
+		}
+		defer response.Close()
 	}
 
-	resp := s.handler(s.ctx, req)
-	defer resp.Close()
-
-	_, _ = io.Copy(conn, resp)
+	_, _ = io.Copy(conn, NewResponseReader(response))
 }
 
-func (s *Server) propagateCancel() {
+func (s *server) propagateCancel() {
 	go func() {
 		defer s.wg.Done()
 
@@ -141,7 +137,7 @@ func (s *Server) propagateCancel() {
 	}()
 }
 
-func (s *Server) closed() bool {
+func (s *server) Closed() bool {
 	select {
 	case <-s.ctx.Done():
 		return true