From fbacc2cea25916c617d66ac4be75f4d7ddcdd4d8 Mon Sep 17 00:00:00 2001 From: Rajesh Date: Mon, 17 Aug 2026 16:11:26 +0530 Subject: [PATCH] feat: encrypted session resumption with zero-downtime connection migration --- .env.example | 13 + README.md | 19 + src/index.js | 94 ++- src/rate-limiter.js | 2 + src/room-manager.js | 145 +++- src/server.js | 532 ++++++++++++++- src/session-manager.js | 545 +++++++++++++++ src/validator.js | 8 + tests/session-manager.test.js | 659 +++++++++++++++++++ tests/session-resumption-integration.test.js | 506 ++++++++++++++ 10 files changed, 2472 insertions(+), 51 deletions(-) create mode 100644 src/session-manager.js create mode 100644 tests/session-manager.test.js create mode 100644 tests/session-resumption-integration.test.js diff --git a/.env.example b/.env.example index 483fbbe..73c0483 100644 --- a/.env.example +++ b/.env.example @@ -14,3 +14,16 @@ MAX_MESSAGES_PER_SECOND=100 # Maximum new WebSocket connections allowed per IP address per minute CONN_RATE_LIMIT=30 + +# Session resumption. Leave SESSION_ENCRYPTION_KEY empty to disable it. +# One 32-byte base64 key: openssl rand -base64 32 +# For rotation, use a JSON key map. Every listed key opens old blobs; new blobs +# are sealed with the "v1" entry, which must be present. +# SESSION_ENCRYPTION_KEY={"v1":"","v2":""} +SESSION_ENCRYPTION_KEY= + +# Sliding TTL of a stored session (ms) +SESSION_TTL_MS=3600000 + +# Identity published in the GW_AFFINITY cookie. Defaults to a per-process uuid. +INSTANCE_ID= 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 ea9e74b..6857707 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,15 +45,30 @@ 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 sessionManager; +let instanceId; +let saveAllSessions; try { - ({ wss } = createServer(config)); + ({ wss, 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. @@ -50,28 +79,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) @@ -128,8 +186,12 @@ export function shutdown(wss, signal) { }); } -process.on("SIGTERM", () => shutdown(wss, "SIGTERM")); -process.on("SIGINT", () => shutdown(wss, "SIGINT")); +// Without a session manager there is nothing to persist, so the shared +// broadcast path stays in use. +const shutdownOptions = sessionManager ? { saveAllSessions } : {}; + +process.on("SIGTERM", () => shutdown(wss, "SIGTERM", shutdownOptions)); +process.on("SIGINT", () => shutdown(wss, "SIGINT", shutdownOptions)); process.on("uncaughtException", (err) => { logger.error("Uncaught exception", { error: err.message }); 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 865f352..3679c5b 100644 --- a/src/room-manager.js +++ b/src/room-manager.js @@ -1,5 +1,10 @@ import { WebSocket } from "ws"; +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; + /** * Manages room membership and message broadcasting for connected WebSocket clients. * @@ -7,14 +12,66 @@ import { WebSocket } from "ws"; * `clientId → WebSocket` so broadcasts are O(members). A reverse index * (`_clientRooms`) enables O(1) lookup of all rooms a client belongs to, * which is used during disconnection cleanup. + * + * Each room also keeps a monotonic sequence number and a bounded ring buffer of + * recent broadcasts so reconnecting clients can replay what they missed. */ export class RoomManager { - constructor({ maxRoomSize = Infinity } = {}) { + /** + * @param {object} [options] + * @param {number} [options.maxRoomSize] - Legacy per-room member cap. When set, + * `join()` reports `{ ok }` results instead of an error frame. + * @param {number} [options.maxRoomsPerClient] - Max rooms a single client may join. + * @param {number} [options.maxMembersPerRoom] - Max members allowed in one room. + * @param {number} [options.maxRooms] - Ceiling on the number of live rooms. + * @param {{ enabled?: boolean, memoryThresholdBytes?: number, recoveryThresholdBytes?: number }} [options.circuitBreaker] + * Heap-pressure breaker: opens above `memoryThresholdBytes`, closes below `recoveryThresholdBytes`. + * @param {number} [options.ringBufferSize] - Replay entries kept per room. + * @param {number} [options.maxBufferBytes] - Byte ceiling for one room's replay buffer. + * @param {number} [options.deduplicationWindowMs] - TTL of a `location_update` dedup key. + * @param {number} [options.maxDedupEntries] - Max dedup keys retained. + */ + constructor({ + maxRoomSize, + maxRoomsPerClient = Infinity, + maxMembersPerRoom = Infinity, + maxRooms = Infinity, + circuitBreaker = {}, + ringBufferSize = DEFAULT_RING_BUFFER_SIZE, + maxBufferBytes = DEFAULT_MAX_BUFFER_BYTES, + deduplicationWindowMs = DEFAULT_DEDUP_WINDOW_MS, + maxDedupEntries = DEFAULT_MAX_DEDUP_ENTRIES, + } = {}) { /** @type {Map>} */ this._rooms = new Map(); /** @type {Map>} */ this._clientRooms = new Map(); - this._maxRoomSize = maxRoomSize; + + this._maxRoomSize = maxRoomSize ?? Infinity; + this._legacyJoinResult = maxRoomSize != null; + this._maxRoomsPerClient = maxRoomsPerClient; + this._maxMembersPerRoom = maxMembersPerRoom; + this._maxRooms = maxRooms; + this._totalMembers = 0; + + this._circuitBreakerEnabled = circuitBreaker.enabled === true; + this._memoryThresholdBytes = circuitBreaker.memoryThresholdBytes ?? Infinity; + this._recoveryThresholdBytes = circuitBreaker.recoveryThresholdBytes ?? 0; + this._circuitBreakerState = "CLOSED"; + + /** @type {Map} roomId → last assigned sequence number */ + this._roomSeq = new Map(); + /** @type {Map>} */ + this._roomBuffers = new Map(); + /** @type {Map} roomId → approximate buffer size in bytes */ + this._roomBufferBytes = new Map(); + this._ringBufferSize = ringBufferSize; + this._maxBufferBytes = maxBufferBytes; + + /** @type {Map} dedup key → time the message was first seen (ms) */ + this._dedupCache = new Map(); + this._deduplicationWindowMs = deduplicationWindowMs; + this._maxDedupEntries = maxDedupEntries; } /** @private */ @@ -49,6 +106,34 @@ export class RoomManager { } } + /** @private */ + _rejection(code, message) { + return { type: "error", payload: { code, message } }; + } + + /** + * @private + * Re-evaluates heap pressure and returns true while the breaker rejects joins. + */ + _circuitBreakerOpen() { + if (!this._circuitBreakerEnabled) return false; + + const { heapUsed } = process.memoryUsage(); + if (this._circuitBreakerState === "OPEN") { + if (heapUsed < this._recoveryThresholdBytes) { + this._circuitBreakerState = "CLOSED"; + return false; + } + return true; + } + + if (heapUsed > this._memoryThresholdBytes) { + this._circuitBreakerState = "OPEN"; + return true; + } + return false; + } + /** @private */ _isDuplicate(roomId, message, excludeClientId) { let parsed = message; @@ -83,19 +168,65 @@ 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"); if (ws == null) throw new TypeError("ws is required"); - const room = this._ensureRoom(roomId); - if (!room.has(clientId) && room.size >= this._maxRoomSize) { - return { ok: false, reason: 'ROOM_FULL' }; + if (this._circuitBreakerOpen()) { + return this._rejection( + "CIRCUIT_BREAKER_OPEN", + "Circuit breaker is open due to high resource pressure", + ); } - room.set(clientId, ws); + const room = this._rooms.get(roomId); + const isMember = room?.has(clientId) === true; + + if (!isMember) { + const joinedRooms = this._clientRooms.get(clientId)?.size ?? 0; + if (joinedRooms >= this._maxRoomsPerClient) { + return this._rejection( + "ROOM_LIMIT_EXCEEDED", + `Client room limit exceeded (${this._maxRoomsPerClient})`, + ); + } + if (!room && this._rooms.size >= this._maxRooms) { + return this._rejection( + "MAX_ROOMS_REACHED", + `Maximum room count ceiling reached (${this._maxRooms})`, + ); + } + const members = room?.size ?? 0; + if (members >= this._maxMembersPerRoom) { + return this._rejection( + "ROOM_FULL", + `Room member limit reached (${this._maxMembersPerRoom})`, + ); + } + if (members >= this._maxRoomSize) { + return { ok: false, reason: "ROOM_FULL" }; + } + } + + this._ensureRoom(roomId).set(clientId, ws); this._ensureClientRooms(clientId).add(roomId); - return { ok: true }; + if (!isMember) this._totalMembers++; + + return this._legacyJoinResult ? { ok: true } : undefined; } /** diff --git a/src/server.js b/src/server.js index 32d9b60..1632a17 100644 --- a/src/server.js +++ b/src/server.js @@ -1,28 +1,124 @@ 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"; import { createRateLimiter } from "./rate-limiter.js"; import { createConnRateLimiter } from "./conn-rate-limiter.js"; -import { VALIDATION_ERROR } from "./errors.js"; +import { ROOM_FULL, VALIDATION_ERROR } from "./errors.js"; -export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit, maxConnectionsPerIp } = {}) { -import { createRateLimiter } from "./rate-limiter.js"; +const AFFINITY_COOKIE = "GW_AFFINITY"; +const AFFINITY_MAX_AGE_S = 3600; +const DEFAULT_SESSION_TTL_MS = 3600000; +const RATE_WINDOW_MS = 1000; +const MAX_LOCAL_SESSIONS = 10000; +const MAX_CLOSE_REASON_BYTES = 123; +const MIGRATE_CLOSE_CODE = 4100; +const MIGRATE_PATH = /^\/admin\/v1\/clients\/([^/]+)\/migrate$/; +const PROTOCOL_VERSION = 1; + +/** + * Reads one cookie out of a raw `Cookie` header. + * + * @param {unknown} header - Value of the `Cookie` request header. + * @param {string} name - Cookie name to look for. + * @returns {string|null} The value, or null when absent. + */ +function readCookie(header, name) { + if (typeof header !== "string") return null; + for (const pair of header.split(";")) { + const eq = pair.indexOf("="); + if (eq === -1) continue; + if (pair.slice(0, eq).trim() === name) return pair.slice(eq + 1).trim(); + } + return null; +} +/** + * Reads the `sid` claim of an already-verified token. `auth.js` verified the + * signature, so decoding without re-verification is safe here. + * + * @param {string|null|undefined} token + * @returns {string|null} The session blob, or null when the claim is absent. + */ +function sessionIdFromToken(token) { + if (typeof token !== "string" || token.length === 0) return null; + let payload; + try { + payload = jwt.decode(token); + } catch { + return null; + } + const sid = payload && typeof payload === "object" ? payload.sid : null; + return typeof sid === "string" && sid.length > 0 ? sid : null; +} + +/** + * Creates the HTTP health endpoint and the WebSocket gateway on top of it. + * + * 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, rooms: RoomManager, ipConnectionCount: Map, rateLimiter: object, sessionManager: SessionManager|null, instanceId: string, saveAllSessions: () => Promise> }} + */ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit, maxConnectionsPerIp, - ringBufferSize: _ringBufferSize, - deduplicationWindowMs: _deduplicationWindowMs, - maxBufferBytes: _maxBufferBytes, - maxDedupEntries: _maxDedupEntries, + maxMessagesPerSecond, + ringBufferSize, + deduplicationWindowMs, + maxBufferBytes, + maxDedupEntries, + sessionManager, + sessionEncryptionKey, + sessionTtlMs, + instanceId, + redis, } = {}) { + const sessionKey = sessionEncryptionKey ?? process.env.SESSION_ENCRYPTION_KEY ?? null; + const sessionTtl = sessionTtlMs ?? (Number(process.env.SESSION_TTL_MS) || DEFAULT_SESSION_TTL_MS); + const ownsSessions = sessionManager == null && sessionKey != null; + /** @type {SessionManager|null} */ + const sessions = ownsSessions + ? new SessionManager({ + redis: redis ?? null, + encryptionKey: sessionKey, + ttlMs: sessionTtl, + logger, + }) + : (sessionManager ?? null); + const resolvedInstanceId = instanceId ?? process.env.INSTANCE_ID ?? uuid(); + + /** @type {Map} clientId → live connection context */ + const liveClients = new Map(); + /** @type {Map} clientId → last known state, for the sticky-cookie path */ + const localSessions = new Map(); + const server = http.createServer((req, res) => { let url; try { @@ -39,6 +135,32 @@ export function createServer({ return; } + if (req.method === "GET" && url.pathname === "/metrics" && sessions) { + res.writeHead(200, { "Content-Type": "application/json" }); + res.end(JSON.stringify(sessions.metrics)); + return; + } + + 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; + } + res.writeHead(404, { "Content-Type": "application/json" }); res.end(JSON.stringify({ error: "Not Found" })); }); @@ -50,23 +172,302 @@ export function createServer({ server.listen(port ?? 8080); - const rooms = new RoomManager(); - const rateLimiter = createRateLimiter(); + const rooms = new RoomManager({ + maxRoomSize: Number(process.env.MAX_ROOM_SIZE) || 500, + ringBufferSize, + deduplicationWindowMs, + maxBufferBytes, + maxDedupEntries, + }); + const rateLimiter = createRateLimiter(maxMessagesPerSecond); 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; } - wss.on("connection", (ws, req) => { + if (sessions) { + wss.on("headers", (headers) => { + headers.push( + `Set-Cookie: ${AFFINITY_COOKIE}=${resolvedInstanceId}; Path=/; HttpOnly; SameSite=Lax; Max-Age=${AFFINITY_MAX_AGE_S}` + ); + }); + } + + wss.on("connection", async (ws, req) => { const clientId = uuid(); ws.isAlive = true; + ws._lastPongAt = Date.now(); const ip = req.socket.remoteAddress; @@ -104,25 +505,38 @@ export function createServer({ return; } - const actualClientId = authResult.clientId ?? clientId; + let actualClientId = authResult.clientId ?? clientId; ws._clientId = actualClientId; logger.info("Client connected", { clientId: actualClientId, ip }); + const ctx = sessions ? createContext(actualClientId, req) : null; + if (ctx) { + ctx.ws = ws; + liveClients.set(actualClientId, ctx); + } + /** @type {Promise|null} Frames wait for the handshake so replayed state lands first. */ + let pendingResume = null; + ws.on("pong", heartbeat); - ws.on("message", (raw) => { + ws.on("message", async (raw) => { + if (pendingResume) { + await pendingResume; + pendingResume = null; + } + if (!rateLimiter.check(actualClientId)) { logger.warn("Message rate limit exceeded", { clientId: actualClientId }); - ws.send(JSON.stringify({ type: "error", payload: { message: "Rate limit exceeded" } })); + safeSend(ws, { type: "error", payload: { message: "Rate limit exceeded" } }); return; } + if (ctx) ctx.messageWindow = pruneWindow([...ctx.messageWindow, Date.now()]); const validation = validateMessage(raw.toString()); if (!validation.ok) { logger.warn("Validation failed", { clientId: actualClientId, error: validation.error }); sendError(ws, validation.error, validation.code ?? VALIDATION_ERROR); - safeSend(ws, { type: "error", payload: { message: validation.error } }); return; } @@ -131,29 +545,51 @@ export function createServer({ switch (msg.type) { case "join_room": { const joinResult = rooms.join(actualClientId, msg.roomId, ws); - if (!joinResult.ok && joinResult.reason === 'ROOM_FULL') { + if (joinResult?.type === "error") { + logger.warn("Join rejected", { + clientId: actualClientId, + roomId: msg.roomId, + code: joinResult.payload.code, + }); + safeSend(ws, joinResult); + break; + } + if (joinResult?.ok === false) { logger.warn("Room is full", { clientId: actualClientId, roomId: msg.roomId }); - ws.send(JSON.stringify({ type: "error", payload: { message: "Room is full", code: "ROOM_FULL" } })); + safeSend(ws, { + type: "error", + payload: { message: "Room is full", code: joinResult.reason ?? ROOM_FULL }, + }); break; } logger.info("Client joined room", { clientId: actualClientId, roomId: msg.roomId }); safeSend(ws, { type: "room_joined", payload: { roomId: msg.roomId } }); + touchSession(ctx); break; } case "leave_room": { rooms.leave(actualClientId, msg.roomId); + if (ctx) { + ctx.ackedSeq.delete(msg.roomId); + ctx.geofence.delete(msg.roomId); + } logger.info("Client left room", { clientId: actualClientId, roomId: msg.roomId }); safeSend(ws, { type: "room_left", payload: { roomId: msg.roomId } }); + touchSession(ctx); break; } case "reconnect": { const clientRooms = rooms.getClientRooms(actualClientId); if (!clientRooms.has(msg.roomId)) { - ws.send(JSON.stringify({ type: "error", payload: { message: "Must join room before reconnecting" } })); + safeSend(ws, { + type: "error", + payload: { message: "Must join room before reconnecting" }, + }); break; } - const replayResult = rooms.handleReconnect(msg.roomId, msg.lastSeq); - ws.send(JSON.stringify(replayResult)); + if (ctx) ctx.ackedSeq.set(msg.roomId, msg.lastSeq); + safeSend(ws, rooms.handleReconnect(msg.roomId, msg.lastSeq)); + touchSession(ctx); break; } case "location_update": { @@ -164,15 +600,26 @@ export function createServer({ payload: { clientId: actualClientId, ...msg.payload }, }, actualClientId); } + touchSession(ctx); break; } case "token_refresh": { const result = await verifyConnection(msg.token); if (result.ok) { + const previousClientId = actualClientId; actualClientId = result.clientId; - ws.send(JSON.stringify({ type: "token_refresh_ok" })); + ws._clientId = actualClientId; + if (ctx && previousClientId !== actualClientId) { + if (liveClients.get(previousClientId) === ctx) liveClients.delete(previousClientId); + ctx.clientId = actualClientId; + liveClients.set(actualClientId, ctx); + } + logger.info("Token refreshed", { clientId: actualClientId }); + safeSend(ws, { type: "token_refresh_ok" }); + touchSession(ctx); } else { - ws.send(JSON.stringify({ type: "error", payload: { message: result.error } })); + logger.warn("Token refresh failed", { clientId: actualClientId, reason: result.error }); + safeSend(ws, { type: "error", payload: { message: result.error } }); } break; } @@ -180,6 +627,18 @@ export function createServer({ }); ws.on("close", (code, reason) => { + if (ctx) { + // Snapshot before teardown so a pending save cannot persist empty rooms. + ctx.frozenState = captureState(ctx); + rememberLocal(ctx.clientId, ctx.frozenState); + if (liveClients.get(ctx.clientId) === ctx) liveClients.delete(ctx.clientId); + sessions.flush(ctx.clientId).catch((err) => { + logger.error("Failed to flush session on close", { + clientId: ctx.clientId, + error: err.message, + }); + }); + } rooms.disconnect(actualClientId); rateLimiter.remove(actualClientId); const trackedIp = ws._trackedIp; @@ -202,25 +661,42 @@ export function createServer({ ws.on("error", (err) => { logger.error("WebSocket error", { clientId: actualClientId, error: err.message }); }); + + if (ctx) { + pendingResume = resumeSession(ws, req, url, ctx, token).catch((err) => { + logger.error("Session resumption failed", { clientId: ctx.clientId, error: err.message }); + }); + } }); const interval = 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); wss.on("close", () => { clearInterval(interval); + if (ownsSessions) sessions.close(); server.close(); }); - return { wss, server, rooms, ipConnectionCount, rateLimiter }; + return { + wss, + server, + rooms, + ipConnectionCount, + rateLimiter, + sessionManager: sessions, + instanceId: resolvedInstanceId, + saveAllSessions, + }; } diff --git a/src/session-manager.js b/src/session-manager.js new file mode 100644 index 0000000..2d1806f --- /dev/null +++ b/src/session-manager.js @@ -0,0 +1,545 @@ +/** + * @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. + */ + +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} 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. + */ + +/** + * 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] - 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] + */ + +/** + * 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 */ + +/** + * 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 {any} timer + * @returns {any} The same timer. + */ +function unrefTimer(timer) { + if (timer && typeof timer.unref === "function") timer.unref(); + return timer; +} + +/** + * Decodes one base64 master key and stretches it into an AES-256 key. + * + * @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 deriveKey(keyId, value) { + if (typeof value !== "string" || value.trim().length === 0) { + throw new TypeError(`encryptionKey["${keyId}"] must be a non-empty base64 string`); + } + 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 Buffer.from(hkdfSync("sha256", master, "", HKDF_INFO, KEY_BYTES)); +} + +/** + * 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 {unknown} encryptionKey + * @param {string} keyId - Key used to seal new blobs; must exist in the result. + * @returns {Map} + */ +function parseKeyMap(encryptionKey, keyId) { + if (encryptionKey == null) { + throw new TypeError( + "encryptionKey is required: a 32-byte base64 key, or a { keyId: base64key } map" + ); + } + + 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 entries = Object.entries(source); + if (entries.length === 0) throw new TypeError("encryptionKey holds no keys"); + + 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; +} + +/** + * In-memory stand-in for Redis with per-entry expiry timestamps. + * + * @implements {SessionStore} + */ +class MemoryStore { + constructor() { + /** @type {Map} */ + this._entries = new Map(); + } + + /** + * @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 }); + } + + /** + * @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; + } + + /** @param {string} key */ + async del(key) { + this._entries.delete(key); + } + + /** 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 {number} */ + get size() { + return this._entries.size; + } +} + +/** + * Seals, stores and restores per-client session state. + */ +export class SessionManager { + /** + * @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 = 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; + } + + /** + * @private + * @returns {SessionStore} + */ + _storage() { + return this._redis ?? this._memory; + } + + /** + * Seals `state` and stores it under `session:`, refreshing the TTL. + * + * @param {string} clientId + * @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) { + 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"); + } + + 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` + ); + } + + await this._storage().set(KEY_PREFIX + clientId, blob, { PX: this._ttlMs }); + return blob; + } + + /** + * 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; + } + + /** + * 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. + * + * @private + * @param {unknown} sessionId + * @returns {SessionState|null} + */ + _open(sessionId) { + if (typeof sessionId !== "string" || sessionId.length === 0) return null; + + const parts = sessionId.split("."); + if (parts.length !== 4) 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 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 { + // Wrong key or corrupt payload — fall through to the next candidate. + } + } + return null; + } + + /** + * 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 + * @param {() => SessionState} stateProvider + * @returns {Promise} Resolves once the save is scheduled, not stored. + */ + 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) + ) + ); + } + + /** + * 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} stateProvider + * @returns {Promise} The blob, or null on failure. + */ + _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; + }); + } + + /** + * @private + * @param {string} clientId + * @returns {(() => SessionState)|null} The cancelled provider, if any. + */ + _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; + } + + /** + * Fires a client's pending debounced save immediately. + * + * @param {string} clientId + * @returns {Promise} The blob, or null when nothing was pending. + */ + async flush(clientId) { + if (!this._debounceTimers.has(clientId)) { + this._cancelDebounce(clientId); + return null; + } + const provider = this._cancelDebounce(clientId); + return provider ? this._runSave(clientId, provider) : null; + } + + /** + * Fires every pending save — call before closing connections on shutdown. + * + * @returns {Promise>} One blob per flushed client. + */ + async flushAll() { + const pending = [...this._debounceTimers.keys()]; + return Promise.all(pending.map((clientId) => this.flush(clientId))); + } + + /** + * Drops a client's stored session and cancels any pending save. + * + * @param {string} clientId + * @returns {Promise} + */ + 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); + } + + /** + * Clears every timer so the process (or a test run) can exit. Idempotent. + * + * @returns {Promise} + */ + async close() { + for (const clientId of [...this._debounceTimers.keys()]) { + this._cancelDebounce(clientId); + } + if (this._sweepTimer) { + clearInterval(this._sweepTimer); + this._sweepTimer = null; + } + } + + /** + * Counts a resumption outcome decided outside this class, such as an + * identity `mismatch` or a `new_session`. + * + * @param {ResumptionResult} result + */ + recordResumption(result) { + if (!Object.hasOwn(this._counters, result)) { + throw new RangeError( + `unknown resumption result "${result}" (expected: ${Object.keys(this._counters).join("|")})` + ); + } + this._counters[result]++; + } + + /** + * Snapshot of the issue-18 metrics. Mutating it does not affect the manager. + * + * @type {{ session_resumption_total: Record, session_state_size_bytes: number }} + */ + 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 19528ba..8835127 100644 --- a/src/validator.js +++ b/src/validator.js @@ -24,6 +24,10 @@ const reconnectSchema = z.object({ lastSeq: z.number().min(0), }); +const tokenRefreshSchema = z.object({ + token: z.string().min(1), +}); + const messageSchema = z.discriminatedUnion("type", [ z.object({ type: z.literal("location_update"), @@ -41,6 +45,10 @@ const messageSchema = z.discriminatedUnion("type", [ type: z.literal("reconnect"), ...reconnectSchema.shape, }), + z.object({ + type: z.literal("token_refresh"), + ...tokenRefreshSchema.shape, + }), ]); const MESSAGE_SIZE_LIMITS = { diff --git a/tests/session-manager.test.js b/tests/session-manager.test.js new file mode 100644 index 0000000..0104f7f --- /dev/null +++ b/tests/session-manager.test.js @@ -0,0 +1,659 @@ +/** + * @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; +} + +/** + * 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-123", iss: "fleet-auth" }, + rooms: [ + { + roomId: "fleet-alpha", + highestAckedSeq: 42, + highestReceivedSeq: 45, + geofenceInsideSet: ["fence-1", "fence-3"], + }, + ], + rateLimitState: { messageWindow: [1001, 1002, 1003], connectionWindow: [900] }, + metadata: { + ip: "10.0.0.1", + userAgent: "FleetApp/2.3", + connectedAt: 1700000000000, + lastActivityAt: 1700000009000, + }, + ...overrides, + }; +} + +/** + * 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)); +} + +afterEach(async () => { + for (const manager of managers) await manager.close(); + managers = []; + vi.useRealTimers(); + vi.restoreAllMocks(); +}); + +// ─── key material ───────────────────────────────────────────────────────────── + +describe("SessionManager key material", () => { + it("throws when encryptionKey is missing", () => { + expect(() => new SessionManager({})).toThrow(/encryptionKey is required/); + }); + + 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/ + ); + }); + + 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("throws when keyId is absent from the key map", () => { + expect(() => new SessionManager({ encryptionKey: { v1: KEY_A }, keyId: "v9" })).toThrow( + /keyId "v9" is absent/ + ); + }); + + it("throws when keyId contains the blob separator", () => { + expect(() => new SessionManager({ encryptionKey: KEY_A, keyId: "v.1" })).toThrow( + /must not contain/ + ); + }); + + it("throws on malformed JSON key material", () => { + expect(() => new SessionManager({ encryptionKey: '{"v1": ' })).toThrow(/does not parse/); + }); + + 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()); + }); +}); + +// ─── 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); + }); + + 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("."); + + 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); + }); + + it("never stores plaintext state", async () => { + const redis = makeFakeRedis(); + const manager = newManager({ redis }); + const blob = await manager.save("client-001", makeState()); + + expect(blob).not.toContain("fleet-alpha"); + expect(redis.entries.get("session:client-001").value).not.toContain("device-123"); + }); + + 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()); + + 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/); + }); +}); + +// ─── load failure paths ─────────────────────────────────────────────────────── + +describe("SessionManager load failures", () => { + it("returns null for a corrupted ciphertext byte", async () => { + const manager = newManager(); + const blob = await manager.save("client-001", makeState()); + + await expect(manager.load(corruptCiphertext(blob))).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.decrypt_failed).toBe(1); + }); + + 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(); + }); + + 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("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(); + } + }); + + 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 }); + + 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); + }); + + it("returns null once the TTL has expired", async () => { + const manager = newManager({ ttlMs: 60 }); + const blob = await manager.save("client-001", makeState()); + + await sleep(90); + + 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); + }); + + 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); + }); + + 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")); + + await expect(manager.load(blob)).resolves.toBeNull(); + expect(manager.metrics.session_resumption_total.expired).toBe(1); + }); +}); + +// ─── sliding TTL ────────────────────────────────────────────────────────────── + +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()); + + await vi.advanceTimersByTimeAsync(700); + const refreshed = await manager.save("client-001", makeState()); + await vi.advanceTimersByTimeAsync(700); + + // 1400 ms after the first save, but only 700 ms after the second. + await expect(manager.load(refreshed)).resolves.toEqual(makeState()); + }); + + 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" })); + + 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" }); + + 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"); + }); + + 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("counts each load outcome", async () => { + const manager = newManager({ ttlMs: 60 }); + const blob = await manager.save("client-001", makeState()); + + await manager.load(blob); + await manager.load(blob); + await manager.load("garbage"); + await sleep(90); + await manager.load(blob); + + 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 new file mode 100644 index 0000000..0948c17 --- /dev/null +++ b/tests/session-resumption-integration.test.js @@ -0,0 +1,506 @@ +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-session-resumption"; +const KEY_V1 = Buffer.alloc(32, 7).toString("base64"); +const KEY_V2 = Buffer.alloc(32, 9).toString("base64"); + +/** Signs the HS256 token the gateway authenticates with. */ +function makeToken(clientId, claims = {}) { + return jwt.sign({ sub: clientId, ...claims }, TEST_SECRET, { expiresIn: 60 }); +} + +/** 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); + }, + }; +} + +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, + }); +} + +function waitClose(ws) { + return new Promise((resolve) => { + ws.once("close", (code, reason) => resolve({ code, reason: reason?.toString() ?? "" })); + }); +} + +/** 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, + }; +} + +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; + } + + /** 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); + }); + } + + /** 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" }); + } + + /** 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; + }); + + afterEach(async () => { + 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(); + } + gateways.length = 0; + delete process.env.AUTH_SECRET; + }); + + it("resumes rooms and sequence numbers, and the client is a member again", async () => { + const gateway = startGateway(); + + 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 b = await connect(gateway.port, makeToken("client-b")); + b.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(b, "room_joined"); + + b.send(JSON.stringify({ type: "location_update", payload: { latitude: 1, longitude: 2 } })); + await waitForFrame(a, "location_update"); + + a.send(JSON.stringify({ type: "reconnect", roomId: "fleet-1", lastSeq: 1 })); + await waitForFrame(a, "replay_complete"); + + 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", + }); + + 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("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 = 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("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"]); + }); + + 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("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 ws = await connect(gateway.port, makeToken("client-exp"), { sessionId: blob }); + await sleep(150); + + 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); + + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(ws, "room_joined"); + }); + + 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 ws = await connect(gateway.port, makeToken("client-bad"), { sessionId: tampered }); + await sleep(150); + + expect(ws.frames.find((f) => f.type === "session_resumed")).toBeUndefined(); + expect(gateway.manager.metrics.session_resumption_total.decrypt_failed).toBe(1); + + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-1" })); + await waitForFrame(ws, "room_joined"); + }); + + 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"])); + + 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("persists a debounced save to the store within 1s of join_room", async () => { + const gateway = startGateway({ debounceMs: 500 }); + + 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 startedAt = Date.now(); + const { state } = await waitForSavedState(gateway, "client-debounce"); + + expect(Date.now() - startedAt).toBeLessThan(1000); + expect(state.rooms.map((room) => room.roomId)).toEqual(["fleet-9"]); + expect(state.authIdentity).toEqual({ sub: "client-debounce" }); + }); + + 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"); + + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-sticky" })); + await waitForFrame(first, "room_joined"); + await disconnect(gateway, first); + + 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"); + + 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("migrates a client on demand and resumes from the handed-over blob", async () => { + const gateway = startGateway(); + + 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 every live session on saveAllSessions()", async () => { + const gateway = startGateway({ debounceMs: 5000 }); + + const first = await connect(gateway.port, makeToken("client-s1")); + first.send(JSON.stringify({ type: "join_room", roomId: "fleet-s1" })); + await waitForFrame(first, "room_joined"); + + 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); + const metrics = await res.json(); + expect(metrics.session_resumption_total.new_session).toBe(1); + expect(metrics.session_resumption_total.success).toBe(1); + expect(metrics.session_resumption_total.mismatch).toBe(0); + expect(metrics.session_state_size_bytes).toBeGreaterThan(0); + }); + + it("keeps /metrics a 404 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(404); + expect(await res.json()).toEqual({ error: "Not Found" }); + }); +});