job-search/internal/websocket/tcpserver/server.go

// Package tcpserver runs a TCP server whose clients join a websocket.Hub
// directly, without an HTTP upgrade: each connection carries newline-delimited
// JSON envelopes (see websocket.Hub.ServeTCP). It can serve a hub on its own,
// or provide a hub's clients in tests.
package tcpserver

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/netip"
	"strconv"
	"time"

	"git.avsdev.uk/Go/job-search/internal/config"
	"git.avsdev.uk/Go/job-search/internal/middleware"
	"git.avsdev.uk/Go/job-search/internal/websocket"
)

// rejectTimeout caps how long telling a rate limited client why it is being
// turned away may take.
const rejectTimeout = time.Second

// Server listens on the configured host and port and runs each connection as
// a client of its hub. New connections from each client IP are limited by
// websocket.connection_limit.
type Server struct {
	conf *config.JSCServer
	log  *slog.Logger
	hub  *websocket.Hub
}

// Option configures a Server.
type Option func(*Server)

// WithLogger makes the server log to logger instead of slog.Default().
func WithLogger(logger *slog.Logger) Option {
	return func(s *Server) {
		s.log = logger
	}
}

// New builds a server for hub's plain TCP clients. Run starts it; it shuts
// hub down when it stops, so hub should not be shared with another server.
// Push to hub to reach the server's clients. It logs to slog.Default()
// unless given WithLogger.
func New(conf *config.JSCServer, hub *websocket.Hub, opts ...Option) *Server {
	server := &Server{
		conf: conf,
		hub:  hub,
	}
	for _, opt := range opts {
		opt(server)
	}
	return server
}

// logger returns the logger the server logs to, tagged with its component.
func (s *Server) logger() *slog.Logger {
	log := s.log
	if log == nil {
		log = slog.Default()
	}
	return log.With("component", "tcp")
}

// Run listens on the configured host and port and serves plain TCP clients
// until ctx is cancelled, then shuts down gracefully (see serve).
func (s *Server) Run(ctx context.Context) error {
	addr := net.JoinHostPort(s.conf.Host, strconv.Itoa(s.conf.Port))
	listen, err := net.Listen("tcp", addr)
	if err != nil {
		return fmt.Errorf("listen %s: %w", addr, err)
	}
	return s.serve(ctx, listen)
}

// serve accepts connections on listen until ctx is cancelled or listen fails,
// then shuts down: it stops accepting, then closes the hub's clients
// gracefully, waiting up to the shutdown timeout (for ever if it is 0 or
// less). It returns nil after a shutdown caused by ctx, or the error that
// stopped it accepting, joined with any error from the hub's shutdown.
func (s *Server) serve(ctx context.Context, listen net.Listener) error {
	// Clients outlive ctx until the hub has closed them, so the shutdown can
	// tell them why; cancelling ctx itself would drop them without a word.
	connCtx, cancelConns := context.WithCancel(context.WithoutCancel(ctx))
	defer cancelConns()

	stop := context.AfterFunc(ctx, func() {
		listen.Close() //nolint:errcheck,gosec // Accept returns net.ErrClosed, which ends the loop
	})
	defer stop()

	s.logger().Info("listening", "addr", listen.Addr().String())
	acceptErr := s.accept(ctx, connCtx, listen)
	if ctx.Err() != nil {
		acceptErr = nil // cancelling ctx is how the server is stopped
	}
	listen.Close() //nolint:errcheck,gosec // closed already when ctx was cancelled

	s.logger().Info("shutting down")
	shutdownCtx := context.Background()
	if s.conf.Timeouts.ShutdownTimeout > 0 {
		var cancel context.CancelFunc
		shutdownCtx, cancel = context.WithTimeout(
			shutdownCtx, s.conf.Timeouts.ShutdownTimeout,
		)
		defer cancel()
	}

	errs := []error{acceptErr}
	if err := s.hub.Shutdown(shutdownCtx); err != nil {
		errs = append(errs, fmt.Errorf("websocket shutdown: %w", err))
	}
	return errors.Join(errs...)
}

// accept runs each connection on listen as a hub client, with connCtx, until
// ctx is cancelled (it returns ctx's error) or listen is closed some other way
// (it returns the error). Other accept errors are retried after a back off.
// Connections over the connection limit are turned away (see reject).
func (s *Server) accept(
	ctx, connCtx context.Context, listen net.Listener,
) error {
	limiter := s.connectionLimiter(ctx)

	var delay time.Duration
	for {
		newConn, err := listen.Accept()
		if err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			if errors.Is(err, net.ErrClosed) {
				return fmt.Errorf("accept: %w", err)
			}

			// Back off and retry e.g. when out of file descriptors.
			delay = min(max(2*delay, 5*time.Millisecond), time.Second)
			s.logger().Warn("tcp accept failed", "error", err, "retry_in", delay)

			select {
			case <-time.After(delay):
				continue
			case <-ctx.Done():
				return ctx.Err()
			}
		}
		delay = 0

		if limiter != nil {
			ip := remoteIP(newConn)
			if allowed, _ := limiter.Allow(ip); !allowed {
				s.logger().Warn("tcp connection rate limited", "ip", ip)
				go reject(newConn, "rate limited")
				continue
			}
		}

		go s.hub.ServeTCP(connCtx, newConn)
	}
}

// connectionLimiter returns a limiter of new connections per client IP for
// websocket.connection_limit, or nil if it is off. Its idle entries are
// evicted after rate_limits.idle_ttl until ctx is cancelled.
func (s *Server) connectionLimiter(ctx context.Context) middleware.Limiter {
	rule := s.conf.WebSocket.ConnectionLimit
	if !rule.Enabled || rule.PerSecond <= 0 {
		return nil
	}
	return middleware.NewKeyedLimiter(
		ctx, float64(rule.PerSecond), max(rule.Burst, 1), s.conf.RateLimits.IdleTTL,
	)
}

// reject tells a client why it is being turned away, with an error envelope
// as the hub would send, then closes its connection.
func reject(conn net.Conn, reason string) {
	defer conn.Close() //nolint:errcheck // turned away; nothing to do if closing fails

	message, err := json.Marshal(websocket.Envelope{
		Type:  websocket.TypeError,
		Error: reason,
	})
	if err != nil {
		return
	}

	conn.SetWriteDeadline(time.Now().Add(rejectTimeout)) //nolint:errcheck,gosec // best effort
	conn.Write(append(message, '\n'))                    //nolint:errcheck,gosec // best effort
}

// remoteIP returns the address of the connection's peer, or "" if it has no
// IP address.
func remoteIP(conn net.Conn) string {
	addrPort, err := netip.ParseAddrPort(conn.RemoteAddr().String())
	if err != nil {
		return ""
	}
	return addrPort.Addr().Unmap().String()
}