diff --git a/src/main/runtime/relay/desktop-relay-service.ts b/src/main/runtime/relay/desktop-relay-service.ts index a08121ac8c1..d2430d3c6fd 100644 --- a/src/main/runtime/relay/desktop-relay-service.ts +++ b/src/main/runtime/relay/desktop-relay-service.ts @@ -65,8 +65,7 @@ export class DesktopRelayService { mobileSocketWiring, isCurrent, refreshAccessToken, - onStatus: options.onStatus, - onResolveDirector: () => this.coordinator.restart() + onStatus: options.onStatus }) void this.flushRevokeOutbox(broker) return broker diff --git a/src/main/runtime/relay/relay-auth-coordinator.ts b/src/main/runtime/relay/relay-auth-coordinator.ts index 0bdd99c6687..7d244afe76f 100644 --- a/src/main/runtime/relay/relay-auth-coordinator.ts +++ b/src/main/runtime/relay/relay-auth-coordinator.ts @@ -63,11 +63,6 @@ export class RelayAuthCoordinator { this.options.onStatus('offline') } - restart(): void { - this.fenceAndCloseNow() - this.reconcile() - } - getActiveBroker(): CoordinatedRelayBroker | null { return this.ownership?.valid ? this.ownership.broker : null } diff --git a/src/main/runtime/relay/relay-control-client.ts b/src/main/runtime/relay/relay-control-client.ts index d431c13d54a..44f3d65798c 100644 --- a/src/main/runtime/relay/relay-control-client.ts +++ b/src/main/runtime/relay/relay-control-client.ts @@ -106,6 +106,10 @@ export class RelayControlClient { }) } + get pendingRequestCount(): number { + return this.requests.size + } + refreshAuthorization(relayJwt: string): void { this.sendActive({ type: 'auth-refresh', relayJwt }) } diff --git a/src/main/runtime/relay/relay-control-origin.ts b/src/main/runtime/relay/relay-control-origin.ts new file mode 100644 index 00000000000..ef607c9862b --- /dev/null +++ b/src/main/runtime/relay/relay-control-origin.ts @@ -0,0 +1,243 @@ +import type WebSocket from 'ws' +import type { E2EEKeypair } from '../e2ee-keypair' +import { CloudRelayTransport } from '../rpc/relay-transport' +import type { MobileSocketWiring } from '../rpc/mobile-socket-wiring' +import { RelayControlClient } from './relay-control-client' +import type { + RelayConnectionOpenMessage, + RelayDrainMessage, + RelayHostHelloAckMessage +} from './relay-control-protocol' +import type { RelayIdentity } from './relay-session-broker-contract' +import type { RelayAssignment } from './relay-http-client' + +type RelayControlOriginOptions = { + assignment: RelayAssignment + relayJwt: string + relayHostId: string + identity: RelayIdentity + keypair: E2EEKeypair + appVersion: string + mobileSocketWiring: MobileSocketWiring + createControlSocket?: (url: string, relayJwt: string) => WebSocket + createDataSocket?: (url: string) => WebSocket + onConnectionOwned: (connectionId: string, origin: RelayControlOrigin) => void + onConnectionReleased: (connectionId: string, origin: RelayControlOrigin) => void + onDrain: (origin: RelayControlOrigin, message: RelayDrainMessage) => void + onClose: (origin: RelayControlOrigin, code: number) => void +} + +export class RelayControlOrigin { + readonly assignment: RelayAssignment + readonly transport: CloudRelayTransport + private readonly options: RelayControlOriginOptions + private readonly controls = new Set() + private readonly retiredControlTimers = new Map< + RelayControlClient, + ReturnType + >() + private activeControl: RelayControlClient | null = null + private generation = 0 + private controlResumeSecret: string | null = null + private leaseExpiresAt = 0 + private acceptingConnections = true + private closed = false + + constructor(options: RelayControlOriginOptions) { + this.options = options + this.assignment = options.assignment + this.transport = new CloudRelayTransport({ + cellUrl: options.assignment.cellUrl, + relayHostId: options.relayHostId, + generation: 0, + createSocket: options.createDataSocket, + onConnectionClosed: (connectionId) => options.onConnectionReleased(connectionId, this) + }) + options.mobileSocketWiring.attachTransport(this.transport, (ws) => + this.transport.metadataFor(ws) + ) + } + + get control(): RelayControlClient { + if (!this.activeControl) { + throw new Error('relay_control_not_active') + } + return this.activeControl + } + + get availableControl(): RelayControlClient | null { + return this.activeControl + } + + get cellUrl(): string { + return this.assignment.cellUrl + } + + get assignmentEpoch(): number { + return this.assignment.assignmentEpoch + } + + get controlLeaseExpiresAt(): number { + return this.leaseExpiresAt + } + + get pendingRequestCount(): number { + let count = 0 + for (const control of this.controls) { + count += control.pendingRequestCount + } + return count + } + + async open(): Promise { + await this.transport.start() + const { control, ack } = await this.openControl() + this.activate(control, ack) + } + + async rebind(relayJwt: string, assignment: RelayAssignment): Promise { + if (assignment.cellUrl !== this.cellUrl || !this.controlResumeSecret || this.generation <= 0) { + throw new Error('relay_control_rebind_origin_mismatch') + } + const previous = this.activeControl + const { control, ack } = await this.openControl({ + relayJwt, + assignmentEpoch: assignment.assignmentEpoch, + previousGeneration: this.generation, + controlResumeSecret: this.controlResumeSecret + }) + this.activate(control, ack) + this.acceptingConnections = true + // Why: the resumed control owns the same server generation and splices; + // the predecessor remains only long enough for any idempotent reply in flight. + if (previous && previous.pendingRequestCount === 0) { + this.closeRetiredControl(previous) + } else if (previous) { + // Why: basis-bound requests keep their original control through its + // bounded request deadline; afterward the resumed control is sole owner. + this.retiredControlTimers.set( + previous, + setTimeout(() => this.closeRetiredControl(previous), 10_100) + ) + } + } + + markDraining(): void { + // The relay changes the control's protocol state when it sends drain. This + // marker exists for the broker's ownership policy, not a second wire event. + this.acceptingConnections = false + } + + refreshAuthorization(relayJwt: string): void { + for (const control of this.controls) { + try { + control.refreshAuthorization(relayJwt) + } catch { + // A closing drain-only origin cannot block refresh on the active target. + } + } + } + + async close(): Promise { + if (this.closed) { + return + } + this.closed = true + for (const timer of this.retiredControlTimers.values()) { + clearTimeout(timer) + } + this.retiredControlTimers.clear() + for (const control of this.controls) { + control.closeNow() + } + this.controls.clear() + this.activeControl = null + await this.transport.stop() + } + + closeNow(): void { + void this.close() + } + + private async openControl(overrides?: { + relayJwt: string + assignmentEpoch: number + previousGeneration: number + controlResumeSecret: string + }): Promise<{ control: RelayControlClient; ack: RelayHostHelloAckMessage }> { + let control!: RelayControlClient + control = new RelayControlClient({ + cellUrl: this.cellUrl, + relayJwt: overrides?.relayJwt ?? this.options.relayJwt, + relayHostId: this.options.relayHostId, + assignmentEpoch: overrides?.assignmentEpoch ?? this.assignmentEpoch, + identity: this.options.identity, + keypair: this.options.keypair, + appVersion: this.options.appVersion, + ...(overrides + ? { + previousGeneration: overrides.previousGeneration, + controlResumeSecret: overrides.controlResumeSecret + } + : {}), + onConnectionOpen: (message) => this.openConnection(message), + onDrain: (message) => this.options.onDrain(this, message), + onClose: (code) => { + this.controls.delete(control) + const timer = this.retiredControlTimers.get(control) + if (timer) { + clearTimeout(timer) + this.retiredControlTimers.delete(control) + } + if (this.activeControl === control) { + this.activeControl = null + this.options.onClose(this, code) + } + }, + createSocket: this.options.createControlSocket + }) + this.controls.add(control) + try { + return { control, ack: await control.connect() } + } catch (error) { + this.controls.delete(control) + control.closeNow() + throw error + } + } + + private closeRetiredControl(control: RelayControlClient): void { + const timer = this.retiredControlTimers.get(control) + if (timer) { + clearTimeout(timer) + this.retiredControlTimers.delete(control) + } + if (this.activeControl !== control && this.controls.delete(control)) { + control.closeNow() + } + } + + private activate(control: RelayControlClient, ack: RelayHostHelloAckMessage): void { + if (ack.generation <= 0) { + throw new Error('invalid_relay_generation') + } + this.transport.setGeneration(ack.generation) + this.generation = ack.generation + this.controlResumeSecret = ack.controlResumeSecret + this.leaseExpiresAt = ack.leaseExpiresAt + this.activeControl = control + for (const connectionId of ack.activeConnIds) { + this.options.onConnectionOwned(connectionId, this) + } + } + + private openConnection(message: RelayConnectionOpenMessage): void { + if (!this.acceptingConnections) { + return + } + this.options.onConnectionOwned(message.connId, this) + void this.transport.openConnection(message).catch(() => { + this.options.onConnectionReleased(message.connId, this) + }) + } +} diff --git a/src/main/runtime/relay/relay-control-requests.ts b/src/main/runtime/relay/relay-control-requests.ts index e67b7b090bf..2a96b94e3ec 100644 --- a/src/main/runtime/relay/relay-control-requests.ts +++ b/src/main/runtime/relay/relay-control-requests.ts @@ -25,6 +25,10 @@ export type DeviceCredentialInstallAuthorization = export class RelayControlRequests { private readonly pending = new Map() + get size(): number { + return this.pending.size + } + createInvite( reqId: string, relayDeviceId: string, diff --git a/src/main/runtime/relay/relay-origin-pool.ts b/src/main/runtime/relay/relay-origin-pool.ts new file mode 100644 index 00000000000..082f8d0778c --- /dev/null +++ b/src/main/runtime/relay/relay-origin-pool.ts @@ -0,0 +1,291 @@ +import type WebSocket from 'ws' +import type { E2EEKeypair } from '../e2ee-keypair' +import type { MobileSocketWiring } from '../rpc/mobile-socket-wiring' +import { RelayControlOrigin } from './relay-control-origin' +import type { RelayControlClient } from './relay-control-client' +import type { RelayDrainMessage } from './relay-control-protocol' +import { requestRelayAssignment, type RelayAssignment } from './relay-http-client' +import type { RelayBrokerStatus, RelayIdentity } from './relay-session-broker-contract' + +type RelayOriginPoolOptions = { + directorUrl: string + relayHostId: string + identity: RelayIdentity + keypair: E2EEKeypair + appVersion: string + mobileSocketWiring: MobileSocketWiring + isCurrent: () => boolean + onStatus: (status: RelayBrokerStatus) => void + fetch?: typeof globalThis.fetch + createControlSocket?: (url: string, relayJwt: string) => WebSocket + createDataSocket?: (url: string) => WebSocket + random?: () => number + now?: () => number +} + +export class RelayOriginPool { + private readonly options: RelayOriginPoolOptions + private activeOrigin: RelayControlOrigin | null = null + private readonly origins = new Set() + private readonly drainingOrigins = new Set() + private readonly basisOrigins = new Map() + private readonly drainTimers = new Map>() + private assignment: RelayAssignment | null = null + private relayJwt: string | null = null + private rotationTimer: ReturnType | null = null + private rotationPromise: Promise | null = null + private closed = false + + constructor(options: RelayOriginPoolOptions) { + this.options = options + } + + get activeAssignment(): RelayAssignment | null { + return this.assignment + } + + get activeControl(): RelayControlClient | null { + return this.activeOrigin?.availableControl ?? null + } + + controlForBasis(basisConnId: string): RelayControlClient | null { + return this.basisOrigins.get(basisConnId)?.availableControl ?? null + } + + async openInitial(assignment: RelayAssignment, relayJwt: string): Promise { + this.assignment = assignment + this.relayJwt = relayJwt + const origin = this.createOrigin(assignment, relayJwt) + this.origins.add(origin) + await origin.open() + this.assertCurrent() + this.activeOrigin = origin + this.scheduleControlRotation() + } + + refreshAuthorization(relayJwt: string): void { + this.relayJwt = relayJwt + for (const origin of this.origins) { + origin.refreshAuthorization(relayJwt) + } + } + + closeNow(): void { + if (this.closed) { + return + } + this.closed = true + if (this.rotationTimer) { + clearTimeout(this.rotationTimer) + this.rotationTimer = null + } + for (const timer of this.drainTimers.values()) { + clearTimeout(timer) + } + this.drainTimers.clear() + for (const origin of this.origins) { + origin.closeNow() + } + this.origins.clear() + this.drainingOrigins.clear() + this.basisOrigins.clear() + this.activeOrigin = null + } + + private createOrigin(assignment: RelayAssignment, relayJwt: string): RelayControlOrigin { + return new RelayControlOrigin({ + assignment, + relayJwt, + relayHostId: this.options.relayHostId, + identity: this.options.identity, + keypair: this.options.keypair, + appVersion: this.options.appVersion, + mobileSocketWiring: this.options.mobileSocketWiring, + createControlSocket: this.options.createControlSocket, + createDataSocket: this.options.createDataSocket, + onConnectionOwned: (connectionId, origin) => { + if (this.isCurrent() && this.origins.has(origin)) { + this.basisOrigins.set(connectionId, origin) + } + }, + onConnectionReleased: (connectionId, origin) => { + if (this.basisOrigins.get(connectionId) === origin) { + this.basisOrigins.delete(connectionId) + } + this.maybeCloseDrainedOrigin(origin) + }, + onDrain: (origin, message) => this.handleDrain(origin, message), + onClose: (origin) => { + if (origin === this.activeOrigin && this.isCurrent()) { + this.options.onStatus('offline') + this.handleDrain(origin, { + type: 'drain', + graceMs: 0, + recovery: 'resolve-director' + }) + } + } + }) + } + + private handleDrain(origin: RelayControlOrigin, message: RelayDrainMessage): void { + if (!this.isCurrent() || origin !== this.activeOrigin) { + return + } + origin.markDraining() + this.drainingOrigins.add(origin) + this.options.onStatus('draining') + if (!this.rotationPromise) { + this.rotationPromise = this.resolveDrainTarget(origin, message).finally(() => { + this.rotationPromise = null + }) + } + } + + private async resolveDrainTarget( + origin: RelayControlOrigin, + message: RelayDrainMessage + ): Promise { + try { + if (!this.relayJwt) { + throw new Error('relay_authorization_unavailable') + } + // Why: only the configured director can choose a migration target. + const assignment = await requestRelayAssignment({ + directorUrl: this.options.directorUrl, + relayToken: this.relayJwt, + relayHostId: this.options.relayHostId, + fetch: this.options.fetch + }) + this.assertCurrent() + if (assignment.cellUrl === origin.cellUrl) { + await origin.rebind(this.relayJwt, assignment) + this.assertCurrent() + this.activeOrigin = origin + this.assignment = assignment + this.drainingOrigins.delete(origin) + } else { + await this.activateTarget(origin, assignment, this.relayJwt, message.graceMs) + } + this.options.onStatus('registered') + this.scheduleControlRotation() + } catch { + if (this.isCurrent()) { + const random = this.options.random ?? Math.random + setTimeout(() => this.handleDrain(origin, message), 250 + Math.floor(random() * 751)) + } + } + } + + private async activateTarget( + origin: RelayControlOrigin, + assignment: RelayAssignment, + relayJwt: string, + graceMs: number + ): Promise { + const target = this.createOrigin(assignment, relayJwt) + this.origins.add(target) + try { + await target.open() + this.assertCurrent() + } catch (error) { + this.origins.delete(target) + target.closeNow() + throw error + } + this.activeOrigin = target + this.assignment = assignment + this.scheduleDrainDeadline(origin, graceMs) + this.maybeCloseDrainedOrigin(origin) + } + + private scheduleControlRotation(): void { + if (this.rotationTimer) { + clearTimeout(this.rotationTimer) + } + const origin = this.activeOrigin + if (!origin || this.closed) { + this.rotationTimer = null + return + } + const now = (this.options.now ?? Date.now)() + const random = this.options.random ?? Math.random + const earlyMs = 60_000 + Math.floor(random() * 60_001) + const delay = Math.max(0, origin.controlLeaseExpiresAt - earlyMs - now) + this.rotationTimer = setTimeout(() => void this.rebindActiveControl(origin), delay) + } + + private async rebindActiveControl(origin: RelayControlOrigin): Promise { + this.rotationTimer = null + if (!this.isCurrent() || origin !== this.activeOrigin || this.rotationPromise) { + return + } + if (!this.relayJwt || !this.assignment) { + return + } + try { + await origin.rebind(this.relayJwt, this.assignment) + this.assertCurrent() + this.scheduleControlRotation() + } catch { + if (this.isCurrent() && origin === this.activeOrigin) { + const random = this.options.random ?? Math.random + this.rotationTimer = setTimeout( + () => void this.rebindActiveControl(origin), + 5_000 + Math.floor(random() * 10_001) + ) + } + } + } + + private scheduleDrainDeadline(origin: RelayControlOrigin, graceMs: number): void { + const existing = this.drainTimers.get(origin) + if (existing) { + clearTimeout(existing) + } + this.drainTimers.set( + origin, + setTimeout(() => this.closeOrigin(origin), graceMs) + ) + } + + private maybeCloseDrainedOrigin(origin: RelayControlOrigin): void { + if ( + !this.drainingOrigins.has(origin) || + origin.pendingRequestCount > 0 || + [...this.basisOrigins.values()].includes(origin) + ) { + return + } + this.closeOrigin(origin) + } + + private closeOrigin(origin: RelayControlOrigin): void { + if (origin === this.activeOrigin) { + return + } + const timer = this.drainTimers.get(origin) + if (timer) { + clearTimeout(timer) + this.drainTimers.delete(origin) + } + for (const [connectionId, owner] of this.basisOrigins) { + if (owner === origin) { + this.basisOrigins.delete(connectionId) + } + } + this.drainingOrigins.delete(origin) + this.origins.delete(origin) + origin.closeNow() + } + + private assertCurrent(): void { + if (!this.isCurrent()) { + throw new Error('stale_relay_origin_pool') + } + } + + private isCurrent(): boolean { + return !this.closed && this.options.isCurrent() + } +} diff --git a/src/main/runtime/relay/relay-session-broker-contract.ts b/src/main/runtime/relay/relay-session-broker-contract.ts index c6e8f03bafd..3c7a707e5bd 100644 --- a/src/main/runtime/relay/relay-session-broker-contract.ts +++ b/src/main/runtime/relay/relay-session-broker-contract.ts @@ -22,7 +22,6 @@ export type RelaySessionBrokerOptions = { isCurrent: () => boolean refreshAccessToken: () => Promise onStatus: (status: RelayBrokerStatus) => void - onResolveDirector: () => void fetch?: typeof globalThis.fetch createControlSocket?: (url: string, relayJwt: string) => WebSocket createDataSocket?: (url: string) => WebSocket diff --git a/src/main/runtime/relay/relay-session-broker.test.ts b/src/main/runtime/relay/relay-session-broker.test.ts index ccd46ecf87a..6223a497357 100644 --- a/src/main/runtime/relay/relay-session-broker.test.ts +++ b/src/main/runtime/relay/relay-session-broker.test.ts @@ -6,14 +6,30 @@ import type * as RelayHttpClientModule from './relay-http-client' const fakes = vi.hoisted(() => ({ controls: [] as { + options: { + onConnectionOpen(message: { + connId: string + connTicket: string + kind: 'invite' | 'resume' + relayDeviceId: string + attachDeadlineMs: number + }): void + onDrain(message: { type: 'drain'; graceMs: number; recovery: 'resolve-director' }): void + previousGeneration?: number + controlResumeSecret?: string + } connect: ReturnType closeNow: ReturnType + confirmResume: ReturnType + installCredential: ReturnType + pendingRequestCount: number }[], transports: [] as { start: ReturnType stop: ReturnType setGeneration: ReturnType metadataFor: ReturnType + openConnection: ReturnType }[], controlConnect: vi.fn(), exchange: vi.fn(), @@ -30,8 +46,26 @@ vi.mock('./relay-control-client', () => ({ RelayControlClient: class { connect = fakes.controlConnect closeNow = vi.fn() + confirmResume = vi.fn().mockResolvedValue({ + type: 'device-resume-confirmed', + v: 1, + reqId: 'confirm-1', + currentVersion: 1, + acceptedAs: 'current', + renewed: true, + resumeExpiresAt: 100_000 + }) + installCredential = vi.fn().mockResolvedValue({ + type: 'device-credential-installed', + v: 1, + reqId: 'install-1', + authorizationMode: 'relay-basis', + currentVersion: 1, + resumeExpiresAt: 100_000 + }) + pendingRequestCount = 0 - constructor() { + constructor(readonly options: (typeof fakes.controls)[number]['options']) { fakes.controls.push(this) } } @@ -43,6 +77,7 @@ vi.mock('../rpc/relay-transport', () => ({ stop = vi.fn().mockResolvedValue(undefined) setGeneration = vi.fn() metadataFor = vi.fn() + openConnection = vi.fn().mockResolvedValue(undefined) constructor() { fakes.transports.push(this) @@ -65,7 +100,7 @@ describe('RelaySessionBroker lifecycle ownership', () => { fakes.controls.length = 0 fakes.transports.length = 0 fakes.controlConnect.mockReset() - fakes.exchange.mockReset().mockResolvedValue({ relayToken: 'relay-jwt', expiresAt: 60_000 }) + fakes.exchange.mockReset().mockResolvedValue({ relayToken: 'relay-jwt', expiresAt: 1_000_000 }) fakes.assign.mockReset().mockResolvedValue({ cellUrl: 'https://relay.example.test', assignmentEpoch: 1, @@ -94,8 +129,7 @@ describe('RelaySessionBroker lifecycle ownership', () => { mobileSocketWiring: { attachTransport: vi.fn() } as never, isCurrent: () => current, refreshAccessToken: async () => null, - onStatus: (status) => statuses.push(status), - onResolveDirector: vi.fn() + onStatus: (status) => statuses.push(status) }) await vi.waitFor(() => expect(fakes.controls).toHaveLength(1)) current = false @@ -104,7 +138,7 @@ describe('RelaySessionBroker lifecycle ownership', () => { v: 1, generation: 1, controlResumeSecret: 'A'.repeat(43), - leaseExpiresAt: 60_000, + leaseExpiresAt: 1_000_000, activeConnIds: [], pendingConns: [] }) @@ -114,4 +148,134 @@ describe('RelaySessionBroker lifecycle ownership', () => { expect(fakes.transports[0]!.stop).toHaveBeenCalledOnce() expect(statuses).toEqual(['connecting']) }) + + it('activates a new origin while keeping basis-bound work on the drained origin', async () => { + const firstAck: RelayHostHelloAckMessage = { + type: 'host-hello-ack', + v: 1, + generation: 1, + controlResumeSecret: 'A'.repeat(43), + leaseExpiresAt: 1_000_000, + activeConnIds: [], + pendingConns: [] + } + fakes.controlConnect.mockResolvedValueOnce(firstAck).mockResolvedValueOnce({ + ...firstAck, + generation: 2, + controlResumeSecret: 'B'.repeat(43) + }) + fakes.assign + .mockResolvedValueOnce({ + cellUrl: 'https://relay-c1.example.test', + assignmentEpoch: 1, + leaseExpiresAt: 1_000_000 + }) + .mockResolvedValueOnce({ + cellUrl: 'https://relay-c2.example.test', + assignmentEpoch: 2, + leaseExpiresAt: 2_000_000 + }) + const broker = await RelaySessionBroker.connect(brokerOptions({ onStatus: vi.fn() })) + fakes.controls[0]!.options.onConnectionOpen({ + connId: 'old-basis', + connTicket: 'T'.repeat(43), + kind: 'resume', + relayDeviceId: 'device-1', + attachDeadlineMs: 1_000 + }) + expect(brokerBasisIds(broker)).toEqual(['old-basis']) + fakes.controls[0]!.options.onDrain({ + type: 'drain', + graceMs: 30_000, + recovery: 'resolve-director' + }) + await vi.waitFor(() => expect(fakes.controls).toHaveLength(2)) + + expect(broker.endpoint?.cellUrl).toBe('https://relay-c2.example.test') + expect(fakes.transports[0]!.openConnection).toHaveBeenCalledOnce() + expect(brokerBasisIds(broker)).toEqual(['old-basis']) + expect(fakes.controls[0]!.closeNow).not.toHaveBeenCalled() + expect(fakes.transports[0]!.stop).not.toHaveBeenCalled() + await broker.confirmResume('old-basis', 'confirm-1') + expect(fakes.controls[0]!.confirmResume).toHaveBeenCalledWith('old-basis', 'confirm-1') + await broker.installCredential( + 'device-1', + { reqId: 'install-1', newResumeTokenHash: 'H'.repeat(43) }, + { mode: 'relay-basis', basisConnId: 'old-basis' } + ) + expect(fakes.controls[0]!.installCredential).toHaveBeenCalledOnce() + }) + + it('rebinds the same process generation with its control resume secret', async () => { + const ack: RelayHostHelloAckMessage = { + type: 'host-hello-ack', + v: 1, + generation: 7, + controlResumeSecret: 'R'.repeat(43), + leaseExpiresAt: 1_000_000, + activeConnIds: ['existing-basis'], + pendingConns: [] + } + fakes.controlConnect.mockResolvedValueOnce(ack).mockResolvedValueOnce({ + ...ack, + leaseExpiresAt: 2_000_000 + }) + fakes.assign + .mockResolvedValueOnce({ + cellUrl: 'https://relay.example.test', + assignmentEpoch: 1, + leaseExpiresAt: 1_000_000 + }) + .mockResolvedValueOnce({ + cellUrl: 'https://relay.example.test', + assignmentEpoch: 1, + leaseExpiresAt: 2_000_000 + }) + const broker = await RelaySessionBroker.connect(brokerOptions()) + expect(brokerBasisIds(broker)).toEqual(['existing-basis']) + fakes.controls[0]!.options.onDrain({ + type: 'drain', + graceMs: 5_000, + recovery: 'resolve-director' + }) + await vi.waitFor(() => expect(fakes.controls).toHaveLength(2)) + + expect(fakes.transports).toHaveLength(1) + expect(fakes.controls[1]!.options.previousGeneration).toBe(7) + expect(fakes.controls[1]!.options.controlResumeSecret).toBe('R'.repeat(43)) + await vi.waitFor(() => expect(brokerBasisIds(broker)).toEqual(['existing-basis'])) + await broker.confirmResume('existing-basis', 'confirm-1') + expect(fakes.controls[1]!.confirmResume).toHaveBeenCalledOnce() + }) }) + +function brokerBasisIds(broker: RelaySessionBroker): string[] { + const pool = (broker as unknown as { originPool: unknown }).originPool + return [...(pool as { basisOrigins: Map }).basisOrigins.keys()] +} + +function brokerOptions( + overrides: Partial[0]> = {} +): Parameters[0] { + const keypair = nacl.box.keyPair() + return { + authConfig: { + relayTokenEndpoint: 'https://auth.example.test/v1/relay-token', + relayDirectorUrl: 'https://relay.example.test' + } as OrcaCloudAuthConfig, + accessToken: 'access-token', + identity: { userId: 'user-1', profileId: 'profile-1', organizationId: 'org-1' }, + keypair: { + ...keypair, + publicKeyB64: Buffer.from(keypair.publicKey).toString('base64') + }, + appVersion: '1.0.0', + mobileSocketWiring: { attachTransport: vi.fn() } as never, + isCurrent: () => true, + refreshAccessToken: async () => null, + onStatus: vi.fn(), + now: () => 0, + random: () => 0, + ...overrides + } +} diff --git a/src/main/runtime/relay/relay-session-broker.ts b/src/main/runtime/relay/relay-session-broker.ts index 9f0bc5c13c7..920bd7b57e9 100644 --- a/src/main/runtime/relay/relay-session-broker.ts +++ b/src/main/runtime/relay/relay-session-broker.ts @@ -6,8 +6,6 @@ import type { MobileRelayEndpoint, PairingProvisionRelayParams } from '../../../shared/mobile-relay-credential-contract' -import { CloudRelayTransport } from '../rpc/relay-transport' -import { RelayControlClient } from './relay-control-client' import type { DeviceCredentialInstallAuthorization } from './relay-control-requests' import { deriveRelayHostId, @@ -16,6 +14,7 @@ import { type RelayAuthorization, type RelayAssignment } from './relay-http-client' +import { RelayOriginPool } from './relay-origin-pool' import type { RelayBrokerStatus, RelaySessionBrokerOptions } from './relay-session-broker-contract' export type { RelayBrokerStatus } from './relay-session-broker-contract' @@ -29,16 +28,29 @@ export class StaleRelayBrokerError extends Error { export class RelaySessionBroker { private readonly options: RelaySessionBrokerOptions private readonly relayHostId: string - private control: RelayControlClient | null = null - private transport: CloudRelayTransport | null = null + private readonly originPool: RelayOriginPool private authorization: RelayAuthorization | null = null - private assignment: RelayAssignment | null = null private refreshTimer: ReturnType | null = null private closed = false private constructor(options: RelaySessionBrokerOptions) { this.options = options this.relayHostId = deriveRelayHostId(options.keypair.publicKey) + this.originPool = new RelayOriginPool({ + directorUrl: options.authConfig.relayDirectorUrl, + relayHostId: this.relayHostId, + identity: options.identity, + keypair: options.keypair, + appVersion: options.appVersion, + mobileSocketWiring: options.mobileSocketWiring, + isCurrent: () => this.isCurrent(), + onStatus: (status) => this.publishStatus(status), + fetch: options.fetch, + createControlSocket: options.createControlSocket, + createDataSocket: options.createDataSocket, + random: options.random, + now: options.now + }) } static async connect(options: RelaySessionBrokerOptions): Promise { @@ -57,7 +69,7 @@ export class RelaySessionBroker { } get currentAssignment(): RelayAssignment | null { - return this.assignment + return this.originPool.activeAssignment } get ownerIdentityKey(): string { @@ -66,7 +78,7 @@ export class RelaySessionBroker { } get endpoint(): MobileRelayEndpoint | null { - const assignment = this.assignment + const assignment = this.originPool.activeAssignment if (!assignment) { return null } @@ -80,19 +92,21 @@ export class RelaySessionBroker { } } - createInvite(relayDeviceId: string): ReturnType { - if (!this.control) { + createInvite(relayDeviceId: string) { + const control = this.originPool.activeControl + if (!control) { return Promise.reject(new Error('relay_control_not_active')) } - return this.control.createInvite(relayDeviceId) + return control.createInvite(relayDeviceId) } async createPairingRelay(relayDeviceId: string): Promise { - const assignment = this.assignment - if (!assignment || !this.control) { + const assignment = this.originPool.activeAssignment + const control = this.originPool.activeControl + if (!assignment || !control) { throw new Error('relay_control_not_active') } - const invite = await this.control.createInvite(relayDeviceId) + const invite = await control.createInvite(relayDeviceId) this.assertCurrent() return { v: 1, @@ -107,10 +121,11 @@ export class RelaySessionBroker { } revokeDevice(relayDeviceId: string, reqId?: string): Promise { - if (!this.control) { + const control = this.originPool.activeControl + if (!control) { return Promise.reject(new Error('relay_control_not_active')) } - return this.control.revokeDevice(relayDeviceId, reqId) + return control.revokeDevice(relayDeviceId, reqId) } async installCredential( @@ -118,10 +133,14 @@ export class RelaySessionBroker { params: PairingProvisionRelayParams, authorization: DeviceCredentialInstallAuthorization ): Promise { - if (!this.control) { + const control = + authorization.mode === 'relay-basis' + ? this.originPool.controlForBasis(authorization.basisConnId) + : this.originPool.activeControl + if (!control) { throw new Error('relay_control_not_active') } - const message = await this.control.installCredential({ + const message = await control.installCredential({ relayDeviceId, authorization, ...params @@ -135,20 +154,22 @@ export class RelaySessionBroker { relayDeviceId: string, reqId: string ): Promise { - if (!this.control) { + const control = this.originPool.activeControl + if (!control) { throw new Error('relay_control_not_active') } - const message = await this.control.credentialInstallStatus(relayDeviceId, reqId) + const message = await control.credentialInstallStatus(relayDeviceId, reqId) this.assertCurrent() const { type: _type, ...result } = message return result } async confirmResume(basisConnId: string, reqId: string): Promise { - if (!this.control) { - throw new Error('relay_control_not_active') + const control = this.originPool.controlForBasis(basisConnId) + if (!control) { + throw new Error('relay_basis_origin_not_found') } - const message = await this.control.confirmResume(basisConnId, reqId) + const message = await control.confirmResume(basisConnId, reqId) this.assertCurrent() const { type: _type, ...result } = message return result @@ -164,10 +185,7 @@ export class RelaySessionBroker { clearTimeout(this.refreshTimer) this.refreshTimer = null } - this.control?.closeNow() - this.control = null - void this.transport?.stop() - this.transport = null + this.originPool.closeNow() if (publishOffline) { this.options.onStatus('offline') } @@ -189,56 +207,16 @@ export class RelaySessionBroker { fetch: this.options.fetch }) this.assertCurrent() - const transport = new CloudRelayTransport({ - cellUrl: assignment.cellUrl, - relayHostId: this.relayHostId, - generation: 0, - createSocket: this.options.createDataSocket - }) - // Why: stale-epoch cleanup must own partially opened resources before the - // next await, even though the broker is not externally active yet. - this.transport = transport - this.options.mobileSocketWiring.attachTransport(transport, (ws) => transport.metadataFor(ws)) - await transport.start() - this.assertCurrent() - const control = new RelayControlClient({ - cellUrl: assignment.cellUrl, - relayJwt: authorization.relayToken, - relayHostId: this.relayHostId, - assignmentEpoch: assignment.assignmentEpoch, - identity: this.options.identity, - keypair: this.options.keypair, - appVersion: this.options.appVersion, - onConnectionOpen: (message) => { - if (this.isCurrent()) { - void transport.openConnection(message).catch(() => {}) - } - }, - onDrain: () => { - if (!this.isCurrent()) { - return - } - this.publishStatus('draining') - // Why: the data plane never chooses recovery targets; every drain - // returns to the configured stable director. - this.options.onResolveDirector() - }, - onClose: () => { - if (this.isCurrent()) { - this.publishStatus('offline') - } - }, - createSocket: this.options.createControlSocket - }) - this.control = control - const ack = await control.connect() - this.assertCurrent() - if (ack.generation <= 0) { - throw new Error('invalid_relay_generation') + try { + await this.originPool.openInitial(assignment, authorization.relayToken) + } catch (error) { + if (!this.isCurrent()) { + throw new StaleRelayBrokerError() + } + throw error } - transport.setGeneration(ack.generation) + this.assertCurrent() this.authorization = authorization - this.assignment = assignment this.publishStatus('registered') this.scheduleRefresh() } @@ -271,7 +249,7 @@ export class RelaySessionBroker { fetch: this.options.fetch }) this.assertCurrent() - this.control?.refreshAuthorization(authorization.relayToken) + this.originPool.refreshAuthorization(authorization.relayToken) this.authorization = authorization this.scheduleRefresh() } catch { diff --git a/src/main/runtime/rpc/relay-transport.test.ts b/src/main/runtime/rpc/relay-transport.test.ts index 38f221b49bc..432503b9a3c 100644 --- a/src/main/runtime/rpc/relay-transport.test.ts +++ b/src/main/runtime/rpc/relay-transport.test.ts @@ -39,6 +39,7 @@ describe('CloudRelayTransport', () => { server.once('connection', (socket, request) => resolve({ socket, path: request.url ?? '' })) }) let clientSocket: WebSocketClient | null = null + const onConnectionClosed = vi.fn() const transport = new CloudRelayTransport({ cellUrl: `http://127.0.0.1:${address.port}`, relayHostId: 'AbCdEf0123_-xyZ9', @@ -46,7 +47,8 @@ describe('CloudRelayTransport', () => { createSocket: (url) => { clientSocket = new WebSocketClient(url, { perMessageDeflate: false }) return clientSocket - } + }, + onConnectionClosed }) transports.push(transport) const received: (string | Uint8Array)[] = [] @@ -87,6 +89,8 @@ describe('CloudRelayTransport', () => { basisConnId: 'conn/with spaces', credentialKind: 'resume' }) + socket.close() + await vi.waitFor(() => expect(onConnectionClosed).toHaveBeenCalledWith('conn/with spaces')) }) it('rejects non-origin cell URLs before opening a socket', () => { diff --git a/src/main/runtime/rpc/relay-transport.ts b/src/main/runtime/rpc/relay-transport.ts index 2275e85d966..1a34afdf849 100644 --- a/src/main/runtime/rpc/relay-transport.ts +++ b/src/main/runtime/rpc/relay-transport.ts @@ -19,6 +19,7 @@ type CloudRelayTransportOptions = { relayHostId: string generation: number createSocket?: (url: string) => WebSocket + onConnectionClosed?: (connectionId: string) => void } function relayWebSocketOrigin(cellUrl: string): string { @@ -41,6 +42,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport private readonly relayHostId: string private generation: number private readonly createSocket: (url: string) => WebSocket + private readonly onConnectionClosed: ((connectionId: string) => void) | undefined private readonly socketsByConnectionId = new Map() private readonly metadataBySocket = new Map() private readonly clientIds = new Map() @@ -52,6 +54,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport this.cellWebSocketOrigin = relayWebSocketOrigin(options.cellUrl) this.relayHostId = options.relayHostId this.generation = options.generation + this.onConnectionClosed = options.onConnectionClosed this.createSocket = options.createSocket ?? ((url) => @@ -81,6 +84,9 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport } setGeneration(generation: number): void { + if (generation === this.generation) { + return + } if ( this.socketsByConnectionId.size > 0 || !Number.isSafeInteger(generation) || @@ -153,6 +159,7 @@ export class CloudRelayTransport implements RpcTransport, MobileSocketTransport this.metadataBySocket.delete(socket) const clientId = this.clientIds.get(socket) ?? null this.clientIds.delete(socket) + this.onConnectionClosed?.(connection.connId) const hasOtherConnections = clientId !== null && [...this.clientIds.values()].includes(clientId) this.closeHandler?.(clientId, socket, hasOtherConnections)