feat(ws): authenticate and meter price streams

docs/websocket-auth.md

@@ -0,0 +1,27 @@
+# Authenticated price WebSocket
+
+`GET /ws` requires a valid Lens API key and validates it before completing the
+WebSocket upgrade. Send the key in `Authorization: Bearer <key>` when the
+client supports custom headers. Browser clients can offer the `lens` protocol
+and a `lens-key.<key>` protocol; Lens selects `lens`, so the credential is
+never returned as the negotiated protocol and is not placed in the URL.
+
+```ts
+const socket = new WebSocket('wss://api.example.com/ws', [
+  'lens',
+  `lens-key.${apiKey}`,
+])
+```
+
+The default maximum is five concurrent price streams per API key and process.
+Set `WS_MAX_CONNECTIONS_PER_KEY` to an integer from 1 to 1000 to change it.
+Connections count against the key's slot until the socket closes. The
+validated key metadata is attached to the Fastify request, allowing the same
+per-key identity and quota context used by HTTP routes to be applied here. Each
+opened stream records 50 cents against the key's daily and monthly usage, which
+matches the existing `$0.50` stream price. A stream is rejected before upgrade
+when that charge would exceed the configured key quota.
+
+The API key is never accepted from a query parameter. The
+`Sec-WebSocket-Protocol` header is redacted from request logs because browser
+clients use it to carry the credential.

src/__tests__/logger.test.ts

@@ -103,6 +103,15 @@ describe('structured logger output', () => {
     expect(output).toContain('[redacted]')
   })
 
