fix: revoking a key now disconnects live connections (#503)

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 })
     },
   )

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<string, number> = {
   ...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<string, string | undefined> | 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
       },
     },
   })

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<string, string | undefined> | 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,17 +112,22 @@ 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' }))
 
     // Register with fan-out manager for backpressure-aware broadcasting
     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

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<FastifyIn
     return reply.send(error)
   })
 
+  app.addHook('onClose', async () => {
+    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')

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