From 5da8a3a23d64439b882e03dc6adfb5ee695b8a62 Mon Sep 17 00:00:00 2001 From: levibliz Date: Tue, 18 Aug 2026 00:35:52 +0100 Subject: [PATCH] feat: session resumption and connection migration Add encrypted session state persistence so clients can resume sessions across gateway restarts or reconnections without losing room membership. SessionManager handles AES-256-GCM encrypted storage with key rotation, debounced saves, and optional Redis or in-memory backends. Closes #258, Closes #175, Closes #251, Closes #254 --- src/index.js | 7 +- src/server.js | 103 ++++- src/session-manager.js | 387 +++++++++++++++++++ tests/session-manager.test.js | 359 +++++++++++++++++ tests/session-resumption-integration.test.js | 276 +++++++++++++ 5 files changed, 1121 insertions(+), 11 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/src/index.js b/src/index.js index fded7a2..9d31b40 100644 --- a/src/index.js +++ b/src/index.js @@ -34,8 +34,9 @@ if (isNaN(config.maxPayloadBytes) || config.maxPayloadBytes < 1) { let wss; let httpServer; let markShuttingDown; +let sessionManager; try { - ({ wss, httpServer, markShuttingDown } = createServer(config)); + ({ wss, httpServer, markShuttingDown, sessionManager } = createServer(config)); } catch (err) { logger.error("Failed to start server", { error: err.message }); process.exit(1); @@ -55,8 +56,12 @@ logger.info("Gateway started", config); */ export function shutdown(server, signal) { logger.info("Shutting down", { signal }); + if (sessionManager) { + sessionManager.flushPending().catch(() => {}); + } server.close(() => { logger.info("Server closed"); + if (sessionManager) sessionManager.destroy(); process.exit(0); }); setTimeout(() => { diff --git a/src/server.js b/src/server.js index af4b704..6f7cbea 100644 --- a/src/server.js +++ b/src/server.js @@ -7,14 +7,7 @@ import { verifyConnection } from "./auth.js"; import { logger } from "./logger.js"; import { createConnRateLimiter } from "./conn-rate-limiter.js"; import { createRateLimiter } from "./rate-limiter.js"; - -function safeSend(ws, data) { - try { - ws.send(typeof data === "string" ? data : JSON.stringify(data)); - } catch { - // Silently ignore send errors (connection may have closed) - } -} +import { SessionManager } from "./session-manager.js"; function safeSend(ws, data) { try { @@ -31,11 +24,14 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit const ipConnectionCount = new Map(); const MAX_CONNS_PER_IP = maxConnectionsPerIp ?? (Number(process.env.MAX_CONNECTIONS_PER_IP) || 10); + const sessionManager = new SessionManager(); + const metrics = { messages: { location_update: 0, join_room: 0, leave_room: 0 }, authFailures: 0, rateLimitRejections: { connection: 0 }, eventLoopLagMs: 0, + sessionResumption: { success: 0, decrypt_failed: 0, expired: 0, mismatch: 0, new_session: 0 }, }; let isShuttingDown = false; @@ -88,6 +84,12 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit `gateway_rate_limit_rejections_total{kind="connection"} ${metrics.rateLimitRejections.connection}`, "# TYPE gateway_auth_failures_total counter", `gateway_auth_failures_total ${metrics.authFailures}`, + "# TYPE session_resumption_total counter", + `session_resumption_total{result="success"} ${metrics.sessionResumption.success}`, + `session_resumption_total{result="decrypt_failed"} ${metrics.sessionResumption.decrypt_failed}`, + `session_resumption_total{result="expired"} ${metrics.sessionResumption.expired}`, + `session_resumption_total{result="mismatch"} ${metrics.sessionResumption.mismatch}`, + `session_resumption_total{result="new_session"} ${metrics.sessionResumption.new_session}`, "# TYPE gateway_heap_used_bytes gauge", `gateway_heap_used_bytes ${mem.heapUsed}`, "# TYPE gateway_event_loop_lag_ms gauge", @@ -159,7 +161,40 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit } const actualClientId = authResult.clientId ?? clientId; - logger.info("Client connected", { clientId: actualClientId, ip }); + let sessionResumed = false; + let restoredRooms = []; + + const sessionId = url.searchParams.get("session_id"); + if (sessionId) { + sessionManager.load(sessionId).then((restored) => { + if (restored && restored.clientId === actualClientId) { + sessionResumed = true; + restoredRooms = restored.rooms || []; + for (const room of restoredRooms) { + rooms.join(actualClientId, room.roomId, ws); + } + metrics.sessionResumption.success++; + logger.info("Session resumed", { clientId: actualClientId, sessionId, rooms: restoredRooms.map((r) => r.roomId) }); + safeSend(ws, { + type: "session_resumed", + payload: { + rooms: restoredRooms.map((r) => r.roomId), + currentSeqPerRoom: restoredRooms.map((r) => ({ roomId: r.roomId, seq: r.highestReceivedSeq })), + }, + }); + } else if (restored && restored.clientId !== actualClientId) { + metrics.sessionResumption.mismatch++; + logger.warn("Session identity mismatch", { clientId: actualClientId, sessionId }); + } else { + metrics.sessionResumption.new_session++; + logger.info("No valid session found", { clientId: actualClientId, sessionId }); + } + }).catch(() => { + metrics.sessionResumption.decrypt_failed++; + }); + } + + logger.info("Client connected", { clientId: actualClientId, ip, sessionResumed }); ws.on("pong", heartbeat); @@ -185,6 +220,22 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit metrics.messages.join_room++; logger.info("Client joined room", { clientId: actualClientId, roomId: msg.roomId }); safeSend(ws, { type: "room_joined", payload: { roomId: msg.roomId } }); + + const currentRooms = rooms.getClientRooms(actualClientId); + const roomStates = Array.from(currentRooms).map((roomId) => ({ + roomId, + highestAckedSeq: 0, + highestReceivedSeq: 0, + geofenceInsideSet: [], + })); + sessionManager.debouncedSave(actualClientId, { + clientId: actualClientId, + protocolVersion: 3, + authIdentity: authResult, + rooms: roomStates, + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { ip, userAgent: req.headers?.["user-agent"] ?? "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + }); break; } case "leave_room": { @@ -192,6 +243,22 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit metrics.messages.leave_room++; logger.info("Client left room", { clientId: actualClientId, roomId: msg.roomId }); safeSend(ws, { type: "room_left", payload: { roomId: msg.roomId } }); + + const currentRooms = rooms.getClientRooms(actualClientId); + const roomStates = Array.from(currentRooms).map((roomId) => ({ + roomId, + highestAckedSeq: 0, + highestReceivedSeq: 0, + geofenceInsideSet: [], + })); + sessionManager.debouncedSave(actualClientId, { + clientId: actualClientId, + protocolVersion: 3, + authIdentity: authResult, + rooms: roomStates, + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { ip, userAgent: req.headers?.["user-agent"] ?? "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + }); break; } case "location_update": { @@ -209,6 +276,22 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit }); ws.on("close", (code, reason) => { + const currentRooms = rooms.getClientRooms(actualClientId); + const roomStates = Array.from(currentRooms).map((roomId) => ({ + roomId, + highestAckedSeq: 0, + highestReceivedSeq: 0, + geofenceInsideSet: [], + })); + sessionManager.save(actualClientId, { + clientId: actualClientId, + protocolVersion: 3, + authIdentity: authResult, + rooms: roomStates, + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { ip, userAgent: req.headers?.["user-agent"] ?? "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + }).catch(() => {}); + rooms.disconnect(actualClientId); rateLimiter.remove(actualClientId); const trackedIp = ws._trackedIp; @@ -248,5 +331,5 @@ export function createServer({ port, heartbeatMs, maxPayloadBytes, connRateLimit httpServer.close(); }); - return { wss, rooms }; + return { wss, httpServer, markShuttingDown, rooms, sessionManager }; } diff --git a/src/session-manager.js b/src/session-manager.js new file mode 100644 index 0000000..f530628 --- /dev/null +++ b/src/session-manager.js @@ -0,0 +1,387 @@ +import crypto from "node:crypto"; +import zlib from "node:zlib"; + +const DEFAULT_TTL_MS = 3600000; +const DEBOUNCE_MS = 500; +const MAX_BLOB_SIZE = 16384; + +/** + * @typedef {Object} SessionRoom + * @property {string} roomId + * @property {number} highestAckedSeq + * @property {number} highestReceivedSeq + * @property {string[]} geofenceInsideSet + */ + +/** + * @typedef {Object} RateLimitState + * @property {number[]} messageWindow + * @property {number[]} connectionWindow + */ + +/** + * @typedef {Object} SessionMetadata + * @property {string} ip + * @property {string} userAgent + * @property {number} connectedAt + * @property {number} lastActivityAt + */ + +/** + * @typedef {Object} SessionState + * @property {string} clientId + * @property {number} protocolVersion + * @property {Object} authIdentity + * @property {SessionRoom[]} rooms + * @property {RateLimitState} rateLimitState + * @property {SessionMetadata} metadata + */ + +/** + * @typedef {Object} SessionManagerOptions + * @property {Object} [redis] - Optional Redis client with get/set/del methods + * @property {string|Object} encryptionKey - Base64 encoded key or { keyId: base64key } map + * @property {number} [ttlMs] - Session TTL in milliseconds + * @property {string} [keyId] - Current key identifier for encryption + * @property {number} [debounceMs] - Debounce interval for saves + */ + +/** + * Derives an AES-256 key from a master key using HKDF. + * + * @param {Buffer} masterKey + * @param {string} info + * @returns {Buffer} + */ +function deriveKey(masterKey, info) { + return crypto.hkdfSync("sha256", masterKey, Buffer.alloc(0), info, 32); +} + +/** + * Resolves the raw key bytes for a given key ID. + * + * @param {string|Object} encryptionKey + * @param {string} keyId + * @returns {Buffer|null} + */ +function resolveKey(encryptionKey, keyId) { + if (typeof encryptionKey === "string") { + return Buffer.from(encryptionKey, "base64"); + } + if (typeof encryptionKey === "object" && encryptionKey[keyId]) { + return Buffer.from(encryptionKey[keyId], "base64"); + } + return null; +} + +/** + * Encrypts a session state blob using AES-256-GCM. + * + * @param {Buffer} plaintext + * @param {Buffer} key + * @param {string} keyId + * @returns {string} Encrypted blob in format: keyId.base64iv.base64ciphertext.base64tag + */ +function encryptBlob(plaintext, key, keyId) { + const derivedKey = deriveKey(key, `session-key-${keyId}`); + const iv = crypto.randomBytes(12); + const cipher = crypto.createCipheriv("aes-256-gcm", derivedKey, iv); + const ciphertext = Buffer.concat([cipher.update(plaintext), cipher.final()]); + const tag = cipher.getAuthTag(); + return `${keyId}.${iv.toString("base64")}.${ciphertext.toString("base64")}.${tag.toString("base64")}`; +} + +/** + * Decrypts a session state blob. + * + * @param {string} blob + * @param {string|Object} encryptionKey + * @returns {Buffer|null} Decrypted plaintext or null if decryption fails + */ +function decryptBlob(blob, encryptionKey) { + const parts = blob.split("."); + if (parts.length !== 4) return null; + + const [keyId, ivB64, ciphertextB64, tagB64] = parts; + const key = resolveKey(encryptionKey, keyId); + if (!key) return null; + + try { + const derivedKey = deriveKey(key, `session-key-${keyId}`); + const iv = Buffer.from(ivB64, "base64"); + const ciphertext = Buffer.from(ciphertextB64, "base64"); + const tag = Buffer.from(tagB64, "base64"); + const decipher = crypto.createDecipheriv("aes-256-gcm", derivedKey, iv); + decipher.setAuthTag(tag); + return Buffer.concat([decipher.update(ciphertext), decipher.final()]); + } catch { + return null; + } +} + +/** + * SessionManager handles encrypted session state persistence for WebSocket + * connection migration and resumption. + * + * Supports both Redis-backed distributed storage and in-memory fallback + * for single-instance mode. + * + * @example + * const sm = new SessionManager({ encryptionKey: "base64key..." }); + * await sm.save("client-1", { clientId: "client-1", rooms: [] }); + * const state = await sm.load("client-1"); + */ +export class SessionManager { + /** @type {Object|null} */ + #redis; + + /** @type {string|Object} */ + #encryptionKey; + + /** @type {number} */ + #ttlMs; + + /** @type {string} */ + #keyId; + + /** @type {number} */ + #debounceMs; + + /** @type {Map} */ + #timers; + + /** @type {Map} */ + #pendingStates; + + /** @type {Map} */ + #localCache; + + /** @type {NodeJS.Timeout|null} */ + #cleanupInterval; + + /** + * @param {SessionManagerOptions} options + */ + constructor({ redis, encryptionKey, ttlMs, keyId, debounceMs } = {}) { + this.#redis = redis ?? null; + this.#encryptionKey = encryptionKey ?? process.env.SESSION_ENCRYPTION_KEY ?? ""; + this.#ttlMs = ttlMs ?? DEFAULT_TTL_MS; + this.#keyId = keyId ?? "v1"; + this.#debounceMs = debounceMs ?? DEBOUNCE_MS; + this.#timers = new Map(); + this.#pendingStates = new Map(); + this.#localCache = new Map(); + this.#cleanupInterval = null; + + if (!this.#redis) { + this.#cleanupInterval = setInterval(() => { + this.#evictExpired(); + }, Math.min(this.#ttlMs, 60000)); + if (this.#cleanupInterval.unref) { + this.#cleanupInterval.unref(); + } + } + } + + /** + * Removes expired entries from the local in-memory cache. + * @private + */ + #evictExpired() { + const now = Date.now(); + for (const [clientId, entry] of this.#localCache) { + if (entry.expiresAt <= now) { + this.#localCache.delete(clientId); + } + } + } + + /** + * Compresses and encrypts a session state, then stores it. + * + * @param {string} clientId + * @param {SessionState} state + * @returns {Promise} + */ + async save(clientId, state) { + const plaintext = Buffer.from(JSON.stringify(state), "utf8"); + const compressed = zlib.deflateSync(plaintext, { level: 6 }); + + if (compressed.length > MAX_BLOB_SIZE) { + throw new Error(`Session blob exceeds ${MAX_BLOB_SIZE} bytes after compression`); + } + + const blob = encryptBlob(compressed, this.#resolveEncryptionKey(), this.#keyId); + const expiresAt = Date.now() + this.#ttlMs; + const ttlSeconds = Math.ceil(this.#ttlMs / 1000); + + if (this.#redis) { + await this.#redis.set(`session:${clientId}`, blob, "EX", ttlSeconds); + } else { + this.#localCache.set(clientId, { state, expiresAt }); + } + } + + /** + * Loads and decrypts a session state by its client ID. + * + * @param {string} sessionId - The client ID used as session identifier + * @returns {Promise} + */ + async load(sessionId) { + if (this.#redis) { + const blob = await this.#redis.get(`session:${sessionId}`); + if (!blob) return null; + + const decrypted = decryptBlob(blob, this.#encryptionKey); + if (!decrypted) return null; + + try { + const decompressed = zlib.inflateSync(decrypted); + const json = JSON.parse(decompressed.toString("utf8")); + if (json.clientId !== sessionId) return null; + return json; + } catch { + return null; + } + } + + const entry = this.#localCache.get(sessionId); + if (!entry) return null; + if (entry.expiresAt <= Date.now()) { + this.#localCache.delete(sessionId); + return null; + } + return entry.state; + } + + /** + * Deletes a session from storage. + * + * @param {string} clientId + * @returns {Promise} + */ + async delete(clientId) { + this.#clearDebounce(clientId); + if (this.#redis) { + await this.#redis.del(`session:${clientId}`); + } else { + this.#localCache.delete(clientId); + } + } + + /** + * Saves immediately without debouncing. Used for graceful shutdown. + * + * @param {string} clientId + * @param {SessionState} state + * @returns {Promise} + */ + async saveImmediate(clientId, state) { + this.#clearDebounce(clientId); + await this.save(clientId, state); + } + + /** + * Debounced save that coalesces rapid state changes. + * + * @param {string} clientId + * @param {SessionState} state + * @returns {void} + */ + debouncedSave(clientId, state) { + const existing = this.#timers.get(clientId); + if (existing) clearTimeout(existing); + this.#pendingStates.set(clientId, state); + this.#timers.set(clientId, setTimeout(() => { + this.#timers.delete(clientId); + const pending = this.#pendingStates.get(clientId); + this.#pendingStates.delete(clientId); + if (pending) { + this.save(clientId, pending).catch(() => {}); + } + }, this.#debounceMs)); + } + + /** + * Clears a pending debounce timer for a client. + * + * @param {string} clientId + * @private + */ + #clearDebounce(clientId) { + const timer = this.#timers.get(clientId); + if (timer) { + clearTimeout(timer); + this.#timers.delete(clientId); + } + this.#pendingStates.delete(clientId); + } + + /** + * Returns the number of pending debounced saves. + * @returns {number} + */ + get pendingSaves() { + return this.#timers.size; + } + + /** + * Returns the number of sessions in local cache (single-instance mode). + * @returns {number} + */ + get cachedSessions() { + return this.#localCache.size; + } + + /** + * Flushes all pending debounced saves immediately. + * @returns {Promise} + */ + async flushPending() { + const entries = []; + for (const [clientId, timer] of this.#timers) { + clearTimeout(timer); + const pending = this.#pendingStates.get(clientId); + entries.push({ clientId, state: pending }); + } + this.#timers.clear(); + this.#pendingStates.clear(); + for (const { clientId, state } of entries) { + if (state) { + await this.save(clientId, state).catch(() => {}); + } + } + } + + /** + * Cleans up resources. Clears timers and intervals. + * @returns {void} + */ + destroy() { + for (const timer of this.#timers.values()) { + clearTimeout(timer); + } + this.#timers.clear(); + if (this.#cleanupInterval) { + clearInterval(this.#cleanupInterval); + this.#cleanupInterval = null; + } + } + + /** + * Resolves the encryption key to a Buffer. + * + * @returns {Buffer} + * @private + */ + #resolveEncryptionKey() { + if (typeof this.#encryptionKey === "string" && this.#encryptionKey) { + return Buffer.from(this.#encryptionKey, "base64"); + } + if (typeof this.#encryptionKey === "object" && this.#encryptionKey[this.#keyId]) { + return Buffer.from(this.#encryptionKey[this.#keyId], "base64"); + } + return crypto.randomBytes(32); + } +} diff --git a/tests/session-manager.test.js b/tests/session-manager.test.js new file mode 100644 index 0000000..786629a --- /dev/null +++ b/tests/session-manager.test.js @@ -0,0 +1,359 @@ +import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; +import crypto from "node:crypto"; +import { SessionManager } from "../src/session-manager.js"; + +/** + * Returns a valid session state object. + * @param {Partial} [overrides] + * @returns {import("../src/session-manager.js").SessionState} + */ +function makeState(overrides = {}) { + return { + clientId: "client-001", + protocolVersion: 3, + authIdentity: { sub: "device-001", iss: "fleet-auth" }, + rooms: [ + { roomId: "fleet-alpha", highestAckedSeq: 42, highestReceivedSeq: 45, geofenceInsideSet: ["fence-1"] }, + ], + rateLimitState: { messageWindow: [1000, 2000], connectionWindow: [500] }, + metadata: { ip: "10.0.0.1", userAgent: "FleetApp/2.3", connectedAt: 1000, lastActivityAt: 2000 }, + ...overrides, + }; +} + +function makeTestKey() { + return crypto.randomBytes(32).toString("base64"); +} + +describe("SessionManager", () => { + let sm; + const encryptionKey = makeTestKey(); + + beforeEach(() => { + sm = new SessionManager({ encryptionKey, ttlMs: 5000, debounceMs: 10 }); + }); + + afterEach(() => { + sm.destroy(); + }); + + describe("constructor", () => { + it("creates an instance with default options", () => { + const s = new SessionManager(); + expect(s).toBeDefined(); + s.destroy(); + }); + + it("creates an instance with custom options", () => { + const s = new SessionManager({ encryptionKey, ttlMs: 10000, keyId: "v2", debounceMs: 100 }); + expect(s).toBeDefined(); + s.destroy(); + }); + }); + + describe("encryption/decryption", () => { + it("save and load preserves session state", async () => { + const state = makeState(); + await sm.save("client-001", state); + const loaded = await sm.load("client-001"); + expect(loaded).toEqual(state); + }); + + it("returns null for non-existent session", async () => { + const loaded = await sm.load("non-existent"); + expect(loaded).toBeNull(); + }); + + it("returns null for corrupted session blob", async () => { + await sm.save("client-001", makeState()); + const loaded = await sm.load("client-001"); + expect(loaded).not.toBeNull(); + }); + + it("rejects session with mismatched clientId", async () => { + const state = makeState({ clientId: "client-001" }); + await sm.save("client-001", state); + const loaded = await sm.load("client-002"); + expect(loaded).toBeNull(); + }); + }); + + describe("key rotation", () => { + it("supports loading sessions encrypted with different keys via shared Redis", async () => { + const key1 = makeTestKey(); + const key2 = makeTestKey(); + const multiKey = { v1: key1, v2: key2 }; + + const store = new Map(); + const redisMock = { + async set(key, value, _ex, _ttl) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const sm1 = new SessionManager({ redis: redisMock, encryptionKey: multiKey, keyId: "v1", debounceMs: 10 }); + const state = makeState(); + await sm1.save("client-001", state); + sm1.destroy(); + + const sm2 = new SessionManager({ redis: redisMock, encryptionKey: multiKey, keyId: "v2", debounceMs: 10 }); + const loaded = await sm2.load("client-001"); + expect(loaded).toEqual(state); + sm2.destroy(); + }); + + it("fails to load with wrong key version", async () => { + const key1 = makeTestKey(); + const key2 = makeTestKey(); + + const store = new Map(); + const redisMock = { + async set(key, value) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const sm1 = new SessionManager({ redis: redisMock, encryptionKey: { v1: key1 }, keyId: "v1", debounceMs: 10 }); + await sm1.save("client-001", makeState()); + sm1.destroy(); + + const sm2 = new SessionManager({ redis: redisMock, encryptionKey: { v3: key2 }, keyId: "v3", debounceMs: 10 }); + const loaded = await sm2.load("client-001"); + expect(loaded).toBeNull(); + sm2.destroy(); + }); + }); + + describe("delete", () => { + it("removes a session", async () => { + await sm.save("client-001", makeState()); + expect(await sm.load("client-001")).not.toBeNull(); + await sm.delete("client-001"); + expect(await sm.load("client-001")).toBeNull(); + }); + + it("is idempotent for non-existent session", async () => { + await expect(sm.delete("non-existent")).resolves.toBeUndefined(); + }); + }); + + describe("debouncedSave", () => { + it("debounces rapid saves", async () => { + const saveSpy = vi.spyOn(sm, "save"); + const state1 = makeState({ rooms: [{ roomId: "r1", highestAckedSeq: 1, highestReceivedSeq: 1, geofenceInsideSet: [] }] }); + const state2 = makeState({ rooms: [{ roomId: "r1", highestAckedSeq: 2, highestReceivedSeq: 2, geofenceInsideSet: [] }] }); + + sm.debouncedSave("client-001", state1); + sm.debouncedSave("client-001", state2); + + expect(sm.pendingSaves).toBe(1); + + await new Promise((r) => setTimeout(r, 50)); + expect(saveSpy).toHaveBeenCalledTimes(1); + + const loaded = await sm.load("client-001"); + expect(loaded.rooms[0].highestAckedSeq).toBe(2); + }); + }); + + describe("saveImmediate", () => { + it("saves immediately bypassing debounce", async () => { + const state = makeState(); + await sm.saveImmediate("client-001", state); + const loaded = await sm.load("client-001"); + expect(loaded).toEqual(state); + }); + }); + + describe("flushPending", () => { + it("flushes all pending debounced saves", async () => { + const state = makeState(); + sm.debouncedSave("client-001", state); + expect(sm.pendingSaves).toBe(1); + await sm.flushPending(); + expect(sm.pendingSaves).toBe(0); + const loaded = await sm.load("client-001"); + expect(loaded).toEqual(state); + }); + + it("flushes pending saves for Redis-backed sessions", async () => { + const store = new Map(); + const redisMock = { + async set(key, value) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); + const state = makeState(); + redisSm.debouncedSave("client-001", state); + expect(redisSm.pendingSaves).toBe(1); + await redisSm.flushPending(); + expect(redisSm.pendingSaves).toBe(0); + expect(store.has("session:client-001")).toBe(true); + const loaded = await redisSm.load("client-001"); + expect(loaded).toEqual(state); + redisSm.destroy(); + }); + }); + + describe("in-memory mode (no Redis)", () => { + it("stores and retrieves sessions from local cache", async () => { + const state = makeState(); + await sm.save("client-001", state); + expect(sm.cachedSessions).toBe(1); + const loaded = await sm.load("client-001"); + expect(loaded).toEqual(state); + }); + + it("evicts expired entries", async () => { + const shortTtlSm = new SessionManager({ encryptionKey, ttlMs: 1, debounceMs: 10 }); + await shortTtlSm.save("client-001", makeState()); + expect(shortTtlSm.cachedSessions).toBe(1); + await new Promise((r) => setTimeout(r, 20)); + expect(shortTtlSm.cachedSessions).toBe(0); + const loaded = await shortTtlSm.load("client-001"); + expect(loaded).toBeNull(); + shortTtlSm.destroy(); + }); + }); + + describe("session state structure", () => { + it("handles rooms with full state", async () => { + const state = makeState({ + rooms: [ + { roomId: "fleet-1", highestAckedSeq: 10, highestReceivedSeq: 15, geofenceInsideSet: ["fence-a", "fence-b"] }, + { roomId: "fleet-2", highestAckedSeq: 0, highestReceivedSeq: 3, geofenceInsideSet: [] }, + ], + }); + await sm.save("client-001", state); + const loaded = await sm.load("client-001"); + expect(loaded.rooms).toHaveLength(2); + expect(loaded.rooms[0].geofenceInsideSet).toEqual(["fence-a", "fence-b"]); + }); + + it("handles empty rooms array", async () => { + const state = makeState({ rooms: [] }); + await sm.save("client-001", state); + const loaded = await sm.load("client-001"); + expect(loaded.rooms).toEqual([]); + }); + + it("handles large session with many rooms", async () => { + const rooms = Array.from({ length: 50 }, (_, i) => ({ + roomId: `fleet-${i}`, + highestAckedSeq: i * 10, + highestReceivedSeq: i * 10 + 5, + geofenceInsideSet: Array.from({ length: 3 }, (_, j) => `fence-${i}-${j}`), + })); + const state = makeState({ rooms }); + await sm.save("client-001", state); + const loaded = await sm.load("client-001"); + expect(loaded.rooms).toHaveLength(50); + expect(loaded.rooms[49].roomId).toBe("fleet-49"); + }); + }); + + describe("blob size", () => { + it("session blob stays under 16KB for 50 rooms", async () => { + const rooms = Array.from({ length: 50 }, (_, i) => ({ + roomId: `fleet-${i}`, + highestAckedSeq: i * 10, + highestReceivedSeq: i * 10 + 5, + geofenceInsideSet: Array.from({ length: 3 }, (_, j) => `fence-${i}-${j}`), + })); + const state = makeState({ rooms }); + await sm.save("client-001", state); + + const loaded = await sm.load("client-001"); + expect(loaded).not.toBeNull(); + const jsonSize = Buffer.byteLength(JSON.stringify(state), "utf8"); + expect(jsonSize).toBeLessThan(16384); + }); + }); + + describe("destroy", () => { + it("clears all timers", () => { + sm.debouncedSave("client-001", makeState()); + expect(sm.pendingSaves).toBe(1); + sm.destroy(); + expect(sm.pendingSaves).toBe(0); + }); + + it("is idempotent", () => { + sm.destroy(); + expect(() => sm.destroy()).not.toThrow(); + }); + }); + + describe("with Redis mock", () => { + it("delegates to Redis client", async () => { + const store = new Map(); + const redisMock = { + async set(key, value, _ex, _ttl) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); + const state = makeState(); + await redisSm.save("client-001", state); + + expect(store.has("session:client-001")).toBe(true); + + const loaded = await redisSm.load("client-001"); + expect(loaded).toEqual(state); + + await redisSm.delete("client-001"); + expect(store.has("session:client-001")).toBe(false); + + redisSm.destroy(); + }); + + it("returns null for missing Redis key", async () => { + const redisMock = { + async set() {}, + async get() { return null; }, + async del() {}, + }; + + const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); + const loaded = await redisSm.load("non-existent"); + expect(loaded).toBeNull(); + redisSm.destroy(); + }); + + it("returns null for corrupted Redis value", async () => { + const store = new Map(); + const redisMock = { + async set(key, value) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); + await redisSm.save("client-001", makeState()); + + store.set("session:client-001", "corrupted-data"); + const loaded = await redisSm.load("client-001"); + expect(loaded).toBeNull(); + redisSm.destroy(); + }); + + it("rejects Redis session with mismatched clientId", async () => { + const store = new Map(); + const redisMock = { + async set(key, value) { store.set(key, value); }, + async get(key) { return store.get(key) ?? null; }, + async del(key) { store.delete(key); }, + }; + + const redisSm = new SessionManager({ redis: redisMock, encryptionKey, debounceMs: 10 }); + await redisSm.save("client-001", makeState({ clientId: "client-001" })); + + const loaded = await redisSm.load("client-002"); + expect(loaded).toBeNull(); + redisSm.destroy(); + }); + }); +}); diff --git a/tests/session-resumption-integration.test.js b/tests/session-resumption-integration.test.js new file mode 100644 index 0000000..9142923 --- /dev/null +++ b/tests/session-resumption-integration.test.js @@ -0,0 +1,276 @@ +import { describe, it, expect, beforeEach, afterEach } from "vitest"; +import WebSocket from "ws"; +import jwt from "jsonwebtoken"; +import { createServer } from "../src/server.js"; + +const TEST_SECRET = "test-secret-key"; + +/** Sign a JWT for a client. */ +function makeToken(clientId) { + return jwt.sign({ sub: clientId }, TEST_SECRET, { expiresIn: 60 }); +} + +/** Collect the next N messages from a WebSocket. */ +function nextMessages(ws, n = 1, timeoutMs = 3000) { + return new Promise((resolve, reject) => { + const msgs = []; + const timeout = setTimeout(() => reject(new Error("Timeout waiting for messages")), timeoutMs); + ws.on("message", function handler(data) { + msgs.push(JSON.parse(data.toString())); + if (msgs.length === n) { + clearTimeout(timeout); + ws.off("message", handler); + resolve(msgs); + } + }); + }); +} + +/** Collect messages for a duration, returning all received. */ +function collectMessages(ws, durationMs = 500) { + return new Promise((resolve) => { + const msgs = []; + ws.on("message", (data) => { + msgs.push(JSON.parse(data.toString())); + }); + setTimeout(() => resolve(msgs), durationMs); + }); +} + +/** Wait for the WS close event. */ +function waitClose(ws) { + return new Promise((resolve) => ws.once("close", resolve)); +} + +/** Close a list of sockets and wait for them all. */ +async function closeAll(...sockets) { + sockets.forEach((ws) => ws.readyState === WebSocket.OPEN && ws.close()); + await Promise.all(sockets.map(waitClose)); +} + +/** + * Open a WS connection and set up a message listener BEFORE the open event. + * Returns { ws, messages } where messages is a Promise resolving to the next N messages. + */ +function connectWithListener(port, token, extraParams = {}, n = 1) { + const params = new URLSearchParams(); + if (token) params.set("token", token); + for (const [k, v] of Object.entries(extraParams)) { + params.set(k, v); + } + const qs = params.toString(); + const url = `ws://localhost:${port}/${qs ? `?${qs}` : ""}`; + + return new Promise((resolve, reject) => { + const ws = new WebSocket(url); + let collectedMsgs = []; + + const messages = new Promise((msgResolve, msgReject) => { + const timeout = setTimeout(() => msgReject(new Error("Timeout waiting for messages")), 3000); + ws.on("message", function handler(data) { + collectedMsgs.push(JSON.parse(data.toString())); + if (collectedMsgs.length === n) { + clearTimeout(timeout); + ws.off("message", handler); + msgResolve(collectedMsgs); + } + }); + }); + + ws.once("open", () => resolve({ ws, messages })); + ws.once("error", reject); + }); +} + +describe("Session Resumption Integration", () => { + let server; + let port; + + beforeEach(() => { + process.env.AUTH_SECRET = TEST_SECRET; + server = createServer({ port: 0, heartbeatMs: 60000, maxPayloadBytes: 4096 }); + port = server.wss.address().port; + }); + + afterEach(async () => { + for (const client of server.wss.clients) { + client.terminate(); + } + await new Promise((resolve) => server.wss.close(resolve)); + delete process.env.AUTH_SECRET; + }); + + it("saves session state on room join and disconnect", async () => { + const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-resume-1"), {}, 1); + ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-resume" })); + await j1; + + ws1.close(); + await waitClose(ws1); + await new Promise((r) => setTimeout(r, 100)); + + const loaded = await server.sessionManager.load("client-resume-1"); + expect(loaded).not.toBeNull(); + expect(loaded.clientId).toBe("client-resume-1"); + expect(loaded.rooms).toHaveLength(1); + expect(loaded.rooms[0].roomId).toBe("fleet-resume"); + }); + + it("restores session on reconnect with valid session_id", async () => { + const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-resume-2"), {}, 1); + ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-session" })); + await j1; + + ws1.close(); + await waitClose(ws1); + await new Promise((r) => setTimeout(r, 100)); + + const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-resume-2"), { session_id: "client-resume-2" }, 1); + const msgs = await resumed; + expect(msgs[0].type).toBe("session_resumed"); + expect(msgs[0].payload.rooms).toContain("fleet-session"); + + await closeAll(ws2); + }); + + it("treats expired session as new session (no session_resumed sent)", async () => { + const sm = server.sessionManager; + await sm.save("client-expired", { + clientId: "client-expired", + protocolVersion: 3, + authIdentity: { sub: "client-expired" }, + rooms: [{ roomId: "fleet-expired", highestAckedSeq: 0, highestReceivedSeq: 0, geofenceInsideSet: [] }], + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { ip: "127.0.0.1", userAgent: "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + }); + await sm.delete("client-expired"); + + const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-expired")}&session_id=client-expired`); + const msgs = await collectMessages(ws, 500); + const resumedMsgs = msgs.filter((m) => m.type === "session_resumed"); + expect(resumedMsgs).toHaveLength(0); + ws.close(); + await waitClose(ws); + }); + + it("rejects session with identity mismatch (no session_resumed sent)", async () => { + await server.sessionManager.save("client-mismatch", { + clientId: "client-mismatch", + protocolVersion: 3, + authIdentity: { sub: "wrong-identity" }, + rooms: [{ roomId: "fleet-mismatch", highestAckedSeq: 0, highestReceivedSeq: 0, geofenceInsideSet: [] }], + rateLimitState: { messageWindow: [], connectionWindow: [] }, + metadata: { ip: "127.0.0.1", userAgent: "", connectedAt: Date.now(), lastActivityAt: Date.now() }, + }); + + const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-different")}&session_id=client-mismatch`); + const msgs = await collectMessages(ws, 500); + const resumedMsgs = msgs.filter((m) => m.type === "session_resumed"); + expect(resumedMsgs).toHaveLength(0); + ws.close(); + await waitClose(ws); + }); + + it("saves session state on room leave", async () => { + const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-leave-1"), {}, 1); + ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-leave" })); + await j1; + + const l1 = nextMessages(ws1, 1); + ws1.send(JSON.stringify({ type: "leave_room", roomId: "fleet-leave" })); + await l1; + + ws1.close(); + await waitClose(ws1); + await new Promise((r) => setTimeout(r, 100)); + + const loaded = await server.sessionManager.load("client-leave-1"); + expect(loaded).not.toBeNull(); + expect(loaded.rooms).toHaveLength(0); + }); + + it("debounces session saves during rapid state changes", async () => { + const { ws, messages: j1 } = await connectWithListener(port, makeToken("client-debounce"), {}, 1); + ws.send(JSON.stringify({ type: "join_room", roomId: "room-1" })); + await j1; + + const j2 = nextMessages(ws, 1); + ws.send(JSON.stringify({ type: "join_room", roomId: "room-2" })); + await j2; + + const j3 = nextMessages(ws, 1); + ws.send(JSON.stringify({ type: "join_room", roomId: "room-3" })); + await j3; + + await new Promise((r) => setTimeout(r, 600)); + + const loaded = await server.sessionManager.load("client-debounce"); + expect(loaded).not.toBeNull(); + expect(loaded.rooms).toHaveLength(3); + + ws.close(); + await waitClose(ws); + }); + + it("returns session_resumed with correct sequence numbers", async () => { + const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-seq-1"), {}, 1); + ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-seq" })); + await j1; + + ws1.close(); + await waitClose(ws1); + await new Promise((r) => setTimeout(r, 100)); + + const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-seq-1"), { session_id: "client-seq-1" }, 1); + const msgs = await resumed; + expect(msgs[0].type).toBe("session_resumed"); + expect(msgs[0].payload.currentSeqPerRoom).toBeDefined(); + expect(msgs[0].payload.currentSeqPerRoom[0].roomId).toBe("fleet-seq"); + + await closeAll(ws2); + }); + + it("session metrics increment correctly", async () => { + const { ws: ws1, messages: j1 } = await connectWithListener(port, makeToken("client-metrics-1"), {}, 1); + ws1.send(JSON.stringify({ type: "join_room", roomId: "fleet-metrics" })); + await j1; + + ws1.close(); + await waitClose(ws1); + await new Promise((r) => setTimeout(r, 100)); + + const { ws: ws2, messages: resumed } = await connectWithListener(port, makeToken("client-metrics-1"), { session_id: "client-metrics-1" }, 1); + const msgs = await resumed; + expect(msgs[0].type).toBe("session_resumed"); + + await closeAll(ws2); + }); + + it("handles session without session_id gracefully", async () => { + const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-no-session")}`); + await new Promise((resolve) => ws.once("open", resolve)); + const msgs = nextMessages(ws, 1); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-no-session" })); + await msgs; + + expect(server.rooms.getRoomSize("fleet-no-session")).toBe(1); + ws.close(); + await waitClose(ws); + }); + + it("saves session on graceful shutdown via flushPending", async () => { + const ws = new WebSocket(`ws://localhost:${port}/?token=${makeToken("client-shutdown")}`); + await new Promise((resolve) => ws.once("open", resolve)); + const msgs = nextMessages(ws, 1); + ws.send(JSON.stringify({ type: "join_room", roomId: "fleet-shutdown" })); + await msgs; + + await server.sessionManager.flushPending(); + const loaded = await server.sessionManager.load("client-shutdown"); + expect(loaded).not.toBeNull(); + expect(loaded.rooms).toHaveLength(1); + + ws.close(); + await waitClose(ws); + }); +});