From 909d87560fbc7c076bae446b3aa9dce4cb4a2ef0 Mon Sep 17 00:00:00 2001 From: 404oops Date: Sun, 1 Nov 2026 18:45:18 +0100 Subject: [PATCH] fix(realtime): evict revoked credentials and recheck paid eligibility --- src/backend/clients/event/types.ts | 6 + .../OpenAIRealtimeRelayService.test.ts | 223 ++++++++++++++++++ .../realtime/OpenAIRealtimeRelayService.ts | 166 +++++++++++-- src/backend/services/realtime/README.md | 5 + src/docs/src/rate-limits-and-quotas.md | 2 +- 5 files changed, 375 insertions(+), 27 deletions(-) diff --git a/src/backend/clients/event/types.ts b/src/backend/clients/event/types.ts index 03fa20e55..4948735a8 100644 --- a/src/backend/clients/event/types.ts +++ b/src/backend/clients/event/types.ts @@ -826,6 +826,12 @@ export type EventMap = { /** One access token was revoked; only its own connections should drop. */ 'auth.access-token.revoked': { token_uid: string }; + /** Same-cluster eviction for raw realtime WebSockets. */ + 'pubsub.realtime.auth-revoked': { + sessionUids?: string[]; + tokenUid?: string; + }; + /** * A grant was withdrawn, so whatever was standing on it has to be settled * rather than left to fail its next check. diff --git a/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts b/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts index 812b57dd1..d6354bdf5 100644 --- a/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts +++ b/src/backend/services/realtime/OpenAIRealtimeRelayService.test.ts @@ -16,6 +16,9 @@ import { vi, } from 'vitest'; import WebSocket, { WebSocketServer } from 'ws'; +import { EventClient } from '../../clients/event/EventClient.js'; +import { BroadcastService } from '../broadcast/BroadcastService.js'; +import { FULL_API_ACCESS } from '../permission/consts.js'; import { makeActor } from '../../core/actor.js'; import { checkRateLimit } from '../../core/http/middleware/rateLimit.js'; import { withAiCostFactorHook } from '../../drivers/integrationTestUtil'; @@ -1444,6 +1447,226 @@ describe('OpenAIRealtimeRelayService relay', () => { ).toBe(true); recovery.ws.close(); }); + + it('evicts a revoked session through the auth API and preserves another session', async () => { + const row = await server.stores.user.getByUsername(user.username); + const first = await server.services.auth.createSessionToken(row!, {}); + const second = await server.services.auth.createSessionToken(row!, {}); + const a = open(`${relayUrl}?auth_token=${first.token}`); + const b = open(`${relayUrl}?auth_token=${second.token}`); + await Promise.all([ + a.waitFor('session.created'), + b.waitFor('session.created'), + ]); + await server.services.auth.revokeSession(first.session.uuid); + expect(await a.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await a.closed; + expect(kv.get(realtimeAuthKey(first.token))).toBeUndefined(); + b.ws.send(JSON.stringify({ type: 'still.valid' })); + await b.waitForEcho('still.valid'); + b.ws.close(); + }); + + it('detects a downgrade without a response or auth cache expiry', async () => { + const c = open(`${relayUrl}?auth_token=${user.token}`); + await c.waitFor('session.created'); + planId = 'user_free'; + const row = await server.stores.user.getByUsername(user.username); + server.services.metering.invalidateActorSubscription(row!.uuid); + await server.services.openaiRealtime.recheckSessions(); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'subscription_required' }, + }); + await c.closed; + }); + + it('evicts only the access token revoked through the auth API', async () => { + const row = await server.stores.user.getByUsername(user.username); + const actor = makeActor({ user: row! }); + const token = await server.services.auth.createAccessToken(actor, [ + [FULL_API_ACCESS], + ]); + const a = open(`${relayUrl}?auth_token=${token}`); + const b = open(`${relayUrl}?auth_token=${user.token}`); + await Promise.all([ + a.waitFor('session.created'), + b.waitFor('session.created'), + ]); + await server.services.auth.revokeAccessToken(actor, token); + expect(await a.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await a.closed; + b.ws.send(JSON.stringify({ type: 'still.valid' })); + await b.waitForEcho('still.valid'); + b.ws.close(); + }); + + it('closes with a retryable error after a sustained verification outage', async () => { + const c = open(`${relayUrl}?auth_token=${user.token}`); + await c.waitFor('session.created'); + vi.spyOn(server.clients.redis, 'get').mockRejectedValue( + new Error('lookup down'), + ); + vi.spyOn(server.clients.db, 'read').mockRejectedValue( + new Error('lookup down'), + ); + vi.spyOn(server.clients.db, 'pread').mockRejectedValue( + new Error('lookup down'), + ); + kv.del(realtimeAuthKey(user.token)); + await server.services.openaiRealtime.recheckSessions(); + expect(c.ws.readyState).toBe(WebSocket.OPEN); + const now = Date.now(); + vi.spyOn(Date, 'now').mockReturnValue(now + 61_000); + await server.services.openaiRealtime.recheckSessions(); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'internal_error' }, + }); + expect((await c.closed).code).toBe(1011); + }); + + it('detects an auth revocation when its event is missed', async () => { + const c = open(`${relayUrl}?auth_token=${user.token}`); + await c.waitFor('session.created'); + const identity = server.services.token.verify<{ session_uid: string }>( + 'auth', + user.token, + ); + const emit = server.clients.event.emit.bind(server.clients.event); + vi.spyOn(server.clients.event, 'emit').mockImplementation((...args) => { + if (args[0] !== 'outer.auth.sessions.revoked') return emit(...args); + }); + await server.services.auth.revokeSession(identity.session_uid); + expect(c.ws.readyState).toBe(WebSocket.OPEN); + vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 61_000); + await server.services.openaiRealtime.recheckSessions(); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await c.closed; + }); + + it('detects a paid downgrade after a missed plan invalidation while free credit remains', async () => { + const holder = await createTestUser(server, { + username: 'missed-plan-event', + password: 'missed-plan-event-password', + }); + const c = open(`${relayUrl}?auth_token=${holder.token}`); + await c.waitFor('session.created'); + planId = 'user_free'; + vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 61_000); + const row = await server.stores.user.getByUsername(holder.username); + expect( + await server.services.metering.getRemainingUsage({ + user: row, + } as never), + ).toBeGreaterThan(0); + await server.services.openaiRealtime.recheckSessions(); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'subscription_required' }, + }); + await c.closed; + }); + + it('bounds an entitlement lookup outage without reporting invalid credentials', async () => { + const c = open(`${relayUrl}?auth_token=${user.token}`); + await c.waitFor('session.created'); + const metering = server.services.metering; + const resolvers = ( + metering as unknown as { + subscriptionResolvers: (() => Promise)[]; + } + ).subscriptionResolvers; + const fail = async (): Promise => { + throw new Error('plan lookup down'); + }; + resolvers.unshift(fail); + const row = await server.stores.user.getByUsername(user.username); + metering.invalidateActorSubscription(row!.uuid); + try { + await server.services.openaiRealtime.recheckSessions(); + expect(c.ws.readyState).toBe(WebSocket.OPEN); + vi.spyOn(Date, 'now').mockReturnValue(Date.now() + 61_000); + await server.services.openaiRealtime.recheckSessions(); + expect(await c.waitFor('error')).toMatchObject({ + error: { code: 'internal_error' }, + }); + expect((await c.closed).code).toBe(1011); + } finally { + resolvers.splice(resolvers.indexOf(fail), 1); + metering.invalidateActorSubscription(row!.uuid); + } + }); + + it('evicts a raw socket on a sibling node through the real cluster pubsub path', async () => { + const config = { + domain: 'puter.localhost', + realtime: { + upstreamUrl: `ws://127.0.0.1:${(upstream.address() as AddressInfo).port}`, + }, + providers: { 'openai-completion': { apiKey: 'sk-test' } }, + broadcast: { + peers: [], + webhook: { + peerId: 'relay-sibling', + secret: 'test-cluster-secret', + }, + }, + } as IConfig; + const event = new EventClient(config); + const remoteClients = { ...server.clients, event }; + const broadcast = new BroadcastService( + config, + remoteClients, + server.stores, + server.services, + ); + const relay = new OpenAIRealtimeRelayService( + config, + remoteClients, + server.stores, + server.services, + ); + const http = createServer(); + relay.attachHttpServer(http); + broadcast.onServerStart(); + relay.onServerStart(); + await new Promise((resolve) => + http.listen(0, '127.0.0.1', resolve), + ); + const row = await server.stores.user.getByUsername(user.username); + const revoked = await server.services.auth.createSessionToken(row!, {}); + const url = (token: string) => + `ws://127.0.0.1:${(http.address() as AddressInfo).port}${REALTIME_PATH}?auth_token=${token}`; + const a = open(url(revoked.token), undefined, { + host: 'api.puter.localhost', + }); + const b = open(url(user.token), undefined, { + host: 'api.puter.localhost', + }); + try { + await Promise.all([ + a.waitFor('session.created'), + b.waitFor('session.created'), + ]); + await server.services.auth.revokeSession(revoked.session.uuid); + expect(await a.waitFor('error')).toMatchObject({ + error: { code: 'invalid_api_key' }, + }); + await a.closed; + b.ws.send(JSON.stringify({ type: 'sibling.valid' })); + await b.waitForEcho('sibling.valid'); + } finally { + a.ws.terminate(); + b.ws.terminate(); + await relay.onServerPrepareShutdown(); + await broadcast.onServerPrepareShutdown(); + await new Promise((resolve) => http.close(() => resolve())); + } + }); it('bounds unverified socket allocation while rejected clients withhold close acknowledgments', async () => { const relay = new OpenAIRealtimeRelayService( { domain: 'puter.localhost', trust_proxy: 'loopback' } as IConfig, diff --git a/src/backend/services/realtime/OpenAIRealtimeRelayService.ts b/src/backend/services/realtime/OpenAIRealtimeRelayService.ts index 24f828b71..0a5ee3713 100644 --- a/src/backend/services/realtime/OpenAIRealtimeRelayService.ts +++ b/src/backend/services/realtime/OpenAIRealtimeRelayService.ts @@ -37,7 +37,10 @@ import { import { readProviderKey } from '../../drivers/util/providerRegistry.js'; import { runWithConcurrencyLimit } from '../../util/concurrency.js'; import { kv } from '../../util/kvSingleton.js'; -import { assertActorHasSubscription } from '../metering/enforcement.js'; +import { + assertActorHasSubscription, + subscriptionEnforcementEnabled, +} from '../metering/enforcement.js'; import type { CreditHold, UsageInput } from '../metering/types.js'; import { decideSocketAuth } from '../socket/SocketService.js'; import type { ServerDrainInfo } from '../../types.js'; @@ -532,6 +535,8 @@ export class OpenAIRealtimeRelayService extends PuterService { close: (code: number, reason: string) => void; heartbeat: () => void; recheck: () => Promise; + revoke: (sessionUids?: ReadonlySet, tokenUid?: string) => void; + planChanged: (userUuid: string) => void; drain: () => Promise; notify: (closeInMs: number | undefined) => void; }>(); @@ -688,6 +693,35 @@ export class OpenAIRealtimeRelayService extends PuterService { /** A host extension can announce a restart before the drain starts. */ override onServerStart(): void { + this.clients.event.on('outer.auth.sessions.revoked', (_key, data) => { + this.clients.event.emit( + 'pubsub.realtime.auth-revoked', + { sessionUids: data.session_uids }, + {}, + ); + }); + this.clients.event.on('auth.access-token.revoked', (_key, data) => { + this.clients.event.emit( + 'pubsub.realtime.auth-revoked', + { tokenUid: data.token_uid }, + {}, + ); + }); + // Raw WebSockets do not participate in socket.io's cluster adapter. + this.clients.event.on('pubsub.realtime.auth-revoked', (_key, data) => { + const revoked = data.sessionUids + ? new Set(data.sessionUids) + : undefined; + for (const session of this.#sessions) + session.revoke(revoked, data.tokenUid); + }); + this.clients.event.on( + 'outer.pubsub.metering.subscription-changed', + (_key, data) => { + for (const session of this.#sessions) + session.planChanged(data.userUuid); + }, + ); this.clients.event.on('serverRestartAnnounced', (_key, data) => this.#announceRestart(data?.closeInMs), ); @@ -762,7 +796,9 @@ export class OpenAIRealtimeRelayService extends PuterService { let settleTimer: ReturnType | null = null; const token = realtimeTokenFrom(req, url); let actor: Actor | null = null; - let rechecking = false; + let rechecking: Promise | undefined; + let identity: AnyTokenPayload | undefined; + let verificationFailedAt: number | undefined; // Responses started upstream and not yet done. A response bills its // usage on `response.done`, so the upstream outlives the client until // then; one the upstream closes on first is billed its hold. @@ -831,6 +867,34 @@ export class OpenAIRealtimeRelayService extends PuterService { }; const session = { + revoke: (sessionUids?: ReadonlySet, tokenUid?: string) => { + const sessionUid = + identity?.session_uid ?? + (identity && + (identity.type === 'session' || identity.type === 'gui') + ? identity.uuid + : undefined); + const parentUid = + identity?.type === 'app-under-user' + ? identity.session + : undefined; + if (!( + sessionUids?.has(sessionUid ?? '') || + sessionUids?.has(parentUid ?? '') || + (tokenUid && + identity?.type === 'access-token' && + identity.token_uid === tokenUid) + )) + return; + if (token) { + kv.del(realtimeAuthKey(token)); + this.#authChecks.delete(realtimeAuthKey(token)); + } + reject('invalid_api_key', 'Credentials are no longer valid'); + }, + planChanged: (userUuid: string) => { + if (actor?.user?.uuid === userUuid) void session.recheck(); + }, close: (code: number, reason: string) => { if (closed) return; closed = true; @@ -861,30 +925,77 @@ export class OpenAIRealtimeRelayService extends PuterService { }, // Called every AUTH_SWEEP_MS; re-authenticates once the token's kv // entry has expired. A failed lookup is retried on the next sweep. - recheck: async () => { - if (!actor || !token || closed || rechecking) return; - if (kv.get(realtimeAuthKey(token)) !== undefined) return; - rechecking = true; - try { - const next = await this.#authenticateCached( - token, - req, - clientIp, - ); - if (closed) return; - if (!next || next.user?.uuid !== actor.user?.uuid) { - reject( - 'invalid_api_key', - 'Credentials are no longer valid', - ); - return; + recheck: () => { + if (!actor || !token || closed) return Promise.resolve(); + if (rechecking) return rechecking; + rechecking = (async () => { + try { + const next = + kv.get(realtimeAuthKey(token)) !== undefined + ? actor + : await this.#authenticateCached( + token, + req, + clientIp, + ); + if (closed) return; + if (!next || next.user?.uuid !== actor.user?.uuid) { + reject( + 'invalid_api_key', + 'Credentials are no longer valid', + ); + return; + } + actor = next; + if (subscriptionEnforcementEnabled(this.config)) { + if (await planProvisional()) + throw new Error('subscription unavailable'); + try { + await assertActorHasSubscription( + this.services.metering, + actor, + true, + this.config, + { + events: this.clients.event, + surface: 'realtime.openai', + }, + ); + } catch (error) { + if ( + [402, 403].includes( + (error as { statusCode?: number }) + .statusCode ?? 0, + ) && + !(await planProvisional()) + ) { + reject( + 'subscription_required', + 'Realtime conversations require a subscription', + ); + return; + } + throw error; + } + } + verificationFailedAt = undefined; + } catch { + verificationFailedAt ??= Date.now(); + if ( + Date.now() - verificationFailedAt >= + AUTH_RECHECK_SEC * 1000 + ) { + reject( + 'internal_error', + 'Could not verify continued access', + 1011, + ); + } + } finally { + rechecking = undefined; } - actor = next; - } catch (err) { - console.warn('[realtime] credential recheck failed', err); - } finally { - rechecking = false; - } + })(); + return rechecking; }, drain: async () => { if (inFlight > 0 && !closed) { @@ -1280,7 +1391,10 @@ export class OpenAIRealtimeRelayService extends PuterService { try { if (!token || Buffer.byteLength(token) > MAX_TOKEN_BYTES) throw new Error('invalid token'); - this.services.token.verify('auth', token); + identity = this.services.token.verify( + 'auth', + token, + ); } catch { reject('invalid_api_key', 'Authentication failed'); return; diff --git a/src/backend/services/realtime/README.md b/src/backend/services/realtime/README.md index d92eb77fd..a23e6d7d7 100644 --- a/src/backend/services/realtime/README.md +++ b/src/backend/services/realtime/README.md @@ -81,6 +81,11 @@ 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. +Session and access-token revocation events close affected sockets and invalidate +their auth cache entries. Subscription eligibility is checked independently of +authentication and after subscription invalidations. A sustained verification +outage gets a 60-second retry grace from the first failed lookup, then closes +with a retryable internal error. Periodic checks remain the missed-event fallback. When the drain starts, open sessions get a `puter.reconnect` event and new sessions receive `server_restarting` and close with 1012. Open sessions can finish an in-flight response for up to five seconds at shutdown, then close diff --git a/src/docs/src/rate-limits-and-quotas.md b/src/docs/src/rate-limits-and-quotas.md index 243a66621..b83b9cb26 100644 --- a/src/docs/src/rate-limits-and-quotas.md +++ b/src/docs/src/rate-limits-and-quotas.md @@ -110,7 +110,7 @@ Realtime sessions ([`puter.ai.convo()`](/AI/convo/), `/puterai/openai/v1/realtim | Session setup timeout | 15 seconds | Error `connection_timeout`, close 1011. | | Upstream handshake timeout | 10 seconds | Close 1011. | | Heartbeat interval | 30 seconds | A missing pong terminates the connection on the next heartbeat. | -| Credential cache and recheck | Cache 60 seconds, sweep every 15 seconds, at most 16 checks in parallel | Invalid credentials close with 1008. A check that fails at connection gets `internal_error` and closes with 1011. A failed recheck keeps the session and tries again on the next sweep. Overlapping sweeps are skipped. | +| Credential cache and recheck | Cache 60 seconds, sweep every 15 seconds, at most 16 checks in parallel | Invalid credentials close with 1008. A check that fails at connection gets `internal_error` and closes with 1011. A failed recheck retries on the next sweep for up to 60 seconds from the first failed lookup, then closes with `internal_error` and 1011. Subscription eligibility is checked independently on each sweep and after a plan invalidation. Revocation events close the affected credentials. Overlapping sweeps are skipped. | | Session attempts per signature-verified credential | 60 per minute | Error `rate_limit_exceeded`, close 1008. | | Pre-auth attempts | 120 per client IP and 600 per node per minute | Local memory admission runs before the WebSocket handshake and signature verification. Exceeding either rate returns HTTP 429. Client IP follows the configured proxy trust policy. The shared memory limiter caps its state at 10,000 keys with a 1-hour retention floor and a 60-second sweep. | | Unauthenticated upgrade sockets, including closing sockets | 64 per node | Further upgrades receive HTTP 503. The slot is retained until successful admission or transport closure. |