feat: session resumption and connection migration with encrypted state transfer

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://<host>:<port>/?token=<jwt>
 
 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://<host>:<port>/?token=<jwt>&session_id=<url-encoded blob>
+```
+
+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=<instanceId>` 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
 
 ```

src/index.js

@@ -5,13 +5,27 @@ import { logger } from "./logger.js";
 /**
  * Parses environment variables into server configuration with integer coercion.
  *
- * @returns {{ port: number, heartbeatMs: number, maxPayloadBytes: number }}
+ * Session-resumption keys are only present when their variables are set, so an
+ * unconfigured deployment keeps the historical three-key shape.
+ *
+ * @returns {{ port: number, heartbeatMs: number, maxPayloadBytes: number, sessionEncryptionKey?: string, sessionTtlMs?: number, instanceId?: string }}
  */
 export function parseConfig() {
   const port = parseInt(process.env.PORT ?? "8080", 10);
   const heartbeatMs = parseInt(process.env.WS_HEARTBEAT_MS ?? "30000", 10);
   const maxPayloadBytes = parseInt(process.env.MAX_PAYLOAD_BYTES ?? "1024", 10);
-  return { port, heartbeatMs, maxPayloadBytes };
+  const config = { port, heartbeatMs, maxPayloadBytes };
+
+  if (process.env.SESSION_ENCRYPTION_KEY) {
+    config.sessionEncryptionKey = process.env.SESSION_ENCRYPTION_KEY;
+  }
+  if (process.env.SESSION_TTL_MS) {
+    config.sessionTtlMs = parseInt(process.env.SESSION_TTL_MS, 10);
+  }
+  if (process.env.INSTANCE_ID) {
+    config.instanceId = process.env.INSTANCE_ID;
+  }
+  return config;
 }
 
 const config = parseConfig();
@@ -31,18 +45,31 @@ if (isNaN(config.maxPayloadBytes) || config.maxPayloadBytes < 1) {
   process.exit(1);
 }
 
+if (config.sessionTtlMs !== undefined && (isNaN(config.sessionTtlMs) || config.sessionTtlMs < 1)) {
+  logger.error("Invalid SESSION_TTL_MS value", { SESSION_TTL_MS: process.env.SESSION_TTL_MS });
+  process.exit(1);
+}
+
 let wss;
-let httpServer;
 let markShuttingDown;
 let sessionManager;
+let instanceId;
+let saveAllSessions;
 try {
-  ({ wss, httpServer, markShuttingDown, sessionManager } = createServer(config));
+  ({ wss, markShuttingDown, sessionManager, instanceId, saveAllSessions } = createServer(config));
 } catch (err) {
   logger.error("Failed to start server", { error: err.message });
   process.exit(1);
 }
 
-logger.info("Gateway started", config);
+// The encryption key never reaches the logs.
+logger.info("Gateway started", {
+  port: config.port,
+  heartbeatMs: config.heartbeatMs,
+  maxPayloadBytes: config.maxPayloadBytes,
+  instanceId,
+  sessionResumption: sessionManager != null,
+});
 
 /**
  * Initiates a multi-phase graceful shutdown of the WebSocket server.
@@ -53,28 +80,57 @@ logger.info("Gateway started", config);
  * Phase 4 (4000ms):    Close connections with WebSocket code 1001 "Going Away".
  * Phase 5 (>5000ms):   Force exit.
  *
+ * When session resumption is active, phase 2 first persists every live session
+ * and hands each client its own fresh blob as `session_id`, so a client can
+ * resume on another instance immediately.
+ *
  * @param {object} wss - The WebSocket server instance.
  * @param {string} signal - The OS signal that triggered the shutdown (e.g. "SIGTERM").
+ * @param {object} [options]
+ * @param {() => Promise<Map<string, string>>} [options.saveAllSessions] - Persists live sessions.
  * @returns {void}
  */
-export function shutdown(wss, signal) {
+export function shutdown(wss, signal, { saveAllSessions } = {}) {
   logger.info("shutdown: stopping accept", { signal });
 
   const clientCount = wss.clients ? wss.clients.size : 0;
 
+  /** Sends `server_shutting_down`, adding a per-client blob when one exists. */
+  function notifyClients(blobs) {
+    const shared = JSON.stringify({ type: "server_shutting_down", payload: { reconnectIn: 5 } });
+    for (const client of wss.clients) {
+      const sessionId = blobs?.get(client._clientId);
+      try {
+        client.send(
+          sessionId
+            ? JSON.stringify({
+                type: "server_shutting_down",
+                payload: { reconnectIn: 5, session_id: sessionId },
+              })
+            : shared
+        );
+      } catch {
+        // Client may already be disconnected
+      }
+    }
+  }
+
   // Phase 2 — Notify clients (100ms)
   setTimeout(() => {
     logger.info("shutdown: notifying N clients", { clientCount });
-    if (wss.clients) {
-      const notification = JSON.stringify({ type: "server_shutting_down", payload: { reconnectIn: 5 } });
-      for (const client of wss.clients) {
-        try {
-          client.send(notification);
-        } catch {
-          // Client may already be disconnected
-        }
-      }
+    if (!wss.clients) return;
+
+    if (typeof saveAllSessions !== "function") {
+      notifyClients(null);
+      return;
     }
+
+    saveAllSessions()
+      .then((blobs) => notifyClients(blobs))
+      .catch((err) => {
+        logger.error("shutdown: session save failed", { error: err.message });
+        notifyClients(null);
+      });
   }, 100);
 
   // Phase 3 — Drain pending sends (500ms–4000ms)
@@ -131,15 +187,20 @@ export function shutdown(wss, signal) {
   });
 }
 
+// Without a session manager there is nothing to persist, so the shared
+// broadcast path stays in use.
+const shutdownOptions = sessionManager ? { saveAllSessions } : {};
+
+// Flip /healthz and /readyz to 503 first so load balancers stop routing here
+// while the drain phases run; closing the WebSocket server also closes the
+// co-located HTTP server.
 process.on("SIGTERM", () => {
   markShuttingDown();
-  wss.close();
-  shutdown(httpServer, "SIGTERM");
+  shutdown(wss, "SIGTERM", shutdownOptions);
 });
 process.on("SIGINT", () => {
   markShuttingDown();
-  wss.close();
-  shutdown(httpServer, "SIGINT");
+  shutdown(wss, "SIGINT", shutdownOptions);
 });
 
 process.on("uncaughtException", (err) => {

src/rate-limiter.js

@@ -83,3 +83,5 @@ export function createRateLimiter(maxPerSecond) {
     get size() {
       return windows.size;
     },
+  };
+}

src/room-manager.js

@@ -1,6 +1,11 @@
 import { WebSocket } from "ws";
 import { v7 as uuidv7 } from "uuid";
 
+const DEFAULT_RING_BUFFER_SIZE = 100;
+const DEFAULT_MAX_BUFFER_BYTES = 1024 * 1024;
+const DEFAULT_DEDUP_WINDOW_MS = 5000;
+const DEFAULT_MAX_DEDUP_ENTRIES = 10_000;
+
 /**
  * @typedef {Object} BackpressureOptions
  * @property {boolean} [enabled=false] - Enable backpressure-aware broadcasting
@@ -195,6 +200,19 @@ export class RoomManager {
     return false;
   }
 
+  /**
+   * Subscribes a client to a room, enforcing the configured DoS limits.
+   *
+   * Re-joining a room the client already occupies only replaces the stored
+   * socket, so it is never rejected by a limit.
+   *
+   * @param {string} clientId - Unique identifier for the client.
+   * @param {string} roomId - Identifier of the room to join.
+   * @param {import("ws").WebSocket} ws - Socket to register for broadcasts.
+   * @returns {undefined | { ok: boolean, reason?: string } | { type: "error", payload: { code: string, message: string } }}
+   *   `{ ok }` when `maxRoomSize` is configured, an error frame when a limit or the
+   *   circuit breaker rejects the join, otherwise `undefined`.
+   */
   join(clientId, roomId, ws) {
     if (clientId == null) throw new TypeError("clientId is required");
     if (roomId == null) throw new TypeError("roomId is required");