Skip to content

fix(api): globalize WS connection/subscription caps via SharedStore [BUG-005] #214

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Morenikeoa wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(api): globalize WS connection/subscription caps via SharedStore [BUG-005] #214

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions src/middleware/shared-store.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,15 @@ export interface SharedStore {
*/
decrementConnectionCount(key: string): Promise<void>;

/**
* 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<void>;

/**
* Record an auth failure for an IP.
* Returns the updated record so callers can decide whether to ban.
Expand Down Expand Up @@ -189,6 +198,15 @@ export class InMemoryStore implements SharedStore {
}
}

async addConnectionCount(key: string, delta: number): Promise<void> {
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,
Expand Down Expand Up @@ -416,6 +434,29 @@ export class UpstashStore implements SharedStore {
}
}

async addConnectionCount(key: string, delta: number): Promise<void> {
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:<ip> → { count, windowStart, bannedUntil }
Expand Down
4 changes: 2 additions & 2 deletions src/routes/health.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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) });
Expand Down
128 changes: 99 additions & 29 deletions src/routes/ws.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<any> {
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,
Expand Down Expand Up @@ -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 });
Comment on lines +628 to +631

@coderabbitai coderabbitai Bot Jun 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

The fleet-wide connection cap is still check-then-act.

Line 629 reads GLOBAL_CONN_KEY, and Line 707 increments it later. Concurrent upgrades on different replicas can all observe the same pre-cap value and all pass, so MAX_WS_CONNECTIONS can still be exceeded under burst load. This needs an atomic reserve/rollback path in SharedStore, not a read followed by a later increment.

Also applies to: 707-709

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/routes/ws.ts` around lines 628 - 631, The fleet-wide WS connection limit
is still vulnerable to check-then-act because
`getConnectionCount(GLOBAL_CONN_KEY)` is read before the later increment path,
so concurrent upgrades can all pass the cap. Update `src/routes/ws.ts` to use an
atomic reserve/rollback flow through `SharedStore` instead of reading
`GLOBAL_CONN_KEY` and incrementing later in separate steps. The fix should be
centered around the global connection check near `MAX_WS_CONNECTIONS` and the
increment logic around the later connection setup path, ensuring only one
replica can reserve a slot when the cap is reached and that any failed upgrade
releases the reservation.

ws.close(1008, "Connection limit reached");
return;
}
Expand Down Expand Up @@ -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", {
Comment on lines +707 to +709

@coderabbitai coderabbitai Bot Jun 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Register close before awaiting the shared counter write.

Line 707 can block on the shared store after the client has already been added locally and after the per-IP counter was incremented, but the close listener is only attached much later. If the peer disconnects during that await, cleanup never runs, so clients, the per-IP count, and potentially the global count leak until restart.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/routes/ws.ts` around lines 707 - 709, The WebSocket connection setup in
ws.ts should attach the close cleanup handler before awaiting
getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY). Move the socket
close listener registration earlier in the connection flow, before the shared
counter write, so a disconnect during that await still triggers cleanup of
clients, the per-IP count, and the global count. Keep the fix localized around
the connection setup logic near logger.info and the shared store increment.

ip: clientIp,
authenticated,
totalClients: clients.size
Expand Down Expand Up @@ -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;

Comment on lines +895 to +907

@coderabbitai coderabbitai Bot Jun 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

The fleet-wide subscription cap can still overshoot under concurrency.

Both subscribe paths snapshot GLOBAL_SUB_KEY once and apply the delta afterward. Multiple clients/replicas can therefore all start below the cap and all commit, so MAX_GLOBAL_SUBSCRIPTIONS is not actually enforced fleet-wide. This needs an atomic capped-add/reservation in the shared store before mutating client.subscriptions.

Also applies to: 954-984, 1086-1106

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/routes/ws.ts` around lines 895 - 907, The global subscription cap logic
in the websocket subscribe flow still uses a stale snapshot of GLOBAL_SUB_KEY,
so concurrent replicas can oversubscribe. Update the subscribe paths in ws.ts
around the current getSharedStore/globalSubsAtStart handling to use an atomic
reservation or capped increment in the shared store before mutating
client.subscriptions, and apply the same fix to the other subscribe blocks
referenced in the comment so MAX_GLOBAL_SUBSCRIPTIONS is enforced fleet-wide.

for (const channel of msg.channels) {
if (typeof channel !== "string") continue;
const safeChannel = channel.slice(0, 100);
Expand Down Expand Up @@ -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 });

Expand Down Expand Up @@ -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);
}
Expand All @@ -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) {
Expand All @@ -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 });
}
Expand All @@ -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 });
}
Expand Down Expand Up @@ -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;

Expand Down
45 changes: 45 additions & 0 deletions tests/middleware/shared-store.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
// ---------------------------------------------------------------------------
Expand Down Expand Up @@ -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"));
Expand Down
Loading