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

@@ -0,0 +1,300 @@
+package websocket
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/coder/websocket"
+	"golang.org/x/time/rate"
+)
+
+// clientSendBuffer is how many messages are queued for a client before it
+// counts as a slow consumer and is closed.
+const clientSendBuffer = 256
+
+// logMessageLimit is how much of a message is logged; longer messages are
+// cut short and logged as truncated.
+const logMessageLimit = 1 << 10
+
+// msgConn is the connection a client reads and writes messages on: a
+// *websocket.Conn, or a tcpConn for plain TCP clients (see tcp.go).
+type msgConn interface {
+	Read(ctx context.Context) (websocket.MessageType, []byte, error)
+	Write(ctx context.Context, messageType websocket.MessageType, payload []byte) error
+	Ping(ctx context.Context) error
+	Close(code websocket.StatusCode, reason string) error
+	CloseNow() error
+}
+
+// client is one connection registered with the hub. Its read loop handles
+// incoming envelopes, its write loop sends the queued messages and its ping
+// loop checks an idle connection is still alive.
+type client struct {
+	hub       *Hub
+	conn      msgConn
+	send      chan []byte // queued messages, sent by writeLoop
+	id        string
+	ip        string
+	lastSeen  atomic.Int64 // unix nanoseconds; set by the read and ping loops
+	cancel    context.CancelFunc
+	limiter   *rate.Limiter // limits the messages the client sends
+	closeOnce sync.Once
+}
+
+// errUnsupportedData is returned by readLoop (and a tcpConn's Write) for a
+// binary message.
+var errUnsupportedData = errors.New("binary frames are not supported")
+
+// newClient creates a client for conn. cancel cancels the context the client
+// runs with; the client calls it when it finishes or is closed. The client's
+// messages are limited by websocket.message_rate_limit if it is enabled.
+func newClient(
+	h *Hub, conn msgConn, ip string, cancel context.CancelFunc,
+) *client {
+	limiter := rate.NewLimiter(rate.Inf, 0)
+	if rule := h.conf.MessageRateLimit; rule.Enabled && rule.PerSecond > 0 {
+		limiter = rate.NewLimiter(
+			rate.Limit(rule.PerSecond),
+			max(rule.Burst, 1),
+		)
+	}
+
+	c := &client{
+		hub:     h,
+		conn:    conn,
+		id:      rand.Text()[:12],
+		ip:      ip,
+		send:    make(chan []byte, clientSendBuffer),
+		cancel:  cancel,
+		limiter: limiter,
+	}
+	c.lastSeen.Store(time.Now().UnixNano())
+
+	return c
+}
+
+// run starts the writer and keepalive goroutines and reads until the
+// connection ends. It returns once the connection is finished.
+func (c *client) run(ctx context.Context) {
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() { defer wg.Done(); c.writeLoop(ctx) }()
+	go func() { defer wg.Done(); c.pingLoop(ctx) }()
+
+	c.sendEnvelope(Envelope{
+		Type: TypeWelcome,
+		Payload: mustRaw(map[string]any{
+			"clientId": c.id,
+		}),
+	})
+
+	err := c.readLoop(ctx)
+	switch {
+	case err == nil:
+	case websocket.CloseStatus(err) != -1:
+		// The peer closed the connection, or close() closed it.
+	case errors.Is(err, errUnsupportedData):
+		c.close(websocket.StatusUnsupportedData, "text frames only")
+	case errors.Is(err, context.Canceled):
+		// The client was cancelled: by close() (shutdown or a slow
+		// consumer), after a failed write or ping, or by the caller.
+		c.conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed
+	default:
+		c.hub.log.Debug("read error", "client", c.id, "error", err)
+		c.conn.CloseNow() //nolint:errcheck,gosec // the connection has failed anyway
+	}
+
+	c.cancel()
+	wg.Wait()
+}
+
+// readLoop reads messages until the connection ends, and returns why.
+// Messages over the rate limit or that aren't valid JSON envelopes get an
+// error reply; the rest are passed to handle.
+func (c *client) readLoop(ctx context.Context) error {
+	for {
+		messageType, data, err := c.conn.Read(ctx)
+		if err != nil {
+			return err
+		}
+		c.logMessage(ctx, "message received", data)
+		c.lastSeen.Store(time.Now().UnixNano())
+		if messageType != websocket.MessageText {
+			return errUnsupportedData
+		}
+		// Check if the client is rate limited for early discard
+		if !c.limiter.Allow() {
+			c.sendError("", "rate limited")
+			continue
+		}
+		// Unpack the envelope and pass to handle to process if valid
+		var env Envelope
+		if err := json.Unmarshal(data, &env); err != nil {
+			c.sendError("", "invalid JSON envelope")
+			continue
+		}
+		c.handle(env)
+	}
+}
+
+// handle answers one envelope from the client: ping and echo, or an error
+// for any other type.
+func (c *client) handle(env Envelope) {
+	switch env.Type {
+	case TypePing:
+		c.sendEnvelope(Envelope{
+			Type:    TypePong,
+			ID:      env.ID,
+			Payload: mustRaw(map[string]any{"time": time.Now().UTC()}),
+		})
+
+	case TypeEcho:
+		c.sendEnvelope(Envelope{
+			Type:    TypeEcho,
+			ID:      env.ID,
+			Payload: env.Payload,
+		})
+
+	default:
+		c.sendError(env.ID, "unknown message type")
+	}
+}
+
+// writeLoop sends queued messages, each within the write timeout (no limit if
+// it is 0 or less), until ctx is cancelled. A failed write cancels the client.
+func (c *client) writeLoop(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case msg := <-c.send:
+			writeCtx, cancel := withTimeout(ctx, c.hub.conf.WriteTimeout)
+			err := c.conn.Write(writeCtx, websocket.MessageText, msg)
+			cancel()
+			if err != nil {
+				c.hub.log.Debug("write failed", "client", c.id, "error", err)
+				c.cancel()
+				return
+			}
+			c.logMessage(ctx, "message sent", msg)
+		}
+	}
+}
+
+// pingLoop pings the connection when it has been idle for the ping interval;
+// no pong within the pong timeout (no limit if it is 0 or less) drops the
+// connection. Ping needs the concurrent reader in readLoop to see pongs.
+func (c *client) pingLoop(ctx context.Context) {
+	pingInterval := c.hub.conf.PingInterval
+	pongTimeout := c.hub.conf.PongTimeout
+
+	if pingInterval <= 0 {
+		return
+	}
+
+	ticker := time.NewTicker(pingInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			// Don't ping if the lastSeen is more recent than the ping interval
+			lastSeen := time.Unix(0, c.lastSeen.Load())
+			if time.Now().Before(lastSeen.Add(pingInterval)) {
+				continue
+			}
+			pingCtx, cancel := withTimeout(ctx, pongTimeout)
+			err := c.conn.Ping(pingCtx)
+			cancel()
+			if err != nil {
+				c.hub.log.Debug("ping failed", "client", c.id, "error", err)
+				c.conn.CloseNow() //nolint:errcheck,gosec // dropping the client anyway
+				c.cancel()
+				return
+			}
+			c.lastSeen.Store(time.Now().UnixNano())
+		}
+	}
+}
+
+// trySend queues msg without blocking; false means the buffer is full.
+func (c *client) trySend(msg []byte) bool {
+	select {
+	case c.send <- msg:
+		return true
+	default:
+		return false
+	}
+}
+
+// sendEnvelope queues env for the client, closing the client as a slow
+// consumer if its queue is full.
+func (c *client) sendEnvelope(env Envelope) {
+	encoded, err := json.Marshal(env)
+	if err != nil {
+		c.hub.log.Error("marshal", "client", c.id, "error", err)
+		return
+	}
+	if !c.trySend(encoded) {
+		c.hub.log.Warn("dropping slow client", "client", c.id, "ip", c.ip)
+		go c.close(websocket.StatusPolicyViolation, "slow consumer")
+	}
+}
+
+// sendError queues an error envelope replying to the message with id (none
+// if empty).
+func (c *client) sendError(id, msg string) {
+	c.sendEnvelope(Envelope{
+		Type:  TypeError,
+		ID:    id,
+		Error: msg,
+	})
+}
+
+// close performs a graceful close handshake once, then cancels the client.
+// It may block for the handshake, so call it from its own goroutine unless
+// the read loop has already exited.
+func (c *client) close(code websocket.StatusCode, reason string) {
+	c.closeOnce.Do(func() {
+		c.conn.Close(code, reason) //nolint:errcheck,gosec // the client is cancelled anyway
+		c.cancel()
+	})
+}
+
+// logMessage logs a message sent to or received from the client at debug
+// level, with its size; at most logMessageLimit bytes of it are logged. It
+// does nothing unless debug logging is enabled.
+func (c *client) logMessage(ctx context.Context, event string, data []byte) {
+	if !c.hub.log.Enabled(ctx, slog.LevelDebug) {
+		return
+	}
+	logged, truncated := data, false
+	if len(logged) > logMessageLimit {
+		logged, truncated = logged[:logMessageLimit], true
+	}
+	c.hub.log.LogAttrs(ctx, slog.LevelDebug, event,
+		slog.String("client", c.id),
+		slog.Int("bytes", len(data)),
+		slog.String("message", string(logged)),
+		slog.Bool("truncated", truncated),
+	)
+}
+
+// withTimeout returns ctx with the timeout applied, or ctx itself if the
+// timeout is 0 or less (no limit).
+func withTimeout(
+	ctx context.Context, timeout time.Duration,
+) (context.Context, context.CancelFunc) {
+	if timeout <= 0 {
+		return ctx, func() {}
+	}
+	return context.WithTimeout(ctx, timeout)
+}

