diff --git a/src/middleware/shared-store.ts b/src/middleware/shared-store.ts index d8763aa..bb8ccec 100644 --- a/src/middleware/shared-store.ts +++ b/src/middleware/shared-store.ts @@ -85,6 +85,15 @@ export interface SharedStore { */ decrementConnectionCount(key: string): Promise; + /** + * Atomically add `delta` (positive or negative) to the counter for `key`. + * For batch updates where calling increment/decrement N times would cost + * N round-trips under Upstash (e.g. N subscriptions acquired/released in + * one WebSocket message). Floors at 0 and deletes the key, same as + * decrementConnectionCount. + */ + addConnectionCount(key: string, delta: number): Promise; + /** * Record an auth failure for an IP. * Returns the updated record so callers can decide whether to ban. @@ -189,6 +198,15 @@ export class InMemoryStore implements SharedStore { } } + async addConnectionCount(key: string, delta: number): Promise { + const next = (this.connCounts.get(key) ?? 0) + delta; + if (next <= 0) { + this.connCounts.delete(key); + } else { + this.connCounts.set(key, next); + } + } + async recordAuthFailure( ip: string, windowMs: number, @@ -416,6 +434,29 @@ export class UpstashStore implements SharedStore { } } + async addConnectionCount(key: string, delta: number): Promise { + try { + // Same 24h safety-net TTL as increment/decrementConnectionCount. + const script = ` + local val = redis.call("INCRBY", KEYS[1], ARGV[1]) + if val <= 0 then + redis.call("DEL", KEYS[1]) + else + redis.call("EXPIRE", KEYS[1], 86400) + end + return val + `; + await this.eval(script, [this.connKey(key)], [delta]); + } catch (err) { + logger.warn("UpstashStore.addConnectionCount failed, using fallback", { + key, + delta, + error: err instanceof Error ? err.message : String(err), + }); + await this.fallback.addConnectionCount(key, delta); + } + } + /** * Auth-failure record stored as a Redis HASH: * pcl:af: → { count, windowStart, bannedUntil } diff --git a/src/routes/health.ts b/src/routes/health.ts index 92050f3..be282d3 100644 --- a/src/routes/health.ts +++ b/src/routes/health.ts @@ -56,7 +56,7 @@ export function healthRoutes(): Hono { // Check WebSocket subsystem — saturated WS means new clients can't connect try { - const wsMetrics = getWebSocketMetrics(); + const wsMetrics = await getWebSocketMetrics(); const utilization = wsMetrics.totalConnections / wsMetrics.limits.maxGlobalConnections; checks.ws = utilization < 0.95; // degraded if >95% of connection slots used } catch { @@ -83,7 +83,7 @@ export function healthRoutes(): Hono { app.get("/ws/stats", requireApiKey(), async (c) => { try { - const metrics = getWebSocketMetrics(); + const metrics = await getWebSocketMetrics(); return c.json(metrics); } catch (err) { logger.error("Failed to get WebSocket metrics", { error: truncateErrorMessage(err instanceof Error ? err.message : err, 120) }); diff --git a/src/routes/ws.ts b/src/routes/ws.ts index b2ee32e..37c070d 100644 --- a/src/routes/ws.ts +++ b/src/routes/ws.ts @@ -110,8 +110,13 @@ interface WsClient { msgWindowStart: number; // Start of current rate-limit window } -// Track global subscription count across all clients -let globalSubscriptionCount = 0; +// Global connection/subscription caps are enforced via the SharedStore (same +// backend as the per-IP counters below) so the limits hold across replicas +// under horizontal scaling. Without this, MAX_WS_CONNECTIONS/ +// MAX_GLOBAL_SUBSCRIPTIONS are each enforced independently per replica — the +// true fleet-wide cap becomes N×limit instead of limit, for N replicas. +const GLOBAL_CONN_KEY = "ws:global-connections"; +const GLOBAL_SUB_KEY = "ws:global-subscriptions"; // Auth failure rate limiting per IP (issue #839: connection flood from repeat auth failures) // Tracks recent auth failures to temporarily ban repeat offenders. @@ -450,19 +455,34 @@ function flushPriceUpdate(slabAddress: string): void { } /** - * Get WebSocket metrics for /ws/stats endpoint + * Get WebSocket metrics for /ws/stats endpoint. + * + * totalConnections and totalSubscriptions are read from the SharedStore and + * reflect the true fleet-wide total across all replicas (see GLOBAL_CONN_KEY/ + * GLOBAL_SUB_KEY). connectionsPerSlab, messagesPerSec, and bytesPerSec remain + * local to this replica — they're inherently per-process observability data + * (this replica's own observed traffic/distribution), not safety caps, so + * aggregating them across replicas is a different and much larger problem + * than fixing the cap-enforcement bug this fixes; left out of scope. */ -export function getWebSocketMetrics(): any { +export async function getWebSocketMetrics(): Promise { const now = Date.now(); const elapsedSec = (now - metrics.lastResetTime) / 1000 || 1; + const store = getSharedStore(); + const [totalConnections, totalSubscriptions] = await Promise.all([ + store.getConnectionCount(GLOBAL_CONN_KEY), + store.getConnectionCount(GLOBAL_SUB_KEY), + ]); return { - totalConnections: metrics.totalConnections, + totalConnections, + totalSubscriptions, connectionsPerSlab: Object.fromEntries(metrics.connectionsPerSlab), messagesPerSec: parseFloat((metrics.messagesReceived / elapsedSec).toFixed(2)), bytesPerSec: parseInt((metrics.bytesSent / elapsedSec).toFixed(0), 10), limits: { maxGlobalConnections: MAX_WS_CONNECTIONS, + maxGlobalSubscriptions: MAX_GLOBAL_SUBSCRIPTIONS, maxConnectionsPerSlab: MAX_CONNECTIONS_PER_SLAB, maxConnectionsPerIp: MAX_CONNECTIONS_PER_IP, maxUnauthConnectionsPerIp: MAX_UNAUTHENTICATED_CONNECTIONS_PER_IP, @@ -605,9 +625,10 @@ export function setupWebSocket(server: Server): WebSocketServer { return; } - // H2: Reject if at max connections - if (clients.size >= MAX_WS_CONNECTIONS) { - logger.warn("Max global WS connections reached", { ip: clientIp }); + // H2: Reject if at max connections (global across all replicas, via SharedStore) + const globalConnCount = await getSharedStore().getConnectionCount(GLOBAL_CONN_KEY); + if (globalConnCount >= MAX_WS_CONNECTIONS) { + logger.warn("Max global WS connections reached", { ip: clientIp, count: globalConnCount }); ws.close(1008, "Connection limit reached"); return; } @@ -683,8 +704,9 @@ export function setupWebSocket(server: Server): WebSocketServer { }; clients.add(client); metrics.totalConnections = clients.size; - - logger.info("WebSocket connection established", { + await getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY); + + logger.info("WebSocket connection established", { ip: clientIp, authenticated, totalClients: clients.size @@ -870,7 +892,19 @@ export function setupWebSocket(server: Server): WebSocketServer { const subscribed: string[] = []; const errors: string[] = []; - + + // Check the global subscription cap once per message (not once + // per channel) and track local accepted-count as a delta, then + // flush a single batched update after the loop. This avoids up to + // MAX_CHANNELS_PER_MESSAGE (50) shared-store round-trips per + // message. The existing per-IP caps in this file are already + // non-atomic check-then-act, so the small intra-message overshoot + // window this introduces (bounded by MAX_CHANNELS_PER_MESSAGE) is + // not a new class of weaker guarantee. + const subStore = getSharedStore(); + const globalSubsAtStart = await subStore.getConnectionCount(GLOBAL_SUB_KEY); + let newSubsThisMessage = 0; + for (const channel of msg.channels) { if (typeof channel !== "string") continue; const safeChannel = channel.slice(0, 100); @@ -917,31 +951,37 @@ export function setupWebSocket(server: Server): WebSocketServer { continue; } - // Cap global subscriptions to prevent DoS - if (globalSubscriptionCount >= MAX_GLOBAL_SUBSCRIPTIONS) { + // Cap global subscriptions to prevent DoS (checked against the + // count as of this message's start plus any accepted so far in + // this same loop — see comment above the loop). + if (globalSubsAtStart + newSubsThisMessage >= MAX_GLOBAL_SUBSCRIPTIONS) { errors.push("Server subscription limit reached"); break; } - + // Cap subscriptions per client if (client.subscriptions.size >= MAX_SUBSCRIPTIONS_PER_CLIENT) { errors.push("Subscription limit per connection reached"); break; } - + // Check per-slab connection limit const slabClients = connectionsPerSlab.get(sanitized); if (slabClients && slabClients.size >= MAX_CONNECTIONS_PER_SLAB) { errors.push("Connection limit for this market reached"); continue; } - + client.subscriptions.add(fullChannel); - globalSubscriptionCount++; + newSubsThisMessage++; addClientToSlab(client, sanitized); subscribed.push(fullChannel); } - + + if (newSubsThisMessage > 0) { + await subStore.addConnectionCount(GLOBAL_SUB_KEY, newSubsThisMessage); + } + if (subscribed.length > 0) { safeSend(ws, { type: "subscribed", channels: subscribed }); @@ -1043,17 +1083,27 @@ export function setupWebSocket(server: Server): WebSocketServer { message: "Please use channels array. Subscribing to all channels for this slab.", }); + // Same batched cap-check pattern as the modern subscribe path + // above (at most 3 channels here, but kept consistent). + const legacySubStore = getSharedStore(); + const legacyGlobalSubsAtStart = await legacySubStore.getConnectionCount(GLOBAL_SUB_KEY); + let legacyNewSubs = 0; + const subscribed: string[] = []; for (const channel of channels) { if (client.subscriptions.has(channel)) continue; - if (globalSubscriptionCount >= MAX_GLOBAL_SUBSCRIPTIONS) break; + if (legacyGlobalSubsAtStart + legacyNewSubs >= MAX_GLOBAL_SUBSCRIPTIONS) break; if (client.subscriptions.size >= MAX_SUBSCRIPTIONS_PER_CLIENT) break; client.subscriptions.add(channel); - globalSubscriptionCount++; + legacyNewSubs++; subscribed.push(channel); } + if (legacyNewSubs > 0) { + await legacySubStore.addConnectionCount(GLOBAL_SUB_KEY, legacyNewSubs); + } + if (subscribed.length > 0) { addClientToSlab(client, sanitized); } @@ -1070,12 +1120,13 @@ export function setupWebSocket(server: Server): WebSocketServer { } const unsubscribed: string[] = []; - + let removedSubs = 0; + for (const channel of msg.channels) { if (client.subscriptions.delete(channel)) { - globalSubscriptionCount--; + removedSubs++; unsubscribed.push(channel); - + // Extract slab and remove from slab tracking if no more subs for this slab const slab = extractSlabFromChannel(channel); if (slab) { @@ -1088,7 +1139,11 @@ export function setupWebSocket(server: Server): WebSocketServer { } } } - + + if (removedSubs > 0) { + await getSharedStore().addConnectionCount(GLOBAL_SUB_KEY, -removedSubs); + } + if (unsubscribed.length > 0) { safeSend(ws, { type: "unsubscribed", channels: unsubscribed }); } @@ -1099,14 +1154,19 @@ export function setupWebSocket(server: Server): WebSocketServer { if (sanitized) { const channels = [`price:${sanitized}`, `trades:${sanitized}`, `funding:${sanitized}`]; const unsubscribed: string[] = []; - + let legacyRemovedSubs = 0; + for (const channel of channels) { if (client.subscriptions.delete(channel)) { - globalSubscriptionCount--; + legacyRemovedSubs++; unsubscribed.push(channel); } } - + + if (legacyRemovedSubs > 0) { + await getSharedStore().addConnectionCount(GLOBAL_SUB_KEY, -legacyRemovedSubs); + } + removeClientFromSlab(client, sanitized); safeSend(ws, { type: "unsubscribed", slabAddress: sanitized, channels: unsubscribed }); } @@ -1153,8 +1213,18 @@ export function setupWebSocket(server: Server): WebSocketServer { } // H2: O(1) removal with Set - // Decrement global subscription count for all client subscriptions - globalSubscriptionCount = Math.max(0, globalSubscriptionCount - client.subscriptions.size); + // Decrement global subscription + connection counts via the shared + // store (fire-and-forget — this handler is synchronous, matching the + // per-IP decrement pattern above). + const subsToRelease = client.subscriptions.size; + if (subsToRelease > 0) { + closeStore.addConnectionCount(GLOBAL_SUB_KEY, -subsToRelease).catch((err) => + logger.warn("addConnectionCount error on close (subscriptions)", { ip: client.ip, error: String(err) }) + ); + } + closeStore.decrementConnectionCount(GLOBAL_CONN_KEY).catch((err) => + logger.warn("decrementConnectionCount error on close (global connections)", { ip: client.ip, error: String(err) }) + ); clients.delete(client); metrics.totalConnections = clients.size; diff --git a/tests/middleware/shared-store.test.ts b/tests/middleware/shared-store.test.ts index 9deec45..ef1ce29 100644 --- a/tests/middleware/shared-store.test.ts +++ b/tests/middleware/shared-store.test.ts @@ -109,6 +109,41 @@ describe("InMemoryStore — connection counts", () => { }); }); +// --------------------------------------------------------------------------- +// InMemoryStore — addConnectionCount (batched delta, BUG-005) +// --------------------------------------------------------------------------- + +describe("InMemoryStore — addConnectionCount", () => { + it("adds a positive delta in one call", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 5); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(5); + }); + + it("adds a negative delta to release a batch", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 5); + await store.addConnectionCount("ws:global-subscriptions", -3); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(2); + }); + + it("floors at 0 and removes the key rather than going negative", async () => { + const store = new InMemoryStore(); + await store.addConnectionCount("ws:global-subscriptions", 2); + await store.addConnectionCount("ws:global-subscriptions", -10); + expect(await store.getConnectionCount("ws:global-subscriptions")).toBe(0); + }); + + it("is interoperable with incrementConnectionCount/decrementConnectionCount on the same key", async () => { + const store = new InMemoryStore(); + await store.incrementConnectionCount("ws:global-connections"); + await store.addConnectionCount("ws:global-connections", 4); + expect(await store.getConnectionCount("ws:global-connections")).toBe(5); + await store.decrementConnectionCount("ws:global-connections"); + expect(await store.getConnectionCount("ws:global-connections")).toBe(4); + }); +}); + // --------------------------------------------------------------------------- // InMemoryStore — auth failure ban tests // --------------------------------------------------------------------------- @@ -233,6 +268,16 @@ describe("UpstashStore — graceful fallback", () => { globalThis.fetch = originalFetch; }); + it("falls back gracefully for addConnectionCount", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network unreachable")); + + const store = new UpstashStore("https://test.upstash.io", "token"); + await expect(store.addConnectionCount("ws:global-subscriptions", 5)).resolves.toBeUndefined(); + + globalThis.fetch = originalFetch; + }); + it("falls back gracefully for auth ban checks", async () => { const originalFetch = globalThis.fetch; globalThis.fetch = vi.fn().mockRejectedValue(new Error("Network unreachable")); diff --git a/tests/routes/ws-global-caps.test.ts b/tests/routes/ws-global-caps.test.ts new file mode 100644 index 0000000..c4e0b21 --- /dev/null +++ b/tests/routes/ws-global-caps.test.ts @@ -0,0 +1,216 @@ +/** + * Regression for BUG-005: MAX_WS_CONNECTIONS and MAX_GLOBAL_SUBSCRIPTIONS + * must be enforced via the SharedStore so they hold fleet-wide across + * replicas, not just against this process's own (possibly empty) local + * `clients` Set / subscription counter. + * + * These tests simulate "other replicas already at the cap" by bumping the + * shared store directly, then prove THIS replica — whose own local state is + * empty — still rejects new connections/subscriptions, which would not have + * been true before this fix. + */ +import { describe, it, expect, afterEach, vi } from "vitest"; +import http from "node:http"; +import WebSocket from "ws"; + +vi.mock("@percolator/shared", () => ({ + createLogger: vi.fn(() => ({ + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + })), + eventBus: { on: vi.fn() }, + getSupabase: vi.fn(() => ({ + from: vi.fn(() => ({ + select: vi.fn(() => ({ + eq: vi.fn(() => ({ + single: vi.fn(() => Promise.resolve({ data: null, error: null })), + })), + })), + })), + })), + sanitizeSlabAddress: vi.fn((s: string) => s), + sendInfoAlert: vi.fn(), +})); + +function waitForOpen(ws: WebSocket): Promise { + return new Promise((resolve, reject) => { + if (ws.readyState === WebSocket.OPEN) return resolve(); + ws.once("open", resolve); + ws.once("error", reject); + }); +} + +function waitForClose(ws: WebSocket): Promise { + return new Promise((resolve) => { + if (ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) { + return resolve(ws.readyState); + } + ws.once("close", (code) => resolve(code)); + }); +} + +function waitForMessage(ws: WebSocket, predicate: (msg: any) => boolean, timeoutMs = 1500): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => reject(new Error("timeout waiting for message")), timeoutMs); + const onMsg = (raw: unknown) => { + const parsed = JSON.parse(String(raw)); + if (predicate(parsed)) { + clearTimeout(timer); + ws.off("message", onMsg); + resolve(parsed); + } + }; + ws.on("message", onMsg); + }); +} + +interface TestServer { + server: http.Server; + port: number; + sharedStoreModule: typeof import("../../src/middleware/shared-store.js"); +} + +async function startServer(env: Record): Promise { + Object.assign(process.env, env); + vi.resetModules(); + const { setupWebSocket } = await import("../../src/routes/ws.js"); + // Import shared-store AFTER ws.js has resolved it, so this test gets the + // same module instance (and same _store singleton) ws.ts uses internally — + // required to seed/inspect the global counters it reads from. + const sharedStoreModule = await import("../../src/middleware/shared-store.js"); + const server = http.createServer((_req, res) => { + res.writeHead(200); + res.end("ok"); + }); + setupWebSocket(server as unknown as import("node:http").Server); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const { port } = server.address() as { port: number }; + return { server, port, sharedStoreModule }; +} + +async function stopServer(ts: TestServer): Promise { + await new Promise((resolve, reject) => + ts.server.close((err) => (err ? reject(err) : resolve())), + ); +} + +describe("WS global connection cap is shared across replicas (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.MAX_WS_CONNECTIONS; + delete process.env.WS_AUTH_REQUIRED; + }); + + it("rejects a new connection once the global count is at the cap, even though this replica's own local client set is empty", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + MAX_WS_CONNECTIONS: "2", + }); + + // Simulate 2 connections already established on OTHER replicas by + // bumping the shared store directly — this replica's own `clients` Set + // is still empty, so the pre-fix (local-only) check would have allowed + // a new connection here. + const store = ts.sharedStoreModule.getSharedStore(); + await store.incrementConnectionCount("ws:global-connections"); + await store.incrementConnectionCount("ws:global-connections"); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + const code = await waitForClose(ws); + expect(code).toBe(1008); + }); + + it("allows a connection when under the global cap", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + MAX_WS_CONNECTIONS: "2", + }); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + expect(ws.readyState).toBe(WebSocket.OPEN); + ws.close(); + await waitForClose(ws); + }); +}); + +describe("WS global subscription cap is shared across replicas (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.WS_AUTH_REQUIRED; + }); + + it("rejects a new subscription once the global subscription count is at the cap (MAX_GLOBAL_SUBSCRIPTIONS = 1000, not env-overridable)", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + // Simulate 1000 subscriptions already held across OTHER replicas. + const store = ts.sharedStoreModule.getSharedStore(); + await store.addConnectionCount("ws:global-subscriptions", 1000); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + + ws.send(JSON.stringify({ type: "subscribe", channels: ["price:SLAB-GLOBAL-CAP"] })); + const errorMsg = await waitForMessage(ws, (m) => m.type === "error"); + expect(errorMsg.message).toMatch(/server subscription limit reached/i); + + ws.close(); + await waitForClose(ws); + }); + + it("allows a subscription when under the global cap", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + const ws = new WebSocket(`ws://127.0.0.1:${ts.port}/`); + await waitForOpen(ws); + + ws.send(JSON.stringify({ type: "subscribe", channels: ["price:SLAB-UNDER-CAP"] })); + const subscribedMsg = await waitForMessage(ws, (m) => m.type === "subscribed"); + expect(subscribedMsg.channels).toContain("price:SLAB-UNDER-CAP"); + + ws.close(); + await waitForClose(ws); + }); +}); + +describe("getWebSocketMetrics reports the global (shared-store) count, not just this replica's local count (BUG-005)", () => { + let ts: TestServer; + + afterEach(async () => { + if (ts) await stopServer(ts); + delete process.env.WS_AUTH_REQUIRED; + }); + + it("totalConnections reflects connections established on other replicas, not seen locally", async () => { + ts = await startServer({ + NODE_ENV: "test", + WS_AUTH_REQUIRED: "false", + }); + + const { getWebSocketMetrics } = await import("../../src/routes/ws.js"); + const store = ts.sharedStoreModule.getSharedStore(); + + // Simulate 7 connections established on other replicas — this replica's + // own `clients` Set is still empty. + for (let i = 0; i < 7; i++) { + await store.incrementConnectionCount("ws:global-connections"); + } + + const metrics = await getWebSocketMetrics(); + expect(metrics.totalConnections).toBe(7); + }); +});