diff --git a/docs/websocket-auth.md b/docs/websocket-auth.md new file mode 100644 index 00000000..100cfef8 --- /dev/null +++ b/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 ` when the +client supports custom headers. Browser clients can offer the `lens` protocol +and a `lens-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. diff --git a/src/__tests__/logger.test.ts b/src/__tests__/logger.test.ts index ed18bf8a..6d41f77e 100644 --- a/src/__tests__/logger.test.ts +++ b/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({ diff --git a/src/api/websocket.ts b/src/api/websocket.ts index acd8dada..7e3ed67f 100644 --- a/src/api/websocket.ts +++ b/src/api/websocket.ts @@ -1,4 +1,4 @@ -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' @@ -6,16 +6,120 @@ import { X402_NETWORK_LABEL, paymentAddressFor, getX402ResourceServer } from '.. 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() +const connectionLeases = new WeakMap 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 } -} \ No newline at end of file +} diff --git a/src/logger.ts b/src/logger.ts index f0674ff3..ac6ba3e7 100644 --- a/src/logger.ts +++ b/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', diff --git a/src/metrics.ts b/src/metrics.ts index ac8ac646..b87555f9 100644 --- a/src/metrics.ts +++ b/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', diff --git a/src/x402/metering.ts b/src/x402/metering.ts index a0a8f5d7..0f0a5ecc 100644 --- a/src/x402/metering.ts +++ b/src/x402/metering.ts @@ -59,10 +59,16 @@ export async function getQuotaConfig(keyId: string): Promise { } } -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 } diff --git a/tests/integration/ws.test.ts b/tests/integration/ws.test.ts index 57f25946..1c70aca9 100644 --- a/tests/integration/ws.test.ts +++ b/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 +const VALID_KEY = 'lens_0123456789abcdef0123456789abcdef0123456789abcdef' +const VALID_KEY_HASH = 'unused-in-this-mock' describe('WebSocket subscription full cycle', () => { let app: ReturnType 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(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(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(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((resolve, reject) => { + ws.once('open', resolve) + ws.once('error', reject) + }) + expect(ws.protocol).toBe('lens') + const closed = new Promise(resolve => ws.once('close', () => resolve())) ws.close() - await new Promise((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((resolve, reject) => { + first.once('open', resolve) + first.once('error', reject) + }) + + const rejected = await new Promise(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(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((resolve, reject) => { + afterClose.once('open', resolve) + afterClose.once('error', reject) + }) + const afterCloseClosed = new Promise(resolve => afterClose.once('close', () => resolve())) + afterClose.close() + await afterCloseClosed }) -}) \ No newline at end of file +})