gus/gemini/serve.go

package gemini

import (
	"bufio"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"

	"tildegit.org/tjp/gus"
	"tildegit.org/tjp/gus/logging"
)

type titanRequestBodyKey struct{}

// TitanRequestBody is the key set in a handler's context for titan requests.
//
// When this key is present in the context (request.URL.Scheme will be "titan"), the
// corresponding value is a *bufio.Reader from which the request body can be read.
var TitanRequestBody = titanRequestBodyKey{}

type server struct {
	ctx      context.Context
	errorLog logging.Logger
	network  string
	address  string
	cancel   context.CancelFunc
	wg       *sync.WaitGroup
	listener net.Listener
	handler  gus.Handler
}

// NewServer builds a gemini server.
func NewServer(
	ctx context.Context,
	errorLog logging.Logger,
	tlsConfig *tls.Config,
	network string,
	address string,
	handler gus.Handler,
) (gus.Server, error) {
	listener, err := net.Listen(network, address)
	if err != nil {
		return nil, err
	}

	addr := listener.Addr()

	s := &server{
		ctx:      ctx,
		errorLog: errorLog,
		network:  addr.Network(),
		address:  addr.String(),
		wg:       &sync.WaitGroup{},
		listener: tls.NewListener(listener, tlsConfig),
		handler:  handler,
	}

	return s, nil
}

// Serve starts the server and blocks until it is closed.
//
// This function will allocate resources which are not cleaned up until
// Close() is called.
//
// 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.
//
// On titan protocol requests it sets a key/value pair in the context. The
// key is TitanRequestBody, and the value is a *bufio.Reader from which the
// request body can be read.
func (s *server) Serve() error {
	s.wg.Add(1)
	defer s.wg.Done()

	s.ctx, s.cancel = context.WithCancel(s.ctx)

	s.wg.Add(1)
	s.propagateCancel()

	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if s.Closed() {
				err = nil
			} else {
				_ = s.errorLog.Log("msg", "accept error", "error", err)
			}

			return err
		}

		s.wg.Add(1)
		go s.handleConn(conn)
	}
}

func (s *server) Close() {
	s.cancel()
	s.wg.Wait()
}

func (s *server) Network() string {
	return s.network
}

func (s *server) Address() string {
	return s.address
}

func (s *server) Hostname() string {
	host, _, _ := net.SplitHostPort(s.address)
	return host
}

func (s *server) Port() string {
	_, portStr, _ := net.SplitHostPort(s.address)
	return portStr
}

func (s *server) handleConn(conn net.Conn) {
	defer s.wg.Done()
	defer conn.Close()

	buf := bufio.NewReader(conn)

	var response *gus.Response
	req, err := ParseRequest(buf)
	if err != nil {
		response = BadRequest(err.Error())
	} else {
		req.Server = s
		req.RemoteAddr = conn.RemoteAddr()
		if tlsconn, ok := conn.(*tls.Conn); ok {
			state := tlsconn.ConnectionState()
			req.TLSState = &state
		}

		ctx := s.ctx
		if req.Scheme == "titan" {
			len, err := sizeParam(req.Path)
			if err == nil {
				ctx = context.WithValue(
					ctx,
					TitanRequestBody,
					io.LimitReader(buf, int64(len)),
				)
			}
		}

		defer func() {
			if r := recover(); r != nil {
				err := fmt.Errorf("%s", r)
				_ = s.errorLog.Log("msg", "panic in handler", "err", err)
				_, _ = io.Copy(conn, NewResponseReader(Failure(err)))
			}
		}()

		response = s.handler(ctx, req)
		if response == nil {
			response = NotFound("Resource does not exist.")
		}
	}

	defer response.Close()
	_, _ = io.Copy(conn, NewResponseReader(response))
}

func (s *server) propagateCancel() {
	go func() {
		defer s.wg.Done()

		<-s.ctx.Done()
		_ = s.listener.Close()
	}()
}

func (s *server) Closed() bool {
	select {
	case <-s.ctx.Done():
		return true
	default:
		return false
	}
}

func sizeParam(path string) (int, error) {
	_, rest, found := strings.Cut(path, ";")
	if !found {
		return 0, errors.New("no params in path")
	}

	for _, piece := range strings.Split(rest, ";") {
		key, val, _ := strings.Cut(piece, "=")
		if key == "size" {
			return strconv.Atoi(val)
		}
	}

	return 0, errors.New("no size param found")
}