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/tcp_test.go

@@ -0,0 +1,405 @@
+package websocket
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"log/slog"
+	"net"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+	ws "github.com/coder/websocket"
+)
+
+var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil))
+
+// tcpHub starts a hub serving plain TCP clients, via ServeTCP, on a local
+// port and returns it with the address. It is shut down when the test ends.
+func tcpHub(t *testing.T, mutate func(*config.JSCServerWebSocket)) (*Hub, string) {
+	t.Helper()
+	conf := config.JSCServerWebSocket{
+		MaxMessageSize: 1 << 10,
+		WriteTimeout:   time.Second,
+	}
+	if mutate != nil {
+		mutate(&conf)
+	}
+	h := NewHub(conf, discardLog)
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go h.ServeTCP(ctx, conn)
+		}
+	}()
+	t.Cleanup(func() {
+		ln.Close()
+		sctx, scancel := context.WithTimeout(context.Background(), time.Second)
+		defer scancel()
+		h.Shutdown(sctx)
+		cancel()
+	})
+	return h, ln.Addr().String()
+}
+
+// tcpClient is a plain TCP test client reading newline-delimited envelopes.
+type tcpClient struct {
+	t    *testing.T
+	conn net.Conn
+	r    *bufio.Reader
+}
+
+func dialTCP(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)}
+}
+
+func (c *tcpClient) send(line string) {
+	c.t.Helper()
+	if _, err := io.WriteString(c.conn, line); err != nil {
+		c.t.Fatal(err)
+	}
+}
+
+// read returns the next envelope, failing the test after a second.
+func (c *tcpClient) read() 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 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: the server closing with unread data
+// from the client resets the connection.
+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)
+	}
+}
+
+func waitCount(t *testing.T, h *Hub, want int) {
+	t.Helper()
+	for deadline := time.Now().Add(time.Second); h.Count() != want; {
+		if time.Now().After(deadline) {
+			t.Fatalf("Count = %d; want %d", h.Count(), want)
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+}
+
+// ---------- ServeTCP ----------
+
+func TestTCPMessages(t *testing.T) {
+	h, addr := tcpHub(t, nil)
+	c := dialTCP(t, addr)
+
+	if e := c.read(); e.Type != TypeWelcome {
+		t.Fatalf("want welcome, got %+v", e)
+	}
+	waitCount(t, h, 1)
+
+	// Blank lines are skipped and a trailing \r is ignored.
+	c.send("\n  \r\n{\"type\":\"ping\",\"id\":\"1\"}\r\n")
+	if e := c.read(); e.Type != TypePong || e.ID != "1" {
+		t.Fatalf("want pong 1, got %+v", e)
+	}
+	c.send(`{"type":"echo","id":"2","payload":{"a":1}}` + "\n")
+	if e := c.read(); e.Type != TypeEcho || e.ID != "2" || string(e.Payload) != `{"a":1}` {
+		t.Fatalf("want echo 2, got %+v", e)
+	}
+	c.send("not json\n")
+	if e := c.read(); e.Type != TypeError || e.Error != "invalid JSON envelope" {
+		t.Fatalf("want invalid JSON error, got %+v", e)
+	}
+	c.send(`{"type":"bogus","id":"3"}` + "\n")
+	if e := c.read(); e.Type != TypeError || e.ID != "3" {
+		t.Fatalf("want unknown type error, got %+v", e)
+	}
+
+	// Pushed messages reach TCP clients.
+	h.Push("item.created", map[string]string{"id": "x"})
+	if e := c.read(); e.Type != "item.created" || string(e.Payload) != `{"id":"x"}` {
+		t.Fatalf("want item.created, got %+v", e)
+	}
+}
+
+// With no write timeout, messages are still sent: a zero timeout must not
+// mean one that has already expired.
+func TestTCPNoWriteTimeout(t *testing.T) {
+	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.WriteTimeout = 0 })
+	c := dialTCP(t, addr)
+	if e := c.read(); e.Type != TypeWelcome {
+		t.Fatalf("want welcome, got %+v", e)
+	}
+	c.send(`{"type":"ping","id":"1"}` + "\n")
+	if e := c.read(); e.Type != TypePong {
+		t.Fatalf("want pong, got %+v", e)
+	}
+}
+
+func TestTCPClientDisconnect(t *testing.T) {
+	h, addr := tcpHub(t, nil)
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	waitCount(t, h, 1)
+	c.conn.Close()
+	waitCount(t, h, 0)
+}
+
+// A line longer than max_message_size drops the connection.
+func TestTCPLineTooLong(t *testing.T) {
+	h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.MaxMessageSize = 64 })
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	c.send(`{"type":"echo","payload":"` + strings.Repeat("x", 100) + `"}` + "\n")
+	c.waitClosed()
+	waitCount(t, h, 0)
+}
+
+// Shutdown tells TCP clients why they are being closed, then closes them,
+// and refuses new ones.
+func TestTCPShutdown(t *testing.T) {
+	h, addr := tcpHub(t, nil)
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	waitCount(t, h, 1)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := h.Shutdown(ctx); err != nil {
+		t.Fatalf("Shutdown: %v", err)
+	}
+	if e := c.read(); e.Type != TypeError || e.Error != "server shutting down" {
+		t.Fatalf("want shutdown error, got %+v", e)
+	}
+	c.waitClosed()
+
+	late := dialTCP(t, addr)
+	if e := late.read(); e.Type != TypeError || e.Error != "server shutting down" {
+		t.Fatalf("late client: want shutdown error, got %+v", e)
+	}
+	late.waitClosed()
+}
+
+// Ping writes a blank line to an idle client. The client's messages, which
+// also count as activity, can arrive while pings run.
+func TestTCPPing(t *testing.T) {
+	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
+		c.PingInterval = 20 * time.Millisecond
+		c.PongTimeout = time.Second
+	})
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	c.conn.SetReadDeadline(time.Now().Add(time.Second))
+	line, err := c.r.ReadString('\n')
+	if err != nil || line != "\n" {
+		t.Fatalf("want a blank ping line, got %q, %v", line, err)
+	}
+	for i := range 5 {
+		c.send(`{"type":"ping"}` + "\n")
+		for {
+			c.conn.SetReadDeadline(time.Now().Add(time.Second))
+			line, err := c.r.ReadString('\n')
+			if err != nil {
+				t.Fatalf("message %d: %v", i, err)
+			}
+			if line != "\n" {
+				break
+			}
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}
+
+// A client that doesn't read is dropped as a slow consumer once its queue is
+// full. net.Pipe has no buffer, so writes to it block until the peer reads.
+func TestTCPSlowConsumer(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{WriteTimeout: 50 * time.Millisecond}, discardLog)
+	server, peer := net.Pipe()
+	t.Cleanup(func() { server.Close(); peer.Close() })
+	done := make(chan struct{})
+	go func() { h.ServeTCP(context.Background(), server); close(done) }()
+	waitCount(t, h, 1)
+
+	for range clientSendBuffer + 1 {
+		h.Push("item.created", nil)
+	}
+	select {
+	case <-done:
+	case <-time.After(2 * time.Second):
+		t.Fatal("slow client not dropped")
+	}
+	waitCount(t, h, 0)
+}
+
+// A pong timeout of 0 means no timeout, not one that has already expired:
+// idle clients are pinged and kept.
+func TestTCPNoPongTimeout(t *testing.T) {
+	h, addr := tcpHub(t, func(c *config.JSCServerWebSocket) {
+		c.PingInterval = 10 * time.Millisecond
+		c.PongTimeout = 0
+	})
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	for range 3 {
+		c.conn.SetReadDeadline(time.Now().Add(time.Second))
+		if line, err := c.r.ReadString('\n'); err != nil || line != "\n" {
+			t.Fatalf("want a blank ping line, got %q, %v", line, err)
+		}
+	}
+	waitCount(t, h, 1)
+}
+
+// A max_message_size of 0 uses the default limit.
+func TestTCPDefaultMessageSize(t *testing.T) {
+	_, addr := tcpHub(t, func(c *config.JSCServerWebSocket) { c.MaxMessageSize = 0 })
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	c.send(`{"type":"echo","id":"1","payload":"` + strings.Repeat("x", 4096) + `"}` + "\n")
+	if e := c.read(); e.Type != TypeEcho || e.ID != "1" {
+		t.Fatalf("want echo, got %+v", e)
+	}
+}
+
+// ---------- tcpConn ----------
+
+// pipeConn returns a tcpConn on one end of an in-memory pipe, and the other
+// end. Writes on a pipe block until the other end reads.
+func pipeConn(t *testing.T, closeTimeout time.Duration) (*tcpConn, net.Conn) {
+	t.Helper()
+	server, peer := net.Pipe()
+	t.Cleanup(func() { server.Close(); peer.Close() })
+	return newTCPConn(server, 0, closeTimeout), peer
+}
+
+func TestTCPConnRead(t *testing.T) {
+	tc, peer := pipeConn(t, time.Second)
+	go func() {
+		io.WriteString(peer, "\r\n{\"a\":1}\r\n")
+		peer.Close()
+	}()
+	typ, data, err := tc.Read(context.Background())
+	if err != nil || typ != ws.MessageText || string(data) != `{"a":1}` {
+		t.Fatalf("Read = %v, %q, %v", typ, data, err)
+	}
+	// The peer closing is a normal closure.
+	_, _, err = tc.Read(context.Background())
+	if ws.CloseStatus(err) != ws.StatusNormalClosure {
+		t.Fatalf("Read at EOF = %v; want a normal closure", err)
+	}
+}
+
+func TestTCPConnReadCancelled(t *testing.T) {
+	tc, _ := pipeConn(t, time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
+	time.AfterFunc(10*time.Millisecond, cancel)
+	if _, _, err := tc.Read(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatalf("Read = %v; want context.Canceled", err)
+	}
+}
+
+func TestTCPConnWriteBinary(t *testing.T) {
+	tc, _ := pipeConn(t, time.Second)
+	err := tc.Write(context.Background(), ws.MessageBinary, []byte("x"))
+	if !errors.Is(err, errUnsupportedData) {
+		t.Fatalf("Write(binary) = %v; want errUnsupportedData", err)
+	}
+}
+
+// Close sends its reason and later reads report its status; a second Close
+// or a CloseNow does nothing more.
+func TestTCPConnClose(t *testing.T) {
+	tc, peer := pipeConn(t, time.Second)
+	got := make(chan string, 1)
+	go func() {
+		line, _ := bufio.NewReader(peer).ReadString('\n')
+		got <- line
+	}()
+	if err := tc.Close(ws.StatusGoingAway, "bye"); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	if line := <-got; line != `{"type":"error","error":"bye"}`+"\n" {
+		t.Errorf("peer got %q", line)
+	}
+	_, _, err := tc.Read(context.Background())
+	var ce ws.CloseError
+	if !errors.As(err, &ce) || ce.Code != ws.StatusGoingAway || ce.Reason != "bye" {
+		t.Errorf("Read after Close = %v; want going away, bye", err)
+	}
+	if err := tc.Close(ws.StatusNormalClosure, "again"); err != nil {
+		t.Errorf("second Close: %v", err)
+	}
+	if err := tc.CloseNow(); err != nil {
+		t.Errorf("CloseNow after Close: %v", err)
+	}
+}
+
+// CloseNow doesn't wait for a Close stuck sending its message to a peer that
+// isn't reading, and cuts it short.
+func TestTCPConnCloseNowDuringClose(t *testing.T) {
+	tc, _ := pipeConn(t, 0) // no timeout: Close would wait for ever
+	closed := make(chan error, 1)
+	go func() { closed <- tc.Close(ws.StatusGoingAway, "bye") }()
+	for !tc.isClosed() {
+		time.Sleep(time.Millisecond)
+	}
+
+	done := make(chan error, 1)
+	go func() { done <- tc.CloseNow() }()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Errorf("CloseNow: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("CloseNow waited for Close")
+	}
+	select {
+	case <-closed:
+	case <-time.After(time.Second):
+		t.Fatal("Close not cut short by CloseNow")
+	}
+}
+
+// Close gives up sending its message after the close timeout.
+func TestTCPConnCloseTimeout(t *testing.T) {
+	tc, _ := pipeConn(t, 20*time.Millisecond)
+	done := make(chan error, 1)
+	go func() { done <- tc.Close(ws.StatusGoingAway, "bye") }()
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Errorf("Close: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Close ignored its timeout")
+	}
+}