Added websocket hub for handling websocket clients, with a wsserve command

which runs a plain TCP server for receiving websocket hub clients, with
server config, CSRF and rate limit middleware.

- internal/websocket: Hub pushing messages to WebSocket clients (ServeHTTP,
  example.com/ws upgrade with optional CSRF check) and plain TCP clients (ServeTCP,
  newline-delimited JSON envelopes); ping/pong keepalive, per-client message
  rate limit, slow consumer handling, graceful Shutdown, debug logging of
  messages sent and received
- internal/websocket/tcpserver: standalone TCP server for a hub, with accept
  back-off, a per-IP connection limit (server.websocket.connection_limit) and
  graceful shutdown
- internal/cli/wsserve: `wsserve` command running the TCP server, with
  --listen/--host/--port overrides
- internal/middleware: keyed token bucket rate limiter, real client IP from
  trusted proxies
- internal/httpx: RFC 9457 problem responses, request IDs in context
- internal/config: server section (host, port, shutdown timeout, rate limiter
  idle TTL, websocket); util.ByteSize settings (e.g. "32KiB"); lists of
  sections as single keys; empty lists stored as nil; isExplicitlySet for
  printLoadedConfig; JobSearchConfigSource renamed to JSCSourcesSource;
  TestSections checks every section's methods
- internal/util: ByteSize, parsed and formatted with go-humanize

internal/websocket/tcpserver/server.go

@@ -0,0 +1,182 @@
+// 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
+}
+
+// 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.
+func New(conf *config.JSCServer, hub *websocket.Hub, log *slog.Logger) *Server {
+	return &Server{
+		conf: conf,
+		log:  log,
+		hub:  hub,
+	}
+}
+
+// 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.log.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.log.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.log.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.log.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()
+}

internal/websocket/tcpserver/server_test.go