internal/websocket/client_test.go

@@ -0,0 +1,258 @@
+package websocket
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"log/slog"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+	ws "github.com/coder/websocket"
+)
+
+// fakeConn is a msgConn that records how it was closed. Read blocks until its
+// context is cancelled.
+type fakeConn struct {
+	closed chan ws.StatusCode
+}
+
+func newFakeConn() *fakeConn { return &fakeConn{closed: make(chan ws.StatusCode, 1)} }
+
+func (f *fakeConn) Read(ctx context.Context) (ws.MessageType, []byte, error) {
+	<-ctx.Done()
+	return 0, nil, ctx.Err()
+}
+
+func (f *fakeConn) Write(context.Context, ws.MessageType, []byte) error { return nil }
+
+func (f *fakeConn) Ping(context.Context) error { return nil }
+
+func (f *fakeConn) Close(code ws.StatusCode, _ string) error {
+	f.closed <- code
+	return nil
+}
+
+func (f *fakeConn) CloseNow() error { return nil }
+
+// newTestClient returns a client of a new hub on f, with no loops running,
+// and its context.
+func newTestClient(f *fakeConn) (*client, context.Context) {
+	h := NewHub(config.JSCServerWebSocket{}, discardLog)
+	ctx, cancel := context.WithCancel(context.Background())
+	return newClient(h, f, "192.0.2.1", cancel), ctx
+}
+
+// An envelope that can't be marshalled is dropped, not queued.
+func TestSendEnvelopeBadPayload(t *testing.T) {
+	c, _ := newTestClient(newFakeConn())
+	c.sendEnvelope(Envelope{Type: TypeEcho, Payload: json.RawMessage("{")})
+	if len(c.send) != 0 {
+		t.Errorf("queued %d messages; want none", len(c.send))
+	}
+}
+
+// A client whose queue is full is closed as a slow consumer.
+func TestSendEnvelopeSlowConsumer(t *testing.T) {
+	f := newFakeConn()
+	c, ctx := newTestClient(f)
+	for range clientSendBuffer {
+		c.sendEnvelope(Envelope{Type: TypePong})
+	}
+	select {
+	case code := <-f.closed:
+		t.Fatalf("closed with %v before the queue was full", code)
+	default:
+	}
+
+	c.sendEnvelope(Envelope{Type: TypePong})
+	select {
+	case code := <-f.closed:
+		if code != ws.StatusPolicyViolation {
+			t.Errorf("closed with %v; want policy violation", code)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("slow consumer not closed")
+	}
+	select {
+	case <-ctx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("slow consumer not cancelled")
+	}
+}
+
+// close closes the connection once, however often it is called.
+func TestClientCloseOnce(t *testing.T) {
+	f := newFakeConn()
+	c, _ := newTestClient(f)
+	c.close(ws.StatusGoingAway, "bye")
+	c.close(ws.StatusNormalClosure, "again")
+	if code := <-f.closed; code != ws.StatusGoingAway {
+		t.Errorf("closed with %v; want going away", code)
+	}
+	select {
+	case code := <-f.closed:
+		t.Errorf("closed a second time with %v", code)
+	default:
+	}
+}
+
+// The message rate limit applies only when it is enabled and has a rate.
+func TestClientMessageRateLimit(t *testing.T) {
+	tests := []struct {
+		rule    config.JSCRateLimitRule
+		limited bool
+	}{
+		{config.JSCRateLimitRule{Enabled: false, PerSecond: 1, Burst: 1}, false},
+		{config.JSCRateLimitRule{Enabled: true, PerSecond: 0, Burst: 1}, false},
+		{config.JSCRateLimitRule{Enabled: true, PerSecond: 1, Burst: 1}, true},
+	}
+	for _, tt := range tests {
+		h := NewHub(config.JSCServerWebSocket{MessageRateLimit: tt.rule}, discardLog)
+		c := newClient(h, newFakeConn(), "", func() {})
+		c.limiter.Allow()
+		if limited := !c.limiter.Allow(); limited != tt.limited {
+			t.Errorf("rule %+v: second message limited = %v; want %v", tt.rule, limited, tt.limited)
+		}
+	}
+}
+
+func TestMustRaw(t *testing.T) {
+	raw := json.RawMessage(`{"a":1}`)
+	tests := []struct {
+		name string
+		v    any
+		want string
+	}{
+		{"nil", nil, ""},
+		{"raw message as is", raw, `{"a":1}`},
+		{"marshalled", map[string]int{"b": 2}, `{"b":2}`},
+		{"unmarshallable", func() {}, ""},
+	}
+	for _, tt := range tests {
+		if got := string(mustRaw(tt.v)); got != tt.want {
+			t.Errorf("%s: mustRaw = %q; want %q", tt.name, got, tt.want)
+		}
+	}
+}
+
+func TestWithTimeout(t *testing.T) {
+	ctx, cancel := withTimeout(context.Background(), 0)
+	defer cancel()
+	if _, ok := ctx.Deadline(); ok {
+		t.Error("timeout 0 set a deadline")
+	}
+	ctx, cancel = withTimeout(context.Background(), time.Minute)
+	defer cancel()
+	if _, ok := ctx.Deadline(); !ok {
+		t.Error("timeout 1m set no deadline")
+	}
+}
+
+// syncBuffer is a bytes.Buffer safe for the client's goroutines to log to.
+type syncBuffer struct {
+	mu  sync.Mutex
+	buf bytes.Buffer
+}
+
+func (b *syncBuffer) Write(p []byte) (int, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buf.Write(p)
+}
+
+// entries returns the JSON log entries written so far.
+func (b *syncBuffer) entries(t *testing.T) []map[string]any {
+	t.Helper()
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	var entries []map[string]any
+	for line := range strings.Lines(b.buf.String()) {
+		var e map[string]any
+		if err := json.Unmarshal([]byte(line), &e); err != nil {
+			t.Fatalf("log line %q: %v", line, err)
+		}
+		entries = append(entries, e)
+	}
+	return entries
+}
+
+// jsonLog returns a logger writing JSON entries at level and above to a
+// buffer.
+func jsonLog(level slog.Level) (*slog.Logger, *syncBuffer) {
+	buf := &syncBuffer{}
+	return slog.New(slog.NewJSONHandler(buf, &slog.HandlerOptions{Level: level})), buf
+}
+
+// At debug level, messages received from and sent to a client are logged.
+func TestLogMessages(t *testing.T) {
+	log, buf := jsonLog(slog.LevelDebug)
+	h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, log)
+	peer, _ := pipeClient(t, h)
+
+	ping := `{"type":"ping","id":"1"}`
+	go io.WriteString(peer, ping+"\n")
+	peer.SetReadDeadline(time.Now().Add(time.Second))
+	pong, err := bufio.NewReader(peer).ReadString('\n')
+	if err != nil {
+		t.Fatalf("pong: %v", err)
+	}
+	pong = strings.TrimSuffix(pong, "\n")
+
+	// The write is logged after it returns, so wait for its entry.
+	want := map[string]string{
+		"message received": ping,
+		"message sent":     pong,
+	}
+	deadline := time.Now().Add(time.Second)
+	for {
+		found := map[string]bool{}
+		for _, e := range buf.entries(t) {
+			if msg, ok := e["msg"].(string); ok && want[msg] == e["message"] {
+				found[msg] = true
+				if e["bytes"] != float64(len(want[msg])) || e["truncated"] != false {
+					t.Errorf("entry %v: bad bytes or truncated", e)
+				}
+			}
+		}
+		if len(found) == len(want) {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("logged %v; want received ping and sent pong", buf.entries(t))
+		}
+		time.Sleep(5 * time.Millisecond)
+	}
+}
+
+// A long message is logged cut to logMessageLimit bytes; nothing is logged
+// unless debug logging is on.
+func TestLogMessageLimit(t *testing.T) {
+	log, buf := jsonLog(slog.LevelDebug)
+	h := NewHub(config.JSCServerWebSocket{}, log)
+	c := newClient(h, newFakeConn(), "", func() {})
+	long := strings.Repeat("x", logMessageLimit+100)
+	c.logMessage(context.Background(), "message sent", []byte(long))
+
+	entries := buf.entries(t)
+	if len(entries) != 1 {
+		t.Fatalf("logged %d entries; want 1", len(entries))
+	}
+	e := entries[0]
+	if e["message"] != long[:logMessageLimit] || e["bytes"] != float64(len(long)) ||
+		e["truncated"] != true || e["client"] != c.id {
+		t.Errorf("entry = %v", e)
+	}
+
+	log, buf = jsonLog(slog.LevelInfo)
+	c.hub = NewHub(config.JSCServerWebSocket{}, log)
+	c.logMessage(context.Background(), "message sent", []byte("x"))
+	if entries := buf.entries(t); len(entries) != 0 {
+		t.Errorf("logged %v at info level", entries)
+	}
+}

