feat(gateway): add websocket ingress rate limiting
This commit is contained in:
@@ -70,6 +70,14 @@ export interface GatewayServerConfig {
|
||||
getTokenUsage?: () => TokenUsageEntry[];
|
||||
/** Maximum allowed request body size for inbound HTTP POST bodies. */
|
||||
maxRequestBodyBytes?: number;
|
||||
/** Per-connection WebSocket ingress rate limiting. */
|
||||
wsRateLimit?: {
|
||||
enabled?: boolean;
|
||||
capacity?: number;
|
||||
refillPerSec?: number;
|
||||
maxViolations?: number;
|
||||
violationWindowMs?: number;
|
||||
};
|
||||
/** Optional pairing manager for DM pairing code management via gateway. */
|
||||
pairingManager?: PairingManager;
|
||||
memoryStore?: MemoryStore;
|
||||
@@ -80,6 +88,13 @@ export interface GatewayServerConfig {
|
||||
|
||||
export class GatewayServer {
|
||||
private static readonly DEFAULT_MAX_REQUEST_BODY_BYTES = 1_048_576; // 1 MiB
|
||||
private static readonly DEFAULT_WS_RATE_LIMIT = {
|
||||
enabled: true,
|
||||
capacity: 30,
|
||||
refillPerSec: 15,
|
||||
maxViolations: 8,
|
||||
violationWindowMs: 10_000,
|
||||
} as const;
|
||||
private wss: WebSocketServer | null = null;
|
||||
private httpServer: HttpServer | null = null;
|
||||
private router: Router;
|
||||
@@ -87,6 +102,12 @@ export class GatewayServer {
|
||||
private laneQueue: LaneQueue;
|
||||
private metrics: MetricsCollector;
|
||||
private connectionMap: Map<WebSocket, string> = new Map();
|
||||
private connectionRateMap: Map<string, {
|
||||
tokens: number;
|
||||
lastRefillMs: number;
|
||||
violations: number;
|
||||
windowStartMs: number;
|
||||
}> = new Map();
|
||||
private config: GatewayServerConfig;
|
||||
private startTime: number = Date.now();
|
||||
|
||||
@@ -291,8 +312,22 @@ export class GatewayServer {
|
||||
const connectionId = randomUUID();
|
||||
this.sessionBridge.connect(connectionId);
|
||||
this.connectionMap.set(ws, connectionId);
|
||||
this.connectionRateMap.set(connectionId, {
|
||||
tokens: this.getWsRateConfig().capacity,
|
||||
lastRefillMs: Date.now(),
|
||||
violations: 0,
|
||||
windowStartMs: Date.now(),
|
||||
});
|
||||
|
||||
ws.on('message', async (data) => {
|
||||
const limit = this.consumeConnectionRateToken(connectionId);
|
||||
if (!limit.allowed) {
|
||||
this.send(ws, makeError(0, ErrorCode.InternalError, `Rate limit exceeded. Retry in ${limit.retryMs}ms.`));
|
||||
if (limit.close) {
|
||||
ws.close(4008, 'Rate limit exceeded');
|
||||
}
|
||||
return;
|
||||
}
|
||||
const raw = data.toString();
|
||||
await this.handleMessage(ws, connectionId, raw);
|
||||
});
|
||||
@@ -300,6 +335,7 @@ export class GatewayServer {
|
||||
ws.on('close', () => {
|
||||
this.sessionBridge.disconnect(connectionId);
|
||||
this.connectionMap.delete(ws);
|
||||
this.connectionRateMap.delete(connectionId);
|
||||
});
|
||||
|
||||
ws.on('error', (err) => {
|
||||
@@ -307,6 +343,55 @@ export class GatewayServer {
|
||||
});
|
||||
}
|
||||
|
||||
private getWsRateConfig(): Required<NonNullable<GatewayServerConfig['wsRateLimit']>> {
|
||||
const raw = this.config.wsRateLimit ?? {};
|
||||
return {
|
||||
enabled: raw.enabled ?? GatewayServer.DEFAULT_WS_RATE_LIMIT.enabled,
|
||||
capacity: raw.capacity ?? GatewayServer.DEFAULT_WS_RATE_LIMIT.capacity,
|
||||
refillPerSec: raw.refillPerSec ?? GatewayServer.DEFAULT_WS_RATE_LIMIT.refillPerSec,
|
||||
maxViolations: raw.maxViolations ?? GatewayServer.DEFAULT_WS_RATE_LIMIT.maxViolations,
|
||||
violationWindowMs: raw.violationWindowMs ?? GatewayServer.DEFAULT_WS_RATE_LIMIT.violationWindowMs,
|
||||
};
|
||||
}
|
||||
|
||||
private consumeConnectionRateToken(connectionId: string): { allowed: boolean; close: boolean; retryMs: number } {
|
||||
const cfg = this.getWsRateConfig();
|
||||
if (!cfg.enabled) {
|
||||
return { allowed: true, close: false, retryMs: 0 };
|
||||
}
|
||||
|
||||
const now = Date.now();
|
||||
const state = this.connectionRateMap.get(connectionId);
|
||||
if (!state) {
|
||||
return { allowed: true, close: false, retryMs: 0 };
|
||||
}
|
||||
|
||||
const elapsedMs = Math.max(0, now - state.lastRefillMs);
|
||||
state.tokens = Math.min(cfg.capacity, state.tokens + (elapsedMs / 1000) * cfg.refillPerSec);
|
||||
state.lastRefillMs = now;
|
||||
|
||||
if (state.tokens >= 1) {
|
||||
state.tokens -= 1;
|
||||
this.connectionRateMap.set(connectionId, state);
|
||||
return { allowed: true, close: false, retryMs: 0 };
|
||||
}
|
||||
|
||||
if (now - state.windowStartMs > cfg.violationWindowMs) {
|
||||
state.windowStartMs = now;
|
||||
state.violations = 0;
|
||||
}
|
||||
state.violations += 1;
|
||||
this.connectionRateMap.set(connectionId, state);
|
||||
|
||||
const deficit = 1 - state.tokens;
|
||||
const retryMs = Math.max(1, Math.ceil((deficit / cfg.refillPerSec) * 1000));
|
||||
return {
|
||||
allowed: false,
|
||||
close: state.violations >= cfg.maxViolations,
|
||||
retryMs,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle incoming HTTP requests.
|
||||
* Optionally applies auth (when authHttp is enabled and a token is configured).
|
||||
|
||||
Reference in New Issue
Block a user