diff --git a/packages/server/src/broadcast.test.ts b/packages/server/src/broadcast.test.ts new file mode 100644 index 0000000..4177003 --- /dev/null +++ b/packages/server/src/broadcast.test.ts @@ -0,0 +1,287 @@ +// Integration test for P4.6: move-intent validation + fact-delta broadcast. +// +// This suite runs under vitest on Node, so we can't spin up a real +// Bun.serve WebSocket server (Bun global is absent). Instead we drive +// `handleMessage` directly with mock ServerWebSocket objects that +// capture every outbound frame. That gives us the same integration +// surface (two clients in a room, cross-client broadcast, error-only- +// to-sender) while keeping the tests runtime-agnostic. +import type { ServerWebSocket } from "bun"; + +import { beforeEach, describe, expect, it } from "vitest"; + +import { + handleMessage, + registerConnection, + roomRegistry, + sessionRegistry, + unregisterConnection, + type ClientData, +} from "./broadcast.js"; +import { PROTOCOL_VERSION, type ClientMessage } from "./protocol.js"; + +// --------------------------------------------------------------------------- +// Mock ServerWebSocket +// --------------------------------------------------------------------------- + +interface MockWs extends ServerWebSocket { + /** Captured outbound frames, parsed back from JSON for easy assertions. */ + readonly sent: unknown[]; + /** Read-only: true once the server called `ws.close()` on this mock. */ + readonly closed: boolean; +} + +function makeMockWs(clientId: string): MockWs { + const sent: unknown[] = []; + // `closedFlag` is a 1-element box so the `close()` method can flip a + // value without needing a `this` binding (which TypeScript types as + // the bare method receiver under strict mode). The resulting mock is + // then assembled and exposed as a MockWs. + const closedFlag = { value: false }; + // We type the mock as the full ServerWebSocket via assertion — we only + // implement the methods `handleMessage` actually uses (`send`, `close`, + // `data`). Accessing any other field in the handlers would throw at + // runtime, which is exactly the contract we want enforced. + const ws = { + data: { clientId } as ClientData, + sent, + get closed(): boolean { + return closedFlag.value; + }, + send(msg: string | Buffer): number { + const str = typeof msg === "string" ? msg : msg.toString("utf8"); + sent.push(JSON.parse(str)); + return str.length; + }, + close(): void { + closedFlag.value = true; + }, + } as unknown as MockWs; + return ws; +} + +function nextMsgOfType( + ws: MockWs, + type: string, +): { payload: Record } { + const idx = ws.sent.findIndex( + (m) => + typeof m === "object" && + m !== null && + (m as { type: unknown }).type === type, + ); + if (idx < 0) { + const types = ws.sent.map((m) => (m as { type?: string }).type); + throw new Error( + `no message of type "${type}" in inbox (got ${JSON.stringify(types)})`, + ); + } + const msg = ws.sent[idx] as { payload: Record }; + ws.sent.splice(idx, 1); + return msg; +} + +function sendClient( + ws: MockWs, + type: ClientMessage["type"], + payload: unknown, + seq = 1, +): void { + handleMessage( + ws, + JSON.stringify({ + v: PROTOCOL_VERSION, + seq, + ts: Date.now(), + type, + payload, + }), + ); +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe("broadcast — room.create / room.join / game.move", () => { + // Each test creates its own rooms with freshly-minted codes, so + // cross-test leakage is bounded — no explicit reset needed. + beforeEach(() => { + // Sanity: the registries are process-global. We touch them here + // just to keep them visible to static analysis as a public surface + // under test. + void roomRegistry; + void sessionRegistry; + }); + + it("end-to-end: two clients exchange a legal move and both receive game.delta", () => { + const a = makeMockWs("A"); + const b = makeMockWs("B"); + registerConnection(a); + registerConnection(b); + + // 1) A creates a room. + sendClient(a, "room.create", { rulesetIds: [] }); + const created = nextMsgOfType(a, "room.created"); + expect(created.payload["color"]).toBe("white"); + expect(created.payload["code"]).toMatch(/^[A-Z0-9]{6}$/); + const code = created.payload["code"] as string; + + // 2) B joins — A and B should both end up with a game.state frame. + sendClient(b, "room.join", { code }); + const joined = nextMsgOfType(b, "room.joined"); + expect(joined.payload["color"]).toBe("black"); + const aState = nextMsgOfType(a, "game.state"); + const bState = nextMsgOfType(b, "game.state"); + expect(aState.payload["turn"]).toBe("white"); + expect(bState.payload["turn"]).toBe("white"); + expect(Array.isArray(aState.payload["facts"])).toBe(true); + expect((aState.payload["facts"] as unknown[]).length).toBeGreaterThan(0); + + // 3) A plays e2->e4. + sendClient(a, "game.move", { from: "e2", to: "e4" }, 2); + const aDelta = nextMsgOfType(a, "game.delta"); + const bDelta = nextMsgOfType(b, "game.delta"); + expect(aDelta.payload["turn"]).toBe("black"); + expect(bDelta.payload["turn"]).toBe("black"); + expect(aDelta.payload["moveNotation"]).toBe("e2e4"); + expect(aDelta.payload["gameOver"]).toBeNull(); + + const inserted = aDelta.payload["inserted"] as Array<{ + attr: string; + value: unknown; + }>; + const retracted = aDelta.payload["retracted"] as Array<{ + attr: string; + value: unknown; + }>; + // e4 = file 4 + rank 3*8 = 28; e2 = file 4 + rank 1*8 = 12. + const e4Fact = inserted.find( + (f) => f.attr === "Position" && f.value === 28, + ); + const e2Fact = retracted.find( + (f) => f.attr === "Position" && f.value === 12, + ); + expect(e4Fact).toBeDefined(); + expect(e2Fact).toBeDefined(); + + sendClient(a, "room.leave", {}, 3); + unregisterConnection(a); + unregisterConnection(b); + }); + + it("illegal move: sender gets ILLEGAL_MOVE error, opponent receives nothing", () => { + const a = makeMockWs("A"); + const b = makeMockWs("B"); + registerConnection(a); + registerConnection(b); + + sendClient(a, "room.create", { rulesetIds: [] }); + const code = nextMsgOfType(a, "room.created").payload["code"] as string; + sendClient(b, "room.join", { code }); + nextMsgOfType(b, "room.joined"); + nextMsgOfType(a, "game.state"); + nextMsgOfType(b, "game.state"); + + // A plays a legal move so it becomes B's turn. + sendClient(a, "game.move", { from: "e2", to: "e4" }, 2); + nextMsgOfType(a, "game.delta"); + nextMsgOfType(b, "game.delta"); + + // B tries an illegal move: pawn can't teleport across the board. + sendClient(b, "game.move", { from: "e7", to: "e2" }, 2); + const err = nextMsgOfType(b, "error"); + expect(err.payload["code"]).toBe("ILLEGAL_MOVE"); + expect(err.payload["fatal"]).toBe(false); + + // A must have NOTHING new in its inbox — no broadcast happened. + expect(a.sent).toHaveLength(0); + + sendClient(a, "room.leave", {}, 3); + unregisterConnection(a); + unregisterConnection(b); + }); + + it("NOT_YOUR_TURN: black trying to move first gets turn error", () => { + const a = makeMockWs("A"); + const b = makeMockWs("B"); + registerConnection(a); + registerConnection(b); + + sendClient(a, "room.create", { rulesetIds: [] }); + const code = nextMsgOfType(a, "room.created").payload["code"] as string; + sendClient(b, "room.join", { code }); + nextMsgOfType(b, "room.joined"); + nextMsgOfType(a, "game.state"); + nextMsgOfType(b, "game.state"); + + // Black tries to move before white — protocol distinguishes + // NOT_YOUR_TURN from ILLEGAL_MOVE so clients can message accordingly. + sendClient(b, "game.move", { from: "e7", to: "e5" }, 2); + const err = nextMsgOfType(b, "error"); + expect(err.payload["code"]).toBe("NOT_YOUR_TURN"); + expect(a.sent).toHaveLength(0); + + sendClient(a, "room.leave", {}, 2); + unregisterConnection(a); + unregisterConnection(b); + }); + + it("malformed JSON: sender gets fatal INVALID_MESSAGE and socket is closed", () => { + const a = makeMockWs("A"); + registerConnection(a); + + handleMessage(a, "not json at all {"); + const err = nextMsgOfType(a, "error"); + expect(err.payload["code"]).toBe("INVALID_MESSAGE"); + expect(err.payload["fatal"]).toBe(true); + expect(a.closed).toBe(true); + + unregisterConnection(a); + }); + + it("version mismatch: fatal disconnect with VERSION_MISMATCH error", () => { + const a = makeMockWs("A"); + registerConnection(a); + + handleMessage( + a, + JSON.stringify({ + v: 2, + seq: 1, + ts: Date.now(), + type: "room.create", + payload: {}, + }), + ); + const err = nextMsgOfType(a, "error"); + expect(err.payload["code"]).toBe("VERSION_MISMATCH"); + expect(err.payload["fatal"]).toBe(true); + expect(a.closed).toBe(true); + + unregisterConnection(a); + }); + + it("room.join with unknown code returns ROOM_NOT_FOUND", () => { + const a = makeMockWs("A"); + registerConnection(a); + + sendClient(a, "room.join", { code: "ZZZZZZ" }); + const err = nextMsgOfType(a, "error"); + expect(err.payload["code"]).toBe("ROOM_NOT_FOUND"); + expect(err.payload["fatal"]).toBe(false); + + unregisterConnection(a); + }); + + it("unauthenticated game.move returns BAD_TOKEN", () => { + const a = makeMockWs("A"); + registerConnection(a); + + sendClient(a, "game.move", { from: "e2", to: "e4" }); + const err = nextMsgOfType(a, "error"); + expect(err.payload["code"]).toBe("BAD_TOKEN"); + + unregisterConnection(a); + }); +}); diff --git a/packages/server/src/broadcast.ts b/packages/server/src/broadcast.ts new file mode 100644 index 0000000..94e1cfa --- /dev/null +++ b/packages/server/src/broadcast.ts @@ -0,0 +1,464 @@ +// WebSocket message router + fact-delta broadcaster (P4.6). +// +// All inbound client frames funnel through handleMessage. Per-connection +// state (clientId, roomCode, token, rateLimiter) lives on ws.data; the +// authoritative room and session state lives on module-scoped registries +// shared across the process. The only side-effects are: +// 1. mutating RoomRegistry / GameSessionRegistry +// 2. sending frames to individual sockets or broadcasting to a room +// Nothing in this module touches the filesystem, network I/O beyond WS, +// or the event loop via timers — handlers are synchronous and O(facts). +import type { ServerWebSocket } from "bun"; + +import { GameSessionRegistry } from "./game-session.js"; +import { logger } from "./logger.js"; +import { RateLimiter } from "./middleware.js"; +import { + PROTOCOL_VERSION, + validateMessageString, + type ErrorCode, + type Fact as WireFact, + type GameMovePayload, + type RoomCreatePayload, + type RoomJoinPayload, + type ServerMessage, +} from "./protocol.js"; +import { RoomRegistry } from "./rooms.js"; + +// --------------------------------------------------------------------------- +// Per-connection data carried on ws.data +// --------------------------------------------------------------------------- + +export interface ClientData { + /** Process-unique connection id (randomUUID), minted at upgrade time. */ + clientId: string; + /** Room the socket is currently authenticated into, if any. */ + roomCode?: string; + /** Player token within `roomCode`. Authenticates subsequent messages. */ + token?: string; + /** Lazy-initialised on first message; one bucket per connection. */ + rateLimiter?: RateLimiter; +} + +// --------------------------------------------------------------------------- +// Module state — shared across the process +// --------------------------------------------------------------------------- + +/** clientId → socket. Used to resolve room → sockets via token lookup. */ +const connections = new Map>(); + +export const roomRegistry = new RoomRegistry(); +export const sessionRegistry = new GameSessionRegistry(); + +/** + * Monotonic server-side sequence counter. Every outbound frame carries + * `seq = nextSeq()` so clients can order messages even if the transport + * reorders (it shouldn't over a single WS, but clients reconnect). + */ +let serverSeq = 0; +function nextSeq(): number { + serverSeq += 1; + return serverSeq; +} + +// --------------------------------------------------------------------------- +// Connection lifecycle +// --------------------------------------------------------------------------- + +export function registerConnection(ws: ServerWebSocket): void { + connections.set(ws.data.clientId, ws); +} + +export function unregisterConnection(ws: ServerWebSocket): void { + connections.delete(ws.data.clientId); + // If the socket was bound to a room, mark the slot disconnected so the + // reconnect grace window (handled elsewhere) can run. We intentionally + // do NOT leaveRoom here — disconnect ≠ leave per PROTOCOL.md. + const { roomCode, token } = ws.data; + if (roomCode !== undefined && token !== undefined) { + roomRegistry.markDisconnected(roomCode, token); + // Broadcast game.end to the remaining player so they know their + // opponent dropped. v1 treats disconnect as immediate game end; the + // 60-second reconnect grace is a follow-up (see PROTOCOL.md). + broadcastGameEnd(roomCode, token, "player_left"); + } +} + +// --------------------------------------------------------------------------- +// Send helpers +// --------------------------------------------------------------------------- + +function sendTo(ws: ServerWebSocket, msg: ServerMessage): void { + ws.send(JSON.stringify(msg)); +} + +function getConnectionsInRoom( + code: string, +): ServerWebSocket[] { + const room = roomRegistry.getRoom(code); + if (!room) return []; + const sockets: ServerWebSocket[] = []; + // Build a token → socket lookup from the flat connections map once per + // broadcast. Rooms have ≤ 2 players and connections are ≤ ~few hundred + // in v1, so a linear scan is fine and avoids a secondary index that + // would need cache-invalidation on every auth change. + const byToken = new Map>(); + for (const ws of connections.values()) { + if (ws.data.token !== undefined) byToken.set(ws.data.token, ws); + } + for (const token of room.players.keys()) { + const ws = byToken.get(token); + if (ws !== undefined) sockets.push(ws); + } + return sockets; +} + +function broadcastToRoom(code: string, msg: ServerMessage): void { + const json = JSON.stringify(msg); + for (const ws of getConnectionsInRoom(code)) { + ws.send(json); + } +} + +// --------------------------------------------------------------------------- +// Envelope builders +// --------------------------------------------------------------------------- + +function envelope( + type: T, + payload: P, +): ServerMessage { + // The zod union enforces payload-type correlation at parse time; here + // we trust the caller (our own handlers) to pass the right payload for + // the declared type. Cast is narrow — no `any`, just a type assertion + // to the discriminated union branch shape. + return { + v: PROTOCOL_VERSION, + seq: nextSeq(), + ts: Date.now(), + type, + payload, + } as ServerMessage; +} + +function errorMessage( + code: ErrorCode, + message: string, + fatal: boolean, +): ServerMessage { + return envelope("error", { code, message, fatal }); +} + +// --------------------------------------------------------------------------- +// Message dispatch +// --------------------------------------------------------------------------- + +/** + * Entry point for every inbound WS frame. Order of checks mirrors + * PROTOCOL.md §Error Handling: framing → size → parse → dispatch. + * Size check is handled by the caller (middleware) before we see the + * frame; here we only handle parse + routing. + */ +export function handleMessage( + ws: ServerWebSocket, + raw: string | Buffer, +): void { + const str = typeof raw === "string" ? raw : raw.toString("utf8"); + + const result = validateMessageString(str); + if (!result.ok) { + // VERSION_MISMATCH is fatal per PROTOCOL.md; other parse failures are + // still fatal in v1 because we have no way to resync on a malformed + // stream. Closing the socket prevents a misbehaving client from + // flooding us with garbage while consuming rate-limit tokens. + const code: ErrorCode = result.error.startsWith("VERSION_MISMATCH") + ? "VERSION_MISMATCH" + : "INVALID_MESSAGE"; + sendTo(ws, errorMessage(code, result.error, true)); + ws.close(); + return; + } + + const msg = result.data; + // Server-originated messages arriving from a client are protocol errors + // — we never expect to see them inbound. The union includes them so the + // single Schema can round-trip; here we gate them out. + switch (msg.type) { + case "room.create": + handleRoomCreate(ws, msg.payload); + break; + case "room.join": + handleRoomJoin(ws, msg.payload); + break; + case "room.leave": + handleRoomLeave(ws); + break; + case "game.move": + handleGameMove(ws, msg.payload); + break; + case "room.created": + case "room.joined": + case "game.state": + case "game.delta": + case "game.end": + case "error": + sendTo( + ws, + errorMessage( + "INVALID_MESSAGE", + `server-only message type "${msg.type}" received from client`, + false, + ), + ); + break; + } +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +function handleRoomCreate( + ws: ServerWebSocket, + payload: RoomCreatePayload, +): void { + if (ws.data.roomCode !== undefined) { + sendTo( + ws, + errorMessage( + "INVALID_MESSAGE", + "connection already bound to a room", + false, + ), + ); + return; + } + const rulesetIds = payload.rulesetIds ?? []; + const { code, token, color } = roomRegistry.createRoom([...rulesetIds]); + sessionRegistry.create(code, rulesetIds); + ws.data.roomCode = code; + ws.data.token = token; + logger + .child({ clientId: ws.data.clientId, roomCode: code }) + .info("room.create"); + sendTo(ws, envelope("room.created", { code, token, color })); +} + +function handleRoomJoin( + ws: ServerWebSocket, + payload: RoomJoinPayload, +): void { + if (ws.data.roomCode !== undefined) { + sendTo( + ws, + errorMessage( + "INVALID_MESSAGE", + "connection already bound to a room", + false, + ), + ); + return; + } + const result = roomRegistry.joinRoom(payload.code); + if ("error" in result) { + sendTo( + ws, + errorMessage(result.error, `room ${payload.code}: ${result.error}`, false), + ); + return; + } + ws.data.roomCode = payload.code; + ws.data.token = result.token; + logger + .child({ clientId: ws.data.clientId, roomCode: payload.code }) + .info("room.join"); + + sendTo( + ws, + envelope("room.joined", { + code: payload.code, + token: result.token, + color: result.color, + activeRules: result.activeRules, + }), + ); + + // Second player joined → send game.state to BOTH players so the UI can + // render the starting position. The session was created on room.create + // so it must exist; defensive lookup logs and bails if it doesn't. + const session = sessionRegistry.get(payload.code); + if (!session) { + logger + .child({ roomCode: payload.code }) + .error("room.join: missing session for existing room"); + return; + } + broadcastToRoom( + payload.code, + envelope("game.state", { + facts: session.getAllFacts(), + turn: session.getTurn(), + lastSeq: 0, + moveHistory: [], + activeRules: [...result.activeRules], + // fen is a UI convenience for v1; we haven't wired FEN generation + // on the server yet, so we send an empty string. Clients that need + // FEN can derive it from `facts`. + fen: "", + }), + ); +} + +function handleRoomLeave(ws: ServerWebSocket): void { + const { roomCode, token } = ws.data; + if (roomCode === undefined || token === undefined) { + sendTo( + ws, + errorMessage( + "INVALID_MESSAGE", + "room.leave: not currently in a room", + false, + ), + ); + return; + } + broadcastGameEnd(roomCode, token, "player_left"); + roomRegistry.leaveRoom(roomCode, token); + // If the room is now empty, reap its session too so a recycled code + // doesn't accidentally resume a stale game. + if (roomRegistry.getRoom(roomCode) === undefined) { + sessionRegistry.delete(roomCode); + } + delete ws.data.roomCode; + delete ws.data.token; +} + +function handleGameMove( + ws: ServerWebSocket, + payload: GameMovePayload, +): void { + const { roomCode, token } = ws.data; + if (roomCode === undefined || token === undefined) { + sendTo( + ws, + errorMessage("BAD_TOKEN", "not authenticated into a room", false), + ); + return; + } + const player = roomRegistry.getPlayerByToken(roomCode, token); + if (!player) { + sendTo(ws, errorMessage("BAD_TOKEN", "unknown token for room", false)); + return; + } + const session = sessionRegistry.get(roomCode); + if (!session) { + // Session missing but room exists: internal inconsistency. Surface + // as INVALID_MESSAGE so the client disconnects and re-joins clean. + sendTo( + ws, + errorMessage( + "INVALID_MESSAGE", + "internal error: missing game session", + true, + ), + ); + ws.close(); + return; + } + + // Turn gate — applyMove would catch this too (the engine only offers + // moves for the current side), but we return NOT_YOUR_TURN specifically + // so clients can distinguish "you're premoving" from "that move is + // illegal for the side to move". + if (player.color !== session.getTurn()) { + sendTo( + ws, + errorMessage("NOT_YOUR_TURN", "it is not your turn", false), + ); + return; + } + + const moveResult = session.applyMove( + payload.from, + payload.to, + payload.promoteTo, + ); + if (!moveResult.ok) { + const code: ErrorCode = moveResult.error; + sendTo( + ws, + errorMessage( + code, + code === "GAME_OVER" + ? "game is already over" + : `illegal move ${payload.from}->${payload.to}`, + false, + ), + ); + return; + } + + // v1 move notation is compact `fromto[promotion]`. PROTOCOL.md example + // uses "e2e4"; promotion is appended as the first letter lower-case + // (e.g. "e7e8q"). Full SAN is a follow-up. + const moveNotation = + payload.from + + payload.to + + (payload.promoteTo ? payload.promoteTo[0] ?? "" : ""); + + broadcastToRoom( + roomCode, + envelope("game.delta", { + inserted: moveResult.inserted satisfies WireFact[], + retracted: moveResult.retracted satisfies WireFact[], + moveNotation, + turn: moveResult.turn, + gameOver: moveResult.gameOver, + }), + ); + + // Terminal positions also get an explicit game.end for clarity per + // PROTOCOL.md §game.end. `finalFen` is empty for v1 (see note above). + if (moveResult.gameOver !== null) { + broadcastToRoom( + roomCode, + envelope("game.end", { + winner: moveResult.gameOver.winner, + reason: moveResult.gameOver.reason, + finalFen: "", + }), + ); + } +} + +/** + * Broadcast a game.end to everyone in `code` EXCEPT the player whose + * token is `leaverToken`. Used when a player disconnects or leaves + * voluntarily — they don't need to hear about their own departure. + */ +function broadcastGameEnd( + code: string, + leaverToken: string, + reason: "player_left", +): void { + const room = roomRegistry.getRoom(code); + if (!room) return; + // Winner is the *other* player if we can identify them; otherwise draw. + let winner: "white" | "black" | "draw" = "draw"; + for (const p of room.players.values()) { + if (p.token !== leaverToken) { + winner = p.color; + break; + } + } + const msg = envelope("game.end", { + winner, + reason, + finalFen: "", + }); + const json = JSON.stringify(msg); + for (const ws of getConnectionsInRoom(code)) { + if (ws.data.token === leaverToken) continue; + ws.send(json); + } +} diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 251d02c..835d9a4 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -1,10 +1,19 @@ // @paratype/chess-server — authoritative Bun WebSocket server import { randomUUID } from "node:crypto"; -import { logger } from "./logger.js"; -interface ClientData { - clientId: string; -} +import { + handleMessage, + registerConnection, + unregisterConnection, + type ClientData, +} from "./broadcast.js"; +import { logger } from "./logger.js"; +import { + checkMessageSize, + checkOrigin, + MAX_MESSAGE_BYTES, + RateLimiter, +} from "./middleware.js"; const port = parseInt(process.env["PORT"] ?? "7357", 10); @@ -14,8 +23,20 @@ export const server = Bun.serve({ const url = new URL(req.url); if (url.pathname === "/ws") { + // Origin check runs before upgrade so browser clients from + // disallowed origins never see a 101 response. Tests bypass the + // check by setting ALLOWED_ORIGINS or by omitting the header + // (Bun's WebSocket client sends no Origin by default, which is + // treated as same-origin and accepted in test environments). + const origin = req.headers.get("origin"); + if (origin !== null && !checkOrigin(origin)) { + return new Response("Forbidden origin", { status: 403 }); + } + const upgraded = srv.upgrade(req, { - data: { clientId: randomUUID() }, + data: { + clientId: randomUUID(), + } satisfies ClientData, }); if (upgraded) return undefined; return new Response("Upgrade failed", { status: 500 }); @@ -29,17 +50,57 @@ export const server = Bun.serve({ }, websocket: { open(ws) { - const log = logger.child({ clientId: ws.data.clientId }); - log.info("ws connected"); + registerConnection(ws); + logger.child({ clientId: ws.data.clientId }).info("ws connected"); }, message(ws, msg) { - const log = logger.child({ clientId: ws.data.clientId }); - const length = typeof msg === "string" ? msg.length : msg.byteLength; - log.info({ length }, "ws message"); + // Size gate — we must enforce this before parsing to avoid + // allocating large JSON trees from a hostile client. + if (!checkMessageSize(msg)) { + ws.send( + JSON.stringify({ + v: 1, + seq: 0, + ts: Date.now(), + type: "error", + payload: { + code: "MSG_TOO_LARGE", + message: `frame exceeds ${String(MAX_MESSAGE_BYTES)} bytes`, + fatal: true, + }, + }), + ); + ws.close(); + return; + } + + // Rate limit — lazily create the bucket so reconnecting clients + // start with a fresh burst allowance. + if (ws.data.rateLimiter === undefined) { + ws.data.rateLimiter = new RateLimiter(); + } + if (!ws.data.rateLimiter.consume()) { + ws.send( + JSON.stringify({ + v: 1, + seq: 0, + ts: Date.now(), + type: "error", + payload: { + code: "RATE_LIMIT", + message: "rate limit exceeded", + fatal: false, + }, + }), + ); + return; + } + + handleMessage(ws, msg); }, close(ws, code, reason) { - const log = logger.child({ clientId: ws.data.clientId }); - log.info({ code, reason }, "ws closed"); + unregisterConnection(ws); + logger.child({ clientId: ws.data.clientId }).info({ code, reason }, "ws closed"); }, }, });