diff --git a/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts b/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts index 3faa8523c..812b57dd1 100644 --- a/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts +++ b/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts @@ -3,7 +3,7 @@ import { type IncomingMessage, type Server as HttpServer, } from 'node:http'; -import type { AddressInfo } from 'node:net'; +import { connect, type AddressInfo, type Socket } from 'node:net'; import { PassThrough } from 'node:stream'; import { afterAll, @@ -16,6 +16,7 @@ import { vi, } from 'vitest'; import WebSocket, { WebSocketServer } from 'ws'; +import { makeActor } from '../../core/actor.js'; import { checkRateLimit } from '../../core/http/middleware/rateLimit.js'; import { withAiCostFactorHook } from '../../drivers/integrationTestUtil'; import type { IConfig } from '../../types'; @@ -199,6 +200,10 @@ describe('OpenAIRealtimeRelayService relay', () => { let upstream: WebSocketServer; let user: TestUserCredentials; let relayUrl: string; + let paidPolicy: NonNullable< + ReturnType + >; + let planId = 'test_paid'; const upstreamRequests: IncomingMessage[] = []; const upstreamEvents: { type: string; @@ -216,13 +221,12 @@ describe('OpenAIRealtimeRelayService relay', () => { await server.services.auth.createSessionToken(row!, {}) ).token; const metering = server.services.metering; + planId = 'test_paid'; + paidPolicy.monthUsageAllowance = 1_000_000_000; + metering.invalidateActorSubscription(row!.uuid); const plan = metering.getActorSubscription.bind(metering); - vi.spyOn(metering, 'getActorSubscription').mockImplementation( - async (actor) => ({ - ...(await plan(actor)), - id: 'test_paid', - monthUsageAllowance: 1_000_000_000, - }), + vi.spyOn(metering, 'getActorSubscription').mockImplementation((actor) => + plan(actor), ); }); @@ -473,6 +477,20 @@ describe('OpenAIRealtimeRelayService relay', () => { username: 'realtime-user', password: 'realtime-user-password', }); + const userRow = await server.stores.user.getByUsername(user.username); + paidPolicy = { + ...(await server.services.metering.getActorSubscription( + makeActor({ user: userRow! }), + )), + id: 'test_paid', + monthUsageAllowance: 1_000_000_000, + free: false, + noAllowance: false, + }; + server.services.metering.registerPolicy(paidPolicy); + server.services.metering.registerSubscriptionResolver( + async () => planId, + ); relayUrl = `ws://api.puter.localhost:${port}${REALTIME_PATH}`; }, 60_000); @@ -482,8 +500,12 @@ describe('OpenAIRealtimeRelayService relay', () => { }, 60_000); /** Open a client and collect every event until it closes. */ - const open = (url: string, protocols?: string[]) => { - const ws = new WebSocket(url, protocols); + const open = ( + url: string, + protocols?: string[], + headers?: Record, + ) => { + const ws = new WebSocket(url, protocols, { headers }); clients.add(ws); ws.once('close', () => clients.delete(ws)); const events: Record[] = []; @@ -496,7 +518,10 @@ describe('OpenAIRealtimeRelayService relay', () => { vi.waitFor( () => { const e = events.find((x) => x.type === type); - if (!e) throw new Error(`no ${type} yet`); + if (!e) + throw new Error( + `no ${type} yet: ${JSON.stringify(events)}`, + ); return e; }, { timeout: 5_000 }, @@ -1073,7 +1098,9 @@ describe('OpenAIRealtimeRelayService relay', () => { }); it('refuses accounts without a subscription', async () => { - // The test user's own plan is a free one. + planId = 'user_free'; + const row = await server.stores.user.getByUsername(user.username); + server.services.metering.invalidateActorSubscription(row!.uuid); vi.mocked(server.services.metering.getActorSubscription).mockRestore(); const c = open(`${relayUrl}?auth_token=${user.token}`); expect(await c.waitFor('error')).toMatchObject({ @@ -1099,13 +1126,15 @@ describe('OpenAIRealtimeRelayService relay', () => { const metering = server.services.metering; const resolvers = (metering as unknown as Resolvers) .subscriptionResolvers; - resolvers.splice(resolvers.indexOf(failing), 1); + const index = resolvers.indexOf(failing); + if (index >= 0) resolvers.splice(index, 1); // The fallback stays cached for its window otherwise. const u = await server.stores.user.getByUsername(user.username); metering.invalidateActorSubscription(u!.uuid); }); it('asks the client to retry instead of reporting no subscription', async () => { + planId = 'user_free'; const u = await server.stores.user.getByUsername(user.username); vi.mocked( server.services.metering.getActorSubscription, @@ -1235,6 +1264,252 @@ describe('OpenAIRealtimeRelayService relay', () => { ws.close(); }); + it('applies the IP gate before opening an upstream connection', async () => { + const before = upstreamRequests.length; + const deny = (_key: string, data: { allow: boolean }) => { + data.allow = false; + }; + server.clients.event.on('ip.validate', deny); + try { + const c = open(`${relayUrl}?auth_token=${user.token}`); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'permission_denied' }, + }); + await c.closed; + expect(upstreamRequests).toHaveLength(before); + } finally { + server.clients.event.off('ip.validate', deny); + } + }); + + it('does not allocate credential limiter buckets for many unique invalid signatures', async () => { + const before = upstreamRequests.length; + const auth = vi.spyOn(server.services.auth, 'authenticate'); + const redis = vi.spyOn(server.clients.redis, 'multi'); + const relay = new OpenAIRealtimeRelayService( + { domain: 'puter.localhost', trust_proxy: 'loopback' } as IConfig, + server.clients, + server.stores, + server.services, + ); + const http = createServer(); + relay.attachHttpServer(http); + await new Promise((resolve) => + http.listen(0, '127.0.0.1', resolve), + ); + try { + for (let i = 0; i < 64; i++) { + const c = open( + `ws://127.0.0.1:${(http.address() as AddressInfo).port}${REALTIME_PATH}?auth_token=invalid-${i}`, + undefined, + { + host: 'api.puter.localhost', + 'x-forwarded-for': `203.0.113.${i + 1}`, + }, + ); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await c.closed; + } + expect(auth).not.toHaveBeenCalled(); + expect(redis).not.toHaveBeenCalled(); + expect(upstreamRequests).toHaveLength(before); + } finally { + await relay.onServerPrepareShutdown(); + await new Promise((resolve) => http.close(() => resolve())); + } + }); + + it('uses non-touching authentication on connection and recheck', async () => { + const auth = vi.spyOn(server.services.auth, 'authenticate'); + const c = open(`${relayUrl}?auth_token=${user.token}`); + await c.waitFor('session.created'); + kv.del(realtimeAuthKey(user.token)); + await server.services.openaiRealtime.recheckSessions(); + expect(auth.mock.calls.length).toBeGreaterThanOrEqual(2); + expect(auth.mock.calls.every(([, ctx]) => ctx?.touch === false)).toBe( + true, + ); + c.ws.close(); + }); + + it.each([ + [false, '203.0.113.9, 10.0.0.1', '127.0.0.1'], + [['loopback', '10.0.0.0/8'], '203.0.113.9, 10.0.0.1', '203.0.113.9'], + [['loopback'], '203.0.113.9, 198.51.100.5', '198.51.100.5'], + ])( + 'resolves forwarded chains using trust policy %j', + async (trustProxy, forwarded, expected) => { + const relay = new OpenAIRealtimeRelayService( + { + domain: 'puter.localhost', + trust_proxy: trustProxy, + } as IConfig, + server.clients, + server.stores, + server.services, + ); + const http = createServer(); + relay.attachHttpServer(http); + await new Promise((resolve) => + http.listen(0, '127.0.0.1', resolve), + ); + let seen = ''; + const deny = ( + _key: string, + data: { allow: boolean; ip: string }, + ) => { + seen = data.ip; + data.allow = false; + }; + server.clients.event.on('ip.validate', deny); + try { + const ws = new WebSocket( + `ws://127.0.0.1:${(http.address() as AddressInfo).port}${REALTIME_PATH}`, + { + headers: { + host: 'api.puter.localhost', + 'x-forwarded-for': forwarded as string, + }, + }, + ); + ws.on('error', () => {}); + await new Promise((resolve) => ws.once('close', resolve)); + expect(seen).toBe(expected); + } finally { + server.clients.event.off('ip.validate', deny); + await relay.onServerPrepareShutdown(); + await new Promise((resolve) => + http.close(() => resolve()), + ); + } + }, + ); + + it.each([undefined, '', 'not-a-jwt', 'x'.repeat(8193)])( + 'rejects missing or malformed credentials before upstream admission', + async (token) => { + const before = upstreamRequests.length; + const c = open( + token === undefined + ? relayUrl + : `${relayUrl}?auth_token=${token}`, + ); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await c.closed; + expect(upstreamRequests).toHaveLength(before); + }, + ); + + it('reports a distributed limiter outage as retryable without authenticating or connecting upstream', async () => { + const before = upstreamRequests.length; + const auth = vi.spyOn(server.services.auth, 'authenticate'); + vi.spyOn(server.clients.redis, 'multi').mockImplementationOnce(() => { + throw new Error('limiter unavailable'); + }); + const c = open(`${relayUrl}?auth_token=${user.token}`); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'internal_error' }, + }); + expect((await c.closed).code).toBe(1011); + expect(auth).not.toHaveBeenCalled(); + expect(upstreamRequests).toHaveLength(before); + }); + + it('keeps recovery attempt capacity after the shared handover budget is exhausted', async () => { + const holder = await createTestUser(server, { + username: 'handover-attempts', + password: 'handover-attempts-password', + }); + for (let i = 0; i < 20; i++) + await checkRateLimit( + `${realtimeConnectKey(holder.token)}:handover`, + 20, + 60_000, + ); + const probe = open( + `${relayUrl}?auth_token=${holder.token}&purpose=handover`, + ); + expect(await probe.waitFor('error')).toMatchObject({ + error: { code: 'rate_limit_exceeded' }, + }); + await probe.closed; + const recovery = open(`${relayUrl}?auth_token=${holder.token}`); + await recovery.waitFor('session.created'); + expect( + await checkRateLimit(realtimeConnectKey(holder.token), 2, 60_000), + ).toBe(true); + recovery.ws.close(); + }); + it('bounds unverified socket allocation while rejected clients withhold close acknowledgments', async () => { + const relay = new OpenAIRealtimeRelayService( + { domain: 'puter.localhost', trust_proxy: 'loopback' } as IConfig, + server.clients, + server.stores, + server.services, + ); + const http = createServer(); + relay.attachHttpServer(http); + await new Promise((resolve) => + http.listen(0, '127.0.0.1', resolve), + ); + const raw = new Set(); + const accepted: WebSocket[] = []; + const upgrade = WebSocketServer.prototype.handleUpgrade; + vi.spyOn(WebSocketServer.prototype, 'handleUpgrade').mockImplementation( + function (this: WebSocketServer, req, socket, head, done) { + upgrade.call(this, req, socket, head, (ws, request) => { + accepted.push(ws); + done(ws, request); + }); + }, + ); + const request = (index: number) => + new Promise((resolve, reject) => { + const socket = connect( + (http.address() as AddressInfo).port, + '127.0.0.1', + ); + raw.add(socket); + socket.on('error', () => {}); + socket.once('error', reject); + socket.once('data', (data) => resolve(data.toString())); + socket.once('connect', () => + socket.write( + [ + `GET ${REALTIME_PATH}?auth_token=invalid-${index} HTTP/1.1`, + 'Host: api.puter.localhost', + 'Connection: Upgrade', + 'Upgrade: websocket', + 'Sec-WebSocket-Version: 13', + 'Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==', + `X-Forwarded-For: 203.0.113.${index + 1}`, + '', + '', + ].join('\r\n'), + ), + ); + }); + try { + for (let i = 0; i < 64; i++) + expect(await request(i)).toMatch(/^HTTP\/1\.1 101/); + await vi.waitFor(() => + expect( + accepted.every((ws) => ws.readyState === WebSocket.CLOSING), + ).toBe(true), + ); + expect(await request(64)).toMatch(/^HTTP\/1\.1 503/); + expect(accepted).toHaveLength(64); + } finally { + for (const socket of raw) socket.destroy(); + await relay.onServerPrepareShutdown(); + await new Promise((resolve) => http.close(() => resolve())); + } + }); + it('closes a session whose credential stops authenticating', async () => { const relay = server.services.openaiRealtime; const auth = vi.spyOn(server.services.auth, 'authenticate'); diff --git a/src/backend/services/realtime/OpenAIRealtimeRelayService.ts b/src/backend/services/realtime/OpenAIRealtimeRelayService.ts index 2e3f3b92f..24f828b71 100644 --- a/src/backend/services/realtime/OpenAIRealtimeRelayService.ts +++ b/src/backend/services/realtime/OpenAIRealtimeRelayService.ts @@ -20,6 +20,8 @@ import { createHash } from 'node:crypto'; import type { IncomingMessage, Server as HttpServer } from 'node:http'; import type { Duplex } from 'node:stream'; +import express from 'express'; +import type { AnyTokenPayload } from '../auth/types.js'; import WebSocket, { WebSocketServer, type RawData } from 'ws'; import type { Actor } from '../../core/actor.js'; import { @@ -148,8 +150,13 @@ const HEARTBEAT_MS = 30_000; const UPSTREAM_HANDSHAKE_MS = 10_000; const SESSION_SETUP_MS = 15_000; const DEFAULT_MAX_PER_USER = 20; -// Session attempts per credential, checked before authenticating. +// Session attempts per signature-verified credential. const CONNECTS_PER_MINUTE = 60; +const HANDOVER_CONNECTS_PER_MINUTE = 20; +const PREAUTH_PER_IP = 120; +const PREAUTH_PER_NODE = 600; +const MAX_AUTHENTICATING = 64; +const MAX_TOKEN_BYTES = 8192; // Each response holds this much audio, at the session model's rates, from the // account's budget until its usage is recorded. const RESPONSE_HOLD_INPUT_TOKENS = 8_000; @@ -533,6 +540,8 @@ export class OpenAIRealtimeRelayService extends PuterService { // One re-authentication per token at a time, shared by its sessions. #authChecks = new Map>(); #draining = false; + #authenticating = 0; + #pendingUpgrades = 0; #sweeping = false; attachHttpServer(server: HttpServer): void { this.#wss = new WebSocketServer({ @@ -546,6 +555,9 @@ export class OpenAIRealtimeRelayService extends PuterService { handleProtocols: (protocols) => protocols.has('realtime') ? 'realtime' : false, }); + // Use Express itself so upgrades follow the ordinary HTTP trust policy. + const ipApp = express(); + ipApp.set('trust proxy', this.config.trust_proxy ?? false); const domain = (this.config.domain ?? '').toLowerCase().trim(); const apiHost = `api.${domain}`; @@ -553,23 +565,101 @@ export class OpenAIRealtimeRelayService extends PuterService { if (!req.url?.startsWith(REALTIME_PATH)) return; const url = new URL(req.url, 'http://placeholder'); if (url.pathname !== REALTIME_PATH) return; + socket.on('error', () => socket.destroy()); + const refuseUpgrade = ( + status: + | '404 Not Found' + | '429 Too Many Requests' + | '503 Service Unavailable', + ) => { + socket.end( + `HTTP/1.1 ${status}\r\nConnection: close\r\n\r\n`, + () => socket.destroy(), + ); + }; const host = (req.headers.host ?? '').split(':')[0].toLowerCase(); if (!domain || host !== apiHost) { - // Node drops its socket error listener once a request upgrades. - socket.on('error', () => socket.destroy()); - socket.end('HTTP/1.1 404 Not Found\r\n\r\n'); + refuseUpgrade('404 Not Found'); return; } - this.#wss!.handleUpgrade(req, socket, head, (client) => { - this.#runSession(client, req, url).catch((err: unknown) => { - console.warn('[realtime] session setup failed', err); - if (client.readyState !== WebSocket.OPEN) return; - client.send( - errorEvent('internal_error', 'Realtime setup failed'), - ); - client.close(1011, 'internal error'); - }); + if ( + this.#authenticating >= MAX_AUTHENTICATING || + this.#pendingUpgrades >= MAX_AUTHENTICATING + ) { + refuseUpgrade('503 Service Unavailable'); + return; + } + const ipRequest = Object.assign(Object.create(ipApp.request), { + headers: req.headers, + connection: req.socket, + socket: req.socket, }); + const clientIp = ipRequest.ip as string; + this.#pendingUpgrades++; + let pendingUpgrade = true; + const releaseUpgrade = () => { + if (!pendingUpgrade) return; + pendingUpgrade = false; + this.#pendingUpgrades--; + socket.removeListener('close', releaseUpgrade); + }; + socket.once('close', releaseUpgrade); + void (async () => { + try { + if ( + !(await checkRateLimit( + 'realtime:preauth:node', + PREAUTH_PER_NODE, + 60_000, + 'memory', + )) || + !(await checkRateLimit( + `realtime:preauth:ip:${clientIp}`, + PREAUTH_PER_IP, + 60_000, + 'memory', + )) + ) { + refuseUpgrade('429 Too Many Requests'); + return; + } + if (socket.destroyed) return; + if (this.#authenticating >= MAX_AUTHENTICATING) { + refuseUpgrade('503 Service Unavailable'); + return; + } + this.#wss!.handleUpgrade(req, socket, head, (client) => { + this.#authenticating++; + this.#runSession( + client, + req, + url, + clientIp, + releaseUpgrade, + ) + .catch((err: unknown) => { + console.warn( + '[realtime] session setup failed', + err, + ); + if (client.readyState !== WebSocket.OPEN) + return; + client.send( + errorEvent( + 'internal_error', + 'Realtime setup failed', + ), + ); + client.close(1011, 'internal error'); + }) + .finally(() => { + this.#authenticating--; + }); + }); + } catch { + refuseUpgrade('503 Service Unavailable'); + } + })(); }); this.#heartbeatTimer = setInterval(() => { @@ -654,6 +744,8 @@ export class OpenAIRealtimeRelayService extends PuterService { client: WebSocket, req: IncomingMessage, url: URL, + clientIp: string, + onAdmitted: () => void, ): Promise { // Buffer what the client sends while auth and the upstream handshake // run; OpenAI clients send `session.update` immediately on open. @@ -774,7 +866,11 @@ export class OpenAIRealtimeRelayService extends PuterService { if (kv.get(realtimeAuthKey(token)) !== undefined) return; rechecking = true; try { - const next = await this.#authenticateCached(token, req); + const next = await this.#authenticateCached( + token, + req, + clientIp, + ); if (closed) return; if (!next || next.user?.uuid !== actor.user?.uuid) { reject( @@ -1174,12 +1270,44 @@ export class OpenAIRealtimeRelayService extends PuterService { return; } + const ipGate = { allow: true, ip: clientIp }; + await this.clients.event.emitAndWait('ip.validate', ipGate, {}); + if (closed) return; + if (!ipGate.allow) { + reject('permission_denied', 'Connection refused'); + return; + } + try { + if (!token || Buffer.byteLength(token) > MAX_TOKEN_BYTES) + throw new Error('invalid token'); + this.services.token.verify('auth', token); + } catch { + reject('invalid_api_key', 'Authentication failed'); + return; + } if ( - token && + url.searchParams.get('purpose') === 'handover' && !(await checkRateLimit( - realtimeConnectKey(token), + `${realtimeConnectKey(token!)}:handover`, + HANDOVER_CONNECTS_PER_MINUTE, + 60_000, + undefined, + { failOpen: false }, + )) + ) { + reject( + 'rate_limit_exceeded', + 'Too many realtime handover attempts', + ); + return; + } + if ( + !(await checkRateLimit( + realtimeConnectKey(token!), CONNECTS_PER_MINUTE, 60_000, + undefined, + { failOpen: false }, )) ) { reject( @@ -1190,7 +1318,9 @@ export class OpenAIRealtimeRelayService extends PuterService { } if (closed) return; try { - actor = token ? await this.#authenticateCached(token, req) : null; + actor = token + ? await this.#authenticateCached(token, req, clientIp) + : null; } catch (err) { // A lookup that failed says nothing about the credential. console.warn('[realtime] credential lookup failed', err); @@ -1390,6 +1520,7 @@ export class OpenAIRealtimeRelayService extends PuterService { // The holds are released once this charge is recorded. meter(unreportedResponseRows(billingModel, inFlight), ...holds); }); + onAdmitted(); } async #record( @@ -1415,19 +1546,25 @@ export class OpenAIRealtimeRelayService extends PuterService { #authenticateCached( token: string, req: IncomingMessage, + clientIp: string, ): Promise { const key = realtimeAuthKey(token); let check = this.#authChecks.get(key); if (!check) { check = (async () => { try { - const actor = await this.#authenticate(token, req); - if (actor) { + const actor = await this.#authenticate( + token, + req, + clientIp, + ); + if (actor && this.#authChecks.get(key) === check) { kv.set(key, actor.user!.uuid, { EX: AUTH_RECHECK_SEC }); } return actor; } finally { - this.#authChecks.delete(key); + if (this.#authChecks.get(key) === check) + this.#authChecks.delete(key); } })(); this.#authChecks.set(key, check); @@ -1438,10 +1575,12 @@ export class OpenAIRealtimeRelayService extends PuterService { async #authenticate( token: string, req: IncomingMessage, + clientIp: string, ): Promise { const ua = req.headers['user-agent']; const result = await this.services.auth.authenticate(token, { - ip: req.socket.remoteAddress, + ip: clientIp, + touch: false, userAgent: typeof ua === 'string' ? ua : undefined, }); if (result.blocked) return null; diff --git a/src/backend/services/realtime/README.md b/src/backend/services/realtime/README.md index 9d949f2a7..d92eb77fd 100644 --- a/src/backend/services/realtime/README.md +++ b/src/backend/services/realtime/README.md @@ -68,7 +68,17 @@ Each response holds a fixed amount of the account's budget until its usage is recorded, so concurrent sessions can't overdraw it; a session opens only if one hold fits. Input transcription is priced at the costlier of the transcription models in force at commit and at completion. Session attempts -are limited per credential before authentication runs. +are limited per credential after signature verification. Local, bounded IP and +node admission limits protect signature verification and authentication lookups. +Upgrades use Express's configured proxy trust policy and the `ip.validate` gate. +Authentication uses `touch: false` on both initial connection and rechecks, so +opening or maintaining a realtime socket does not update session activity metadata. + +Text, transcripts and raw audio are forwarded without new content logging. +The relay does not invoke `ai.prompt.validate`. A realtime adapter needs the host +listeners' chat payload and denial contracts before that hook can safely apply +to instructions and conversation text. Raw audio and generated transcripts need +separate policy decisions. Credential checks run in batches of at most 16 with no overlapping sweeps. When the drain starts, open sessions get a `puter.reconnect` event and new diff --git a/src/docs/src/AI/realtime.md b/src/docs/src/AI/realtime.md index 27bd884b9..e7a621e29 100644 --- a/src/docs/src/AI/realtime.md +++ b/src/docs/src/AI/realtime.md @@ -30,6 +30,8 @@ Browser WebSocket clients pass `realtime` and `openai-insecure-api-key.