internal/websocket/http.go

@@ -0,0 +1,44 @@
+package websocket
+
+import (
+	"context"
+	"net/http"
+
+	"git.avsdev.uk/Go/job-search/internal/middleware"
+	ws "github.com/coder/websocket"
+)
+
+// ServeHTTP upgrades an http connection and runs the client until it closes.
+func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	mode := ws.CompressionDisabled
+	if h.conf.EnableCompression {
+		mode = ws.CompressionContextTakeover
+	}
+
+	// Accept enforces the Origin check: same host, or one of OriginPatterns.
+	// That is the defence against cross-site WebSocket hijacking.
+	conn, err := ws.Accept(w, r, &ws.AcceptOptions{
+		CompressionMode: mode,
+	})
+	if err != nil {
+		h.log.Debug("upgrade failed", "error", err)
+		return // Accept has already written the HTTP error.
+	}
+	defer conn.CloseNow() //nolint:errcheck // finished with; it may already be closed
+	conn.SetReadLimit(int64(h.maxMessageSize()))
+
+	ip := middleware.ClientIP(r)
+	ctx, cancel := context.WithCancel(r.Context())
+	c := newClient(h, conn, ip, cancel)
+
+	if !h.register(c) {
+		cancel()
+		conn.Close(ws.StatusGoingAway, "server shutting down") //nolint:errcheck,gosec // turned away
+		return
+	}
+	defer h.unregister(c)
+
+	h.log.Info("client connected", "client", c.id, "ip", ip)
+	c.run(ctx)
+	h.log.Info("client disconnected", "client", c.id, "ip", ip)
+}

