diff --git a/src/__tests__/keyRevocationConnections.test.ts b/src/__tests__/keyRevocationConnections.test.ts new file mode 100644 index 00000000..1b793578 --- /dev/null +++ b/src/__tests__/keyRevocationConnections.test.ts @@ -0,0 +1,412 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' +import Fastify, { type FastifyInstance } from 'fastify' +import WebSocket from 'ws' +import type { AddressInfo } from 'net' +import { EventEmitter } from 'events' + +// Stub database and cache for pure test isolation +const mockFindUnique = vi.fn() +const mockUpdate = vi.fn() +const mockCreate = vi.fn() + +vi.mock('../db', () => ({ + prisma: { + apiKey: { + findUnique: (...args: any[]) => mockFindUnique(...args), + update: (...args: any[]) => mockUpdate(...args), + create: (...args: any[]) => mockCreate(...args), + }, + }, + pgPool: { query: vi.fn() }, +})) + +const mockRedisPublish = vi.fn().mockResolvedValue(1) +vi.mock('../redis', () => ({ + redis: { + publish: (...args: any[]) => mockRedisPublish(...args), + on: vi.fn(), + get: vi.fn(), + set: vi.fn(), + }, + getCachedPrice: vi.fn(), +})) + +vi.mock('../aggregator/vwap', () => ({ getAggregatedPrice: vi.fn() })) +vi.mock('../aggregator/bestRoute', () => ({ getBestRoute: vi.fn() })) +vi.mock('../config', () => ({ + config: { + pairs: [], + redis: { url: 'redis://localhost:6379' }, + }, + activeNetwork: 'testnet', +})) + +import { registerGraphQL } from '../api/graphql' +import { registerAdminRoutes } from '../api/admin' +import { registerWebSocket } from '../api/websocket' +import { hashApiKey } from '../api/auth' +import { publishPriceUpdate, priceEmitter } from '../events' +import { + connectionRegistry, + ConnectionRegistry, + publishRevocation, + handleRevocationMessage, + CLOSE_CODE_KEY_REVOKED, + CLOSE_REASON_KEY_REVOKED, + REVOCATION_REDIS_CHANNEL, +} from '../ws/connectionRegistry' + +const SUBPROTOCOL = 'graphql-transport-ws' + +async function buildServer(): Promise<{ app: FastifyInstance; wsUrl: string; httpUrl: string }> { + const app = Fastify({ logger: false }) + await registerAdminRoutes(app) + await registerGraphQL(app) + await registerWebSocket(app) + await app.listen({ port: 0, host: '127.0.0.1' }) + const { port } = app.server.address() as AddressInfo + return { + app, + wsUrl: `ws://127.0.0.1:${port}/graphql`, + httpUrl: `http://127.0.0.1:${port}`, + } +} + +function connectGraphQL(url: string, token?: string): Promise { + return new Promise((resolve, reject) => { + const ws = new WebSocket(url, SUBPROTOCOL, token ? { headers: { authorization: `Bearer ${token}` } } : undefined) + ws.on('error', reject) + ws.on('open', () => { + const initPayload = token ? { type: 'connection_init', payload: { authorization: `Bearer ${token}` } } : { type: 'connection_init' } + ws.send(JSON.stringify(initPayload)) + }) + ws.on('message', function onAck(raw) { + const msg = JSON.parse(raw.toString()) + if (msg.type === 'connection_ack') { + ws.off('message', onAck) + resolve(ws) + } + }) + }) +} + +function waitFor(ws: WebSocket, type: string, timeoutMs = 2000): Promise { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + ws.off('message', onMsg) + reject(new Error(`timed out waiting for "${type}"`)) + }, timeoutMs) + function onMsg(raw: WebSocket.RawData) { + const msg = JSON.parse(raw.toString()) + if (msg.type === type) { + clearTimeout(timer) + ws.off('message', onMsg) + resolve(msg) + } + } + ws.on('message', onMsg) + }) +} + +const SUBSCRIBE = (id: string, pair: string) => ({ + id, + type: 'subscribe', + payload: { + query: `subscription($pair: String!) { + priceUpdated(pair: $pair) { pair price ts network } + }`, + variables: { pair }, + }, +}) + +describe('Issue #503: API Key Revocation Connection Lifecycle', () => { + let app: FastifyInstance + let wsUrl: string + let httpUrl: string + + const adminToken = 'admin-secret-test-token-12345' + const keyPlaintextA = 'lens_0123456789abcdef0123456789abcdef0123456789abcdef' + const keyHashA = hashApiKey(keyPlaintextA) + const keyIdA = 'key-uuid-aaaa-1111' + + const keyPlaintextB = 'lens_fedcba9876543210fedcba9876543210fedcba9876543210' + const keyHashB = hashApiKey(keyPlaintextB) + const keyIdB = 'key-uuid-bbbb-2222' + + beforeEach(async () => { + process.env.ADMIN_TOKEN = adminToken + vi.clearAllMocks() + connectionRegistry.clear() + + // Default mock finding active, unrevoked keys + mockFindUnique.mockImplementation(({ where }: { where: { id?: string; hash?: string } }) => { + if (where.id === keyIdA || where.hash === keyHashA) { + return Promise.resolve({ + id: keyIdA, + hash: keyHashA, + label: 'key-a', + ratePerMin: 100, + ratePerDay: 10000, + revokedAt: null, + expiresAt: null, + }) + } + if (where.id === keyIdB || where.hash === keyHashB) { + return Promise.resolve({ + id: keyIdB, + hash: keyHashB, + label: 'key-b', + ratePerMin: 100, + ratePerDay: 10000, + revokedAt: null, + expiresAt: null, + }) + } + return Promise.resolve(null) + }) + + mockUpdate.mockImplementation(({ where, data }: any) => { + return Promise.resolve({ + id: where.id, + ...data, + }) + }) + + ;({ app, wsUrl, httpUrl } = await buildServer()) + }) + + afterEach(async () => { + connectionRegistry.clear() + await app.close() + delete process.env.ADMIN_TOKEN + }) + + describe('1. Connection Registry unit behaviour', () => { + it('tracks live connections by hashed API key and allows lookup', () => { + const mockConn1 = { close: vi.fn() } + const mockConn2 = { close: vi.fn() } + + const unreg1 = connectionRegistry.register(keyHashA, mockConn1) + const unreg2 = connectionRegistry.register(keyHashA, mockConn2) + + const connections = connectionRegistry.getConnections(keyHashA) + expect(connections).toHaveLength(2) + expect(connections).toContain(mockConn1) + expect(connections).toContain(mockConn2) + + unreg1() + expect(connectionRegistry.getConnections(keyHashA)).toEqual([mockConn2]) + + unreg2() + expect(connectionRegistry.getConnections(keyHashA)).toEqual([]) + }) + + it('deregisters a connection automatically when it closes naturally', () => { + const emitter = new EventEmitter() as any + emitter.close = vi.fn() + + connectionRegistry.register(keyHashA, emitter) + expect(connectionRegistry.getConnections(keyHashA)).toHaveLength(1) + + // Natural socket close emits 'close' + emitter.emit('close') + + expect(connectionRegistry.getConnections(keyHashA)).toHaveLength(0) + }) + + it('closes connections belonging to the specified key hash and leaves others unaffected', () => { + const connA1 = { close: vi.fn() } + const connA2 = { close: vi.fn() } + const connB1 = { close: vi.fn() } + + connectionRegistry.register(keyHashA, connA1) + connectionRegistry.register(keyHashA, connA2) + connectionRegistry.register(keyHashB, connB1) + + expect(connectionRegistry.totalConnections).toBe(3) + + const closed = connectionRegistry.closeConnectionsForKey(keyHashA, CLOSE_CODE_KEY_REVOKED, CLOSE_REASON_KEY_REVOKED) + expect(closed).toBe(2) + + expect(connA1.close).toHaveBeenCalledWith(4401, 'key_revoked') + expect(connA2.close).toHaveBeenCalledWith(4401, 'key_revoked') + expect(connB1.close).not.toHaveBeenCalled() + + expect(connectionRegistry.getConnections(keyHashA)).toHaveLength(0) + expect(connectionRegistry.getConnections(keyHashB)).toHaveLength(1) + }) + }) + + describe('2. Revocation drops live subscriptions with distinct code and reason', () => { + it('proves a GraphQL subscription stops and closes with 4401 "key_revoked" on key revocation', async () => { + const ws = await connectGraphQL(wsUrl, keyPlaintextA) + ws.send(JSON.stringify(SUBSCRIBE('sub-1', 'XLM/USDC'))) + + await new Promise(r => setTimeout(r, 100)) + + // Verify connection is actively streaming + publishPriceUpdate({ pair: 'XLM/USDC', price: 0.1234, ts: new Date().toISOString(), network: 'testnet' }) + const first = await waitFor(ws, 'next') + expect(first.payload.data.priceUpdated.price).toBe(0.1234) + + // Set up close listener + const closePromise = new Promise<{ code: number; reason: string }>((resolve) => { + ws.on('close', (code, reason) => { + resolve({ code, reason: reason.toString() }) + }) + }) + + // Revoke key A via admin API + const res = await app.inject({ + method: 'DELETE', + url: `/admin/keys/${keyIdA}`, + headers: { 'x-admin-token': adminToken }, + }) + expect(res.statusCode).toBe(200) + expect(res.json()).toMatchObject({ id: keyIdA, revoked: true }) + + // Assert connection closes immediately with 4401 and "key_revoked" + const { code, reason } = await closePromise + expect(code).toBe(4401) + expect(reason).toBe('key_revoked') + expect(ws.readyState).toBe(WebSocket.CLOSED) + + // Verify published event to Redis channel + expect(mockRedisPublish).toHaveBeenCalledWith( + REVOCATION_REDIS_CHANNEL, + JSON.stringify({ keyHash: keyHashA }) + ) + }) + + it('proves a WebSocket price stream (/ws) closes with 4401 "key_revoked" when revoked', async () => { + const port = (app.server.address() as AddressInfo).port + const ws = new WebSocket(`ws://127.0.0.1:${port}/ws?apiKey=${keyPlaintextA}`) + + await new Promise((resolve, reject) => { + ws.on('open', () => resolve()) + ws.on('error', reject) + }) + + const closePromise = new Promise<{ code: number; reason: string }>((resolve) => { + ws.on('close', (code, reason) => { + resolve({ code, reason: reason.toString() }) + }) + }) + + // Admin revokes key A + await app.inject({ + method: 'DELETE', + url: `/admin/keys/${keyIdA}`, + headers: { 'x-admin-token': adminToken }, + }) + + const { code, reason } = await closePromise + expect(code).toBe(4401) + expect(reason).toBe('key_revoked') + expect(ws.readyState).toBe(WebSocket.CLOSED) + }) + }) + + describe('3. Multi-key isolation: other keys are unaffected', () => { + it('asserts other keys’ connections are unaffected when one key is revoked', async () => { + const wsA = await connectGraphQL(wsUrl, keyPlaintextA) + const wsB = await connectGraphQL(wsUrl, keyPlaintextB) + + wsA.send(JSON.stringify(SUBSCRIBE('sub-a', 'XLM/USDC'))) + wsB.send(JSON.stringify(SUBSCRIBE('sub-b', 'XLM/USDC'))) + + await new Promise(r => setTimeout(r, 100)) + + let aClosed = false + let aCloseCode: number | null = null + let aCloseReason: string | null = null + + let bClosed = false + + wsA.on('close', (code, reason) => { + aClosed = true + aCloseCode = code + aCloseReason = reason.toString() + }) + + wsB.on('close', () => { + bClosed = true + }) + + // Revoke Key A only + const res = await app.inject({ + method: 'DELETE', + url: `/admin/keys/${keyIdA}`, + headers: { 'x-admin-token': adminToken }, + }) + expect(res.statusCode).toBe(200) + + await new Promise(r => setTimeout(r, 150)) + + // Assert Key A connection closed with distinct code and reason + expect(aClosed).toBe(true) + expect(aCloseCode).toBe(4401) + expect(aCloseReason).toBe('key_revoked') + + // Assert Key B connection remains open and functional + expect(bClosed).toBe(false) + expect(wsB.readyState).toBe(WebSocket.OPEN) + + // Publish update and verify Key B still receives it + publishPriceUpdate({ pair: 'XLM/USDC', price: 0.9999, ts: new Date().toISOString(), network: 'testnet' }) + const nextB = await waitFor(wsB, 'next') + expect(nextB.payload.data.priceUpdated.price).toBe(0.9999) + + wsB.close() + }) + }) + + describe('4. Cross-instance propagation via Redis pub/sub', () => { + it('asserts receiving a revocation event from Redis drops connections on the subscriber instance', async () => { + // Connect client for Key A + const ws = await connectGraphQL(wsUrl, keyPlaintextA) + ws.send(JSON.stringify(SUBSCRIBE('sub-cross', 'XLM/USDC'))) + await new Promise(r => setTimeout(r, 100)) + + expect(connectionRegistry.getConnections(keyHashA)).toHaveLength(1) + + const closePromise = new Promise<{ code: number; reason: string }>((resolve) => { + ws.on('close', (code, reason) => { + resolve({ code, reason: reason.toString() }) + }) + }) + + // Simulate Redis pub/sub message arriving from another instance + handleRevocationMessage(JSON.stringify({ keyHash: keyHashA })) + + const { code, reason } = await closePromise + expect(code).toBe(4401) + expect(reason).toBe('key_revoked') + expect(ws.readyState).toBe(WebSocket.CLOSED) + expect(connectionRegistry.getConnections(keyHashA)).toHaveLength(0) + }) + + it('handles raw string keyHash messages as well as JSON objects', async () => { + const mockConn = { close: vi.fn() } + connectionRegistry.register(keyHashB, mockConn) + + expect(connectionRegistry.getConnections(keyHashB)).toHaveLength(1) + + // Handle raw string message + handleRevocationMessage(keyHashB) + + expect(mockConn.close).toHaveBeenCalledWith(4401, 'key_revoked') + expect(connectionRegistry.getConnections(keyHashB)).toHaveLength(0) + }) + + it('publishRevocation publishes to the correct Redis channel with the key hash', async () => { + await publishRevocation(keyHashA) + + expect(mockRedisPublish).toHaveBeenCalledWith( + 'key:revoked', + JSON.stringify({ keyHash: keyHashA }), + ) + }) + }) +}) + diff --git a/src/api/admin.ts b/src/api/admin.ts index 77a3b4a4..b9bfba4a 100644 --- a/src/api/admin.ts +++ b/src/api/admin.ts @@ -6,6 +6,12 @@ import { admin_auth_failures_total } from '../metrics' import { formatValidationError } from '../validation' import { hashApiKey, isAdminAuthorized, checkAdminAuth, invalidateApiKey } from './auth' import { copyUsageHistory } from '../x402/metering' +import { + connectionRegistry, + publishRevocation, + CLOSE_CODE_KEY_REVOKED, + CLOSE_REASON_KEY_REVOKED, +} from '../ws/connectionRegistry' /** Generates a new opaque API key with a recognizable `lens_` prefix. */ export function generateApiKey(): string { @@ -314,6 +320,17 @@ export async function registerAdminRoutes(app: FastifyInstance) { // instance (Issue #495). Other instances pick it up within the cache TTL. invalidateApiKey(existing.hash) + // Tear down any live connections this key holds on this instance, then + // tell the other instances to do the same (#503). + if (existing.hash) { + connectionRegistry.closeConnectionsForKey( + existing.hash, + CLOSE_CODE_KEY_REVOKED, + CLOSE_REASON_KEY_REVOKED, + ) + await publishRevocation(existing.hash) + } + return reply.status(200).send({ id, revoked: true }) }, ) diff --git a/src/api/graphql.ts b/src/api/graphql.ts index 675a070f..f741ebbe 100644 --- a/src/api/graphql.ts +++ b/src/api/graphql.ts @@ -9,7 +9,7 @@ import { type OperationDefinitionNode, type FragmentDefinitionNode, } from 'graphql' -import { hasScope, extractBearer, lookupApiKey } from './auth' +import { hasScope, extractBearer, lookupApiKey, hashApiKey } from './auth' import { randomUUID } from 'node:crypto' import { getCachedPrice } from '../redis' import { getAggregatedPrice } from '../aggregator/vwap' @@ -37,6 +37,7 @@ import { priceEmitter, PRICE_PUBLISHED, type PricePublishedEvent } from '../even import { HISTORY_INTERVAL_SECONDS, MAX_HISTORY_POINTS } from './history' import { getMaxGraphQLQueryDepth, queryDepthValidationRules } from './queryDepth' import { formatPrice, formatVolume } from '../pricing/precision' +import { connectionRegistry } from '../ws/connectionRegistry' const WINDOW_SECONDS: Record = { ...HISTORY_INTERVAL_SECONDS, @@ -959,6 +960,34 @@ export const resolvers = { }, } +function extractTokenFromConnection(connContext: any, data: any): string | null { + const payload = data?.payload + if (payload) { + if (typeof payload.authorization === 'string') { + const b = extractBearer(payload.authorization) + return b ?? payload.authorization.trim() + } + if (typeof payload.headers?.authorization === 'string') { + const b = extractBearer(payload.headers.authorization) + return b ?? payload.headers.authorization.trim() + } + if (typeof payload.apiKey === 'string') return payload.apiKey.trim() + if (typeof payload.key === 'string') return payload.key.trim() + if (typeof payload.token === 'string') return payload.token.trim() + } + const req = connContext?.request + if (req) { + const authHeader = req.headers?.['authorization'] + const b = extractBearer(authHeader) + if (b) return b + if (typeof req.headers?.['x-api-key'] === 'string') return req.headers['x-api-key'].trim() + const query = req.query as Record | undefined + if (query?.apiKey) return query.apiKey.trim() + if (query?.key) return query.key.trim() + if (query?.token) return query.token.trim() + } + return null +} export async function registerGraphQL(app: FastifyInstance) { const maxQueryDepth = getMaxGraphQLQueryDepth() @@ -985,17 +1014,25 @@ export async function registerGraphQL(app: FastifyInstance) { apiKey: (req as any)?.apiKey, } }, - onConnect: async (data: any) => { - if (!data?.payload) return {} - const auth = data.payload.authorization ?? data.payload.Authorization ?? data.payload.apiKey - if (typeof auth === 'string') { - const token = extractBearer(auth) ?? auth - const keyCtx = await lookupApiKey(token) - if (keyCtx) { - return { apiKey: keyCtx } + onConnect: async function (this: any, data: any) { + const token = extractTokenFromConnection(this.context, data) + if (token) { + const keyContext = await lookupApiKey(token) + if (!keyContext) { + return false + } + const keyHash = hashApiKey(token) + if (this.socket) { + connectionRegistry.register(keyHash, this.socket) } + return { apiKey: keyContext, keyHash } + } + return true + }, + onDisconnect: async function (this: any, context: any) { + if (context?.keyHash && this?.socket) { + connectionRegistry.deregister(context.keyHash, this.socket) } - return data.payload }, }, }) diff --git a/src/api/websocket.ts b/src/api/websocket.ts index 268fa987..b4b9c98d 100644 --- a/src/api/websocket.ts +++ b/src/api/websocket.ts @@ -5,16 +5,29 @@ import { activeNetwork, config, type NetworkName } from '../config' import { X402_NETWORK_LABEL, paymentAddressFor, getX402ResourceServer } from '../x402/network' import { resolveNetworkName } from '../middleware/network' import { fanOutManager } from '../ws/fanout' +import { hashApiKey, lookupApiKey, extractBearer } from './auth' +import { connectionRegistry, CLOSE_CODE_KEY_REVOKED, CLOSE_REASON_KEY_REVOKED } from '../ws/connectionRegistry' import { v4 as uuid } from 'uuid' import { createInternalErrorBody, getCorrelationId } from '../errors' const FACILITATOR_URL = process.env.X402_FACILITATOR_URL ?? 'https://x402.org/facilitator' +function extractTokenFromRequest(req: FastifyRequest): string | null { + const authHeader = req.headers['authorization'] + const bearer = extractBearer(authHeader) + if (bearer) return bearer + if (typeof req.headers['x-api-key'] === 'string') return req.headers['x-api-key'].trim() + const query = req.query as Record | undefined + if (query?.apiKey) return query.apiKey.trim() + if (query?.key) return query.key.trim() + if (query?.token) return query.token.trim() + return null +} export async function registerWebSocket(app: FastifyInstance) { await app.register(websocket) // @fastify/websocket v11: handler receives (socket, req) directly — no connection wrapper - app.get('/ws', { websocket: true, config: { public: true } }, (socket: any, req: FastifyRequest) => { + app.get('/ws', { websocket: true, config: { public: true } }, async (socket: any, req: FastifyRequest) => { app.log.info('[ws] New connection attempt') const rawNetwork = (req.query as any)?.network ?? (req.headers['x-network'] as string | undefined) @@ -26,6 +39,18 @@ export async function registerWebSocket(app: FastifyInstance) { } const network: NetworkName = resolved.network + const rawKey = extractTokenFromRequest(req) + let authenticatedKeyHash: string | null = null + if (rawKey) { + const keyContext = await lookupApiKey(rawKey) + if (!keyContext) { + socket.send(JSON.stringify({ type: 'error', status: 401, message: 'Invalid or revoked API key' })) + socket.close(CLOSE_CODE_KEY_REVOKED, CLOSE_REASON_KEY_REVOKED) + return + } + authenticatedKeyHash = hashApiKey(rawKey) + } + // This process only ingests and streams live prices for `activeNetwork` // (see src/config.ts) — the underlying price events carry no network tag // yet (that's the deeper aggregation-layer work), so a request for any @@ -69,7 +94,7 @@ export async function registerWebSocket(app: FastifyInstance) { socket.send(JSON.stringify({ type: 'error', message: 'Invalid payment' })) socket.close() } else { - setupStream(socket, req) + setupStream(socket, req, authenticatedKeyHash) } }) .catch(err => { @@ -87,10 +112,14 @@ export async function registerWebSocket(app: FastifyInstance) { return } - setupStream(socket, req) + setupStream(socket, req, authenticatedKeyHash) }) - function setupStream(socket: any, req: FastifyRequest) { + function setupStream(socket: any, req: FastifyRequest, keyHash?: string | null) { + if (keyHash) { + const deregisterKey = connectionRegistry.register(keyHash, socket) + socket.on('close', deregisterKey) + } app.log.info('[ws] Connection authorized') socket.send(JSON.stringify({ type: 'status', message: 'Streaming active' })) @@ -98,6 +127,7 @@ export async function registerWebSocket(app: FastifyInstance) { const clientId = uuid() const unsubscribe = fanOutManager.register('*', { id: clientId, + keyHash: keyHash ?? undefined, send: (data: string, cb?: (err?: Error | null) => void) => { if (socket.readyState !== 1) { // OPEN // Tell the fan-out rather than silently swallowing it: a socket that diff --git a/src/app.ts b/src/app.ts index 85de28aa..685e8b60 100644 --- a/src/app.ts +++ b/src/app.ts @@ -17,6 +17,7 @@ import { registerNetworkSelector } from './middleware/network' import { registerHttpMetrics } from './middleware/httpMetrics' import { registerETag } from './middleware/etag' import { registerWebSocket } from './api/websocket' +import { destroyRevocationSubscriber } from './ws/connectionRegistry' import { registerApiKeyAuth, registerAdminAuth, isAdminConfigured } from './api/auth' import { registerAdminRoutes } from './api/admin' import { registerUsageRoutes } from './api/usage' @@ -282,6 +283,10 @@ export async function buildApp(options: BuildAppOptions = {}): Promise { + await destroyRevocationSubscriber() + }) + // Prometheus metrics endpoint (un-gated, public — no API key required) app.get('/metrics', { config: { public: true } }, async (req, reply) => { reply.type('text/plain; version=0.0.4; charset=utf-8') diff --git a/src/index.ts b/src/index.ts index 92965520..f59e8f2c 100644 --- a/src/index.ts +++ b/src/index.ts @@ -14,6 +14,7 @@ import { assertSchemaValid } from './db/schemaVerify' import { applyMigrations } from './db/migrate' import { buildApp } from './app' import { fanOutManager } from './ws/fanout' +import { initRevocationSubscriber } from './ws/connectionRegistry' import { venueAdapters, getEnabledAdapters, startVenueAdapter, startBlendIngester } from './ingest/venues' import { startSnapshotIngester } from './ingest/snapshot' @@ -101,6 +102,14 @@ async function main() { log.warn({ err }, 'websocket fan-out init skipped') } + // ── Key revocation subscriber (requires Redis for multi-instance) ──────── + try { + await initRevocationSubscriber() + log.info('key revocation subscriber initialized') + } catch (err) { + log.warn({ err }, 'key revocation subscriber init skipped') + } + // ── Aggregate refresh worker (non-blocking — requires Redis) ───────────── // BullMQ opens its own ioredis connections, separate from ./redis and // without our error handler on them. If scheduling fails they are still diff --git a/src/ws/connectionRegistry.ts b/src/ws/connectionRegistry.ts new file mode 100644 index 00000000..3d7b00d8 --- /dev/null +++ b/src/ws/connectionRegistry.ts @@ -0,0 +1,243 @@ +import type Redis from 'ioredis' +import { redis } from '../redis' +import { config } from '../config' +import { createLogger } from '../logger' + +const log = createLogger('ws:connectionRegistry') + +/** + * Distinguishable WebSocket close code and reason for revoked API keys. + * A client must be able to tell this apart from a normal server restart so + * reconnect logic does not retry forever. + */ +export const CLOSE_CODE_KEY_REVOKED = 4401 +export const CLOSE_REASON_KEY_REVOKED = 'key_revoked' + +/** + * Redis pub/sub channel for propagating API key revocation events across instances. + */ +export const REVOCATION_REDIS_CHANNEL = 'key:revoked' + +export interface ClosableConnection { + id?: string + close: (code?: number, reason?: string) => void + on?: (event: 'close', listener: () => void) => any + once?: (event: 'close', listener: () => void) => any + off?: (event: 'close', listener: () => void) => any + removeListener?: (event: 'close', listener: () => void) => any +} + +/** + * Tracks live connections keyed by their SHA-256 hashed API key. + * + * Built to prevent revoked keys from continuing to stream over long-lived + * WebSocket or GraphQL subscriptions that authenticated once. + */ +export class ConnectionRegistry { + private connections = new Map>() + + /** + * Registers a connection under its hashed API key. + * If the connection supports close events, automatically deregisters on natural close. + * Returns an unregister function. + */ + register(keyHash: string, conn: ClosableConnection): () => void { + if (!keyHash || !conn) return () => {} + + let set = this.connections.get(keyHash) + if (!set) { + set = new Set() + this.connections.set(keyHash, set) + } + set.add(conn) + + let unregistered = false + const deregister = () => { + if (unregistered) return + unregistered = true + this.deregister(keyHash, conn) + } + + if (typeof conn.once === 'function') { + conn.once('close', deregister) + } else if (typeof conn.on === 'function') { + conn.on('close', deregister) + } + + return deregister + } + + /** + * Deregisters a connection from a specific key hash. + */ + deregister(keyHash: string, conn: ClosableConnection): void { + if (!keyHash || !conn) return + const set = this.connections.get(keyHash) + if (set) { + set.delete(conn) + if (set.size === 0) { + this.connections.delete(keyHash) + } + } + } + + /** + * Looks up all live connections belonging to a specific key hash. + */ + getConnections(keyHash: string): ClosableConnection[] { + if (!keyHash) return [] + const set = this.connections.get(keyHash) + return set ? Array.from(set) : [] + } + + /** + * Closes all live connections belonging to the specified key hash + * with code 4401 and reason "key_revoked". + * + * Connections belonging to any other keys are unaffected. + * Returns the count of closed connections. + */ + closeConnectionsForKey( + keyHash: string, + code: number = CLOSE_CODE_KEY_REVOKED, + reason: string = CLOSE_REASON_KEY_REVOKED, + ): number { + if (!keyHash) return 0 + const set = this.connections.get(keyHash) + if (!set || set.size === 0) return 0 + + const conns = Array.from(set) + this.connections.delete(keyHash) + + for (const conn of conns) { + try { + conn.close(code, reason) + } catch { + // Socket may already be closing or closed + } + } + + log.info({ keyHash, closedCount: conns.length, code, reason }, 'closed connections for revoked key') + return conns.length + } + + /** + * Total number of registered connections across all keys. + */ + get totalConnections(): number { + let count = 0 + for (const set of this.connections.values()) { + count += set.size + } + return count + } + + /** + * Total number of active keys with registered connections. + */ + get activeKeyCount(): number { + return this.connections.size + } + + /** + * Clears all registered connections. + */ + clear(): void { + this.connections.clear() + } +} + +export const connectionRegistry = new ConnectionRegistry() + +/** + * Publishes a revocation event to Redis for cross-instance propagation. + */ +export async function publishRevocation(keyHash: string): Promise { + if (!keyHash) return + try { + const payload = JSON.stringify({ keyHash }) + await redis.publish(REVOCATION_REDIS_CHANNEL, payload) + log.info({ keyHash }, 'published key revocation event to redis') + } catch (err) { + log.warn({ err, keyHash }, 'failed to publish revocation to redis, continuing with local close only') + } +} + +/** + * Handles an incoming revocation message from the Redis channel. + */ +export function handleRevocationMessage(message: string): void { + try { + let keyHash: string | undefined + try { + const data = JSON.parse(message) + keyHash = typeof data === 'string' ? data : data?.keyHash + } catch { + keyHash = message + } + + if (keyHash && typeof keyHash === 'string') { + log.info({ keyHash }, 'received key revocation event via redis pubsub') + connectionRegistry.closeConnectionsForKey( + keyHash, + CLOSE_CODE_KEY_REVOKED, + CLOSE_REASON_KEY_REVOKED, + ) + } + } catch (err) { + log.warn({ err }, 'failed to process revocation message from redis') + } +} + +let subscriberClient: Redis | null = null + +/** + * Subscribes to the Redis revocation channel so this instance drops its + * own connections when a key is revoked on any instance. + */ +export async function initRevocationSubscriber(): Promise<() => Promise> { + if (subscriberClient) { + return async () => destroyRevocationSubscriber() + } + + try { + const IORedis = (await import('ioredis')).default + const client = new IORedis(config.redis.url, { + maxRetriesPerRequest: 1, + lazyConnect: true, + enableOfflineQueue: false, + }) + + await client.connect() + await client.subscribe(REVOCATION_REDIS_CHANNEL) + subscriberClient = client + + client.on('message', (channel: string, message: string) => { + if (channel === REVOCATION_REDIS_CHANNEL) { + handleRevocationMessage(message) + } + }) + + log.info('revocation redis subscriber ready') + } catch (err) { + log.warn({ err }, 'redis revocation pubsub unavailable, running in single-instance mode') + } + + return async () => destroyRevocationSubscriber() +} + +/** + * Cleans up and disconnects the revocation subscriber. + */ +export async function destroyRevocationSubscriber(): Promise { + if (subscriberClient) { + try { + await subscriberClient.unsubscribe(REVOCATION_REDIS_CHANNEL) + } catch { + // Ignore unsubscribe error on shutdown + } + subscriberClient.disconnect() + subscriberClient = null + } +} + diff --git a/src/ws/fanout.ts b/src/ws/fanout.ts index c6cacd6f..7833258e 100644 --- a/src/ws/fanout.ts +++ b/src/ws/fanout.ts @@ -5,6 +5,7 @@ import { Gauge, Counter } from 'prom-client' import { register } from '../metrics' import { createLogger } from '../logger' import { resolveBackpressureLimits, CLOSE_CODE_SLOW_CONSUMER } from './backpressure' +import { connectionRegistry } from './connectionRegistry' const log = createLogger('ws:fanout') @@ -90,6 +91,7 @@ function pairMetricLabel(pairKey: string): string | null { // ─── Client wrapper with backpressure & coalescing ──────────────────────────── export interface FanOutClient { id: string + keyHash?: string /** * Queues `data` on the socket. * @@ -193,6 +195,9 @@ class FanOutManager { } this.isDestroyed = false + if (socket.keyHash) { + connectionRegistry.register(socket.keyHash, socket) + } this.clients.set(id, client) let clients = this.pairClients.get(pairKey) if (!clients) { @@ -219,6 +224,9 @@ class FanOutManager { // the whole point of bounding the buffer is that the bytes go away with it. client.pendingUpdate = null + if (client.socket.keyHash) { + connectionRegistry.deregister(client.socket.keyHash, client.socket) + } this.clients.delete(id) const clients = this.pairClients.get(client.pairKey) if (clients) {