+  it('redacts API keys carried in the WebSocket subprotocol header', () => {
+    const apiKey = 'lens_0123456789abcdef0123456789abcdef0123456789abcdef'
+    const { lines, destination } = capture()
+    const log = buildLogger({ level: 'info', pretty: false, destination })
+    log.info({ req: { headers: { 'sec-websocket-protocol': `lens,lens-key.${apiKey}` } } }, 'request')
+    expect(lines.join('\n')).not.toContain(apiKey)
+    expect(lines.join('\n')).toContain('[redacted]')
+  })
+
   it('renders human-readable lines in development mode', () => {
     const { lines, destination } = capture()
     const log = buildLogger({ level: 'info', pretty: true, destination }).child({

src/api/websocket.ts

@@ -1,21 +1,125 @@
-import type { FastifyInstance, FastifyRequest } from 'fastify'
+import type { FastifyInstance, FastifyReply, FastifyRequest } from 'fastify'
 import websocket from '@fastify/websocket'
 import { priceEmitter, PRICE_UPDATE, PriceUpdateEvent } from '../events'
 import { activeNetwork, type NetworkName } from '../config'
 import { X402_NETWORK_LABEL, paymentAddressFor, getX402ResourceServer } from '../x402/network'
 import { resolveNetworkName } from '../middleware/network'
 import { fanOutManager } from '../ws/fanout'
 import { v4 as uuid } from 'uuid'
+import { extractBearer, lookupApiKey } from './auth'
+import { api_key_auth_failures_total, rate_limit_rejections_total, websocket_connections_active } from '../metrics'
+import { checkQuota, recordUsage } from '../x402/metering'
 
 const FACILITATOR_URL = process.env.X402_FACILITATOR_URL ?? 'https://x402.org/facilitator'
+const STREAM_PRICE_CENTS = 50
+
+const activeConnectionsByKey = new Map<string, number>()
+const connectionLeases = new WeakMap<FastifyRequest, { release: () => void; upgraded: boolean }>()
+
+function maxConnectionsPerKey(): number {
+  const configured = process.env.WS_MAX_CONNECTIONS_PER_KEY
+  if (configured === undefined) return 5
+  if (!/^[1-9]\d*$/.test(configured)) {
+    throw new Error('WS_MAX_CONNECTIONS_PER_KEY must be a positive integer')
+  }
+  const value = Number(configured)
+  if (!Number.isSafeInteger(value) || value > 1000) {
+    throw new Error('WS_MAX_CONNECTIONS_PER_KEY must be between 1 and 1000')
+  }
+  return value
+}
+
+function handshakeApiKey(req: FastifyRequest): { token?: string; malformed: boolean } {
+  const authorization = req.headers.authorization
+  const bearer = extractBearer(authorization)
+  const protocols = (req.headers['sec-websocket-protocol'] ?? '')
+    .split(',')
+    .map(value => value.trim())
+    .filter(Boolean)
+  const offeredKeys = protocols
+    .filter(value => value.startsWith('lens-key.'))
+    .map(value => value.slice('lens-key.'.length))
+  if (offeredKeys.length > 1 || (offeredKeys.length === 1 && !protocols.includes('lens'))) {
+    return { malformed: true }
+  }
+  if (authorization !== undefined && !bearer) return { malformed: true }
+  if (bearer && offeredKeys[0] && bearer !== offeredKeys[0]) return { malformed: true }
+  return { token: bearer ?? offeredKeys[0], malformed: false }
+}
 
 export async function registerWebSocket(app: FastifyInstance) {
-  await app.register(websocket)
+  await app.register(websocket, {
+    options: {
+      // Never select the credential-bearing protocol back to the client.
+      handleProtocols: protocols => protocols.has('lens') ? 'lens' : false,
+    },
+  })
+
+  // Failed upgrades never reach the socket handler, so release a reserved key
+  // slot when Fastify completes an ordinary HTTP error response.
+  app.addHook('onResponse', async req => {
+    const lease = connectionLeases.get(req)
+    if (lease && !lease.upgraded) lease.release()
+  })
 
   // @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,
+    // The regular API-key hook skips this route; authenticate the handshake
+    // here so missing/revoked keys are rejected before the upgrade completes.
+    config: { public: true },
+    preValidation: async (req: FastifyRequest, reply: FastifyReply) => {
+      const supplied = handshakeApiKey(req)
+      if (supplied.malformed) {
+        api_key_auth_failures_total.inc({ reason: 'malformed' })
+        return reply.status(401).send({ error: 'Unauthorized', code: 'malformed' })
+      }
+      if (!supplied.token) {
+        api_key_auth_failures_total.inc({ reason: 'missing' })
+        return reply.status(401).send({ error: 'Unauthorized', code: 'missing' })
+      }
+
+      const apiKey = await lookupApiKey(supplied.token)
+      if (!apiKey) {
+        api_key_auth_failures_total.inc({ reason: 'unknown' })
+        return reply.status(401).send({ error: 'Unauthorized', code: 'unknown' })
+      }
+
+      const quota = await checkQuota(apiKey.id, STREAM_PRICE_CENTS)
+      if (!quota.allowed) {
+        return reply.status(402).send({ error: 'API key quota exceeded', code: 'quota_exceeded' })
+      }
+
+      const active = activeConnectionsByKey.get(apiKey.id) ?? 0
+      if (active >= maxConnectionsPerKey()) {
+        rate_limit_rejections_total.inc({ scope: 'websocket_per_key' })
+        return reply.status(429).send({ error: 'Too many WebSocket connections for this API key' })
+      }
+
+      activeConnectionsByKey.set(apiKey.id, active + 1)
+      websocket_connections_active.inc()
+      ;(req as FastifyRequest & { apiKey?: typeof apiKey }).apiKey = apiKey
+      let released = false
+      connectionLeases.set(req, {
+        upgraded: false,
+        release: () => {
+          if (released) return
+          released = true
+          const remaining = (activeConnectionsByKey.get(apiKey.id) ?? 1) - 1
+          if (remaining === 0) activeConnectionsByKey.delete(apiKey.id)
+          else activeConnectionsByKey.set(apiKey.id, remaining)
+          websocket_connections_active.dec()
+        },
+      })
+    },
+  }, (socket: any, req: FastifyRequest) => {
     app.log.info('[ws] New connection attempt')
 
+    const lease = connectionLeases.get(req)
+    if (lease) lease.upgraded = true
+    socket.once('close', () => lease?.release())
+    socket.once('error', () => lease?.release())
+
     const rawNetwork = (req.query as any)?.network ?? (req.headers['x-network'] as string | undefined)
     const resolved = resolveNetworkName(rawNetwork)
     if (!resolved.ok) {
@@ -68,7 +172,7 @@ export async function registerWebSocket(app: FastifyInstance) {
             socket.send(JSON.stringify({ type: 'error', message: 'Invalid payment' }))
             socket.close()
           } else {
-            setupStream(socket, req)
+          void setupStream(socket, req)
           }
         })
         .catch(err => {
@@ -78,11 +182,20 @@ export async function registerWebSocket(app: FastifyInstance) {
       return
     }
 
-    setupStream(socket, req)
+    void setupStream(socket, req)
   })
 
-  function setupStream(socket: any, req: FastifyRequest) {
-    app.log.info('[ws] Connection authorized')
+  async function setupStream(socket: any, req: FastifyRequest) {
+    // A stream costs the same $0.50 as this route's existing x402 price and is
+    // recorded against the authenticated key's daily/monthly quota counters.
+    try {
+      await recordUsage(req.apiKey!.id, STREAM_PRICE_CENTS)
+    } catch (err) {
+      app.log.error({ err, apiKeyId: req.apiKey?.id }, '[ws] Could not record stream usage')
+      socket.close(1011, 'Unable to start stream')
+      return
+    }
+    app.log.info({ apiKeyId: req.apiKey?.id }, '[ws] Connection authorized')
     socket.send(JSON.stringify({ type: 'status', message: 'Streaming active' }))
 
     // Register with fan-out manager for backpressure-aware broadcasting
@@ -135,4 +248,4 @@ async function verifyPayment(paymentHeader: string, requirements: any, resourceS
   } catch {
     return false
   }
-}
+}

src/logger.ts

@@ -86,12 +86,14 @@ const REDACT_PATHS = [
   'req.headers["x-api-key"]',
   'req.headers["x-admin-token"]',
   'req.headers["x-payment"]',
+  'req.headers["sec-websocket-protocol"]',
   'res.headers["set-cookie"]',
   'headers.authorization',
   'headers.cookie',
   'headers["x-api-key"]',
   'headers["x-admin-token"]',
   'headers["x-payment"]',
+  'headers["sec-websocket-protocol"]',
   'password',
   'secret',
   'token',

src/metrics.ts

@@ -59,6 +59,12 @@ export const admin_auth_failures_total = new Counter({
   registers: [register]
 })
 
+export const websocket_connections_active = new Gauge({
+  name: 'websocket_connections_active',
+  help: 'Number of currently authenticated price WebSocket connections',
+  registers: [register]
+})
+
 export const rate_limit_rejections_total = new Counter({
   name: 'rate_limit_rejections_total',
   help: 'Total number of requests rejected by the per-key or fallback IP rate limiter',

src/x402/metering.ts

@@ -59,10 +59,16 @@ export async function getQuotaConfig(keyId: string): Promise<QuotaConfig> {
   }
 }
 
-export async function checkQuota(keyId: string): Promise<{ allowed: boolean; reason?: string }> {
+export async function checkQuota(keyId: string, additionalCents = 0): Promise<{ allowed: boolean; reason?: string }> {
   const config = await getQuotaConfig(keyId)
   const summary = await getUsageSummaryInternal(keyId, config)
-  if (summary.dailyCents >= summary.dailyQuotaCents || summary.monthlyCents >= summary.monthlyQuotaCents) {
+  const exceedsDaily = additionalCents === 0
+    ? summary.dailyCents >= summary.dailyQuotaCents
+    : summary.dailyCents + additionalCents > summary.dailyQuotaCents
+  const exceedsMonthly = additionalCents === 0
+    ? summary.monthlyCents >= summary.monthlyQuotaCents
+    : summary.monthlyCents + additionalCents > summary.monthlyQuotaCents
+  if (exceedsDaily || exceedsMonthly) {
     return { allowed: false, reason: `${config.overagePolicy} quota exceeded` }
   }
   return { allowed: true }

tests/integration/ws.test.ts

@@ -1,16 +1,42 @@
-import { beforeAll, afterAll, describe, it, expect } from 'vitest'
+import { beforeAll, afterAll, describe, it, expect, vi } from 'vitest'
 import Fastify from 'fastify'
 import WebSocket from 'ws'
 
+vi.mock('../../src/db', () => ({ prisma: { apiKey: { findUnique: vi.fn() } } }))
+vi.mock('../../src/redis', () => {
+  const multi = {
+    incrby: vi.fn(() => multi),
+    expire: vi.fn(() => multi),
+    exec: vi.fn().mockResolvedValue([]),
+  }
+  return {
+    redis: {
+      get: vi.fn().mockResolvedValue(null),
+      multi: vi.fn(() => multi),
+    },
+  }
+})
+
 import { registerWebSocket } from '../../src/api/websocket'
 import { priceEmitter, PRICE_UPDATE } from '../../src/events'
 import { fanOutManager } from '../../src/ws/fanout'
+import { prisma } from '../../src/db'
+
+const mockFindUnique = prisma.apiKey.findUnique as unknown as ReturnType<typeof vi.fn>
+const VALID_KEY = 'lens_0123456789abcdef0123456789abcdef0123456789abcdef'
+const VALID_KEY_HASH = 'unused-in-this-mock'
 
 describe('WebSocket subscription full cycle', () => {
   let app: ReturnType<typeof Fastify>
   let port: number
 
   beforeAll(async () => {
+    process.env.WS_MAX_CONNECTIONS_PER_KEY = '1'
+    mockFindUnique.mockResolvedValue({
+      id: 'ws-key-1', label: 'integration-test', ratePerMin: 60, ratePerDay: 1000,
+      revokedAt: null, expiresAt: null, hash: VALID_KEY_HASH,
+      dailyQuotaCents: 500, monthlyQuotaCents: 10000, overagePolicy: 'block',
+    })
     // Initialize the fan-out manager so it listens for price updates
     // and broadcasts to registered WebSocket clients
     await fanOutManager.initialize()
@@ -32,7 +58,7 @@ describe('WebSocket subscription full cycle', () => {
     const messages: string[] = []
     const waiters: Array<(msg: string) => void> = []
 
-    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`)
+    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`, { headers: { authorization: `Bearer ${VALID_KEY}` } })
 
     // Attach message listener before awaiting open to avoid the race where the
     // server sends the status message synchronously on connect.
@@ -83,7 +109,68 @@ describe('WebSocket subscription full cycle', () => {
     expect(new Date(next.timestamp).toISOString()).toBe(event.timestamp.toISOString())
 
     // 5. Disconnect and verify client cleanup
+    const closed = new Promise<void>(resolve => ws.once('close', () => resolve()))
+    ws.close()
+    await closed
+  })
+
+  it('rejects a missing API key before the websocket upgrade', async () => {
+    const error = await new Promise<Error>(resolve => {
+      const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`)
+      ws.once('error', resolve)
+    })
+    expect(error.message).toContain('401')
+  })
+
+  it('rejects a revoked API key before the websocket upgrade', async () => {
+    mockFindUnique.mockResolvedValueOnce({
+      id: 'ws-key-revoked', label: 'revoked', ratePerMin: 60, ratePerDay: 1000,
+      revokedAt: new Date(), expiresAt: null,
+    })
+    const error = await new Promise<Error>(resolve => {
+      const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`, { headers: { authorization: `Bearer ${VALID_KEY}` } })
+      ws.once('error', resolve)
+    })
+    expect(error.message).toContain('401')
+  })
+
+  it('accepts browser subprotocol credentials without returning the credential', async () => {
+    const url = `ws://127.0.0.1:${port}/ws`
+    expect(url).not.toContain(VALID_KEY)
+    const ws = new WebSocket(url, ['lens', `lens-key.${VALID_KEY}`])
+    await new Promise<void>((resolve, reject) => {
+      ws.once('open', resolve)
+      ws.once('error', reject)
+    })
+    expect(ws.protocol).toBe('lens')
+    const closed = new Promise<void>(resolve => ws.once('close', () => resolve()))
     ws.close()
-    await new Promise<void>((resolve) => ws.on('close', () => resolve()))
+    await closed
+  })
+
+  it('caps concurrent connections per key and frees the slot on close', async () => {
+    const first = new WebSocket(`ws://127.0.0.1:${port}/ws`, { headers: { authorization: `Bearer ${VALID_KEY}` } })
+    await new Promise<void>((resolve, reject) => {
+      first.once('open', resolve)
+      first.once('error', reject)
+    })
+
+    const rejected = await new Promise<Error>(resolve => {
+      const second = new WebSocket(`ws://127.0.0.1:${port}/ws`, { headers: { authorization: `Bearer ${VALID_KEY}` } })
+      second.once('error', resolve)
+    })
+    expect(rejected.message).toContain('429')
+
+    const firstClosed = new Promise<void>(resolve => first.once('close', () => resolve()))
+    first.close()
+    await firstClosed
+    const afterClose = new WebSocket(`ws://127.0.0.1:${port}/ws`, { headers: { authorization: `Bearer ${VALID_KEY}` } })
+    await new Promise<void>((resolve, reject) => {
+      afterClose.once('open', resolve)
+      afterClose.once('error', reject)
+    })
+    const afterCloseClosed = new Promise<void>(resolve => afterClose.once('close', () => resolve()))
+    afterClose.close()
+    await afterCloseClosed
   })
-})
+})