internal/websocket/http_test.go

@@ -0,0 +1,183 @@
+package websocket
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+	ws "github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+)
+
+// wsHub serves h's WebSocket clients on a test HTTP server and returns its
+// ws:// URL.
+func wsHub(t *testing.T, h *Hub) string {
+	t.Helper()
+	srv := httptest.NewServer(h)
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		h.Shutdown(ctx)
+		srv.Close()
+	})
+	return "ws" + strings.TrimPrefix(srv.URL, "http")
+}
+
+// dialWS connects a WebSocket client to url and reads its welcome.
+func dialWS(t *testing.T, ctx context.Context, url string) *ws.Conn {
+	t.Helper()
+	conn, _, err := ws.Dial(ctx, url, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	t.Cleanup(func() { conn.CloseNow() })
+	var e Envelope
+	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypeWelcome {
+		t.Fatalf("want welcome, got %+v, %v", e, err)
+	}
+	return conn
+}
+
+// A client upgraded over HTTP gets replies. A max_message_size of 0 uses the
+// default limit rather than refusing every message.
+func TestHTTPMessages(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, discardLog)
+	url := wsHub(t, h)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	conn := dialWS(t, ctx, url)
+
+	if err := wsjson.Write(ctx, conn, Envelope{Type: TypePing, ID: "1"}); err != nil {
+		t.Fatal(err)
+	}
+	var e Envelope
+	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong || e.ID != "1" {
+		t.Fatalf("want pong 1, got %+v, %v", e, err)
+	}
+}
+
+// One Push reaches WebSocket and plain TCP clients of the same hub.
+func TestPushReachesBothTransports(t *testing.T) {
+	h, addr := tcpHub(t, nil)
+	url := wsHub(t, h)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	conn := dialWS(t, ctx, url)
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	waitCount(t, h, 2)
+
+	h.Push("item.created", map[string]string{"id": "x"})
+	var e Envelope
+	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != "item.created" {
+		t.Fatalf("websocket: want item.created, got %+v, %v", e, err)
+	}
+	if e := c.read(); e.Type != "item.created" {
+		t.Fatalf("tcp: want item.created, got %+v", e)
+	}
+}
+
+// A payload that can't be marshalled is not sent.
+func TestPushBadPayload(t *testing.T) {
+	h, addr := tcpHub(t, nil)
+	c := dialTCP(t, addr)
+	c.read() // welcome
+	waitCount(t, h, 1)
+
+	h.Push("bad", func() {})
+	h.Push("good", nil)
+	if e := c.read(); e.Type != "good" || e.Payload != nil {
+		t.Fatalf("want good with no payload, got %+v", e)
+	}
+}
+
+// A binary frame closes the connection with 1003 (unsupported data).
+func TestHTTPBinaryFrame(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{WriteTimeout: time.Second}, discardLog)
+	url := wsHub(t, h)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	conn := dialWS(t, ctx, url)
+
+	if err := conn.Write(ctx, ws.MessageBinary, []byte{1, 2, 3}); err != nil {
+		t.Fatal(err)
+	}
+	_, _, err := conn.Read(ctx)
+	if ws.CloseStatus(err) != ws.StatusUnsupportedData {
+		t.Fatalf("read after binary frame = %v; want close 1003", err)
+	}
+}
+
+// Messages over the message rate limit get a "rate limited" error and are not
+// handled; the connection stays open.
+func TestHTTPMessageRateLimit(t *testing.T) {
+	conf := config.JSCServerWebSocket{
+		WriteTimeout:     time.Second,
+		MessageRateLimit: config.JSCRateLimitRule{Enabled: true, PerSecond: 1, Burst: 1},
+	}
+	h := NewHub(conf, discardLog)
+	url := wsHub(t, h)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	conn := dialWS(t, ctx, url)
+
+	for _, id := range []string{"1", "2"} {
+		if err := wsjson.Write(ctx, conn, Envelope{Type: TypePing, ID: id}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	var e Envelope
+	if err := wsjson.Read(ctx, conn, &e); err != nil || e.Type != TypePong || e.ID != "1" {
+		t.Fatalf("first: want pong 1, got %+v, %v", e, err)
+	}
+	err := wsjson.Read(ctx, conn, &e)
+	if err != nil || e.Type != TypeError || e.Error != "rate limited" {
+		t.Fatalf("second: want rate limited, got %+v, %v", e, err)
+	}
+	if h.Count() != 1 {
+		t.Errorf("Count = %d; want the client kept", h.Count())
+	}
+}
+
+// A request that isn't a WebSocket upgrade is refused by the upgrade and no
+// client is registered.
+func TestHTTPNotAnUpgrade(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{}, discardLog)
+	srv := httptest.NewServer(h)
+	defer srv.Close()
+	res, err := http.Get(srv.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	res.Body.Close()
+	if res.StatusCode < 400 {
+		t.Errorf("status = %d; want an error", res.StatusCode)
+	}
+	if h.Count() != 0 {
+		t.Errorf("Count = %d; want 0", h.Count())
+	}
+}
+
+// After Shutdown, an upgrade succeeds but the client is closed at once with
+// 1001 (going away).
+func TestHTTPRefusedAfterShutdown(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{}, discardLog)
+	url := wsHub(t, h)
+	if err := h.Shutdown(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	conn, _, err := ws.Dial(ctx, url, nil)
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.CloseNow()
+	if _, _, err := conn.Read(ctx); ws.CloseStatus(err) != ws.StatusGoingAway {
+		t.Fatalf("read = %v; want close 1001", err)
+	}
+}

internal/websocket/hub.go

@@ -0,0 +1,158 @@
+package websocket
+
+import (
+	"context"
+	"encoding/json"
+	"log/slog"
+	"sync"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+	ws "github.com/coder/websocket"
+)
+
+// Hub tracks connected clients, WebSocket (ServeHTTP) and plain TCP
+// (ServeTCP) alike, and sends pushed messages to all of them.
+type Hub struct {
+	conf config.JSCServerWebSocket
+	log  *slog.Logger
+
+	m       sync.RWMutex
+	clients map[*client]struct{}
+	closed  bool
+	wg      sync.WaitGroup
+}
+
+// defaultMaxMessageSize is the largest message read from a client when
+// max_message_size is not set.
+const defaultMaxMessageSize = 32 << 10
+
+// NewHub creates a hub.
+func NewHub(conf config.JSCServerWebSocket, log *slog.Logger) *Hub {
+	return &Hub{
+		conf:    conf,
+		log:     log.With("component", "ws"),
+		clients: make(map[*client]struct{}),
+	}
+}
+
+// Push sends a message of type msgType to every client. The payload is
+// marshalled once and the same bytes are queued for every client; a client
+// whose queue is full is closed as a slow consumer. A payload that can't be
+// marshalled is logged and nothing is sent.
+func (h *Hub) Push(msgType string, payload any) {
+	var rawPayload json.RawMessage
+	if payload != nil {
+		var err error
+		if rawPayload, err = json.Marshal(payload); err != nil {
+			h.log.Error("marshal push payload", "type", msgType, "error", err)
+			return
+		}
+	}
+
+	msg, err := json.Marshal(Envelope{
+		Type:    msgType,
+		Payload: rawPayload,
+	})
+	if err != nil {
+		h.log.Error("marshal push", "type", msgType, "error", err)
+		return
+	}
+
+	var slow []*client
+	h.m.RLock()
+	for client := range h.clients {
+		if !client.trySend(msg) {
+			slow = append(slow, client)
+		}
+	}
+	h.m.RUnlock()
+
+	for _, client := range slow {
+		h.log.Warn("dropping slow client", "client", client.id, "ip", client.ip)
+		go client.close(ws.StatusPolicyViolation, "slow consumer")
+	}
+}
+
+// Count returns the number of connected clients.
+func (h *Hub) Count() int {
+	h.m.RLock()
+	defer h.m.RUnlock()
+	return len(h.clients)
+}
+
+// Shutdown closes every connection with 1001 (going away) and waits for the
+// handlers to return. If ctx expires first, the remaining connections are
+// closed at once and ctx's error is returned. New connections are refused
+// afterwards.
+func (h *Hub) Shutdown(ctx context.Context) error {
+	h.m.Lock()
+	h.closed = true
+	clientShutdowns := make([]*client, 0, len(h.clients))
+	for client := range h.clients {
+		clientShutdowns = append(clientShutdowns, client)
+	}
+	h.m.Unlock()
+
+	if len(clientShutdowns) == 0 {
+		// Nothing to close. Clients that just left are only finishing their
+		// unregister, so this wait is short and ctx doesn't apply.
+		h.wg.Wait()
+		return nil
+	}
+
+	for _, client := range clientShutdowns {
+		go client.close(ws.StatusGoingAway, "server shutting down")
+	}
+
+	done := make(chan struct{})
+	go func() {
+		h.wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return nil
+	case <-ctx.Done():
+		// Both may be ready; a finished shutdown is not a timeout.
+		select {
+		case <-done:
+			return nil
+		default:
+		}
+		for _, client := range clientShutdowns {
+			client.conn.CloseNow() //nolint:errcheck,gosec // out of time; dropping it
+		}
+		return ctx.Err()
+	}
+}
+
+// maxMessageSize is the largest message read from a client: max_message_size,
+// or defaultMaxMessageSize if that is 0 or less.
+func (h *Hub) maxMessageSize() int {
+	if h.conf.MaxMessageSize <= 0 {
+		return defaultMaxMessageSize
+	}
+	return int(h.conf.MaxMessageSize)
+}
+
+// register adds a client to the hub. Hub clients receive all push messages and
+// have their lifecycles tracked by the hub.
+func (h *Hub) register(c *client) bool {
+	h.m.Lock()
+	defer h.m.Unlock()
+	if h.closed {
+		return false
+	}
+	h.clients[c] = struct{}{}
+	h.wg.Add(1)
+	return true
+}
+
+// unregister removes a client from the hub.
+func (h *Hub) unregister(c *client) {
+	h.m.Lock()
+	delete(h.clients, c)
+	h.m.Unlock()
+	h.wg.Done()
+}

internal/websocket/hub_test.go

@@ -0,0 +1,93 @@
+package websocket
+
+import (
+	"bufio"
+	"context"
+	"errors"
+	"net"
+	"testing"
+	"time"
+
+	"git.avsdev.uk/Go/job-search/internal/config"
+)
+
+// pipeClient runs a plain TCP client of h on an in-memory pipe and reads its
+// welcome. It returns the peer's end, which the test stops reading from to
+// make the client's writes block, and a channel closed when ServeTCP returns.
+func pipeClient(t *testing.T, h *Hub) (net.Conn, chan struct{}) {
+	t.Helper()
+	server, peer := net.Pipe()
+	t.Cleanup(func() { server.Close(); peer.Close() })
+	done := make(chan struct{})
+	go func() { h.ServeTCP(context.Background(), server); close(done) }()
+	peer.SetReadDeadline(time.Now().Add(time.Second))
+	if _, err := bufio.NewReader(peer).ReadString('\n'); err != nil {
+		t.Fatalf("welcome: %v", err)
+	}
+	waitCount(t, h, 1)
+	return peer, done
+}
+
+// If ctx expires before the clients have closed, Shutdown closes the rest at
+// once and returns ctx's error.
+func TestShutdownTimeout(t *testing.T) {
+	// No write timeout: the client's goodbye blocks until it is cut short.
+	h := NewHub(config.JSCServerWebSocket{}, discardLog)
+	_, done := pipeClient(t, h)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	defer cancel()
+	if err := h.Shutdown(ctx); !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("Shutdown = %v; want context.DeadlineExceeded", err)
+	}
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("client not closed after the shutdown timeout")
+	}
+}
+
+// Shutdown with no clients returns at once, even with an expired ctx.
+func TestShutdownNoClients(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{}, discardLog)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	for range 20 {
+		if err := h.Shutdown(ctx); err != nil {
+			t.Fatalf("Shutdown = %v; want nil", err)
+		}
+	}
+}
+
+// A ping that can't be sent within the pong timeout drops the client.
+func TestPingFailureDropsClient(t *testing.T) {
+	h := NewHub(config.JSCServerWebSocket{
+		PingInterval: 10 * time.Millisecond,
+		PongTimeout:  20 * time.Millisecond,
+	}, discardLog)
+	_, done := pipeClient(t, h)
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("client not dropped after a failed ping")
+	}
+	waitCount(t, h, 0)
+}
+
+func TestMaxMessageSize(t *testing.T) {
+	tests := []struct {
+		conf config.JSCServerWebSocket
+		want int
+	}{
+		{config.JSCServerWebSocket{}, defaultMaxMessageSize},
+		{config.JSCServerWebSocket{MaxMessageSize: -1}, defaultMaxMessageSize},
+		{config.JSCServerWebSocket{MaxMessageSize: 100}, 100},
+	}
+	for _, tt := range tests {
+		h := NewHub(tt.conf, discardLog)
+		if got := h.maxMessageSize(); got != tt.want {
+			t.Errorf("maxMessageSize(%d) = %d; want %d", tt.conf.MaxMessageSize, got, tt.want)
+		}
+	}
+}

