diff --git a/README.md b/README.md index 7a94df2..20abf4a 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,9 @@ cp .env.example .env | `LOG_LEVEL` | `info` | Minimum log severity (`debug` \| `info` \| `warn` \| `error`) | | `MAX_MESSAGES_PER_SECOND` | `100` | Per-client message rate limit (messages per second) | | `CONN_RATE_LIMIT` | `30` | Max new connections per IP address per minute | +| `SESSION_ENCRYPTION_KEY` | — | 32-byte base64 key (or `{"v1":"…"}` key map) that enables session resumption | +| `SESSION_TTL_MS` | `3600000` | Sliding TTL of a stored session (ms) | +| `INSTANCE_ID` | uuid | Identity published in the `GW_AFFINITY` cookie | #### Tuning rate limits for high-traffic deployments @@ -135,6 +138,22 @@ wss://:/?token= Clients must provide a valid JWT as a query parameter. Connections without a valid token are rejected with a `4001` close code. +### Session Resumption + +Set `SESSION_ENCRYPTION_KEY` to enable it; without the key the gateway behaves exactly as before. + +The gateway seals each client's session state (rooms, sequence numbers, rate-limit window) with AES-256-GCM. The blob is the `session_id`. On reconnect the client presents it as a URL-encoded query parameter, or as the JWT `sid` claim: + +``` +wss://:/?token=&session_id= +``` + +The gateway restores the rooms and replies with `session_resumed`, carrying each room's saved `highestAckedSeq` / `highestReceivedSeq` plus the room's live `currentSeqPerRoom`. The client then sends the usual `reconnect` message for any room that shows a gap. A blob that fails to decrypt, has expired, or belongs to another identity is ignored and the connection continues as a new session. + +Clients that cannot store a blob get the `GW_AFFINITY=` cookie on the handshake response: when they land back on the same instance, the session is restored from that instance's local cache. + +A fresh blob also arrives with `server_shutting_down` on graceful shutdown, and with close code `4100` (in the close reason, or a preceding `migrate` frame) after `POST /admin/v1/clients/{clientId}/migrate`. `GET /metrics` reports the `session_resumption_total` counters. + ### HTTP Health Check ``` diff --git a/src/index.js b/src/index.js index b03bcf9..b738d6e 100644 --- a/src/index.js +++ b/src/index.js @@ -5,13 +5,27 @@ import { logger } from "./logger.js"; /** * Parses environment variables into server configuration with integer coercion. * - * @returns {{ port: number, heartbeatMs: number, maxPayloadBytes: number }} + * Session-resumption keys are only present when their variables are set, so an + * unconfigured deployment keeps the historical three-key shape. + * + * @returns {{ port: number, heartbeatMs: number, maxPayloadBytes: number, sessionEncryptionKey?: string, sessionTtlMs?: number, instanceId?: string }} */ export function parseConfig() { const port = parseInt(process.env.PORT ?? "8080", 10); const heartbeatMs = parseInt(process.env.WS_HEARTBEAT_MS ?? "30000", 10); const maxPayloadBytes = parseInt(process.env.MAX_PAYLOAD_BYTES ?? "1024", 10); - return { port, heartbeatMs, maxPayloadBytes }; + const config = { port, heartbeatMs, maxPayloadBytes }; + + if (process.env.SESSION_ENCRYPTION_KEY) { + config.sessionEncryptionKey = process.env.SESSION_ENCRYPTION_KEY; + } + if (process.env.SESSION_TTL_MS) { + config.sessionTtlMs = parseInt(process.env.SESSION_TTL_MS, 10); + } + if (process.env.INSTANCE_ID) { + config.instanceId = process.env.INSTANCE_ID; + } + return config; } const config = parseConfig(); @@ -31,18 +45,31 @@ if (isNaN(config.maxPayloadBytes) || config.maxPayloadBytes < 1) { process.exit(1); } +if (config.sessionTtlMs !== undefined && (isNaN(config.sessionTtlMs) || config.sessionTtlMs < 1)) { + logger.error("Invalid SESSION_TTL_MS value", { SESSION_TTL_MS: process.env.SESSION_TTL_MS }); + process.exit(1); +} + let wss; -let httpServer; let markShuttingDown; let sessionManager; +let instanceId; +let saveAllSessions; try { - ({ wss, httpServer, markShuttingDown, sessionManager } = createServer(config)); + ({ wss, markShuttingDown, sessionManager, instanceId, saveAllSessions } = createServer(config)); } catch (err) { logger.error("Failed to start server", { error: err.message }); process.exit(1); } -logger.info("Gateway started", config); +// The encryption key never reaches the logs. +logger.info("Gateway started", { + port: config.port, + heartbeatMs: config.heartbeatMs, + maxPayloadBytes: config.maxPayloadBytes, + instanceId, + sessionResumption: sessionManager != null, +}); /** * Initiates a multi-phase graceful shutdown of the WebSocket server. @@ -53,28 +80,57 @@ logger.info("Gateway started", config); * Phase 4 (4000ms): Close connections with WebSocket code 1001 "Going Away". * Phase 5 (>5000ms): Force exit. * + * When session resumption is active, phase 2 first persists every live session + * and hands each client its own fresh blob as `session_id`, so a client can + * resume on another instance immediately. + * * @param {object} wss - The WebSocket server instance. * @param {string} signal - The OS signal that triggered the shutdown (e.g. "SIGTERM"). + * @param {object} [options] + * @param {() => Promise>} [options.saveAllSessions] - Persists live sessions. * @returns {void} */ -export function shutdown(wss, signal) { +export function shutdown(wss, signal, { saveAllSessions } = {}) { logger.info("shutdown: stopping accept", { signal }); const clientCount = wss.clients ? wss.clients.size : 0; + /** Sends `server_shutting_down`, adding a per-client blob when one exists. */ + function notifyClients(blobs) { + const shared = JSON.stringify({ type: "server_shutting_down", payload: { reconnectIn: 5 } }); + for (const client of wss.clients) { + const sessionId = blobs?.get(client._clientId); + try { + client.send( + sessionId + ? JSON.stringify({ + type: "server_shutting_down", + payload: { reconnectIn: 5, session_id: sessionId }, + }) + : shared + ); + } catch { + // Client may already be disconnected + } + } + } + // Phase 2 — Notify clients (100ms) setTimeout(() => { logger.info("shutdown: notifying N clients", { clientCount }); - if (wss.clients) { - const notification = JSON.stringify({ type: "server_shutting_down", payload: { reconnectIn: 5 } }); - for (const client of wss.clients) { - try { - client.send(notification); - } catch { - // Client may already be disconnected - } - } + if (!wss.clients) return; + + if (typeof saveAllSessions !== "function") { + notifyClients(null); + return; } + + saveAllSessions() + .then((blobs) => notifyClients(blobs)) + .catch((err) => { + logger.error("shutdown: session save failed", { error: err.message }); + notifyClients(null); + }); }, 100); // Phase 3 — Drain pending sends (500ms–4000ms) @@ -131,15 +187,20 @@ export function shutdown(wss, signal) { }); } +// Without a session manager there is nothing to persist, so the shared +// broadcast path stays in use. +const shutdownOptions = sessionManager ? { saveAllSessions } : {}; + +// Flip /healthz and /readyz to 503 first so load balancers stop routing here +// while the drain phases run; closing the WebSocket server also closes the +// co-located HTTP server. process.on("SIGTERM", () => { markShuttingDown(); - wss.close(); - shutdown(httpServer, "SIGTERM"); + shutdown(wss, "SIGTERM", shutdownOptions); }); process.on("SIGINT", () => { markShuttingDown(); - wss.close(); - shutdown(httpServer, "SIGINT"); + shutdown(wss, "SIGINT", shutdownOptions); }); process.on("uncaughtException", (err) => { diff --git a/src/rate-limiter.js b/src/rate-limiter.js index fe5b805..8c4f182 100644 --- a/src/rate-limiter.js +++ b/src/rate-limiter.js @@ -83,3 +83,5 @@ export function createRateLimiter(maxPerSecond) { get size() { return windows.size; }, + }; +} diff --git a/src/room-manager.js b/src/room-manager.js index e6d0ae6..6354bd1 100644 --- a/src/room-manager.js +++ b/src/room-manager.js @@ -1,6 +1,11 @@ import { WebSocket } from "ws"; import { v7 as uuidv7 } from "uuid"; +const DEFAULT_RING_BUFFER_SIZE = 100; +const DEFAULT_MAX_BUFFER_BYTES = 1024 * 1024; +const DEFAULT_DEDUP_WINDOW_MS = 5000; +const DEFAULT_MAX_DEDUP_ENTRIES = 10_000; + /** * @typedef {Object} BackpressureOptions * @property {boolean} [enabled=false] - Enable backpressure-aware broadcasting @@ -195,6 +200,19 @@ export class RoomManager { return false; } + /** + * Subscribes a client to a room, enforcing the configured DoS limits. + * + * Re-joining a room the client already occupies only replaces the stored + * socket, so it is never rejected by a limit. + * + * @param {string} clientId - Unique identifier for the client. + * @param {string} roomId - Identifier of the room to join. + * @param {import("ws").WebSocket} ws - Socket to register for broadcasts. + * @returns {undefined | { ok: boolean, reason?: string } | { type: "error", payload: { code: string, message: string } }} + * `{ ok }` when `maxRoomSize` is configured, an error frame when a limit or the + * circuit breaker rejects the join, otherwise `undefined`. + */ join(clientId, roomId, ws) { if (clientId == null) throw new TypeError("clientId is required"); if (roomId == null) throw new TypeError("roomId is required"); diff --git a/src/server.js b/src/server.js index 7f414c9..072da0a 100644 --- a/src/server.js +++ b/src/server.js @@ -1,8 +1,9 @@ import http from "node:http"; -import { WebSocketServer } from "ws"; +import { WebSocket, WebSocketServer } from "ws"; import { v4 as uuid } from "uuid"; import jwt from "jsonwebtoken"; import { RoomManager } from "./room-manager.js"; +import { SessionManager } from "./session-manager.js"; import { validateMessage } from "./validator.js"; import { verifyConnection } from "./auth.js"; import { logger } from "./logger.js"; @@ -11,6 +12,33 @@ import { createConnRateLimiter } from "./conn-rate-limiter.js"; import { VALIDATION_ERROR } from "./errors.js"; import { SessionManager } from "./session-manager.js"; +/** + * Creates the co-located HTTP server (health checks, Prometheus metrics, + * admin migration) and the WebSocket gateway on the same port. + * + * Session resumption is opt-in: it activates when `sessionManager` is injected + * or when an encryption key is available, and stays completely inert otherwise. + * + * @param {object} [options] + * @param {number} [options.port] - TCP port to listen on. Defaults to 8080. + * @param {number} [options.heartbeatMs] - Ping interval used to detect zombies. Defaults to 30000. + * @param {number} [options.maxPayloadBytes] - Max WebSocket frame size. Defaults to 1024. + * @param {number} [options.connRateLimit] - New connections allowed per IP per minute. + * @param {number} [options.maxConnectionsPerIp] - Concurrent connections allowed per IP. + * @param {number} [options.maxMessagesPerSecond] - Per-client message rate limit. + * @param {number} [options.ringBufferSize] - Replay entries kept per room. + * @param {number} [options.deduplicationWindowMs] - TTL of a `location_update` dedup key. + * @param {number} [options.maxBufferBytes] - Byte ceiling for one room's replay buffer. + * @param {number} [options.maxDedupEntries] - Max dedup keys retained. + * @param {SessionManager} [options.sessionManager] - Pre-built manager; takes precedence over `sessionEncryptionKey`. + * @param {string|object} [options.sessionEncryptionKey] - Key (or key map) used to build a manager. + * Falls back to `SESSION_ENCRYPTION_KEY`. + * @param {number} [options.sessionTtlMs] - Session TTL. Falls back to `SESSION_TTL_MS`, then 1 hour. + * @param {string} [options.instanceId] - Value published in the `GW_AFFINITY` cookie. + * Falls back to `INSTANCE_ID`, then a uuid. + * @param {object} [options.redis] - node-redis v4 style client handed to a self-built manager. + * @returns {{ wss: WebSocketServer, server: http.Server, httpServer: http.Server, rooms: RoomManager, ipConnectionCount: Map, rateLimiter: object, metrics: object, markShuttingDown: () => void, sessionManager: SessionManager|null, instanceId: string, saveAllSessions: () => Promise> }} + */ export function createServer({ port, heartbeatMs, @@ -51,26 +79,60 @@ export function createServer({ return; } - const pathname = new URL(req.url, `http://${req.headers.host ?? "localhost"}`).pathname; + const migrate = req.method === "POST" ? MIGRATE_PATH.exec(url.pathname) : null; + if (migrate) { + migrateClient(decodeURIComponent(migrate[1])) + .then((blob) => { + if (!blob) { + res.writeHead(404, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: "Not Found" })); + return; + } + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ ok: true })); + }) + .catch((err) => { + logger.error("Client migration failed", { error: err.message }); + res.writeHead(500, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: "Internal Server Error" })); + }); + return; + } + + if (req.method !== "GET") { + res.writeHead(405, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: "Method Not Allowed" })); + return; + } + + if (url.pathname === "/health") { + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ status: "OK" })); + return; + } - if (pathname === "/health" || pathname === "/healthz") { - if (isShuttingDown && pathname === "/healthz") { + if (url.pathname === "/healthz") { + if (isShuttingDown) { res.writeHead(503, { "Content-Type": "application/json" }); res.end(JSON.stringify({ status: "shutting down" })); return; } res.writeHead(200, { "Content-Type": "application/json" }); - if (pathname === "/healthz") { - res.end(JSON.stringify({ status: "ok", uptime: process.uptime() })); - } else { - res.end(JSON.stringify({ status: "OK" })); - } - } else if (pathname === "/readyz") { + res.end(JSON.stringify({ status: "ok", uptime: process.uptime() })); + return; + } + + if (url.pathname === "/readyz") { if (isShuttingDown) { res.writeHead(503, { "Content-Type": "application/json" }); res.end(JSON.stringify({ status: "not ready", reason: "server is shutting down" })); return; } + if (!isReady) { + res.writeHead(503, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ status: "not ready", reason: "initializing" })); + return; + } res.writeHead(200, { "Content-Type": "application/json" }); res.end(JSON.stringify({ status: "ready", @@ -106,14 +168,13 @@ export function createServer({ `gateway_event_loop_lag_ms ${metrics.eventLoopLagMs}`, ]; res.writeHead(200, { "Content-Type": "text/plain; version=0.0.4; charset=utf-8" }); - res.end(lines.join("\n") + "\n"); - } else { - res.writeHead(404, { "Content-Type": "application/json" }); - res.end(JSON.stringify({ error: "Not Found" })); + res.end(renderMetrics()); + return; } - }); - httpServer.listen(port ?? 8080); + res.writeHead(404, { "Content-Type": "application/json" }); + res.end(JSON.stringify({ error: "Not Found" })); + }); const wss = new WebSocketServer({ server: httpServer, @@ -135,13 +196,285 @@ export function createServer({ const connRateLimiter = createConnRateLimiter(connRateLimit); const ipConnectionCount = new Map(); const MAX_CONNS_PER_IP = maxConnectionsPerIp ?? (Number(process.env.MAX_CONNECTIONS_PER_IP) || 10); + const MAX_MESSAGES_PER_SECOND = + maxMessagesPerSecond ?? (Number(process.env.MAX_MESSAGES_PER_SECOND) || 100); + + const HEARTBEAT_MS = heartbeatMs ?? 30000; + // A client is a zombie once it stays silent for two ping intervals. The floor + // keeps event-loop jitter from reaping healthy clients under tiny intervals. + const PONG_TIMEOUT_MS = Math.max(HEARTBEAT_MS * 2, 1000); function heartbeat() { - this.isAlive = true; + this._lastPongAt = Date.now(); + } + + /** Serialises and sends a frame, skipping sockets that are already going away. */ + function safeSend(ws, message) { + if (ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED) return; + try { + ws.send(typeof message === "string" ? message : JSON.stringify(message)); + } catch (err) { + logger.warn("Failed to send frame", { error: err.message }); + } } function sendError(ws, message, code) { - ws.send(JSON.stringify({ type: "error", payload: { message, code } })); + safeSend(ws, { type: "error", payload: { message, code } }); + } + + /** + * Per-connection state that session capture reads from. + * + * @param {string} clientId + * @param {import("http").IncomingMessage} req + * @returns {object} + */ + function createContext(clientId, req) { + return { + clientId, + ip: req.socket.remoteAddress, + userAgent: req.headers["user-agent"] ?? null, + connectedAt: Date.now(), + lastActivityAt: Date.now(), + /** @type {Map} roomId → highest seq the client ACKed via `reconnect` */ + ackedSeq: new Map(), + /** @type {Map} roomId → opaque geofence set carried across resumptions */ + geofence: new Map(), + /** @type {number[]} timestamps of accepted messages, mirroring the rate-limiter window */ + messageWindow: [], + /** @type {object|null} snapshot taken before teardown wipes room membership */ + frozenState: null, + }; + } + + /** Drops timestamps outside the rate-limit window and caps the retained count. */ + function pruneWindow(timestamps, now = Date.now()) { + const cutoff = now - RATE_WINDOW_MS; + const live = timestamps.filter((ts) => typeof ts === "number" && ts > cutoff); + return live.length > MAX_MESSAGES_PER_SECOND ? live.slice(-MAX_MESSAGES_PER_SECOND) : live; + } + + /** + * Builds the session state for a client from live room and rate-limit state. + * + * @param {object} ctx + * @returns {import("./session-manager.js").SessionState} + */ + function captureState(ctx) { + if (ctx.frozenState) return ctx.frozenState; + const roomIds = [...rooms.getClientRooms(ctx.clientId)]; + return { + clientId: ctx.clientId, + protocolVersion: PROTOCOL_VERSION, + authIdentity: { sub: ctx.clientId }, + rooms: roomIds.map((roomId) => ({ + roomId, + highestAckedSeq: ctx.ackedSeq.get(roomId) ?? 0, + highestReceivedSeq: rooms.getRoomSeq(roomId), + geofenceInsideSet: ctx.geofence.get(roomId) ?? [], + })), + rateLimitState: { + messageWindow: pruneWindow(ctx.messageWindow), + connectionWindow: [], + }, + metadata: { + ip: ctx.ip, + userAgent: ctx.userAgent, + connectedAt: ctx.connectedAt, + lastActivityAt: ctx.lastActivityAt, + }, + }; + } + + /** Keeps the newest state for the sticky-cookie path, bounded by insertion order. */ + function rememberLocal(clientId, state) { + localSessions.delete(clientId); + localSessions.set(clientId, state); + while (localSessions.size > MAX_LOCAL_SESSIONS) { + const oldest = localSessions.keys().next().value; + localSessions.delete(oldest); + } + } + + /** Reads the local cache, dropping entries older than the session TTL. */ + function readLocal(clientId) { + const state = localSessions.get(clientId); + if (!state) return null; + const lastActivityAt = state.metadata?.lastActivityAt ?? 0; + if (Date.now() - lastActivityAt > sessionTtl) { + localSessions.delete(clientId); + return null; + } + return state; + } + + /** + * Marks activity, refreshes the local cache and schedules an encrypted save. + * + * @param {object|null} ctx + */ + function touchSession(ctx) { + if (!sessions || !ctx) return; + ctx.lastActivityAt = Date.now(); + rememberLocal(ctx.clientId, captureState(ctx)); + sessions.debouncedSave(ctx.clientId, () => captureState(ctx)).catch((err) => { + logger.error("Failed to schedule session save", { clientId: ctx.clientId, error: err.message }); + }); + } + + /** + * Re-applies a decrypted session to this instance and confirms it to the client. + * + * @param {import("ws").WebSocket} ws + * @param {object} ctx + * @param {import("./session-manager.js").SessionState} state + */ + function restoreSession(ws, ctx, state) { + const restored = []; + const skipped = []; + + for (const entry of Array.isArray(state.rooms) ? state.rooms : []) { + if (!entry || typeof entry.roomId !== "string") continue; + const joinResult = rooms.join(ctx.clientId, entry.roomId, ws); + if (joinResult?.type === "error" || joinResult?.ok === false) { + skipped.push(entry.roomId); + continue; + } + const highestAckedSeq = Number(entry.highestAckedSeq) || 0; + ctx.ackedSeq.set(entry.roomId, highestAckedSeq); + ctx.geofence.set( + entry.roomId, + Array.isArray(entry.geofenceInsideSet) ? entry.geofenceInsideSet : [] + ); + restored.push({ + roomId: entry.roomId, + highestAckedSeq, + highestReceivedSeq: Number(entry.highestReceivedSeq) || 0, + }); + } + + // Conservative choice: the limiter exposes no window import, so every saved + // in-window timestamp is re-consumed to deny a burst after migration. + ctx.messageWindow = pruneWindow(state.rateLimitState?.messageWindow ?? []); + for (let i = 0; i < ctx.messageWindow.length; i++) rateLimiter.check(ctx.clientId); + + ctx.connectedAt = Number(state.metadata?.connectedAt) || ctx.connectedAt; + + if (skipped.length > 0) { + logger.warn("Session rooms skipped on resume", { clientId: ctx.clientId, rooms: skipped }); + } + + const currentSeqPerRoom = {}; + for (const room of restored) currentSeqPerRoom[room.roomId] = rooms.getRoomSeq(room.roomId); + + safeSend(ws, { type: "session_resumed", payload: { rooms: restored, currentSeqPerRoom } }); + logger.info("Session resumed", { clientId: ctx.clientId, rooms: restored.length }); + touchSession(ctx); + } + + /** + * Runs the resumption handshake: explicit `session_id` first, then the + * sticky-cookie fallback, otherwise a fresh session. + * + * @param {import("ws").WebSocket} ws + * @param {import("http").IncomingMessage} req + * @param {URL} url + * @param {object} ctx + * @param {string|null} token + * @returns {Promise} True when a session was restored. + */ + async function resumeSession(ws, req, url, ctx, token) { + const sessionId = url.searchParams.get("session_id") ?? sessionIdFromToken(token); + + if (sessionId) { + // `load()` counts decrypt_failed / expired; success is recorded here, + // after the identity check, so a mismatch is not also a success. + const state = await sessions.load(sessionId, { countSuccess: false }); + if (!state) { + logger.info("Session not resumable", { clientId: ctx.clientId }); + return false; + } + if (state.clientId !== ctx.clientId) { + sessions.recordResumption("mismatch"); + logger.warn("Session identity mismatch", { + clientId: ctx.clientId, + sessionClientId: state.clientId, + }); + return false; + } + sessions.recordResumption("success"); + restoreSession(ws, ctx, state); + return true; + } + + const affinity = readCookie(req.headers.cookie, AFFINITY_COOKIE); + if (affinity === resolvedInstanceId) { + const cached = readLocal(ctx.clientId); + if (cached) { + sessions.recordResumption("success"); + restoreSession(ws, ctx, cached); + return true; + } + } + + sessions.recordResumption("new_session"); + return false; + } + + /** + * Persists a client's state now and hands the blob over as it disconnects. + * + * @param {string} clientId + * @returns {Promise} The blob, or null when the client is unknown. + */ + async function migrateClient(clientId) { + const ctx = sessions ? liveClients.get(clientId) : null; + if (!ctx) return null; + + await sessions.flush(clientId); + const state = captureState(ctx); + const blob = await sessions.save(clientId, state); + rememberLocal(clientId, state); + + if (Buffer.byteLength(blob, "utf8") <= MAX_CLOSE_REASON_BYTES) { + ctx.ws.close(MIGRATE_CLOSE_CODE, blob); + } else { + safeSend(ctx.ws, { type: "migrate", payload: { session_id: blob } }); + ctx.ws.close(MIGRATE_CLOSE_CODE, "migrated"); + } + logger.info("Client migrated", { clientId }); + return blob; + } + + /** + * Flushes pending saves and persists every live client, for graceful shutdown. + * + * @returns {Promise>} clientId → fresh session blob. + */ + async function saveAllSessions() { + /** @type {Map} */ + const blobs = new Map(); + if (!sessions) return blobs; + + await sessions.flushAll(); + for (const ctx of liveClients.values()) { + const state = captureState(ctx); + rememberLocal(ctx.clientId, state); + try { + blobs.set(ctx.clientId, await sessions.save(ctx.clientId, state)); + } catch (err) { + logger.error("Failed to save session", { clientId: ctx.clientId, error: err.message }); + } + } + return blobs; + } + + if (sessions) { + wss.on("headers", (headers) => { + headers.push( + `Set-Cookie: ${AFFINITY_COOKIE}=${resolvedInstanceId}; Path=/; HttpOnly; SameSite=Lax; Max-Age=${AFFINITY_MAX_AGE_S}` + ); + }); } function safeSend(ws, data) { @@ -157,6 +490,7 @@ export function createServer({ wss.on("connection", (ws, req) => { const clientId = uuid(); ws.isAlive = true; + ws._lastPongAt = Date.now(); const ip = req.socket.remoteAddress; @@ -237,6 +571,7 @@ export function createServer({ } }).catch(() => {}); } + if (ctx) ctx.messageWindow = pruneWindow([...ctx.messageWindow, Date.now()]); ws.on("message", (raw) => { if (!rateLimiter.check(identity.clientId)) { @@ -417,23 +752,48 @@ export function createServer({ }).catch(() => { ws.close(4001, "Authentication failed"); }); + + if (ctx) { + pendingResume = resumeSession(ws, req, url, ctx, token).catch((err) => { + logger.error("Session resumption failed", { clientId: ctx.clientId, error: err.message }); + }); + } }); const heartbeatInterval = setInterval(() => { + const now = Date.now(); wss.clients.forEach((ws) => { - if (ws.isAlive === false) { + const silentMs = now - (ws._lastPongAt ?? now); + if (ws.isAlive === false || silentMs > PONG_TIMEOUT_MS) { logger.warn("Terminating zombie connection", { clientId: ws._clientId ?? ws._trackedIp ?? "unknown", }); return ws.terminate(); } - ws.isAlive = false; ws.ping(); }); - }, heartbeatMs ?? 30000); + }, HEARTBEAT_MS); + + // Sampled event-loop lag: a zero-delay timer that fires late measures how + // far behind the loop is running. + function measureLag() { + const start = Date.now(); + setTimeout(() => { + metrics.eventLoopLagMs = Date.now() - start; + }, 0); + } + const lagInterval = setInterval(measureLag, LAG_SAMPLE_MS); + measureLag(); + + function markShuttingDown() { + isShuttingDown = true; + isReady = false; + } wss.on("close", () => { clearInterval(heartbeatInterval); + clearInterval(lagInterval); + if (ownsSessions) sessions.close(); httpServer.close(); }); diff --git a/src/session-manager.js b/src/session-manager.js index f530628..2d1806f 100644 --- a/src/session-manager.js +++ b/src/session-manager.js @@ -1,387 +1,545 @@ -import crypto from "node:crypto"; -import zlib from "node:zlib"; - -const DEFAULT_TTL_MS = 3600000; -const DEBOUNCE_MS = 500; -const MAX_BLOB_SIZE = 16384; - /** - * @typedef {Object} SessionRoom - * @property {string} roomId - * @property {number} highestAckedSeq - * @property {number} highestReceivedSeq - * @property {string[]} geofenceInsideSet + * @fileoverview Encrypted session state for connection migration and resumption. + * + * A client's session state is serialised, deflated and sealed with AES-256-GCM. + * The resulting blob IS the `session_id` the client presents on reconnect, so a + * different gateway instance can restore room memberships, sequence numbers, + * geofence state and rate-limit windows without extra round-trips. + * + * Storage holds the same blob under `session:` with a sliding TTL. + * A blob only resumes while its storage entry is still live, which makes + * `delete()` and TTL expiry authoritative even though the blob is self-contained. + * + * Redis is optional: without it the manager keeps an in-memory map with expiry + * timestamps, so single-instance deployments need no extra infrastructure. */ -/** - * @typedef {Object} RateLimitState - * @property {number[]} messageWindow - * @property {number[]} connectionWindow - */ +import { createCipheriv, createDecipheriv, hkdfSync, randomBytes } from "node:crypto"; +import { deflateSync, inflateSync } from "node:zlib"; +import { logger as defaultLogger } from "./logger.js"; + +const ALGORITHM = "aes-256-gcm"; +const IV_BYTES = 12; +const TAG_BYTES = 16; +const KEY_BYTES = 32; +const HKDF_INFO = "session-key"; +const MAX_BLOB_BYTES = 16384; +const KEY_PREFIX = "session:"; +const DEFAULT_TTL_MS = 3600000; +const DEFAULT_DEBOUNCE_MS = 500; +const MIN_SWEEP_MS = 1000; +const MAX_SWEEP_MS = 60000; /** - * @typedef {Object} SessionMetadata - * @property {string} ip - * @property {string} userAgent - * @property {number} connectedAt - * @property {number} lastActivityAt + * @typedef {object} SessionRoomState + * @property {string} roomId + * @property {number} highestAckedSeq - Highest sequence number the client ACKed. + * @property {number} highestReceivedSeq - Highest sequence number sent to the client. + * @property {string[]} geofenceInsideSet - Fence IDs the client is currently inside. */ /** - * @typedef {Object} SessionState + * Session state is opaque to this class: it is serialised as-is and never + * validated beyond needing a `clientId` to locate the storage entry. + * + * @typedef {object} SessionState * @property {string} clientId - * @property {number} protocolVersion - * @property {Object} authIdentity - * @property {SessionRoom[]} rooms - * @property {RateLimitState} rateLimitState - * @property {SessionMetadata} metadata + * @property {number} [protocolVersion] + * @property {object} [authIdentity] - JWT claims or mTLS device identity. + * @property {SessionRoomState[]} [rooms] + * @property {{ messageWindow?: number[], connectionWindow?: number[] }} [rateLimitState] + * @property {{ ip?: string, userAgent?: string, connectedAt?: number, lastActivityAt?: number }} [metadata] */ /** - * @typedef {Object} SessionManagerOptions - * @property {Object} [redis] - Optional Redis client with get/set/del methods - * @property {string|Object} encryptionKey - Base64 encoded key or { keyId: base64key } map - * @property {number} [ttlMs] - Session TTL in milliseconds - * @property {string} [keyId] - Current key identifier for encryption - * @property {number} [debounceMs] - Debounce interval for saves + * Minimal subset of the node-redis v4 client used by this module. + * + * @typedef {object} SessionStore + * @property {(key: string, value: string, opts: { PX: number }) => Promise} set + * @property {(key: string) => Promise} get + * @property {(key: string) => Promise} del */ +/** @typedef {"success"|"decrypt_failed"|"expired"|"mismatch"|"new_session"} ResumptionResult */ + /** - * Derives an AES-256 key from a master key using HKDF. + * Marks a timer as non-blocking so a pending save never holds the event loop + * open. Fake timers may omit `unref`, hence the guard. * - * @param {Buffer} masterKey - * @param {string} info - * @returns {Buffer} + * @param {any} timer + * @returns {any} The same timer. */ -function deriveKey(masterKey, info) { - return crypto.hkdfSync("sha256", masterKey, Buffer.alloc(0), info, 32); +function unrefTimer(timer) { + if (timer && typeof timer.unref === "function") timer.unref(); + return timer; } /** - * Resolves the raw key bytes for a given key ID. + * Decodes one base64 master key and stretches it into an AES-256 key. * - * @param {string|Object} encryptionKey - * @param {string} keyId - * @returns {Buffer|null} + * @param {string} keyId - Key name, used only in error messages. + * @param {unknown} value - Base64-encoded 32-byte master key. + * @returns {Buffer} 32-byte derived key. */ -function resolveKey(encryptionKey, keyId) { - if (typeof encryptionKey === "string") { - return Buffer.from(encryptionKey, "base64"); +function deriveKey(keyId, value) { + if (typeof value !== "string" || value.trim().length === 0) { + throw new TypeError(`encryptionKey["${keyId}"] must be a non-empty base64 string`); } - if (typeof encryptionKey === "object" && encryptionKey[keyId]) { - return Buffer.from(encryptionKey[keyId], "base64"); + const master = Buffer.from(value, "base64"); + if (master.length !== KEY_BYTES) { + throw new RangeError( + `encryptionKey["${keyId}"] must decode to ${KEY_BYTES} bytes, got ${master.length}` + ); } - return null; + return Buffer.from(hkdfSync("sha256", master, "", HKDF_INFO, KEY_BYTES)); } /** - * Encrypts a session state blob using AES-256-GCM. + * Normalises the `encryptionKey` option into a `keyId → derived key` map. + * Accepts a bare base64 key (bound to `keyId`), a JSON string, or an object. * - * @param {Buffer} plaintext - * @param {Buffer} key - * @param {string} keyId - * @returns {string} Encrypted blob in format: keyId.base64iv.base64ciphertext.base64tag + * @param {unknown} encryptionKey + * @param {string} keyId - Key used to seal new blobs; must exist in the result. + * @returns {Map} */ -function encryptBlob(plaintext, key, keyId) { - const derivedKey = deriveKey(key, `session-key-${keyId}`); - const iv = crypto.randomBytes(12); - const cipher = crypto.createCipheriv("aes-256-gcm", derivedKey, iv); - const ciphertext = Buffer.concat([cipher.update(plaintext), cipher.final()]); - const tag = cipher.getAuthTag(); - return `${keyId}.${iv.toString("base64")}.${ciphertext.toString("base64")}.${tag.toString("base64")}`; -} +function parseKeyMap(encryptionKey, keyId) { + if (encryptionKey == null) { + throw new TypeError( + "encryptionKey is required: a 32-byte base64 key, or a { keyId: base64key } map" + ); + } -/** - * Decrypts a session state blob. - * - * @param {string} blob - * @param {string|Object} encryptionKey - * @returns {Buffer|null} Decrypted plaintext or null if decryption fails - */ -function decryptBlob(blob, encryptionKey) { - const parts = blob.split("."); - if (parts.length !== 4) return null; + let source = encryptionKey; + if (typeof source === "string") { + const trimmed = source.trim(); + if (trimmed.length === 0) throw new TypeError("encryptionKey must not be empty"); + if (trimmed.startsWith("{")) { + try { + source = JSON.parse(trimmed); + } catch (err) { + throw new TypeError(`encryptionKey looks like JSON but does not parse: ${err.message}`); + } + } else { + source = { [keyId]: trimmed }; + } + } + + if (typeof source !== "object" || Array.isArray(source)) { + throw new TypeError("encryptionKey must be a base64 string, a JSON string, or an object"); + } - const [keyId, ivB64, ciphertextB64, tagB64] = parts; - const key = resolveKey(encryptionKey, keyId); - if (!key) return null; + const entries = Object.entries(source); + if (entries.length === 0) throw new TypeError("encryptionKey holds no keys"); - try { - const derivedKey = deriveKey(key, `session-key-${keyId}`); - const iv = Buffer.from(ivB64, "base64"); - const ciphertext = Buffer.from(ciphertextB64, "base64"); - const tag = Buffer.from(tagB64, "base64"); - const decipher = crypto.createDecipheriv("aes-256-gcm", derivedKey, iv); - decipher.setAuthTag(tag); - return Buffer.concat([decipher.update(ciphertext), decipher.final()]); - } catch { - return null; + const keys = new Map(); + for (const [id, value] of entries) keys.set(id, deriveKey(id, value)); + + if (!keys.has(keyId)) { + throw new RangeError( + `keyId "${keyId}" is absent from encryptionKey (have: ${[...keys.keys()].join(", ")})` + ); } + return keys; } /** - * SessionManager handles encrypted session state persistence for WebSocket - * connection migration and resumption. + * In-memory stand-in for Redis with per-entry expiry timestamps. * - * Supports both Redis-backed distributed storage and in-memory fallback - * for single-instance mode. - * - * @example - * const sm = new SessionManager({ encryptionKey: "base64key..." }); - * await sm.save("client-1", { clientId: "client-1", rooms: [] }); - * const state = await sm.load("client-1"); + * @implements {SessionStore} */ -export class SessionManager { - /** @type {Object|null} */ - #redis; - - /** @type {string|Object} */ - #encryptionKey; - - /** @type {number} */ - #ttlMs; - - /** @type {string} */ - #keyId; +class MemoryStore { + constructor() { + /** @type {Map} */ + this._entries = new Map(); + } - /** @type {number} */ - #debounceMs; + /** + * @param {string} key + * @param {string} value + * @param {{ PX: number }} opts - Millisecond TTL, matching node-redis. + */ + async set(key, value, { PX }) { + this._entries.set(key, { value, expiresAt: Date.now() + PX }); + } - /** @type {Map} */ - #timers; + /** + * @param {string} key + * @returns {Promise} Null when absent or expired. + */ + async get(key) { + const entry = this._entries.get(key); + if (!entry) return null; + if (entry.expiresAt <= Date.now()) { + this._entries.delete(key); + return null; + } + return entry.value; + } - /** @type {Map} */ - #pendingStates; + /** @param {string} key */ + async del(key) { + this._entries.delete(key); + } - /** @type {Map} */ - #localCache; + /** Drops every expired entry so idle sessions do not leak. */ + sweep() { + const now = Date.now(); + for (const [key, entry] of this._entries) { + if (entry.expiresAt <= now) this._entries.delete(key); + } + } - /** @type {NodeJS.Timeout|null} */ - #cleanupInterval; + /** @type {number} */ + get size() { + return this._entries.size; + } +} +/** + * Seals, stores and restores per-client session state. + */ +export class SessionManager { /** - * @param {SessionManagerOptions} options + * @param {object} options + * @param {SessionStore|null} [options.redis] - node-redis v4 style client. + * Omit for the in-memory fallback. + * @param {string|object} options.encryptionKey - 32-byte base64 key, or a + * `{ keyId: base64key }` map (object or JSON string) for key rotation. + * @param {number} [options.ttlMs=3600000] - Sliding TTL of a stored session. + * @param {string} [options.keyId="v1"] - Key that seals new blobs. + * @param {number} [options.debounceMs=500] - Debounce window per client. + * @param {{ error: Function }} [options.logger] - Sink for background failures. */ - constructor({ redis, encryptionKey, ttlMs, keyId, debounceMs } = {}) { - this.#redis = redis ?? null; - this.#encryptionKey = encryptionKey ?? process.env.SESSION_ENCRYPTION_KEY ?? ""; - this.#ttlMs = ttlMs ?? DEFAULT_TTL_MS; - this.#keyId = keyId ?? "v1"; - this.#debounceMs = debounceMs ?? DEBOUNCE_MS; - this.#timers = new Map(); - this.#pendingStates = new Map(); - this.#localCache = new Map(); - this.#cleanupInterval = null; - - if (!this.#redis) { - this.#cleanupInterval = setInterval(() => { - this.#evictExpired(); - }, Math.min(this.#ttlMs, 60000)); - if (this.#cleanupInterval.unref) { - this.#cleanupInterval.unref(); - } + constructor({ + redis = null, + encryptionKey, + ttlMs = DEFAULT_TTL_MS, + keyId = "v1", + debounceMs = DEFAULT_DEBOUNCE_MS, + logger = defaultLogger, + } = {}) { + if (typeof keyId !== "string" || keyId.length === 0) { + throw new TypeError("keyId must be a non-empty string"); + } + if (keyId.includes(".")) { + throw new TypeError('keyId must not contain "." — it separates blob fields'); + } + if (!Number.isFinite(ttlMs) || ttlMs <= 0) { + throw new RangeError("ttlMs must be a positive number of milliseconds"); + } + if (!Number.isFinite(debounceMs) || debounceMs < 0) { + throw new RangeError("debounceMs must be a non-negative number of milliseconds"); + } + + this._keys = parseKeyMap(encryptionKey, keyId); + this._keyId = keyId; + this._ttlMs = ttlMs; + this._debounceMs = debounceMs; + this._logger = logger; + + this._redis = redis ?? null; + this._memory = this._redis ? null : new MemoryStore(); + /** @type {any} */ + this._sweepTimer = null; + if (this._memory) { + const period = Math.max(MIN_SWEEP_MS, Math.min(ttlMs, MAX_SWEEP_MS)); + this._sweepTimer = unrefTimer(setInterval(() => this._memory.sweep(), period)); } + + /** @type {Map} clientId → pending save timer */ + this._debounceTimers = new Map(); + /** @type {Map SessionState>} clientId → latest state provider */ + this._debounceProviders = new Map(); + + this._counters = { + success: 0, + decrypt_failed: 0, + expired: 0, + mismatch: 0, + new_session: 0, + }; + this._stateSizeBytes = 0; } /** - * Removes expired entries from the local in-memory cache. * @private + * @returns {SessionStore} */ - #evictExpired() { - const now = Date.now(); - for (const [clientId, entry] of this.#localCache) { - if (entry.expiresAt <= now) { - this.#localCache.delete(clientId); - } - } + _storage() { + return this._redis ?? this._memory; } /** - * Compresses and encrypts a session state, then stores it. + * Seals `state` and stores it under `session:`, refreshing the TTL. * * @param {string} clientId - * @param {SessionState} state - * @returns {Promise} + * @param {SessionState} state - Opaque session state; must be serialisable. + * @returns {Promise} The blob to hand back as the client's `session_id`. + * @throws {RangeError} When the blob exceeds 16 KB. */ async save(clientId, state) { - const plaintext = Buffer.from(JSON.stringify(state), "utf8"); - const compressed = zlib.deflateSync(plaintext, { level: 6 }); + if (typeof clientId !== "string" || clientId.length === 0) { + throw new TypeError("clientId must be a non-empty string"); + } + if (state == null || typeof state !== "object") { + throw new TypeError("state must be an object"); + } - if (compressed.length > MAX_BLOB_SIZE) { - throw new Error(`Session blob exceeds ${MAX_BLOB_SIZE} bytes after compression`); + const compressed = deflateSync(Buffer.from(JSON.stringify(state), "utf8")); + const iv = randomBytes(IV_BYTES); + const cipher = createCipheriv(ALGORITHM, this._keys.get(this._keyId), iv); + const ciphertext = Buffer.concat([cipher.update(compressed), cipher.final()]); + const blob = [ + this._keyId, + iv.toString("base64"), + ciphertext.toString("base64"), + cipher.getAuthTag().toString("base64"), + ].join("."); + + this._stateSizeBytes = Buffer.byteLength(blob, "utf8"); + if (this._stateSizeBytes > MAX_BLOB_BYTES) { + throw new RangeError( + `session blob for ${clientId} is ${this._stateSizeBytes} bytes, over the ${MAX_BLOB_BYTES} byte limit` + ); } - const blob = encryptBlob(compressed, this.#resolveEncryptionKey(), this.#keyId); - const expiresAt = Date.now() + this.#ttlMs; - const ttlSeconds = Math.ceil(this.#ttlMs / 1000); + await this._storage().set(KEY_PREFIX + clientId, blob, { PX: this._ttlMs }); + return blob; + } - if (this.#redis) { - await this.#redis.set(`session:${clientId}`, blob, "EX", ttlSeconds); - } else { - this.#localCache.set(clientId, { state, expiresAt }); + /** + * Opens a `session_id` blob and confirms the session is still live. + * + * Never throws: every failure (bad format, unknown key, tampering, expiry, + * explicit delete) returns null and bumps the matching counter. A storage + * read error counts as `expired` because liveness cannot be proven. + * + * @param {string} sessionId - The blob returned by `save()`. + * @param {{ countSuccess?: boolean }} [options] - Pass `countSuccess: false` + * when the caller still has its own checks (e.g. identity binding) and + * records the final outcome itself. + * @returns {Promise} + */ + async load(sessionId, { countSuccess = true } = {}) { + const state = this._open(sessionId); + if (!state) { + this._counters.decrypt_failed++; + return null; } + + let stored; + try { + stored = await this._storage().get(KEY_PREFIX + state.clientId); + } catch (err) { + this._logger?.error?.("session store read failed", { + clientId: state.clientId, + error: err.message, + }); + this._counters.expired++; + return null; + } + + if (stored == null) { + this._counters.expired++; + return null; + } + + if (countSuccess) this._counters.success++; + return state; } /** - * Loads and decrypts a session state by its client ID. + * Parses and decrypts a blob. The key named by the blob prefix is tried + * first, then every other key so blobs sealed before a rotation still open. * - * @param {string} sessionId - The client ID used as session identifier - * @returns {Promise} + * @private + * @param {unknown} sessionId + * @returns {SessionState|null} */ - async load(sessionId) { - if (this.#redis) { - const blob = await this.#redis.get(`session:${sessionId}`); - if (!blob) return null; + _open(sessionId) { + if (typeof sessionId !== "string" || sessionId.length === 0) return null; + + const parts = sessionId.split("."); + if (parts.length !== 4) return null; - const decrypted = decryptBlob(blob, this.#encryptionKey); - if (!decrypted) return null; + const [blobKeyId, ivB64, ciphertextB64, tagB64] = parts; + const iv = Buffer.from(ivB64, "base64"); + const ciphertext = Buffer.from(ciphertextB64, "base64"); + const tag = Buffer.from(tagB64, "base64"); + if (iv.length !== IV_BYTES || tag.length !== TAG_BYTES || ciphertext.length === 0) { + return null; + } + + const candidates = []; + if (this._keys.has(blobKeyId)) candidates.push(this._keys.get(blobKeyId)); + for (const [id, key] of this._keys) { + if (id !== blobKeyId) candidates.push(key); + } + for (const key of candidates) { try { - const decompressed = zlib.inflateSync(decrypted); - const json = JSON.parse(decompressed.toString("utf8")); - if (json.clientId !== sessionId) return null; - return json; + const decipher = createDecipheriv(ALGORITHM, key, iv); + decipher.setAuthTag(tag); + const plaintext = Buffer.concat([decipher.update(ciphertext), decipher.final()]); + const state = JSON.parse(inflateSync(plaintext).toString("utf8")); + // An authentic blob without a clientId cannot be matched to storage. + if (state == null || typeof state !== "object" || typeof state.clientId !== "string") { + return null; + } + return state; } catch { - return null; + // Wrong key or corrupt payload — fall through to the next candidate. } } - - const entry = this.#localCache.get(sessionId); - if (!entry) return null; - if (entry.expiresAt <= Date.now()) { - this.#localCache.delete(sessionId); - return null; - } - return entry.state; + return null; } /** - * Deletes a session from storage. + * Coalesces rapid state changes into one save per `debounceMs` per client. + * `stateProvider` runs when the timer fires, so the newest state is captured. * * @param {string} clientId - * @returns {Promise} + * @param {() => SessionState} stateProvider + * @returns {Promise} Resolves once the save is scheduled, not stored. */ - async delete(clientId) { - this.#clearDebounce(clientId); - if (this.#redis) { - await this.#redis.del(`session:${clientId}`); - } else { - this.#localCache.delete(clientId); + async debouncedSave(clientId, stateProvider) { + if (typeof clientId !== "string" || clientId.length === 0) { + throw new TypeError("clientId must be a non-empty string"); } + if (typeof stateProvider !== "function") { + throw new TypeError("stateProvider must be a function returning the latest state"); + } + + this._cancelDebounce(clientId); + this._debounceProviders.set(clientId, stateProvider); + this._debounceTimers.set( + clientId, + unrefTimer( + setTimeout(() => { + this._debounceTimers.delete(clientId); + const provider = this._debounceProviders.get(clientId); + this._debounceProviders.delete(clientId); + if (provider) this._runSave(clientId, provider); + }, this._debounceMs) + ) + ); } /** - * Saves immediately without debouncing. Used for graceful shutdown. + * Runs a pending save. Failures are logged rather than thrown so a + * background timer never produces an unhandled rejection. * + * @private * @param {string} clientId - * @param {SessionState} state - * @returns {Promise} + * @param {() => SessionState} stateProvider + * @returns {Promise} The blob, or null on failure. */ - async saveImmediate(clientId, state) { - this.#clearDebounce(clientId); - await this.save(clientId, state); + _runSave(clientId, stateProvider) { + let state; + try { + state = stateProvider(); + } catch (err) { + this._logger?.error?.("session stateProvider threw", { clientId, error: err.message }); + return Promise.resolve(null); + } + return this.save(clientId, state).catch((err) => { + this._logger?.error?.("session save failed", { clientId, error: err.message }); + return null; + }); } /** - * Debounced save that coalesces rapid state changes. - * + * @private * @param {string} clientId - * @param {SessionState} state - * @returns {void} + * @returns {(() => SessionState)|null} The cancelled provider, if any. */ - debouncedSave(clientId, state) { - const existing = this.#timers.get(clientId); - if (existing) clearTimeout(existing); - this.#pendingStates.set(clientId, state); - this.#timers.set(clientId, setTimeout(() => { - this.#timers.delete(clientId); - const pending = this.#pendingStates.get(clientId); - this.#pendingStates.delete(clientId); - if (pending) { - this.save(clientId, pending).catch(() => {}); - } - }, this.#debounceMs)); + _cancelDebounce(clientId) { + const timer = this._debounceTimers.get(clientId); + if (timer) clearTimeout(timer); + this._debounceTimers.delete(clientId); + const provider = this._debounceProviders.get(clientId) ?? null; + this._debounceProviders.delete(clientId); + return provider; } /** - * Clears a pending debounce timer for a client. + * Fires a client's pending debounced save immediately. * * @param {string} clientId - * @private + * @returns {Promise} The blob, or null when nothing was pending. */ - #clearDebounce(clientId) { - const timer = this.#timers.get(clientId); - if (timer) { - clearTimeout(timer); - this.#timers.delete(clientId); + async flush(clientId) { + if (!this._debounceTimers.has(clientId)) { + this._cancelDebounce(clientId); + return null; } - this.#pendingStates.delete(clientId); + const provider = this._cancelDebounce(clientId); + return provider ? this._runSave(clientId, provider) : null; } /** - * Returns the number of pending debounced saves. - * @returns {number} + * Fires every pending save — call before closing connections on shutdown. + * + * @returns {Promise>} One blob per flushed client. */ - get pendingSaves() { - return this.#timers.size; + async flushAll() { + const pending = [...this._debounceTimers.keys()]; + return Promise.all(pending.map((clientId) => this.flush(clientId))); } /** - * Returns the number of sessions in local cache (single-instance mode). - * @returns {number} + * Drops a client's stored session and cancels any pending save. + * + * @param {string} clientId + * @returns {Promise} */ - get cachedSessions() { - return this.#localCache.size; + async delete(clientId) { + if (typeof clientId !== "string" || clientId.length === 0) { + throw new TypeError("clientId must be a non-empty string"); + } + this._cancelDebounce(clientId); + await this._storage().del(KEY_PREFIX + clientId); } /** - * Flushes all pending debounced saves immediately. + * Clears every timer so the process (or a test run) can exit. Idempotent. + * * @returns {Promise} */ - async flushPending() { - const entries = []; - for (const [clientId, timer] of this.#timers) { - clearTimeout(timer); - const pending = this.#pendingStates.get(clientId); - entries.push({ clientId, state: pending }); + async close() { + for (const clientId of [...this._debounceTimers.keys()]) { + this._cancelDebounce(clientId); } - this.#timers.clear(); - this.#pendingStates.clear(); - for (const { clientId, state } of entries) { - if (state) { - await this.save(clientId, state).catch(() => {}); - } + if (this._sweepTimer) { + clearInterval(this._sweepTimer); + this._sweepTimer = null; } } /** - * Cleans up resources. Clears timers and intervals. - * @returns {void} + * Counts a resumption outcome decided outside this class, such as an + * identity `mismatch` or a `new_session`. + * + * @param {ResumptionResult} result */ - destroy() { - for (const timer of this.#timers.values()) { - clearTimeout(timer); - } - this.#timers.clear(); - if (this.#cleanupInterval) { - clearInterval(this.#cleanupInterval); - this.#cleanupInterval = null; + recordResumption(result) { + if (!Object.hasOwn(this._counters, result)) { + throw new RangeError( + `unknown resumption result "${result}" (expected: ${Object.keys(this._counters).join("|")})` + ); } + this._counters[result]++; } /** - * Resolves the encryption key to a Buffer. + * Snapshot of the issue-18 metrics. Mutating it does not affect the manager. * - * @returns {Buffer} - * @private + * @type {{ session_resumption_total: Record, session_state_size_bytes: number }} */ - #resolveEncryptionKey() { - if (typeof this.#encryptionKey === "string" && this.#encryptionKey) { - return Buffer.from(this.#encryptionKey, "base64"); - } - if (typeof this.#encryptionKey === "object" && this.#encryptionKey[this.#keyId]) { - return Buffer.from(this.#encryptionKey[this.#keyId], "base64"); - } - return crypto.randomBytes(32); + get metrics() { + return { + session_resumption_total: { ...this._counters }, + session_state_size_bytes: this._stateSizeBytes, + }; } } diff --git a/src/validator.js b/src/validator.js index 4fdfb0f..106684c 100644 --- a/src/validator.js +++ b/src/validator.js @@ -40,6 +40,10 @@ const tokenRefreshSchema = z.object({ token: z.string().min(1), }); +const tokenRefreshSchema = z.object({ + token: z.string().min(1), +}); + const messageSchema = z.discriminatedUnion("type", [ z.object({ type: z.literal("location_update"), diff --git a/tests/session-manager.test.js b/tests/session-manager.test.js index 786629a..0104f7f 100644 --- a/tests/session-manager.test.js +++ b/tests/session-manager.test.js @@ -1,359 +1,659 @@ -import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; -import crypto from "node:crypto"; +/** + * @fileoverview Unit tests for the encrypted session resumption store. + * + * Covers AES-256-GCM round-trips, every `load()` rejection path, sliding TTL, + * key rotation, debounced saves, both storage backends, the 16 KB blob budget + * and the issue-18 metrics. + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { randomBytes } from "node:crypto"; import { SessionManager } from "../src/session-manager.js"; +// ─── helpers ────────────────────────────────────────────────────────────────── + +const KEY_A = randomBytes(32).toString("base64"); +const KEY_B = randomBytes(32).toString("base64"); + +/** Keeps expected background failures out of the test output. */ +const silentLogger = { error: () => {} }; + +/** @type {SessionManager[]} */ +let managers = []; + +/** + * Builds a manager and registers it for teardown so no timer outlives a test. + * + * @param {object} [options] + * @returns {SessionManager} + */ +function newManager(options = {}) { + const manager = new SessionManager({ + encryptionKey: KEY_A, + logger: silentLogger, + ...options, + }); + managers.push(manager); + return manager; +} + /** - * Returns a valid session state object. - * @param {Partial} [overrides] + * Fake node-redis v4 client recording every call. + * + * @returns {{ entries: Map, set: Function, get: Function, del: Function }} + */ +function makeFakeRedis() { + const entries = new Map(); + return { + entries, + set: vi.fn(async (key, value, opts) => { + entries.set(key, { value, opts }); + return "OK"; + }), + get: vi.fn(async (key) => entries.get(key)?.value ?? null), + del: vi.fn(async (key) => (entries.delete(key) ? 1 : 0)), + }; +} + +/** + * @param {object} [overrides] * @returns {import("../src/session-manager.js").SessionState} */ function makeState(overrides = {}) { return { clientId: "client-001", protocolVersion: 3, - authIdentity: { sub: "device-001", iss: "fleet-auth" }, + authIdentity: { sub: "device-123", iss: "fleet-auth" }, rooms: [ - { roomId: "fleet-alpha", highestAckedSeq: 42, highestReceivedSeq: 45, geofenceInsideSet: ["fence-1"] }, + { + roomId: "fleet-alpha", + highestAckedSeq: 42, + highestReceivedSeq: 45, + geofenceInsideSet: ["fence-1", "fence-3"], + }, ], - rateLimitState: { messageWindow: [1000, 2000], connectionWindow: [500] }, - metadata: { ip: "10.0.0.1", userAgent: "FleetApp/2.3", connectedAt: 1000, lastActivityAt: 2000 }, + rateLimitState: { messageWindow: [1001, 1002, 1003], connectionWindow: [900] }, + metadata: { + ip: "10.0.0.1", + userAgent: "FleetApp/2.3", + connectedAt: 1700000000000, + lastActivityAt: 1700000009000, + }, ...overrides, }; } -function makeTestKey() { - return crypto.randomBytes(32).toString("base64"); +/** + * Session state for a client subscribed to 50 rooms with full per-room state. + * + * @returns {import("../src/session-manager.js").SessionState} + */ +function makeFiftyRoomState() { + const now = 1700000000000; + const rooms = Array.from({ length: 50 }, (_, i) => ({ + roomId: `fleet-region-${i}-vehicles`, + highestAckedSeq: 100000 + i * 7, + highestReceivedSeq: 100010 + i * 7, + geofenceInsideSet: [`fence-${i}-depot`, `fence-${i}-zone-a`, `fence-${i}-zone-b`], + })); + return { + clientId: "client-fleet-050", + protocolVersion: 3, + authIdentity: { sub: "device-abcdef123456", iss: "fleet-auth", aud: "gateway", exp: now }, + rooms, + rateLimitState: { + messageWindow: Array.from({ length: 100 }, (_, i) => now - i * 9), + connectionWindow: Array.from({ length: 10 }, (_, i) => now - i * 1000), + }, + metadata: { + ip: "10.42.0.17", + userAgent: "FleetApp/2.3 (iOS 17.4)", + connectedAt: now - 60000, + lastActivityAt: now, + }, + }; +} + +/** + * Flips one ciphertext byte, leaving the blob structurally valid. + * + * @param {string} blob + * @returns {string} + */ +function corruptCiphertext(blob) { + const parts = blob.split("."); + const ciphertext = Buffer.from(parts[2], "base64"); + ciphertext[0] ^= 0xff; + parts[2] = ciphertext.toString("base64"); + return parts.join("."); +} + +/** + * @param {number} ms + * @returns {Promise} + */ +function sleep(ms) { + return new Promise((resolve) => setTimeout(resolve, ms)); } -describe("SessionManager", () => { - let sm; - const encryptionKey = makeTestKey(); +afterEach(async () => { + for (const manager of managers) await manager.close(); + managers = []; + vi.useRealTimers(); + vi.restoreAllMocks(); +}); - beforeEach(() => { - sm = new SessionManager({ encryptionKey, ttlMs: 5000, debounceMs: 10 }); +// ─── key material ───────────────────────────────────────────────────────────── + +describe("SessionManager key material", () => { + it("throws when encryptionKey is missing", () => { + expect(() => new SessionManager({})).toThrow(/encryptionKey is required/); }); - afterEach(() => { - sm.destroy(); + it("throws when the base64 key is not 32 bytes", () => { + expect(() => new SessionManager({ encryptionKey: randomBytes(16).toString("base64") })).toThrow( + /must decode to 32 bytes, got 16/ + ); }); - describe("constructor", () => { - it("creates an instance with default options", () => { - const s = new SessionManager(); - expect(s).toBeDefined(); - s.destroy(); - }); + it("throws when a key map entry is not a string", () => { + expect(() => new SessionManager({ encryptionKey: { v1: 42 } })).toThrow( + /must be a non-empty base64 string/ + ); + }); - it("creates an instance with custom options", () => { - const s = new SessionManager({ encryptionKey, ttlMs: 10000, keyId: "v2", debounceMs: 100 }); - expect(s).toBeDefined(); - s.destroy(); - }); + it("throws when keyId is absent from the key map", () => { + expect(() => new SessionManager({ encryptionKey: { v1: KEY_A }, keyId: "v9" })).toThrow( + /keyId "v9" is absent/ + ); }); - describe("encryption/decryption", () => { - it("save and load preserves session state", async () => { - const state = makeState(); - await sm.save("client-001", state); - const loaded = await sm.load("client-001"); - expect(loaded).toEqual(state); - }); + it("throws when keyId contains the blob separator", () => { + expect(() => new SessionManager({ encryptionKey: KEY_A, keyId: "v.1" })).toThrow( + /must not contain/ + ); + }); - it("returns null for non-existent session", async () => { - const loaded = await sm.load("non-existent"); - expect(loaded).toBeNull(); - }); + it("throws on malformed JSON key material", () => { + expect(() => new SessionManager({ encryptionKey: '{"v1": ' })).toThrow(/does not parse/); + }); - it("returns null for corrupted session blob", async () => { - await sm.save("client-001", makeState()); - const loaded = await sm.load("client-001"); - expect(loaded).not.toBeNull(); - }); + it("accepts a JSON string key map", async () => { + const manager = newManager({ encryptionKey: JSON.stringify({ v1: KEY_A, v2: KEY_B }) }); + const blob = await manager.save("client-001", makeState()); + expect(blob.startsWith("v1.")).toBe(true); + await expect(manager.load(blob)).resolves.toEqual(makeState()); + }); +}); - it("rejects session with mismatched clientId", async () => { - const state = makeState({ clientId: "client-001" }); - await sm.save("client-001", state); - const loaded = await sm.load("client-002"); - expect(loaded).toBeNull(); - }); +// ─── save / load round trip ─────────────────────────────────────────────────── + +describe("SessionManager save/load", () => { + it("round-trips state through encryption and compression", async () => { + const manager = newManager(); + const state = makeState(); + const blob = await manager.save("client-001", state); + + await expect(manager.load(blob)).resolves.toEqual(state); }); - describe("key rotation", () => { - it("supports loading sessions encrypted with different keys via shared Redis", async () => { - const key1 = makeTestKey(); - const key2 = makeTestKey(); - const multiKey = { v1: key1, v2: key2 }; - - const store = new Map(); - const redisMock = { - async set(key, value, _ex, _ttl) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; - - const sm1 = new SessionManager({ redis: redisMock, encryptionKey: multiKey, keyId: "v1", debounceMs: 10 }); - const state = makeState(); - await sm1.save("client-001", state); - sm1.destroy(); - - const sm2 = new SessionManager({ redis: redisMock, encryptionKey: multiKey, keyId: "v2", debounceMs: 10 }); - const loaded = await sm2.load("client-001"); - expect(loaded).toEqual(state); - sm2.destroy(); - }); + it("produces a keyId.iv.ciphertext.tag blob", async () => { + const manager = newManager({ keyId: "v7", encryptionKey: { v7: KEY_A } }); + const blob = await manager.save("client-001", makeState()); + const parts = blob.split("."); - it("fails to load with wrong key version", async () => { - const key1 = makeTestKey(); - const key2 = makeTestKey(); - - const store = new Map(); - const redisMock = { - async set(key, value) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; - - const sm1 = new SessionManager({ redis: redisMock, encryptionKey: { v1: key1 }, keyId: "v1", debounceMs: 10 }); - await sm1.save("client-001", makeState()); - sm1.destroy(); - - const sm2 = new SessionManager({ redis: redisMock, encryptionKey: { v3: key2 }, keyId: "v3", debounceMs: 10 }); - const loaded = await sm2.load("client-001"); - expect(loaded).toBeNull(); - sm2.destroy(); - }); + expect(parts).toHaveLength(4); + expect(parts[0]).toBe("v7"); + expect(Buffer.from(parts[1], "base64")).toHaveLength(12); + expect(Buffer.from(parts[2], "base64").length).toBeGreaterThan(0); + expect(Buffer.from(parts[3], "base64")).toHaveLength(16); }); - describe("delete", () => { - it("removes a session", async () => { - await sm.save("client-001", makeState()); - expect(await sm.load("client-001")).not.toBeNull(); - await sm.delete("client-001"); - expect(await sm.load("client-001")).toBeNull(); - }); + it("never stores plaintext state", async () => { + const redis = makeFakeRedis(); + const manager = newManager({ redis }); + const blob = await manager.save("client-001", makeState()); - it("is idempotent for non-existent session", async () => { - await expect(sm.delete("non-existent")).resolves.toBeUndefined(); - }); + expect(blob).not.toContain("fleet-alpha"); + expect(redis.entries.get("session:client-001").value).not.toContain("device-123"); }); - describe("debouncedSave", () => { - it("debounces rapid saves", async () => { - const saveSpy = vi.spyOn(sm, "save"); - const state1 = makeState({ rooms: [{ roomId: "r1", highestAckedSeq: 1, highestReceivedSeq: 1, geofenceInsideSet: [] }] }); - const state2 = makeState({ rooms: [{ roomId: "r1", highestAckedSeq: 2, highestReceivedSeq: 2, geofenceInsideSet: [] }] }); + it("uses a fresh IV for every save", async () => { + const manager = newManager(); + const first = await manager.save("client-001", makeState()); + const second = await manager.save("client-001", makeState()); - sm.debouncedSave("client-001", state1); - sm.debouncedSave("client-001", state2); + expect(first.split(".")[1]).not.toBe(second.split(".")[1]); + await expect(manager.load(first)).resolves.toEqual(makeState()); + }); + + it("rejects a non-object state and an empty clientId", async () => { + const manager = newManager(); + await expect(manager.save("client-001", null)).rejects.toThrow(/state must be an object/); + await expect(manager.save("", makeState())).rejects.toThrow(/clientId/); + }); +}); - expect(sm.pendingSaves).toBe(1); +// ─── load failure paths ─────────────────────────────────────────────────────── - await new Promise((r) => setTimeout(r, 50)); - expect(saveSpy).toHaveBeenCalledTimes(1); +describe("SessionManager load failures", () => { + it("returns null for a corrupted ciphertext byte", async () => { + const manager = newManager(); + const blob = await manager.save("client-001", makeState()); - const loaded = await sm.load("client-001"); - expect(loaded.rooms[0].highestAckedSeq).toBe(2); - }); + await expect(manager.load(corruptCiphertext(blob))).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.decrypt_failed).toBe(1); }); - describe("saveImmediate", () => { - it("saves immediately bypassing debounce", async () => { - const state = makeState(); - await sm.saveImmediate("client-001", state); - const loaded = await sm.load("client-001"); - expect(loaded).toEqual(state); - }); + it("returns null for a tampered auth tag", async () => { + const manager = newManager(); + const parts = (await manager.save("client-001", makeState())).split("."); + const tag = Buffer.from(parts[3], "base64"); + tag[15] ^= 0x01; + parts[3] = tag.toString("base64"); + + await expect(manager.load(parts.join("."))).resolves.toBeNull(); }); - describe("flushPending", () => { - it("flushes all pending debounced saves", async () => { - const state = makeState(); - sm.debouncedSave("client-001", state); - expect(sm.pendingSaves).toBe(1); - await sm.flushPending(); - expect(sm.pendingSaves).toBe(0); - const loaded = await sm.load("client-001"); - expect(loaded).toEqual(state); - }); + it("returns null for truncated and garbage blobs", async () => { + const manager = newManager(); + const blob = await manager.save("client-001", makeState()); + + for (const bad of [ + "", + "garbage", + "v1.only.three", + "v1.a.b.c.d", + blob.split(".").slice(0, 3).join("."), + blob.slice(0, blob.length - 4), + ]) { + await expect(manager.load(bad)).resolves.toBeNull(); + } + }); - it("flushes pending saves for Redis-backed sessions", async () => { - const store = new Map(); - const redisMock = { - async set(key, value) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; - - const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); - const state = makeState(); - redisSm.debouncedSave("client-001", state); - expect(redisSm.pendingSaves).toBe(1); - await redisSm.flushPending(); - expect(redisSm.pendingSaves).toBe(0); - expect(store.has("session:client-001")).toBe(true); - const loaded = await redisSm.load("client-001"); - expect(loaded).toEqual(state); - redisSm.destroy(); - }); + it("returns null instead of throwing for non-string input", async () => { + const manager = newManager(); + for (const bad of [undefined, null, 42, {}, []]) { + await expect(manager.load(bad)).resolves.toBeNull(); + } }); - describe("in-memory mode (no Redis)", () => { - it("stores and retrieves sessions from local cache", async () => { - const state = makeState(); - await sm.save("client-001", state); - expect(sm.cachedSessions).toBe(1); - const loaded = await sm.load("client-001"); - expect(loaded).toEqual(state); - }); + it("returns null when the blob was sealed with a different key", async () => { + const redis = makeFakeRedis(); + const writer = newManager({ redis, encryptionKey: KEY_A }); + const reader = newManager({ redis, encryptionKey: KEY_B }); - it("evicts expired entries", async () => { - const shortTtlSm = new SessionManager({ encryptionKey, ttlMs: 1, debounceMs: 10 }); - await shortTtlSm.save("client-001", makeState()); - expect(shortTtlSm.cachedSessions).toBe(1); - await new Promise((r) => setTimeout(r, 20)); - expect(shortTtlSm.cachedSessions).toBe(0); - const loaded = await shortTtlSm.load("client-001"); - expect(loaded).toBeNull(); - shortTtlSm.destroy(); - }); + const blob = await writer.save("client-001", makeState()); + + await expect(reader.load(blob)).resolves.toBeNull(); + expect(reader.metrics.session_resumption_total.decrypt_failed).toBe(1); + expect(reader.metrics.session_resumption_total.expired).toBe(0); }); - describe("session state structure", () => { - it("handles rooms with full state", async () => { - const state = makeState({ - rooms: [ - { roomId: "fleet-1", highestAckedSeq: 10, highestReceivedSeq: 15, geofenceInsideSet: ["fence-a", "fence-b"] }, - { roomId: "fleet-2", highestAckedSeq: 0, highestReceivedSeq: 3, geofenceInsideSet: [] }, - ], - }); - await sm.save("client-001", state); - const loaded = await sm.load("client-001"); - expect(loaded.rooms).toHaveLength(2); - expect(loaded.rooms[0].geofenceInsideSet).toEqual(["fence-a", "fence-b"]); - }); + it("returns null once the TTL has expired", async () => { + const manager = newManager({ ttlMs: 60 }); + const blob = await manager.save("client-001", makeState()); - it("handles empty rooms array", async () => { - const state = makeState({ rooms: [] }); - await sm.save("client-001", state); - const loaded = await sm.load("client-001"); - expect(loaded.rooms).toEqual([]); - }); + await sleep(90); - it("handles large session with many rooms", async () => { - const rooms = Array.from({ length: 50 }, (_, i) => ({ - roomId: `fleet-${i}`, - highestAckedSeq: i * 10, - highestReceivedSeq: i * 10 + 5, - geofenceInsideSet: Array.from({ length: 3 }, (_, j) => `fence-${i}-${j}`), - })); - const state = makeState({ rooms }); - await sm.save("client-001", state); - const loaded = await sm.load("client-001"); - expect(loaded.rooms).toHaveLength(50); - expect(loaded.rooms[49].roomId).toBe("fleet-49"); - }); + await expect(manager.load(blob)).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.expired).toBe(1); + expect(manager.metrics.session_resumption_total.decrypt_failed).toBe(0); }); - describe("blob size", () => { - it("session blob stays under 16KB for 50 rooms", async () => { - const rooms = Array.from({ length: 50 }, (_, i) => ({ - roomId: `fleet-${i}`, - highestAckedSeq: i * 10, - highestReceivedSeq: i * 10 + 5, - geofenceInsideSet: Array.from({ length: 3 }, (_, j) => `fence-${i}-${j}`), - })); - const state = makeState({ rooms }); - await sm.save("client-001", state); - - const loaded = await sm.load("client-001"); - expect(loaded).not.toBeNull(); - const jsonSize = Buffer.byteLength(JSON.stringify(state), "utf8"); - expect(jsonSize).toBeLessThan(16384); - }); + it("returns null for a deleted session", async () => { + const manager = newManager(); + const blob = await manager.save("client-001", makeState()); + + await manager.delete("client-001"); + + await expect(manager.load(blob)).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.expired).toBe(1); }); - describe("destroy", () => { - it("clears all timers", () => { - sm.debouncedSave("client-001", makeState()); - expect(sm.pendingSaves).toBe(1); - sm.destroy(); - expect(sm.pendingSaves).toBe(0); - }); + it("counts a storage read error as expired", async () => { + const redis = makeFakeRedis(); + const manager = newManager({ redis }); + const blob = await manager.save("client-001", makeState()); + redis.get.mockRejectedValueOnce(new Error("connection lost")); - it("is idempotent", () => { - sm.destroy(); - expect(() => sm.destroy()).not.toThrow(); - }); + await expect(manager.load(blob)).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.expired).toBe(1); }); +}); - describe("with Redis mock", () => { - it("delegates to Redis client", async () => { - const store = new Map(); - const redisMock = { - async set(key, value, _ex, _ttl) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; +// ─── sliding TTL ────────────────────────────────────────────────────────────── - const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); - const state = makeState(); - await redisSm.save("client-001", state); +describe("SessionManager sliding TTL", () => { + it("refreshes the expiry on every save", async () => { + vi.useFakeTimers(); + const manager = newManager({ ttlMs: 1000 }); + await manager.save("client-001", makeState()); - expect(store.has("session:client-001")).toBe(true); + await vi.advanceTimersByTimeAsync(700); + const refreshed = await manager.save("client-001", makeState()); + await vi.advanceTimersByTimeAsync(700); - const loaded = await redisSm.load("client-001"); - expect(loaded).toEqual(state); + // 1400 ms after the first save, but only 700 ms after the second. + await expect(manager.load(refreshed)).resolves.toEqual(makeState()); + }); - await redisSm.delete("client-001"); - expect(store.has("session:client-001")).toBe(false); + it("lets a session that is not re-saved lapse", async () => { + vi.useFakeTimers(); + const manager = newManager({ ttlMs: 1000 }); + const stale = await manager.save("client-stale", makeState({ clientId: "client-stale" })); - redisSm.destroy(); - }); + await vi.advanceTimersByTimeAsync(700); + const fresh = await manager.save("client-001", makeState()); + await vi.advanceTimersByTimeAsync(700); + + await expect(manager.load(stale)).resolves.toBeNull(); + await expect(manager.load(fresh)).resolves.toEqual(makeState()); + }); + + it("sends PX on every redis write", async () => { + const redis = makeFakeRedis(); + const manager = newManager({ redis, ttlMs: 45000 }); + + await manager.save("client-001", makeState()); + await manager.save("client-001", makeState()); + + expect(redis.set).toHaveBeenCalledTimes(2); + for (const call of redis.set.mock.calls) { + expect(call[0]).toBe("session:client-001"); + expect(call[2]).toEqual({ PX: 45000 }); + } + }); +}); + +// ─── key rotation ───────────────────────────────────────────────────────────── + +describe("SessionManager key rotation", () => { + it("loads v1 blobs after rotating to v2 and seals new blobs with v2", async () => { + const redis = makeFakeRedis(); + const before = newManager({ redis, encryptionKey: { v1: KEY_A }, keyId: "v1" }); + const after = newManager({ redis, encryptionKey: { v1: KEY_A, v2: KEY_B }, keyId: "v2" }); + + const oldBlob = await before.save("client-001", makeState()); + expect(oldBlob.startsWith("v1.")).toBe(true); + + await expect(after.load(oldBlob)).resolves.toEqual(makeState()); + + const newBlob = await after.save("client-001", makeState()); + expect(newBlob.startsWith("v2.")).toBe(true); + await expect(after.load(newBlob)).resolves.toEqual(makeState()); + + // The pre-rotation manager has no v2 key, so it cannot open the new blob. + await expect(before.load(newBlob)).resolves.toBeNull(); + }); + + it("falls back to the other keys when the blob prefix is unknown", async () => { + const redis = makeFakeRedis(); + const writer = newManager({ redis, encryptionKey: { v1: KEY_A }, keyId: "v1" }); + const reader = newManager({ redis, encryptionKey: { v2: KEY_B, v1: KEY_A }, keyId: "v2" }); - it("returns null for missing Redis key", async () => { - const redisMock = { - async set() {}, - async get() { return null; }, - async del() {}, - }; - - const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); - const loaded = await redisSm.load("non-existent"); - expect(loaded).toBeNull(); - redisSm.destroy(); + const blob = await writer.save("client-001", makeState()); + const relabelled = ["v99", ...blob.split(".").slice(1)].join("."); + + await expect(reader.load(relabelled)).resolves.toEqual(makeState()); + }); +}); + +// ─── debounced saves ────────────────────────────────────────────────────────── + +describe("SessionManager debouncedSave", () => { + let redis; + let manager; + + beforeEach(() => { + vi.useFakeTimers(); + redis = makeFakeRedis(); + manager = newManager({ redis }); + }); + + it("coalesces rapid calls into a single save of the newest state", async () => { + const provider = vi.fn(() => makeState({ protocolVersion: 9 })); + + await manager.debouncedSave("client-001", () => makeState({ protocolVersion: 1 })); + await manager.debouncedSave("client-001", () => makeState({ protocolVersion: 2 })); + await manager.debouncedSave("client-001", provider); + + expect(redis.set).not.toHaveBeenCalled(); + expect(provider).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(500); + + expect(redis.set).toHaveBeenCalledTimes(1); + expect(provider).toHaveBeenCalledTimes(1); + const blob = redis.set.mock.calls[0][1]; + await expect(manager.load(blob)).resolves.toEqual(makeState({ protocolVersion: 9 })); + }); + + it("does not fire before the debounce window elapses", async () => { + await manager.debouncedSave("client-001", () => makeState()); + + await vi.advanceTimersByTimeAsync(499); + expect(redis.set).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(1); + expect(redis.set).toHaveBeenCalledTimes(1); + }); + + it("honours a custom debounceMs", async () => { + const fast = newManager({ redis, debounceMs: 50 }); + await fast.debouncedSave("client-001", () => makeState()); + + await vi.advanceTimersByTimeAsync(50); + expect(redis.set).toHaveBeenCalledTimes(1); + }); + + it("keeps separate windows per client", async () => { + await manager.debouncedSave("client-a", () => makeState({ clientId: "client-a" })); + await manager.debouncedSave("client-b", () => makeState({ clientId: "client-b" })); + + await vi.advanceTimersByTimeAsync(500); + + expect(redis.set).toHaveBeenCalledTimes(2); + expect(redis.entries.has("session:client-a")).toBe(true); + expect(redis.entries.has("session:client-b")).toBe(true); + }); + + it("flush stores immediately and cancels the pending timer", async () => { + await manager.debouncedSave("client-001", () => makeState()); + + const blob = await manager.flush("client-001"); + + expect(redis.set).toHaveBeenCalledTimes(1); + expect(typeof blob).toBe("string"); + + await vi.advanceTimersByTimeAsync(500); + expect(redis.set).toHaveBeenCalledTimes(1); + }); + + it("flush returns null when nothing is pending", async () => { + await expect(manager.flush("client-001")).resolves.toBeNull(); + expect(redis.set).not.toHaveBeenCalled(); + }); + + it("flushAll stores every pending session", async () => { + for (const id of ["c1", "c2", "c3"]) { + await manager.debouncedSave(id, () => makeState({ clientId: id })); + } + + const blobs = await manager.flushAll(); + + expect(blobs).toHaveLength(3); + expect(redis.set).toHaveBeenCalledTimes(3); + expect(await manager.flushAll()).toEqual([]); + }); + + it("swallows and logs a save failure raised inside the timer", async () => { + await manager.debouncedSave("client-001", () => { + throw new Error("state gone"); }); - it("returns null for corrupted Redis value", async () => { - const store = new Map(); - const redisMock = { - async set(key, value) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; - - const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); - await redisSm.save("client-001", makeState()); - - store.set("session:client-001", "corrupted-data"); - const loaded = await redisSm.load("client-001"); - expect(loaded).toBeNull(); - redisSm.destroy(); + await vi.advanceTimersByTimeAsync(500); + + expect(redis.set).not.toHaveBeenCalled(); + }); + + it("rejects an invalid stateProvider", async () => { + await expect(manager.debouncedSave("client-001", "nope")).rejects.toThrow(/stateProvider/); + }); + + it("delete cancels a pending debounced save", async () => { + await manager.debouncedSave("client-001", () => makeState()); + await manager.delete("client-001"); + + await vi.advanceTimersByTimeAsync(500); + + expect(redis.set).not.toHaveBeenCalled(); + expect(redis.del).toHaveBeenCalledWith("session:client-001"); + }); + + it("close clears pending timers", async () => { + await manager.debouncedSave("client-001", () => makeState()); + await manager.close(); + await manager.close(); + + await vi.advanceTimersByTimeAsync(500); + + expect(redis.set).not.toHaveBeenCalled(); + }); +}); + +// ─── storage backends ───────────────────────────────────────────────────────── + +describe("SessionManager storage backends", () => { + it("works without redis using the in-memory fallback", async () => { + const manager = newManager({ redis: null }); + const blob = await manager.save("client-001", makeState()); + + await expect(manager.load(blob)).resolves.toEqual(makeState()); + + await manager.delete("client-001"); + await expect(manager.load(blob)).resolves.toBeNull(); + }); + + it("sweeps expired in-memory entries on a periodic timer", async () => { + vi.useFakeTimers(); + const manager = newManager({ ttlMs: 1000 }); + await manager.save("client-001", makeState()); + expect(manager._memory.size).toBe(1); + + await vi.advanceTimersByTimeAsync(1100); + + expect(manager._memory.size).toBe(0); + }); + + it("drives a redis client with session: keys and PX TTL", async () => { + const redis = makeFakeRedis(); + const manager = newManager({ redis, ttlMs: 1234 }); + + const blob = await manager.save("client-001", makeState()); + expect(redis.set).toHaveBeenCalledWith("session:client-001", blob, { PX: 1234 }); + + await manager.load(blob); + expect(redis.get).toHaveBeenCalledWith("session:client-001"); + + await manager.delete("client-001"); + expect(redis.del).toHaveBeenCalledWith("session:client-001"); + expect(redis.entries.size).toBe(0); + }); +}); + +// ─── blob size budget ───────────────────────────────────────────────────────── + +describe("SessionManager blob size", () => { + it("keeps a 50-room session under 16 KB", async () => { + const manager = newManager(); + const state = makeFiftyRoomState(); + + const blob = await manager.save(state.clientId, state); + + expect(Buffer.byteLength(blob, "utf8")).toBeLessThan(16384); + expect(manager.metrics.session_state_size_bytes).toBe(Buffer.byteLength(blob, "utf8")); + await expect(manager.load(blob)).resolves.toEqual(state); + }); + + it("rejects a state whose blob exceeds 16 KB", async () => { + const manager = newManager(); + const state = makeState({ metadata: { payload: randomBytes(24 * 1024).toString("base64") } }); + + await expect(manager.save("client-001", state)).rejects.toThrow(/over the 16384 byte limit/); + }); +}); + +// ─── metrics ────────────────────────────────────────────────────────────────── + +describe("SessionManager metrics", () => { + it("starts every counter at zero", () => { + const manager = newManager(); + + expect(manager.metrics).toEqual({ + session_resumption_total: { + success: 0, + decrypt_failed: 0, + expired: 0, + mismatch: 0, + new_session: 0, + }, + session_state_size_bytes: 0, }); + }); - it("rejects Redis session with mismatched clientId", async () => { - const store = new Map(); - const redisMock = { - async set(key, value) { store.set(key, value); }, - async get(key) { return store.get(key) ?? null; }, - async del(key) { store.delete(key); }, - }; + it("counts each load outcome", async () => { + const manager = newManager({ ttlMs: 60 }); + const blob = await manager.save("client-001", makeState()); - const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); - await redisSm.save("client-001", makeState({ clientId: "client-001" })); + await manager.load(blob); + await manager.load(blob); + await manager.load("garbage"); + await sleep(90); + await manager.load(blob); - const loaded = await redisSm.load("client-002"); - expect(loaded).toBeNull(); - redisSm.destroy(); + expect(manager.metrics.session_resumption_total).toMatchObject({ + success: 2, + decrypt_failed: 1, + expired: 1, }); }); + + it("records server-side outcomes through recordResumption", () => { + const manager = newManager(); + + manager.recordResumption("mismatch"); + manager.recordResumption("new_session"); + manager.recordResumption("new_session"); + + expect(manager.metrics.session_resumption_total.mismatch).toBe(1); + expect(manager.metrics.session_resumption_total.new_session).toBe(2); + }); + + it("throws on an unknown resumption result", () => { + const manager = newManager(); + expect(() => manager.recordResumption("bogus")).toThrow(/unknown resumption result/); + }); + + it("tracks the last observed blob size", async () => { + const manager = newManager(); + await manager.save("client-001", makeState()); + const small = manager.metrics.session_state_size_bytes; + + await manager.save("client-001", makeFiftyRoomState()); + + expect(small).toBeGreaterThan(0); + expect(manager.metrics.session_state_size_bytes).toBeGreaterThan(small); + }); + + it("returns a snapshot that cannot mutate internal counters", () => { + const manager = newManager(); + const snapshot = manager.metrics; + + snapshot.session_resumption_total.success = 99; + snapshot.session_state_size_bytes = 99; + + expect(manager.metrics.session_resumption_total.success).toBe(0); + expect(manager.metrics.session_state_size_bytes).toBe(0); + }); }); diff --git a/tests/session-resumption-integration.test.js b/tests/session-resumption-integration.test.js index 9142923..f16bd6b 100644 --- a/tests/session-resumption-integration.test.js +++ b/tests/session-resumption-integration.test.js @@ -2,275 +2,510 @@ import { describe, it, expect, beforeEach, afterEach } from "vitest"; import WebSocket from "ws"; import jwt from "jsonwebtoken"; import { createServer } from "../src/server.js"; +import { SessionManager } from "../src/session-manager.js"; -const TEST_SECRET = "test-secret-key"; +const TEST_SECRET = "test-secret-session-resumption"; +const KEY_V1 = Buffer.alloc(32, 7).toString("base64"); +const KEY_V2 = Buffer.alloc(32, 9).toString("base64"); -/** Sign a JWT for a client. */ -function makeToken(clientId) { - return jwt.sign({ sub: clientId }, TEST_SECRET, { expiresIn: 60 }); +/** Signs the HS256 token the gateway authenticates with. */ +function makeToken(clientId, claims = {}) { + return jwt.sign({ sub: clientId, ...claims }, TEST_SECRET, { expiresIn: 60 }); } -/** Collect the next N messages from a WebSocket. */ -function nextMessages(ws, n = 1, timeoutMs = 3000) { - return new Promise((resolve, reject) => { - const msgs = []; - const timeout = setTimeout(() => reject(new Error("Timeout waiting for messages")), timeoutMs); - ws.on("message", function handler(data) { - msgs.push(JSON.parse(data.toString())); - if (msgs.length === n) { - clearTimeout(timeout); - ws.off("message", handler); - resolve(msgs); +/** In-memory stand-in for node-redis v4 with PX expiry and call counters. */ +function createFakeRedis() { + const entries = new Map(); + const calls = { get: 0, set: 0, del: 0 }; + return { + entries, + calls, + async set(key, value, opts) { + calls.set++; + entries.set(key, { value, expiresAt: Date.now() + (opts?.PX ?? 60000) }); + }, + async get(key) { + calls.get++; + const entry = entries.get(key); + if (!entry) return null; + if (entry.expiresAt <= Date.now()) { + entries.delete(key); + return null; } - }); - }); + return entry.value; + }, + async del(key) { + calls.del++; + entries.delete(key); + }, + }; } -/** Collect messages for a duration, returning all received. */ -function collectMessages(ws, durationMs = 500) { - return new Promise((resolve) => { - const msgs = []; - ws.on("message", (data) => { - msgs.push(JSON.parse(data.toString())); - }); - setTimeout(() => resolve(msgs), durationMs); +const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms)); + +/** Polls `fn` until it returns something truthy. */ +async function waitFor(fn, { timeout = 1000, interval = 10, label = "condition" } = {}) { + const deadline = Date.now() + timeout; + for (;;) { + const result = await fn(); + if (result) return result; + if (Date.now() > deadline) throw new Error(`timed out waiting for ${label}`); + await sleep(interval); + } +} + +/** Waits for the first buffered frame of a given type. */ +function waitForFrame(ws, type, options = {}) { + return waitFor(() => ws.frames.find((frame) => frame.type === type), { + label: `${type} frame`, + ...options, }); } -/** Wait for the WS close event. */ function waitClose(ws) { - return new Promise((resolve) => ws.once("close", resolve)); + return new Promise((resolve) => { + ws.once("close", (code, reason) => resolve({ code, reason: reason?.toString() ?? "" })); + }); } -/** Close a list of sockets and wait for them all. */ -async function closeAll(...sockets) { - sockets.forEach((ws) => ws.readyState === WebSocket.OPEN && ws.close()); - await Promise.all(sockets.map(waitClose)); +/** Session state in the shape the gateway captures. */ +function makeState(clientId, roomIds = [], overrides = {}) { + return { + clientId, + protocolVersion: 1, + authIdentity: { sub: clientId }, + rooms: roomIds.map((roomId) => ({ + roomId, + highestAckedSeq: 0, + highestReceivedSeq: 0, + geofenceInsideSet: [], + })), + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { + ip: "127.0.0.1", + userAgent: "vitest", + connectedAt: Date.now(), + lastActivityAt: Date.now(), + }, + ...overrides, + }; } -/** - * Open a WS connection and set up a message listener BEFORE the open event. - * Returns { ws, messages } where messages is a Promise resolving to the next N messages. - */ -function connectWithListener(port, token, extraParams = {}, n = 1) { - const params = new URLSearchParams(); - if (token) params.set("token", token); - for (const [k, v] of Object.entries(extraParams)) { - params.set(k, v); +describe("session resumption (integration)", () => { + /** @type {Array} */ + const gateways = []; + /** @type {Array} */ + const sockets = []; + + /** + * Starts a gateway on an ephemeral port with a fake-redis backed manager. + * `inspector` shares the store and keys so tests can decrypt saved blobs. + */ + function startGateway({ + encryptionKey = KEY_V1, + keyId = "v1", + ttlMs = 60000, + debounceMs = 50, + instanceId = "instance-1", + redis = createFakeRedis(), + sessionManager, + withSessions = true, + } = {}) { + const manager = + sessionManager ?? + (withSessions ? new SessionManager({ redis, encryptionKey, keyId, ttlMs, debounceMs }) : null); + const inspector = withSessions ? new SessionManager({ redis, encryptionKey, keyId, ttlMs }) : null; + const server = createServer({ + port: 0, + heartbeatMs: 60000, + maxPayloadBytes: 4096, + instanceId, + ...(manager ? { sessionManager: manager } : {}), + }); + const gateway = { server, port: server.wss.address().port, redis, manager, inspector, instanceId }; + gateways.push(gateway); + return gateway; } - const qs = params.toString(); - const url = `ws://localhost:${port}/${qs ? `?${qs}` : ""}`; - - return new Promise((resolve, reject) => { - const ws = new WebSocket(url); - let collectedMsgs = []; - - const messages = new Promise((msgResolve, msgReject) => { - const timeout = setTimeout(() => msgReject(new Error("Timeout waiting for messages")), 3000); - ws.on("message", function handler(data) { - collectedMsgs.push(JSON.parse(data.toString())); - if (collectedMsgs.length === n) { - clearTimeout(timeout); - ws.off("message", handler); - msgResolve(collectedMsgs); - } + + /** Opens a client socket that buffers every frame it receives. */ + function connect(port, token, { sessionId, cookie } = {}) { + let url = `ws://localhost:${port}/?token=${token}`; + if (sessionId) url += `&session_id=${encodeURIComponent(sessionId)}`; + const options = cookie ? { headers: { Cookie: cookie } } : undefined; + + return new Promise((resolve, reject) => { + const ws = new WebSocket(url, options); + ws.frames = []; + ws.upgradeHeaders = null; + sockets.push(ws); + ws.on("upgrade", (res) => { + ws.upgradeHeaders = res.headers; }); + ws.on("message", (data) => ws.frames.push(JSON.parse(data.toString()))); + ws.once("open", () => resolve(ws)); + ws.once("error", reject); }); + } - ws.once("open", () => resolve({ ws, messages })); - ws.once("error", reject); - }); -} + /** Closes a client and waits until the gateway has released the connection. */ + async function disconnect(gateway, ws) { + const before = gateway.server.wss.clients.size; + const closed = waitClose(ws); + ws.close(); + await closed; + await waitFor(() => gateway.server.wss.clients.size < before, { label: "server-side close" }); + } -describe("Session Resumption Integration", () => { - let server; - let port; + /** Waits for a stored blob whose decrypted state satisfies `predicate`. */ + function waitForSavedState(gateway, clientId, predicate = () => true) { + return waitFor( + async () => { + const entry = gateway.redis.entries.get(`session:${clientId}`); + if (!entry) return null; + const state = await gateway.inspector.load(entry.value); + return state && predicate(state) ? { blob: entry.value, state } : null; + }, + { label: `saved session for ${clientId}` } + ); + } beforeEach(() => { process.env.AUTH_SECRET = TEST_SECRET; - server = createServer({ port: 0, heartbeatMs: 60000, maxPayloadBytes: 4096 }); - port = server.wss.address().port; }); afterEach(async () => { - for (const client of server.wss.clients) { - client.terminate(); + for (const ws of sockets) { + if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) ws.terminate(); + } + sockets.length = 0; + + for (const gateway of gateways) { + for (const client of gateway.server.wss.clients) client.terminate(); + await new Promise((resolve) => gateway.server.wss.close(resolve)); + if (gateway.manager) await gateway.manager.close(); + if (gateway.inspector) await gateway.inspector.close(); } - await new Promise((resolve) => server.wss.close(resolve)); + gateways.length = 0; delete process.env.AUTH_SECRET; }); - it("saves session state on room join and disconnect", async () => { - const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-resume-1"), {}, 1); - ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-resume" })); - await j1; + it("resumes rooms and sequence numbers, and the client is a member again", async () => { + const gateway = startGateway(); - ws1.close(); - await waitClose(ws1); - await new Promise((r) => setTimeout(r, 100)); + const a = await connect(gateway.port, makeToken("client-a")); + a.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(a, "room_joined"); - const loaded = await server.sessionManager.load("client-resume-1"); - expect(loaded).not.toBeNull(); - expect(loaded.clientId).toBe("client-resume-1"); - expect(loaded.rooms).toHaveLength(1); - expect(loaded.rooms[0].roomId).toBe("fleet-resume"); - }); + const b = await connect(gateway.port, makeToken("client-b")); + b.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(b, "room_joined"); - it("restores session on reconnect with valid session_id", async () => { - const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-resume-2"), {}, 1); - ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-session" })); - await j1; + b.send(JSON.stringify({ type: "location_update", payload: { latitude: 1, longitude: 2 } })); + await waitForFrame(a, "location_update"); - ws1.close(); - await waitClose(ws1); - await new Promise((r) => setTimeout(r, 100)); + a.send(JSON.stringify({ type: "reconnect", roomId: "fleet-1", lastSeq: 1 })); + await waitForFrame(a, "replay_complete"); - const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-resume-2"), { session_id: "client-resume-2" }, 1); - const msgs = await resumed; - expect(msgs[0].type).toBe("session_resumed"); - expect(msgs[0].payload.rooms).toContain("fleet-session"); + const { blob } = await waitForSavedState( + gateway, + "client-a", + (state) => state.rooms[0]?.highestAckedSeq === 1 + ); + + await disconnect(gateway, a); + await waitFor(() => gateway.server.rooms.getRoomSize("fleet-1") === 1, { + label: "membership cleanup", + }); - await closeAll(ws2); + const resumed = await connect(gateway.port, makeToken("client-a"), { sessionId: blob }); + const frame = await waitForFrame(resumed, "session_resumed"); + + expect(frame.payload.rooms).toEqual([ + { roomId: "fleet-1", highestAckedSeq: 1, highestReceivedSeq: 1 }, + ]); + expect(frame.payload.currentSeqPerRoom).toEqual({ "fleet-1": 1 }); + expect(gateway.server.rooms.getRoomSize("fleet-1")).toBe(2); + + b.send(JSON.stringify({ type: "location_update", payload: { latitude: 3, longitude: 4 } })); + const broadcast = await waitForFrame(resumed, "location_update"); + expect(broadcast.payload.latitude).toBe(3); }); - it("treats expired session as new session (no session_resumed sent)", async () => { - const sm = server.sessionManager; - await sm.save("client-expired", { - clientId: "client-expired", - protocolVersion: 3, - authIdentity: { sub: "client-expired" }, - rooms: [{ roomId: "fleet-expired", highestAckedSeq: 0, highestReceivedSeq: 0, geofenceInsideSet: [] }], - rateLimitState: { messageWindow: [], connectionWindow: [] }, - metadata: { ip: "127.0.0.1", userAgent: "", connectedAt: Date.now(), lastActivityAt: Date.now() }, - }); - await sm.delete("client-expired"); + it("accepts the session blob from the JWT sid claim", async () => { + const gateway = startGateway(); + const blob = await gateway.manager.save("client-sid", makeState("client-sid", ["fleet-sid"])); - const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-expired")}&session_id=client-expired`); - const msgs = await collectMessages(ws, 500); - const resumedMsgs = msgs.filter((m) => m.type === "session_resumed"); - expect(resumedMsgs).toHaveLength(0); - ws.close(); - await waitClose(ws); + const ws = await connect(gateway.port, makeToken("client-sid", { sid: blob })); + const frame = await waitForFrame(ws, "session_resumed"); + + expect(frame.payload.rooms.map((room) => room.roomId)).toEqual(["fleet-sid"]); + expect(gateway.server.rooms.getRoomSize("fleet-sid")).toBe(1); }); - it("rejects session with identity mismatch (no session_resumed sent)", async () => { - await server.sessionManager.save("client-mismatch", { - clientId: "client-mismatch", - protocolVersion: 3, - authIdentity: { sub: "wrong-identity" }, - rooms: [{ roomId: "fleet-mismatch", highestAckedSeq: 0, highestReceivedSeq: 0, geofenceInsideSet: [] }], - rateLimitState: { messageWindow: [], connectionWindow: [] }, - metadata: { ip: "127.0.0.1", userAgent: "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + it("builds its own manager from an encryption key and the injected store", async () => { + const redis = createFakeRedis(); + const server = createServer({ + port: 0, + heartbeatMs: 60000, + maxPayloadBytes: 4096, + sessionEncryptionKey: KEY_V1, + sessionTtlMs: 60000, + instanceId: "instance-own", + redis, }); + const gateway = { + server, + port: server.wss.address().port, + redis, + manager: server.sessionManager, + inspector: new SessionManager({ redis, encryptionKey: KEY_V1, ttlMs: 60000 }), + }; + gateways.push(gateway); + + expect(server.sessionManager).toBeInstanceOf(SessionManager); + expect(server.instanceId).toBe("instance-own"); + + const ws = await connect(gateway.port, makeToken("client-own")); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-own" })); + await waitForFrame(ws, "room_joined"); + + const { state } = await waitForSavedState(gateway, "client-own"); + expect(state.rooms.map((room) => room.roomId)).toEqual(["fleet-own"]); + }); - const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-different")}&session_id=client-mismatch`); - const msgs = await collectMessages(ws, 500); - const resumedMsgs = msgs.filter((m) => m.type === "session_resumed"); - expect(resumedMsgs).toHaveLength(0); - ws.close(); - await waitClose(ws); + it("reads SESSION_ENCRYPTION_KEY from the environment and works without a store", async () => { + process.env.SESSION_ENCRYPTION_KEY = KEY_V1; + let server; + try { + server = createServer({ port: 0, heartbeatMs: 60000, maxPayloadBytes: 4096 }); + } finally { + delete process.env.SESSION_ENCRYPTION_KEY; + } + const gateway = { + server, + port: server.wss.address().port, + redis: null, + manager: server.sessionManager, + inspector: null, + }; + gateways.push(gateway); + + expect(server.sessionManager).toBeInstanceOf(SessionManager); + expect(server.instanceId).toMatch(/^[0-9a-f-]{36}$/); + + const first = await connect(gateway.port, makeToken("client-env")); + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-env" })); + await waitForFrame(first, "room_joined"); + await disconnect(gateway, first); + + const second = await connect(gateway.port, makeToken("client-env"), { + cookie: `GW_AFFINITY=${server.instanceId}`, + }); + const frame = await waitForFrame(second, "session_resumed"); + expect(frame.payload.rooms.map((room) => room.roomId)).toEqual(["fleet-env"]); }); - it("saves session state on room leave", async () => { - const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-leave-1"), {}, 1); - ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-leave" })); - await j1; + it("treats an expired session as a new session", async () => { + const gateway = startGateway({ ttlMs: 60 }); + const blob = await gateway.manager.save("client-exp", makeState("client-exp", ["fleet-1"])); + await sleep(120); - const l1 = nextMessages(ws1, 1); - ws1.send(JSON.stringify({ type: "leave_room", roomId: "fleet-leave" })); - await l1; + const ws = await connect(gateway.port, makeToken("client-exp"), { sessionId: blob }); + await sleep(150); - ws1.close(); - await waitClose(ws1); - await new Promise((r) => setTimeout(r, 100)); + expect(ws.frames.find((f) => f.type === "session_resumed")).toBeUndefined(); + expect(gateway.manager.metrics.session_resumption_total.expired).toBe(1); + expect(gateway.server.rooms.getRoomSize("fleet-1")).toBe(0); - const loaded = await server.sessionManager.load("client-leave-1"); - expect(loaded).not.toBeNull(); - expect(loaded.rooms).toHaveLength(0); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(ws, "room_joined"); }); - it("debounces session saves during rapid state changes", async () => { - const { ws, messages: j1 } = await connectWithListener(port, makeToken("client-debounce"), {}, 1); - ws.send(JSON.stringify({ type: "join_room", roomId: "room-1" })); - await j1; + it("treats a corrupted session blob as a new session", async () => { + const gateway = startGateway(); + const blob = await gateway.manager.save("client-bad", makeState("client-bad", ["fleet-1"])); + const parts = blob.split("."); + parts[2] = (parts[2][0] === "A" ? "B" : "A") + parts[2].slice(1); + const tampered = parts.join("."); - const j2 = nextMessages(ws, 1); - ws.send(JSON.stringify({ type: "join_room", roomId: "room-2" })); - await j2; + const ws = await connect(gateway.port, makeToken("client-bad"), { sessionId: tampered }); + await sleep(150); - const j3 = nextMessages(ws, 1); - ws.send(JSON.stringify({ type: "join_room", roomId: "room-3" })); - await j3; + expect(ws.frames.find((f) => f.type === "session_resumed")).toBeUndefined(); + expect(gateway.manager.metrics.session_resumption_total.decrypt_failed).toBe(1); - await new Promise((r) => setTimeout(r, 600)); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(ws, "room_joined"); + }); - const loaded = await server.sessionManager.load("client-debounce"); - expect(loaded).not.toBeNull(); - expect(loaded.rooms).toHaveLength(3); + it("ignores a session blob issued for another identity", async () => { + const gateway = startGateway(); + const blob = await gateway.manager.save("client-one", makeState("client-one", ["fleet-1"])); - ws.close(); - await waitClose(ws); + const ws = await connect(gateway.port, makeToken("client-two"), { sessionId: blob }); + await sleep(150); + + expect(ws.frames.find((f) => f.type === "session_resumed")).toBeUndefined(); + expect(gateway.manager.metrics.session_resumption_total.mismatch).toBe(1); + expect(gateway.server.rooms.getClientRooms("client-two").size).toBe(0); }); - it("returns session_resumed with correct sequence numbers", async () => { - const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-seq-1"), {}, 1); - ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-seq" })); - await j1; + it("persists a debounced save to the store within 1s of join_room", async () => { + const gateway = startGateway({ debounceMs: 500 }); - ws1.close(); - await waitClose(ws1); - await new Promise((r) => setTimeout(r, 100)); + const ws = await connect(gateway.port, makeToken("client-debounce")); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-9" })); + await waitForFrame(ws, "room_joined"); - const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-seq-1"), { session_id: "client-seq-1" }, 1); - const msgs = await resumed; - expect(msgs[0].type).toBe("session_resumed"); - expect(msgs[0].payload.currentSeqPerRoom).toBeDefined(); - expect(msgs[0].payload.currentSeqPerRoom[0].roomId).toBe("fleet-seq"); + const startedAt = Date.now(); + const { state } = await waitForSavedState(gateway, "client-debounce"); - await closeAll(ws2); + expect(Date.now() - startedAt).toBeLessThan(1000); + expect(state.rooms.map((room) => room.roomId)).toEqual(["fleet-9"]); + expect(state.authIdentity).toEqual({ sub: "client-debounce" }); }); - it("session metrics increment correctly", async () => { - const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-metrics-1"), {}, 1); - ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-metrics" })); - await j1; + it("sets GW_AFFINITY on upgrade and restores from the local cache without a store read", async () => { + const gateway = startGateway({ instanceId: "instance-7" }); + + const first = await connect(gateway.port, makeToken("client-sticky")); + const cookies = first.upgradeHeaders["set-cookie"] ?? []; + expect(cookies.join(";")).toContain("GW_AFFINITY=instance-7"); + expect(cookies.join(";")).toContain("HttpOnly"); + expect(cookies.join(";")).toContain("SameSite=Lax"); - ws1.close(); - await waitClose(ws1); - await new Promise((r) => setTimeout(r, 100)); + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-sticky" })); + await waitForFrame(first, "room_joined"); + await disconnect(gateway, first); - const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-metrics-1"), { session_id: "client-metrics-1" }, 1); - const msgs = await resumed; - expect(msgs[0].type).toBe("session_resumed"); + const readsBefore = gateway.redis.calls.get; + const second = await connect(gateway.port, makeToken("client-sticky"), { + cookie: "GW_AFFINITY=instance-7", + }); + const frame = await waitForFrame(second, "session_resumed"); - await closeAll(ws2); + expect(frame.payload.rooms.map((room) => room.roomId)).toEqual(["fleet-sticky"]); + expect(gateway.redis.calls.get).toBe(readsBefore); + expect(gateway.server.rooms.getRoomSize("fleet-sticky")).toBe(1); }); - it("handles session without session_id gracefully", async () => { - const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-no-session")}`); - await new Promise((resolve) => ws.once("open", resolve)); - const msgs = nextMessages(ws, 1); - ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-no-session" })); - await msgs; + it("migrates a client on demand and resumes from the handed-over blob", async () => { + const gateway = startGateway(); - expect(server.rooms.getRoomSize("fleet-no-session")).toBe(1); - ws.close(); - await waitClose(ws); + const ws = await connect(gateway.port, makeToken("client-migrate")); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-migrate" })); + await waitForFrame(ws, "room_joined"); + + const closed = waitClose(ws); + const res = await fetch(`http://localhost:${gateway.port}/admin/v1/clients/client-migrate/migrate`, { + method: "POST", + }); + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ ok: true }); + + const { code, reason } = await closed; + expect(code).toBe(4100); + + const frame = ws.frames.find((f) => f.type === "migrate"); + const sessionId = frame ? frame.payload.session_id : reason; + expect(sessionId.split(".")).toHaveLength(4); + + await waitFor(() => gateway.server.wss.clients.size === 0, { label: "server-side close" }); + + const resumed = await connect(gateway.port, makeToken("client-migrate"), { sessionId }); + const handshake = await waitForFrame(resumed, "session_resumed"); + expect(handshake.payload.rooms.map((room) => room.roomId)).toEqual(["fleet-migrate"]); + }); + + it("returns 404 from the migrate endpoint for an unknown client", async () => { + const gateway = startGateway(); + const res = await fetch(`http://localhost:${gateway.port}/admin/v1/clients/nobody/migrate`, { + method: "POST", + }); + expect(res.status).toBe(404); + expect(await res.json()).toEqual({ error: "Not Found" }); + }); + + it("opens a v1 blob after rotation and seals new saves with v2", async () => { + const redis = createFakeRedis(); + const legacy = new SessionManager({ redis, encryptionKey: KEY_V1, keyId: "v1" }); + const blob = await legacy.save("client-rot", makeState("client-rot", ["fleet-rot"])); + expect(blob.startsWith("v1.")).toBe(true); + await legacy.close(); + + const gateway = startGateway({ + redis, + encryptionKey: { v1: KEY_V1, v2: KEY_V2 }, + keyId: "v2", + }); + + const ws = await connect(gateway.port, makeToken("client-rot"), { sessionId: blob }); + const frame = await waitForFrame(ws, "session_resumed"); + expect(frame.payload.rooms.map((room) => room.roomId)).toEqual(["fleet-rot"]); + + const rotated = await waitFor( + () => { + const entry = redis.entries.get("session:client-rot"); + return entry && entry.value.startsWith("v2.") ? entry.value : null; + }, + { label: "v2 blob" } + ); + expect(rotated.startsWith("v2.")).toBe(true); }); - it("saves session on graceful shutdown via flushPending", async () => { - const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-shutdown")}`); - await new Promise((resolve) => ws.once("open", resolve)); - const msgs = nextMessages(ws, 1); - ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-shutdown" })); - await msgs; + it("saves every live session on saveAllSessions()", async () => { + const gateway = startGateway({ debounceMs: 5000 }); - await server.sessionManager.flushPending(); - const loaded = await server.sessionManager.load("client-shutdown"); - expect(loaded).not.toBeNull(); - expect(loaded.rooms).toHaveLength(1); + const first = await connect(gateway.port, makeToken("client-s1")); + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-s1" })); + await waitForFrame(first, "room_joined"); - ws.close(); - await waitClose(ws); + const second = await connect(gateway.port, makeToken("client-s2")); + second.send(JSON.stringify({ type: "join_room", roomId: "fleet-s2" })); + await waitForFrame(second, "room_joined"); + + const blobs = await gateway.server.saveAllSessions(); + + expect([...blobs.keys()].sort()).toEqual(["client-s1", "client-s2"]); + expect(gateway.redis.entries.has("session:client-s1")).toBe(true); + expect(gateway.redis.entries.has("session:client-s2")).toBe(true); + + const state = await gateway.inspector.load(blobs.get("client-s2")); + expect(state.rooms.map((room) => room.roomId)).toEqual(["fleet-s2"]); + }); + + it("reports resumption counters on GET /metrics", async () => { + const gateway = startGateway(); + + const first = await connect(gateway.port, makeToken("client-metrics")); + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-metrics" })); + await waitForFrame(first, "room_joined"); + const { blob } = await waitForSavedState(gateway, "client-metrics"); + await disconnect(gateway, first); + + const resumed = await connect(gateway.port, makeToken("client-metrics"), { sessionId: blob }); + await waitForFrame(resumed, "session_resumed"); + + const res = await fetch(`http://localhost:${gateway.port}/metrics`); + expect(res.status).toBe(200); + expect(res.headers.get("content-type")).toBe("text/plain; version=0.0.4; charset=utf-8"); + const body = await res.text(); + expect(body).toContain('session_resumption_total{result="new_session"} 1'); + expect(body).toContain('session_resumption_total{result="success"} 1'); + expect(body).toContain('session_resumption_total{result="mismatch"} 0'); + const sizeMatch = body.match(/session_state_size_bytes (\d+)/); + expect(sizeMatch).not.toBeNull(); + expect(parseInt(sizeMatch[1], 10)).toBeGreaterThan(0); + }); + + it("omits session counters from /metrics when resumption is disabled", async () => { + const gateway = startGateway({ withSessions: false }); + expect(gateway.server.sessionManager).toBeNull(); + + const res = await fetch(`http://localhost:${gateway.port}/metrics`); + expect(res.status).toBe(200); + const body = await res.text(); + expect(body).toContain("gateway_connections_active"); + expect(body).not.toContain("session_resumption_total"); }); });