@@ -0,0 +1,263 @@
+package tcpserver
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"log/slog"
+	"net"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+	"git.avsdev.uk/Go/job-search/internal/websocket"
+)
+
+var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil))
+
+// testConf is a server configuration for tests, changed by mutate if it is
+// not nil.
+func testConf(mutate func(*config.JSCServer)) *config.JSCServer {
+	c := &config.JSCServer{
+		Host:      "127.0.0.1",
+		Port:      0,
+		WebSocket: config.JSCServerWebSocket{WriteTimeout: time.Second},
+		Timeouts:  config.JSCServerTimeouts{ShutdownTimeout: time.Second},
+	}
+	if mutate != nil {
+		mutate(c)
+	}
+	return c
+}
+
+// newServer returns a server for conf with its own hub.
+func newServer(conf *config.JSCServer) (*Server, *websocket.Hub) {
+	h := websocket.NewHub(conf.WebSocket, discardLog)
+	return New(conf, h, discardLog), h
+}
+
+// startServer serves on a local port, returning the server's hub, its
+// address, a function to cancel it and a channel with serve's result.
+func startServer(
+	t *testing.T, conf *config.JSCServer,
+) (*websocket.Hub, string, context.CancelFunc, chan error) {
+	t.Helper()
+	s, h := newServer(conf)
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	done := make(chan error, 1)
+	go func() { done <- s.serve(ctx, ln) }()
+	return h, ln.Addr().String(), cancel, done
+}
+
+// tcpClient is a plain TCP test client reading newline-delimited envelopes.
+type tcpClient struct {
+	t    *testing.T
+	conn net.Conn
+	r    *bufio.Reader
+}
+
+func dial(t *testing.T, addr string) *tcpClient {
+	t.Helper()
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { conn.Close() })
+	return &tcpClient{t: t, conn: conn, r: bufio.NewReader(conn)}
+}
+
+// read returns the next envelope, failing the test after a second.
+func (c *tcpClient) read() websocket.Envelope {
+	c.t.Helper()
+	c.conn.SetReadDeadline(time.Now().Add(time.Second))
+	line, err := c.r.ReadString('\n')
+	if err != nil {
+		c.t.Fatalf("read: %v (got %q)", err, line)
+	}
+	var env websocket.Envelope
+	if err := json.Unmarshal([]byte(line), &env); err != nil {
+		c.t.Fatalf("bad envelope %q: %v", line, err)
+	}
+	return env
+}
+
+// waitClosed checks the server closes the connection, skipping any lines
+// sent first. A reset counts as closed.
+func (c *tcpClient) waitClosed() {
+	c.t.Helper()
+	c.conn.SetReadDeadline(time.Now().Add(time.Second))
+	if _, err := io.ReadAll(c.r); err != nil && !errors.Is(err, syscall.ECONNRESET) {
+		c.t.Fatalf("connection not closed: %v", err)
+	}
+}
+
+// Clients get pushed messages; cancelling the server's context tells them it
+// is shutting down, closes them, stops accepting and returns nil.
+func TestServerShutdown(t *testing.T) {
+	h, addr, cancel, done := startServer(t, testConf(nil))
+	c := dial(t, addr)
+	if e := c.read(); e.Type != websocket.TypeWelcome {
+		t.Fatalf("want welcome, got %+v", e)
+	}
+	for h.Count() != 1 {
+		time.Sleep(time.Millisecond)
+	}
+	h.Push("item.created", map[string]string{"id": "x"})
+	if e := c.read(); e.Type != "item.created" {
+		t.Fatalf("want item.created, got %+v", e)
+	}
+
+	cancel()
+	if e := c.read(); e.Type != websocket.TypeError || e.Error != "server shutting down" {
+		t.Fatalf("want shutdown error, got %+v", e)
+	}
+	c.waitClosed()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Fatalf("serve = %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("serve did not return")
+	}
+	if conn, err := net.Dial("tcp", addr); err == nil {
+		conn.Close()
+		t.Error("listener still accepting")
+	}
+}
+
+// Over websocket.connection_limit, a client is told it is rate limited and
+// closed; other clients are unaffected.
+func TestServerConnectionLimit(t *testing.T) {
+	conf := testConf(func(c *config.JSCServer) {
+		c.WebSocket.ConnectionLimit = config.JSCRateLimitRule{
+			Enabled: true, PerSecond: 1, Burst: 1,
+		}
+		c.RateLimits.IdleTTL = time.Minute
+	})
+	_, addr, cancel, _ := startServer(t, conf)
+	defer cancel()
+
+	first := dial(t, addr)
+	if e := first.read(); e.Type != websocket.TypeWelcome {
+		t.Fatalf("first: want welcome, got %+v", e)
+	}
+	second := dial(t, addr)
+	if e := second.read(); e.Type != websocket.TypeError || e.Error != "rate limited" {
+		t.Fatalf("second: want rate limited, got %+v", e)
+	}
+	second.waitClosed()
+
+	first.conn.Write([]byte(`{"type":"ping","id":"1"}` + "\n"))
+	if e := first.read(); e.Type != websocket.TypePong {
+		t.Fatalf("first: want pong, got %+v", e)
+	}
+}
+
+// With the connection limit off (or no rate), there is no limiter.
+func TestServerConnectionLimitOff(t *testing.T) {
+	for _, rule := range []config.JSCRateLimitRule{
+		{Enabled: false, PerSecond: 1, Burst: 1},
+		{Enabled: true, PerSecond: 0, Burst: 1},
+	} {
+		s, _ := newServer(testConf(func(c *config.JSCServer) {
+			c.WebSocket.ConnectionLimit = rule
+		}))
+		if l := s.connectionLimiter(context.Background()); l != nil {
+			t.Errorf("rule %+v: got a limiter", rule)
+		}
+	}
+}
+
+func TestServerRun(t *testing.T) {
+	s, _ := newServer(testConf(nil))
+	ctx, cancel := context.WithCancel(context.Background())
+	time.AfterFunc(20*time.Millisecond, cancel)
+	if err := s.Run(ctx); err != nil {
+		t.Fatalf("Run = %v", err)
+	}
+}
+
+func TestServerRunListenError(t *testing.T) {
+	s, _ := newServer(testConf(func(c *config.JSCServer) { c.Port = -1 }))
+	if err := s.Run(context.Background()); err == nil || !strings.Contains(err.Error(), "listen") {
+		t.Fatalf("Run = %v; want a listen error", err)
+	}
+}
+
+// flakyListener fails Accept with a temporary error a number of times, then
+// as if closed.
+type flakyListener struct {
+	net.Listener
+	fails int
+}
+
+func (l *flakyListener) Accept() (net.Conn, error) {
+	if l.fails > 0 {
+		l.fails--
+		return nil, errors.New("too many open files")
+	}
+	return nil, net.ErrClosed
+}
+
+// Accept errors other than a closed listener are retried after a back off.
+func TestServerAcceptRetries(t *testing.T) {
+	s, _ := newServer(testConf(nil))
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	fl := &flakyListener{Listener: ln, fails: 3}
+	err = s.accept(context.Background(), context.Background(), fl)
+	if !errors.Is(err, net.ErrClosed) || fl.fails != 0 {
+		t.Fatalf("accept = %v with %d failures left; want net.ErrClosed", err, fl.fails)
+	}
+}
+
+// Cancelling ctx stops accept while it is backing off.
+func TestServerAcceptCancelledInBackOff(t *testing.T) {
+	s, _ := newServer(testConf(nil))
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+	ctx, cancel := context.WithCancel(context.Background())
+	fl := &flakyListener{Listener: ln, fails: 1 << 30}
+	time.AfterFunc(20*time.Millisecond, cancel)
+	if err := s.accept(ctx, ctx, fl); !errors.Is(err, context.Canceled) {
+		t.Fatalf("accept = %v; want context.Canceled", err)
+	}
+}
+
+// A listener closed from elsewhere is an error, after a graceful shutdown.
+func TestServerListenerClosed(t *testing.T) {
+	s, _ := newServer(testConf(nil))
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ln.Close()
+	if err := s.serve(context.Background(), ln); !errors.Is(err, net.ErrClosed) {
+		t.Fatalf("serve = %v; want net.ErrClosed", err)
+	}
+}
+
+func TestRemoteIP(t *testing.T) {
+	a, b := net.Pipe()
+	defer a.Close()
+	defer b.Close()
+	if ip := remoteIP(a); ip != "" {
+		t.Errorf("remoteIP(pipe) = %q; want empty", ip)
+	}
+}