internal/websocket/messages.go

@@ -0,0 +1,64 @@
+// Package websocket implements the hub that pushes messages to connected
+// clients: WebSocket clients upgraded by Hub.ServeHTTP (the /ws endpoint) and
+// plain TCP clients run by Hub.ServeTCP, which send the same messages as
+// newline-delimited JSON (see tcp.go). Package tcpserver runs a TCP server
+// for those clients.
+//
+// Every message is a JSON Envelope: a WebSocket text frame, or one line on
+// TCP.
+//
+// Client -> server:
+//
+//	{"type":"ping","id":"1"}
+//	{"type":"echo","id":"2","payload":{...}}
+//
+// Server -> client:
+//
+//	{"type":"welcome","payload":{"clientId":"..."}}
+//	{"type":"pong","id":"1","payload":{"time":"..."}}
+//	{"type":"echo","id":"2","payload":{...}}
+//	{"type":"error","id":"3","error":"..."}
+//	{"type":"item.created","payload":{...}}
+//
+// Any other type from a client gets an "unknown message type" error. The last
+// message is one sent by Hub.Push: its type is the pushed message type.
+package websocket
+
+import "encoding/json"
+
+// Envelope is the wire format for every message in both directions.
+type Envelope struct {
+	Type    string          `json:"type"`
+	ID      string          `json:"id,omitempty"`
+	Payload json.RawMessage `json:"payload,omitempty"`
+	Error   string          `json:"error,omitempty"`
+}
+
+// Envelope types. TypeEvent, TypeRequest and TypeAck are reserved: nothing
+// sends or handles them yet.
+const (
+	TypeWelcome = "welcome"
+	TypePing    = "ping"
+	TypePong    = "pong"
+	TypeEcho    = "echo"
+	TypeEvent   = "event"
+	TypeRequest = "request"
+	TypeAck     = "ack"
+	TypeError   = "error"
+)
+
+// mustRaw returns v as raw JSON: v itself if it is already json.RawMessage,
+// nil if v is nil or can't be marshalled.
+func mustRaw(v any) json.RawMessage {
+	if v == nil {
+		return nil
+	}
+	if raw, ok := v.(json.RawMessage); ok {
+		return raw
+	}
+	b, err := json.Marshal(v)
+	if err != nil {
+		return nil
+	}
+	return b
+}

internal/websocket/tcp.go

@@ -0,0 +1,261 @@
+package websocket
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/netip"
+	"strings"
+	"sync"
+	"time"
+
+	ws "github.com/coder/websocket"
+)
+
+// Plain TCP clients join the hub as if they were WebSocket clients. Instead
+// of WebSocket frames, each message is one line of JSON (an Envelope)
+// terminated by "\n" in both directions, so a client can be as simple as
+// netcat:
+//
+//	$ nc localhost 9393
+//	{"type":"welcome","payload":{"clientId":"..."}}
+//	{"type":"ping","id":"1"}
+//	{"type":"pong","id":"1","payload":{"time":"..."}}
+//
+// A trailing "\r" is ignored and blank lines are skipped. A line longer than
+// the websocket max_message_size closes the connection, as an oversized frame
+// does. When the server closes a connection (shutdown, slow consumer, a line
+// too long) it sends an error envelope with the reason first.
+//
+// Plain TCP has no ping frames, so liveness uses TCP keepalives probing at the
+// ping interval; Ping itself writes a blank line, which clients ignore, so a
+// dead peer is found when a write fails.
+//
+// There is no HTTP request, so there is no CSRF or Origin check and no proxy
+// handling: the client's address is the TCP peer. Only listen where clients
+// are trusted (e.g. localhost), or wrap the listener (e.g. with PROXY protocol
+// or TLS) before passing its connections to ServeTCP. ServeTCP doesn't limit
+// connections; package tcpserver applies websocket.connection_limit.
+
+// ServeTCP runs a plain TCP connection as a hub client until it closes, as
+// ServeHTTP does for a WebSocket upgrade. Cancelling ctx drops the client
+// without a message; use Shutdown to close clients gracefully. The
+// connection is closed when ServeTCP returns.
+func (h *Hub) ServeTCP(ctx context.Context, newConn net.Conn) {
+	if tcpSocket, ok := newConn.(*net.TCPConn); ok && h.conf.PingInterval > 0 {
+		tcpSocket.SetKeepAliveConfig(net.KeepAliveConfig{ //nolint:errcheck,gosec // best effort
+			Enable:   true,
+			Idle:     h.conf.PingInterval,
+			Interval: h.conf.PingInterval,
+			Count:    3,
+		})
+	}
+
+	ip := ""
+	if addrPort, err := netip.ParseAddrPort(newConn.RemoteAddr().String()); err == nil {
+		ip = addrPort.Addr().Unmap().String()
+	}
+
+	conn := newTCPConn(newConn, h.maxMessageSize(), h.conf.WriteTimeout)
+
+	ctx, cancel := context.WithCancel(ctx)
+	client := newClient(h, conn, ip, cancel)
+
+	if !h.register(client) {
+		cancel()
+		conn.Close(ws.StatusGoingAway, "server shutting down") //nolint:errcheck,gosec // turned away
+		return
+	}
+	defer h.unregister(client)
+
+	h.log.Info("client connected", "client", client.id, "ip", ip, "transport", "tcp")
+	client.run(ctx)
+	conn.CloseNow() //nolint:errcheck,gosec // finished with; it may already be closed
+	h.log.Info("client disconnected", "client", client.id, "ip", ip, "transport", "tcp")
+}
+
+// ---------- tcpConn ----------
+
+// tcpConn adapts a net.Conn carrying newline-delimited JSON to the msgConn a
+// client uses. Read is only called from the client's read loop; writes may
+// come from the write loop, Ping and Close concurrently, so they are
+// serialised.
+type tcpConn struct {
+	conn         net.Conn
+	scan         *bufio.Scanner
+	closeTimeout time.Duration // caps Close's message; 0 or less for none
+
+	wm sync.Mutex // serialises writes
+
+	closeOnce sync.Once     // records the first close
+	closed    chan struct{} // closed by the first Close or CloseNow
+	closeErr  ws.CloseError // why the connection was closed
+}
+
+// newTCPConn wraps conn. maxLine caps the length of a line read
+// (defaultMaxMessageSize if 0 or less) and closeTimeout how long Close spends
+// sending its message.
+func newTCPConn(
+	conn net.Conn, maxLine int, closeTimeout time.Duration,
+) *tcpConn {
+	if maxLine <= 0 {
+		maxLine = defaultMaxMessageSize
+	}
+
+	scan := bufio.NewScanner(conn)
+	// The buffer must hold the line and its "\n".
+	scan.Buffer(make([]byte, 0, min(maxLine+1, 4096)), maxLine+1)
+
+	return &tcpConn{
+		conn:         conn,
+		scan:         scan,
+		closeTimeout: closeTimeout,
+		closed:       make(chan struct{}),
+	}
+}
+
+// Read returns the next non-blank line as a text message. Cancelling ctx
+// interrupts it. The end of the stream is reported as a normal closure, and a
+// read after Close as the close status Close sent, as a *websocket.Conn
+// reports them, so the client treats them as a finished close handshake.
+func (t *tcpConn) Read(ctx context.Context) (ws.MessageType, []byte, error) {
+	stop := context.AfterFunc(ctx, func() {
+		t.conn.SetReadDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the read
+	})
+	defer stop()
+
+	for t.scan.Scan() {
+		line := strings.TrimSuffix(t.scan.Text(), "\r")
+		if strings.TrimSpace(line) == "" {
+			continue
+		}
+		return ws.MessageText, []byte(line), nil
+	}
+
+	err := t.scan.Err()
+	switch {
+	case ctx.Err() != nil:
+		return 0, nil, ctx.Err()
+	case t.isClosed():
+		return 0, nil, t.closeErr
+	case err == nil || errors.Is(err, io.EOF):
+		return 0, nil, ws.CloseError{
+			Code:   ws.StatusNormalClosure,
+			Reason: "connection closed",
+		}
+	case errors.Is(err, bufio.ErrTooLong):
+		// As a *websocket.Conn closes with 1009 for an oversized message.
+		// The client may not see the reason: closing with its unread data
+		// still buffered makes the connection reset.
+		t.Close(ws.StatusMessageTooBig, "message too long") //nolint:errcheck,gosec // dropped anyway
+		return 0, nil, t.closeErr
+	}
+	return 0, nil, err
+}
+
+// Write sends payload as one line. Only text messages are supported.
+func (t *tcpConn) Write(
+	ctx context.Context, messageType ws.MessageType, payload []byte,
+) error {
+	if messageType != ws.MessageText {
+		return errUnsupportedData
+	}
+	line := make([]byte, 0, len(payload)+1)
+	line = append(append(line, payload...), '\n')
+	return t.write(ctx, line)
+}
+
+// Ping writes a blank line, which clients ignore, so a dead peer shows up as a
+// failed write. TCP keepalives (see ServeTCP) do the rest.
+func (t *tcpConn) Ping(ctx context.Context) error {
+	return t.write(ctx, []byte("\n"))
+}
+
+// Close sends an error envelope giving the reason, then closes the
+// connection. Only the first Close or CloseNow does anything; later calls
+// return nil. A CloseNow while Close is sending its message cuts it short.
+func (t *tcpConn) Close(code ws.StatusCode, reason string) error {
+	if !t.markClosed(ws.CloseError{Code: code, Reason: reason}) {
+		return nil
+	}
+
+	ctx := context.Background()
+	if t.closeTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, t.closeTimeout)
+		defer cancel()
+	}
+
+	message, err := json.Marshal(Envelope{
+		Type:  TypeError,
+		Error: reason,
+	})
+	if err == nil {
+		t.write(ctx, append(message, '\n')) //nolint:errcheck,gosec // best effort, closing anyway
+	}
+
+	if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+		return err
+	}
+	return nil
+}
+
+// CloseNow closes the connection without a message. It doesn't wait for a
+// Close that is sending its message.
+func (t *tcpConn) CloseNow() error {
+	t.markClosed(ws.CloseError{Code: ws.StatusAbnormalClosure})
+	if err := t.conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+		return err
+	}
+	return nil
+}
+
+// markClosed records why the connection is being closed and reports whether
+// this is the first close. It never blocks, so a Close sending its message
+// can't hold up CloseNow.
+func (t *tcpConn) markClosed(why ws.CloseError) bool {
+	first := false
+	t.closeOnce.Do(func() {
+		t.closeErr = why
+		close(t.closed)
+		first = true
+	})
+	return first
+}
+
+// isClosed reports whether Close or CloseNow has been called.
+func (t *tcpConn) isClosed() bool {
+	select {
+	case <-t.closed:
+		return true
+	default:
+		return false
+	}
+}
+
+// write writes data whole, giving up at ctx's deadline or when ctx is
+// cancelled.
+func (t *tcpConn) write(ctx context.Context, data []byte) error {
+	t.wm.Lock()
+	defer t.wm.Unlock()
+
+	deadline, _ := ctx.Deadline() // zero means none
+	if err := t.conn.SetWriteDeadline(deadline); err != nil {
+		return err
+	}
+
+	stop := context.AfterFunc(ctx, func() {
+		t.conn.SetWriteDeadline(time.Now()) //nolint:errcheck,gosec // only interrupts the write
+	})
+	defer stop()
+
+	_, err := t.conn.Write(data)
+	if err != nil && ctx.Err() != nil {
+		return ctx.Err()
+	}
+
+	return err
+}

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")
+	}
+}

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)
+	}
+}