From 0bafaaf2e9aacde7c4e007fe37a77f45eb28be1f Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Mon, 13 Jul 2026 00:12:07 -0400 Subject: [PATCH] refactor(runtime): unify mobile socket wiring Co-authored-by: Orca --- .../e2ee-channel-text-backpressure.test.ts | 5 +- src/main/runtime/rpc/e2ee-channel-v2.test.ts | 24 ++- src/main/runtime/rpc/e2ee-channel.test.ts | 11 +- src/main/runtime/rpc/e2ee-channel.ts | 27 ++- src/main/runtime/rpc/e2ee-integration.test.ts | 10 +- .../runtime/rpc/mobile-socket-wiring.test.ts | 183 +++++++++++++++++ src/main/runtime/rpc/mobile-socket-wiring.ts | 188 ++++++++++++++++++ src/main/runtime/runtime-rpc.test.ts | 30 ++- src/main/runtime/runtime-rpc.ts | 129 ++++-------- 9 files changed, 488 insertions(+), 119 deletions(-) create mode 100644 src/main/runtime/rpc/mobile-socket-wiring.test.ts create mode 100644 src/main/runtime/rpc/mobile-socket-wiring.ts diff --git a/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts b/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts index 5d0b2d3da92..4c45e97ea10 100644 --- a/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts +++ b/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts @@ -34,7 +34,10 @@ function setup(overrides?: Partial) { const onError = vi.fn() const channel = new E2EEChannel(ws as unknown as WebSocket, { serverSecretKey: serverKeys.secretKey, - validateToken: (token) => token === 'valid-token', + resolveAuthenticatedDevice: (token) => + token === 'valid-token' + ? { deviceId: 'device-1', deviceToken: token, scope: 'mobile' } + : null, onReady: vi.fn(), onError, ...overrides diff --git a/src/main/runtime/rpc/e2ee-channel-v2.test.ts b/src/main/runtime/rpc/e2ee-channel-v2.test.ts index c25aec8f1e2..00fd0093c4c 100644 --- a/src/main/runtime/rpc/e2ee-channel-v2.test.ts +++ b/src/main/runtime/rpc/e2ee-channel-v2.test.ts @@ -52,16 +52,20 @@ function setup() { const ws = createMockWs() const onReady = vi.fn() const onError = vi.fn() - const validateToken = vi.fn((token: string) => token === 'valid-token') + const resolveAuthenticatedDevice = vi.fn((token: string) => + token === 'valid-token' + ? { deviceId: 'device-1', deviceToken: token, scope: 'mobile' as const } + : null + ) const channel = new E2EEChannel(ws as unknown as WebSocket, { serverSecretKey: server.secretKey, - validateToken, + resolveAuthenticatedDevice, onReady, onError, transportContext: { transport: 'relay', relayHostId: 'AbCdEf0123_-xyZ9' }, requireV2: true }) - return { ws, channel, onReady, onError, validateToken } + return { ws, channel, onReady, onError, resolveAuthenticatedDevice } } function startV2(ctx: ReturnType) { @@ -138,8 +142,12 @@ describe('E2EEChannel v2', () => { const { schedule } = startV2(ctx) authenticate(ctx, schedule) - expect(ctx.validateToken).toHaveBeenCalledOnce() - expect(ctx.onReady).toHaveBeenCalledWith(ctx.channel) + expect(ctx.resolveAuthenticatedDevice).toHaveBeenCalledOnce() + expect(ctx.onReady).toHaveBeenCalledWith(ctx.channel, { + deviceId: 'device-1', + deviceToken: 'valid-token', + scope: 'mobile' + }) const authenticated = openServerFrame(ctx.ws.sent[1]!.data, 'text', schedule, 0n) expect(JSON.parse(new TextDecoder().decode(authenticated!))).toEqual({ type: 'e2ee_authenticated', @@ -171,7 +179,7 @@ describe('E2EEChannel v2', () => { 0n ) ) - expect(ctx.validateToken).not.toHaveBeenCalled() + expect(ctx.resolveAuthenticatedDevice).not.toHaveBeenCalled() }) it('rejects a captured auth frame replayed onto a fresh desktop nonce', () => { @@ -192,7 +200,7 @@ describe('E2EEChannel v2', () => { const secondHandshake = startV2(second) expect(secondHandshake.ready.desktopNonceB64).not.toBe(firstHandshake.ready.desktopNonceB64) second.channel.handleRawMessage(capturedAuth) - expect(second.validateToken).not.toHaveBeenCalled() + expect(second.resolveAuthenticatedDevice).not.toHaveBeenCalled() expect(second.onReady).not.toHaveBeenCalled() }) @@ -222,7 +230,7 @@ describe('E2EEChannel v2', () => { second.channel.onMessage(replayedMutation) second.channel.handleRawMessage(capturedAuth) second.channel.handleRawMessage(capturedMutation) - expect(second.validateToken).not.toHaveBeenCalled() + expect(second.resolveAuthenticatedDevice).not.toHaveBeenCalled() expect(replayedMutation).not.toHaveBeenCalled() }) diff --git a/src/main/runtime/rpc/e2ee-channel.test.ts b/src/main/runtime/rpc/e2ee-channel.test.ts index a1daec13669..846b82903db 100644 --- a/src/main/runtime/rpc/e2ee-channel.test.ts +++ b/src/main/runtime/rpc/e2ee-channel.test.ts @@ -27,7 +27,10 @@ function setup(overrides?: Partial) { const channel = new E2EEChannel(ws as unknown as WebSocket, { serverSecretKey: serverKeys.secretKey, - validateToken: (token) => token === 'valid-token', + resolveAuthenticatedDevice: (token) => + token === 'valid-token' + ? { deviceId: 'device-1', deviceToken: token, scope: 'mobile' } + : null, onReady, onError, ...overrides @@ -63,7 +66,11 @@ describe('E2EEChannel', () => { const ctx = setup() doHandshake(ctx) - expect(ctx.onReady).toHaveBeenCalledWith(ctx.channel) + expect(ctx.onReady).toHaveBeenCalledWith(ctx.channel, { + deviceId: 'device-1', + deviceToken: 'valid-token', + scope: 'mobile' + }) expect(ctx.onError).not.toHaveBeenCalled() expect(ctx.channel.deviceToken).toBe('valid-token') diff --git a/src/main/runtime/rpc/e2ee-channel.ts b/src/main/runtime/rpc/e2ee-channel.ts index cf79f562e2d..cec40fd38ce 100644 --- a/src/main/runtime/rpc/e2ee-channel.ts +++ b/src/main/runtime/rpc/e2ee-channel.ts @@ -26,13 +26,19 @@ const MAX_BINARY_BUFFERED_AMOUNT = 8 * 1024 * 1024 export type E2EEChannelOptions = { serverSecretKey: Uint8Array - validateToken: (token: string) => boolean - onReady: (channel: E2EEChannel) => void + resolveAuthenticatedDevice: (token: string) => E2EEAuthenticatedDevice | null + onReady: (channel: E2EEChannel, device: E2EEAuthenticatedDevice) => void onError: (code: number, reason: string) => void transportContext?: DesktopMobileE2EEV2Context requireV2?: boolean } +export type E2EEAuthenticatedDevice = { + deviceId: string + deviceToken: string + scope: 'mobile' | 'runtime' +} + export class E2EEChannel { private state: ChannelState = 'awaiting_hello' private sharedKey: Uint8Array | null = null @@ -40,8 +46,8 @@ export class E2EEChannel { private handshakeTimer: ReturnType | null = null private readonly ws: WebSocket private readonly serverSecretKey: Uint8Array - private readonly validateToken: (token: string) => boolean - private readonly onReady: (channel: E2EEChannel) => void + private readonly resolveAuthenticatedDevice: (token: string) => E2EEAuthenticatedDevice | null + private readonly onReady: (channel: E2EEChannel, device: E2EEAuthenticatedDevice) => void private readonly onError: (code: number, reason: string) => void private readonly transportContext: DesktopMobileE2EEV2Context private readonly requireV2: boolean @@ -65,11 +71,12 @@ export class E2EEChannel { private textReplyQueue: WsOutboundBackpressureQueue | null = null deviceToken: string | null = null + authenticatedDevice: E2EEAuthenticatedDevice | null = null constructor(ws: WebSocket, options: E2EEChannelOptions) { this.ws = ws this.serverSecretKey = options.serverSecretKey - this.validateToken = options.validateToken + this.resolveAuthenticatedDevice = options.resolveAuthenticatedDevice this.onReady = options.onReady this.onError = options.onError this.transportContext = options.transportContext ?? { transport: 'direct' } @@ -245,13 +252,15 @@ export class E2EEChannel { this.onError(4001, 'Invalid e2ee_auth') return } - if (!this.validateToken(auth.deviceToken)) { + const authenticatedDevice = this.resolveAuthenticatedDevice(auth.deviceToken) + if (!authenticatedDevice || authenticatedDevice.deviceToken !== auth.deviceToken) { this.sendEncryptedControl({ type: 'e2ee_error', error: { code: 'unauthorized' } }) this.onError(4001, 'Unauthorized') return } this.deviceToken = auth.deviceToken + this.authenticatedDevice = authenticatedDevice this.state = 'ready' if (this.handshakeTimer) { @@ -259,6 +268,10 @@ export class E2EEChannel { this.handshakeTimer = null } + // Why: transport-bound identity checks must complete before the peer sees + // authentication success; relay sockets additionally bind this context to + // their immutable relayDeviceId in the resolver. + this.onReady(this, authenticatedDevice) this.sendEncryptedControl( this.v2Session ? { @@ -268,7 +281,6 @@ export class E2EEChannel { } : { type: 'e2ee_authenticated' } ) - this.onReady(this) } private handleV2RawMessage(raw: string | Uint8Array): void { @@ -334,6 +346,7 @@ export class E2EEChannel { this.handshakeTimer = null } this.sharedKey = null + this.authenticatedDevice = null this.v2Session = null this.messageHandler = null this.binaryMessageHandler = null diff --git a/src/main/runtime/rpc/e2ee-integration.test.ts b/src/main/runtime/rpc/e2ee-integration.test.ts index f4029b82473..9065a7b21aa 100644 --- a/src/main/runtime/rpc/e2ee-integration.test.ts +++ b/src/main/runtime/rpc/e2ee-integration.test.ts @@ -44,7 +44,10 @@ describe('E2EE integration (simulated mobile ↔ desktop)', () => { channel = new E2EEChannel(mockWs as unknown as WebSocket, { serverSecretKey: serverKeys.secretKey, - validateToken: (token) => token === 'device-abc', + resolveAuthenticatedDevice: (token) => + token === 'device-abc' + ? { deviceId: 'device-abc', deviceToken: token, scope: 'mobile' } + : null, onReady, onError }) @@ -118,7 +121,10 @@ describe('E2EE integration (simulated mobile ↔ desktop)', () => { const newMobileKeys = generateKeyPair() const newChannel = new E2EEChannel(mockWs as unknown as WebSocket, { serverSecretKey: serverKeys.secretKey, - validateToken: (token) => token === 'device-abc', + resolveAuthenticatedDevice: (token) => + token === 'device-abc' + ? { deviceId: 'device-abc', deviceToken: token, scope: 'mobile' } + : null, onReady, onError }) diff --git a/src/main/runtime/rpc/mobile-socket-wiring.test.ts b/src/main/runtime/rpc/mobile-socket-wiring.test.ts new file mode 100644 index 00000000000..4011e728597 --- /dev/null +++ b/src/main/runtime/rpc/mobile-socket-wiring.test.ts @@ -0,0 +1,183 @@ +import { describe, expect, it, vi } from 'vitest' +import nacl from 'tweetnacl' +import type { WebSocket } from 'ws' +import { + encodeMobileE2EEV2Transcript, + validateMobileE2EEV2Handshake, + type MobileE2EEV2Hello, + type MobileE2EEV2Ready +} from '../../../shared/mobile-e2ee-v2-contract' +import { sealMobileE2EEV2Frame } from '../../../shared/mobile-e2ee-v2-framing' +import type { DeviceRegistry } from '../device-registry' +import { deriveSharedKey, encrypt, generateKeyPair } from './e2ee-crypto' +import { deriveMobileE2EEV2KeySchedule } from './mobile-e2ee-v2-key-schedule' +import { + MobileSocketWiring, + type MobileSocketTransport, + type MobileSocketTransportMetadata +} from './mobile-socket-wiring' + +class FakeSocket { + readonly OPEN = 1 + readyState = this.OPEN + bufferedAmount = 0 + readonly sent: (string | Buffer)[] = [] + readonly send = vi.fn((data: string | Buffer) => this.sent.push(data)) + readonly close = vi.fn() +} + +class FakeTransport implements MobileSocketTransport { + private messageHandler: Parameters[0] | null = null + private closeHandler: Parameters[0] | null = null + readonly setClientId = vi.fn() + readonly terminateClientConnections = vi.fn(() => 0) + + onMessage(handler: Parameters[0]): void { + this.messageHandler = handler + } + + onConnectionClose(handler: Parameters[0]): void { + this.closeHandler = handler + } + + receive(ws: FakeSocket, message: string): void { + this.messageHandler?.(message, vi.fn(), ws as unknown as WebSocket) + } + + disconnect(ws: FakeSocket): void { + this.closeHandler?.(null, ws as unknown as WebSocket, false) + } +} + +function registryFor(deviceId: string, token: string): DeviceRegistry { + return { + validateToken: (candidate: string) => + candidate === token + ? { + deviceId, + token, + name: 'Phone', + scope: 'mobile' as const, + pairedAt: 1, + lastSeenAt: 0 + } + : null, + updateLastSeen: vi.fn() + } as unknown as DeviceRegistry +} + +describe('MobileSocketWiring', () => { + it('preserves the legacy direct handshake, identity, and close cleanup', () => { + const desktop = generateKeyPair() + const phone = generateKeyPair() + const ws = new FakeSocket() + const transport = new FakeTransport() + const onText = vi.fn() + const onClose = vi.fn() + const wiring = new MobileSocketWiring({ + deviceRegistry: registryFor('device-1', 'valid-token'), + e2eeKeypair: { + publicKey: desktop.publicKey, + secretKey: desktop.secretKey, + publicKeyB64: Buffer.from(desktop.publicKey).toString('base64') + }, + onText, + onBinary: vi.fn(), + onClose + }) + wiring.attachTransport(transport) + + transport.receive( + ws, + JSON.stringify({ + type: 'e2ee_hello', + publicKeyB64: Buffer.from(phone.publicKey).toString('base64') + }) + ) + const sharedKey = deriveSharedKey(phone.secretKey, desktop.publicKey) + transport.receive( + ws, + encrypt(JSON.stringify({ type: 'e2ee_auth', deviceToken: 'valid-token' }), sharedKey) + ) + transport.receive(ws, encrypt('{"id":"rpc-1","method":"status.get"}', sharedKey)) + + expect(transport.setClientId).toHaveBeenCalledWith(ws, 'valid-token') + expect(onText).toHaveBeenCalledOnce() + expect(onText.mock.calls[0]?.[0]).toMatchObject({ + device: { deviceId: 'device-1', deviceToken: 'valid-token', scope: 'mobile' }, + transport: { transport: 'direct' } + }) + + transport.disconnect(ws) + expect(onClose).toHaveBeenCalledWith(expect.objectContaining({ ws }), false) + expect(wiring.channelCount).toBe(0) + expect(wiring.connectionCount).toBe(0) + }) + + it('rejects a relay socket whose immutable relayDeviceId differs from E2EE identity', () => { + const desktop = nacl.box.keyPair.fromSecretKey(new Uint8Array(32).fill(1)) + const phone = nacl.box.keyPair.fromSecretKey(new Uint8Array(32).fill(2)) + const ws = new FakeSocket() + const transport = new FakeTransport() + const metadata: MobileSocketTransportMetadata = { + transport: 'relay', + relayHostId: 'AbCdEf0123_-xyZ9', + relayDeviceId: 'outer-device', + basisConnId: 'connection-1' + } + const wiring = new MobileSocketWiring({ + deviceRegistry: registryFor('e2ee-device', 'valid-token'), + e2eeKeypair: { + publicKey: desktop.publicKey, + secretKey: desktop.secretKey, + publicKeyB64: Buffer.from(desktop.publicKey).toString('base64') + }, + onText: vi.fn(), + onBinary: vi.fn(), + onClose: vi.fn() + }) + wiring.attachTransport(transport, () => metadata) + const hello: MobileE2EEV2Hello = { + type: 'e2ee_hello', + v: 2, + clientPublicKeyB64: Buffer.from(phone.publicKey).toString('base64'), + clientNonceB64: Buffer.from(new Uint8Array(32).fill(3)).toString('base64'), + capabilities: { framing: [2], payloadKinds: ['text', 'binary'] }, + context: { + protocol: 'orca-mobile-e2ee', + initiator: 'mobile', + responder: 'desktop', + transport: 'relay', + relayHostId: metadata.relayHostId + } + } + transport.receive(ws, JSON.stringify(hello)) + const ready = JSON.parse(ws.sent[0]!.toString()) as MobileE2EEV2Ready + const handshake = validateMobileE2EEV2Handshake(hello, ready)! + const schedule = deriveMobileE2EEV2KeySchedule({ + sharedSecret: deriveSharedKey(phone.secretKey, desktop.publicKey), + transcript: encodeMobileE2EEV2Transcript(handshake), + clientNonce: handshake.clientNonce, + desktopNonce: handshake.desktopNonce + }) + const auth = sealMobileE2EEV2Frame({ + payload: new TextEncoder().encode( + JSON.stringify({ + type: 'e2ee_auth', + v: 2, + transcriptHashB64: Buffer.from(schedule.transcriptHash).toString('base64'), + deviceToken: 'valid-token' + }) + ), + key: schedule.mobileToDesktopKey, + sessionId: schedule.sessionId, + direction: 'mobile-to-desktop', + payloadKind: 'text', + counter: 0n + }) + transport.receive(ws, Buffer.from(auth).toString('base64')) + + expect(transport.setClientId).not.toHaveBeenCalled() + expect(ws.close).toHaveBeenCalledWith(4001, 'Unauthorized') + }) +}) diff --git a/src/main/runtime/rpc/mobile-socket-wiring.ts b/src/main/runtime/rpc/mobile-socket-wiring.ts new file mode 100644 index 00000000000..9b8b75cc4b4 --- /dev/null +++ b/src/main/runtime/rpc/mobile-socket-wiring.ts @@ -0,0 +1,188 @@ +import { randomBytes } from 'node:crypto' +import type { WebSocket } from 'ws' +import type { DeviceEntry, DeviceRegistry } from '../device-registry' +import type { E2EEKeypair } from '../e2ee-keypair' +import { E2EEChannel, type E2EEAuthenticatedDevice } from './e2ee-channel' + +type MobileSocketPayload = string | Uint8Array + +export type MobileSocketTransportMetadata = + | { transport: 'direct' } + | { + transport: 'relay' + relayHostId: string + relayDeviceId: string + basisConnId: string + } + +export type MobileSocketTransport = { + onMessage( + handler: ( + message: MobileSocketPayload, + reply: (response: string) => void, + ws: WebSocket + ) => void + ): void + onConnectionClose( + handler: (clientId: string | null, ws: WebSocket, hasOtherConnections: boolean) => void + ): void + setClientId(ws: WebSocket, clientId: string): void + terminateClientConnections(clientId: string): number +} + +export type AuthenticatedMobileSocket = { + ws: WebSocket + connectionId: string + device: E2EEAuthenticatedDevice + transport: MobileSocketTransportMetadata +} + +type MobileSocketWiringOptions = { + deviceRegistry: DeviceRegistry + e2eeKeypair: E2EEKeypair + onText: ( + socket: AuthenticatedMobileSocket, + plaintext: string, + reply: (response: string) => void, + sendBinary: (response: Uint8Array) => boolean | void + ) => void + onBinary: (socket: AuthenticatedMobileSocket, bytes: Uint8Array) => void + onClose: (socket: AuthenticatedMobileSocket | null, hasOtherConnections: boolean) => void + onReady?: (socket: AuthenticatedMobileSocket) => void +} + +function toAuthenticatedDevice(device: DeviceEntry): E2EEAuthenticatedDevice { + return { + deviceId: device.deviceId, + deviceToken: device.token, + scope: device.scope + } +} + +export class MobileSocketWiring { + private readonly deviceRegistry: DeviceRegistry + private readonly e2eeKeypair: E2EEKeypair + private readonly onText: MobileSocketWiringOptions['onText'] + private readonly onBinary: MobileSocketWiringOptions['onBinary'] + private readonly onClose: MobileSocketWiringOptions['onClose'] + private readonly onReady: MobileSocketWiringOptions['onReady'] + private readonly channels = new Map() + private readonly connectionIds = new Map() + private readonly authenticatedSockets = new Map() + private readonly transports = new Set() + + constructor(options: MobileSocketWiringOptions) { + this.deviceRegistry = options.deviceRegistry + this.e2eeKeypair = options.e2eeKeypair + this.onText = options.onText + this.onBinary = options.onBinary + this.onClose = options.onClose + this.onReady = options.onReady + } + + attachTransport( + transport: MobileSocketTransport, + getMetadata: (ws: WebSocket) => MobileSocketTransportMetadata = () => ({ + transport: 'direct' + }) + ): void { + this.transports.add(transport) + transport.onMessage((message, _reply, ws) => { + this.handleRawMessage(transport, ws, message, getMetadata(ws)) + }) + transport.onConnectionClose((_clientId, ws) => this.handleClose(ws)) + } + + getConnectionId(ws: WebSocket): string | undefined { + return this.connectionIds.get(ws) + } + + get channelCount(): number { + return this.channels.size + } + + get connectionCount(): number { + return this.connectionIds.size + } + + terminateDeviceConnections(deviceToken: string): number { + let terminated = 0 + for (const transport of this.transports) { + terminated += transport.terminateClientConnections(deviceToken) + } + return terminated + } + + private handleRawMessage( + transport: MobileSocketTransport, + ws: WebSocket, + message: MobileSocketPayload, + metadata: MobileSocketTransportMetadata + ): void { + let channel = this.channels.get(ws) + if (!channel) { + const connectionId = randomBytes(8).toString('hex') + this.connectionIds.set(ws, connectionId) + channel = new E2EEChannel(ws, { + serverSecretKey: this.e2eeKeypair.secretKey, + transportContext: + metadata.transport === 'relay' + ? { transport: 'relay', relayHostId: metadata.relayHostId } + : { transport: 'direct' }, + requireV2: metadata.transport === 'relay', + resolveAuthenticatedDevice: (token) => { + const device = this.deviceRegistry.validateToken(token) + if (!device) { + return null + } + // Why: outer relay authorization cannot choose the local Orca + // identity; E2EE must resolve the same device before readiness. + if (metadata.transport === 'relay' && metadata.relayDeviceId !== device.deviceId) { + return null + } + return toAuthenticatedDevice(device) + }, + onReady: (_channel, device) => { + const socket = { ws, connectionId, device, transport: metadata } + this.authenticatedSockets.set(ws, socket) + transport.setClientId(ws, device.deviceToken) + this.deviceRegistry.updateLastSeen(device.deviceId) + this.onReady?.(socket) + }, + onError: (code, reason) => { + this.channels.get(ws)?.destroy() + this.channels.delete(ws) + ws.close(code, reason) + } + }) + channel.onMessage((plaintext, reply, sendBinary) => { + const socket = this.authenticatedSockets.get(ws) + if (socket) { + this.onText(socket, plaintext, reply, sendBinary) + } + }) + channel.onBinaryMessage((bytes) => { + const socket = this.authenticatedSockets.get(ws) + if (socket) { + this.onBinary(socket, bytes) + } + }) + this.channels.set(ws, channel) + } + channel.handleRawMessage(message) + } + + private handleClose(ws: WebSocket): void { + const socket = this.authenticatedSockets.get(ws) ?? null + this.authenticatedSockets.delete(ws) + this.channels.get(ws)?.destroy() + this.channels.delete(ws) + this.connectionIds.delete(ws) + const hasOtherConnections = + socket !== null && + Array.from(this.authenticatedSockets.values()).some( + (candidate) => candidate.device.deviceToken === socket.device.deviceToken + ) + this.onClose(socket, hasOtherConnections) + } +} diff --git a/src/main/runtime/runtime-rpc.test.ts b/src/main/runtime/runtime-rpc.test.ts index 2d7579cce48..6efe04fb3b8 100644 --- a/src/main/runtime/runtime-rpc.test.ts +++ b/src/main/runtime/runtime-rpc.test.ts @@ -533,12 +533,16 @@ describe('OrcaRuntimeRpcServer', () => { }) ) expect(JSON.parse(await nextWsMessage(ws))).toEqual({ type: 'e2ee_ready' }) - expect(server['e2eeChannels'].size).toBe(1) - expect(server['wsConnectionIds'].size).toBe(1) + expect(server['mobileSocketWiring']?.channelCount).toBe(1) + expect(server['mobileSocketWiring']?.connectionCount).toBe(1) ws.close() await waitForWsClose(ws) - await waitFor(() => server['e2eeChannels'].size === 0 && server['wsConnectionIds'].size === 0) + await waitFor( + () => + server['mobileSocketWiring']?.channelCount === 0 && + server['mobileSocketWiring']?.connectionCount === 0 + ) } finally { await server.stop() } @@ -572,7 +576,11 @@ describe('OrcaRuntimeRpcServer', () => { expect(server.revokeMobileDevice(offer.deviceId)).toBe(true) await Promise.all([waitForWsClose(first), waitForWsClose(second)]) - await waitFor(() => server['e2eeChannels'].size === 0 && server['wsConnectionIds'].size === 0) + await waitFor( + () => + server['mobileSocketWiring']?.channelCount === 0 && + server['mobileSocketWiring']?.connectionCount === 0 + ) expect(disconnectSpy).toHaveBeenCalledTimes(1) } finally { @@ -638,7 +646,11 @@ describe('OrcaRuntimeRpcServer', () => { expect(server.revokeRuntimeAccess(offer.deviceId)).toBe(true) await Promise.all([waitForWsClose(first), waitForWsClose(second)]) - await waitFor(() => server['e2eeChannels'].size === 0 && server['wsConnectionIds'].size === 0) + await waitFor( + () => + server['mobileSocketWiring']?.channelCount === 0 && + server['mobileSocketWiring']?.connectionCount === 0 + ) expect(server.getDeviceRegistry()?.getDevice(offer.deviceId)).toBeNull() } finally { @@ -718,7 +730,9 @@ describe('OrcaRuntimeRpcServer', () => { server['deviceRegistry'] = new DeviceRegistry(userDataPath) const entry = server['deviceRegistry']!.addDevice('runtime-test', 'runtime') const ws = new FakeWebSocket() - server['wsConnectionIds'].set(ws as unknown as WebSocket, 'conn-test') + server['mobileSocketWiring'] = { + getConnectionId: () => 'conn-test' + } as unknown as NonNullable<(typeof server)['mobileSocketWiring']> const replies: Record[] = [] try { @@ -778,7 +792,9 @@ describe('OrcaRuntimeRpcServer', () => { server['deviceRegistry'] = new DeviceRegistry(userDataPath) const entry = server['deviceRegistry']!.addDevice('runtime-test', 'runtime') const ws = new FakeWebSocket() - server['wsConnectionIds'].set(ws as unknown as WebSocket, 'conn-test') + server['mobileSocketWiring'] = { + getConnectionId: () => 'conn-test' + } as unknown as NonNullable<(typeof server)['mobileSocketWiring']> let activeDispatches = 0 ;( server as unknown as { diff --git a/src/main/runtime/runtime-rpc.ts b/src/main/runtime/runtime-rpc.ts index ace969138c4..5d60cb1fd3f 100644 --- a/src/main/runtime/runtime-rpc.ts +++ b/src/main/runtime/runtime-rpc.ts @@ -20,7 +20,7 @@ import { readWsFallbackPort, writeWsFallbackPort } from './rpc/ws-fallback-port- import type { WebSocket } from 'ws' import { DeviceRegistry, type DeviceScope } from './device-registry' import { loadOrCreateE2EEKeypair, type E2EEKeypair } from './e2ee-keypair' -import { E2EEChannel } from './rpc/e2ee-channel' +import { MobileSocketWiring } from './rpc/mobile-socket-wiring' import { encodePairingOffer, PAIRING_OFFER_VERSION } from '../../shared/pairing' import { decodeTerminalStreamFrame, @@ -407,16 +407,9 @@ export class OrcaRuntimeRpcServer { private deviceRegistry: DeviceRegistry | null = null private e2eeKeypair: E2EEKeypair | null = null private tlsFingerprint: string | null = null - private wsTransport: WebSocketTransport | null = null private activeTransports: RpcTransport[] = [] private transports: RuntimeTransportMetadata[] = [] - // Why: each WebSocket connection has its own E2EE channel that manages the - // handshake and encrypt/decrypt lifecycle. Keyed by WebSocket instance. - private e2eeChannels = new Map() - // Why: stable per-WebSocket id used as the cleanup key for streaming - // subscriptions, so the server can reap a closing socket's subscriptions - // without affecting other live sockets that share the same deviceToken. - private wsConnectionIds = new Map() + private mobileSocketWiring: MobileSocketWiring | null = null private readonly binaryStreamHandlers = new Map< string, Map void> @@ -474,7 +467,7 @@ export class OrcaRuntimeRpcServer { if (device?.scope !== 'mobile' || !this.deviceRegistry?.removeDevice(deviceId)) { return false } - this.wsTransport?.terminateClientConnections(device.token) + this.mobileSocketWiring?.terminateDeviceConnections(device.token) return true } @@ -483,7 +476,7 @@ export class OrcaRuntimeRpcServer { if (device?.scope !== 'runtime' || !this.deviceRegistry?.removeDevice(deviceId)) { return false } - this.wsTransport?.terminateClientConnections(device.token) + this.mobileSocketWiring?.terminateDeviceConnections(device.token) return true } @@ -562,7 +555,7 @@ export class OrcaRuntimeRpcServer { } private handleWebSocketBinaryMessage(bytes: Uint8Array, ws: WebSocket): void { - const connectionId = this.wsConnectionIds.get(ws) + const connectionId = this.mobileSocketWiring?.getConnectionId(ws) if (!connectionId) { return } @@ -714,85 +707,37 @@ export class OrcaRuntimeRpcServer { // pin it. ...(this.wsPort !== 0 ? { fallbackPort: readWsFallbackPort(this.userDataPath) } : {}) }) - this.wsTransport = wsTransport - - // Why: each WebSocket connection gets an E2EE channel that handles the - // handshake before any RPC messages are processed. The channel decrypts - // inbound messages and encrypts outbound replies transparently. - wsTransport.onMessage((msg, _reply, ws) => { - let channel = this.e2eeChannels.get(ws) - if (!channel) { - // Why: stable per-ws id used as the cleanup-index key for - // streaming subscriptions, so the server can reap them exactly - // when this socket closes (without affecting other live sockets - // that share the same deviceToken). - this.wsConnectionIds.set(ws, randomBytes(8).toString('hex')) - channel = new E2EEChannel(ws, { - serverSecretKey: this.e2eeKeypair!.secretKey, - validateToken: (token) => this.deviceRegistry?.validateToken(token) != null, - onReady: (ch) => { - if (ch.deviceToken) { - wsTransport.setClientId(ws, ch.deviceToken) - // Why: mark the device as actually connected so it appears - // in the "Paired Devices" list. Devices that were only - // generated as QR codes but never scanned stay hidden. - const device = this.deviceRegistry?.validateToken(ch.deviceToken) - if (device) { - this.deviceRegistry?.updateLastSeen(device.deviceId) - } - } - }, - onError: (code, reason) => { - this.e2eeChannels.get(ws)?.destroy() - this.e2eeChannels.delete(ws) - ws.close(code, reason) - } - }) - channel.onMessage((plaintext, encryptedReply, encryptedBinaryReply) => { - const authenticatedDeviceToken = this.e2eeChannels.get(ws)?.deviceToken ?? null - void this.handleWebSocketMessage( - plaintext, - encryptedReply, - encryptedBinaryReply, - wsTransport, - ws, - authenticatedDeviceToken - ) - }) - channel.onBinaryMessage((bytes) => this.handleWebSocketBinaryMessage(bytes, ws)) - this.e2eeChannels.set(ws, channel) - } - channel.handleRawMessage(msg) - }) - - // Why: when a mobile client disconnects, the runtime must clean up - // connection-scoped state like mobile-fit overrides and the E2EE - // channel to prevent orphaned state. A single paired device can hold - // multiple concurrent sockets (host screen + accounts screen, etc.), - // so destroy the channel for THIS exact ws and skip the per-client - // teardown when other sockets for the same token are still alive. - wsTransport.onConnectionClose((clientId, ws, hasOtherConnections) => { - this.abortWebSocketDispatches(ws) - // Why: sweep streaming subscriptions for THIS ws regardless of - // hasOtherConnections, so per-ws listeners (notifications, - // accounts, terminal) don't leak across reconnects. This is - // independent of the deviceToken-scoped onClientDisconnected. - const connectionId = this.wsConnectionIds.get(ws) - if (connectionId) { - this.runtime.cleanupSubscriptionsForConnection(connectionId) - this.runtime.cancelMobileDictationForConnection(connectionId) - this.binaryStreamHandlers.delete(connectionId) - this.wsConnectionIds.delete(ws) - } - const channel = this.e2eeChannels.get(ws) - if (channel) { - channel.destroy() - this.e2eeChannels.delete(ws) - } - if (clientId && !hasOtherConnections) { - this.runtime.onClientDisconnected(clientId) + const mobileSocketWiring = new MobileSocketWiring({ + deviceRegistry: this.deviceRegistry, + e2eeKeypair: this.e2eeKeypair, + onText: (socket, plaintext, reply, sendBinary) => { + void this.handleWebSocketMessage( + plaintext, + reply, + sendBinary, + undefined, + socket.ws, + socket.device.deviceToken + ) + }, + onBinary: (socket, bytes) => this.handleWebSocketBinaryMessage(bytes, socket.ws), + onClose: (socket, hasOtherConnections) => { + if (!socket) { + return + } + this.abortWebSocketDispatches(socket.ws) + // Why: subscriptions and binary streams are socket-scoped, while + // client disconnect state is device-scoped across both transports. + this.runtime.cleanupSubscriptionsForConnection(socket.connectionId) + this.runtime.cancelMobileDictationForConnection(socket.connectionId) + this.binaryStreamHandlers.delete(socket.connectionId) + if (!hasOtherConnections) { + this.runtime.onClientDisconnected(socket.device.deviceToken) + } } }) + mobileSocketWiring.attachTransport(wsTransport) + this.mobileSocketWiring = mobileSocketWiring await wsTransport.start() if (this.wsPort !== 0 && wsTransport.resolvedPort !== this.wsPort) { @@ -808,7 +753,7 @@ export class OrcaRuntimeRpcServer { // function if it fails to start (e.g., port in use). Log and continue // with Unix socket only. console.error('[runtime] Failed to start WebSocket transport:', error) - this.wsTransport = null + this.mobileSocketWiring = null } } @@ -835,7 +780,7 @@ export class OrcaRuntimeRpcServer { const transports = this.activeTransports this.activeTransports = [] this.transports = [] - this.wsTransport = null + this.mobileSocketWiring = null if (transports.length === 0) { return } @@ -1013,7 +958,7 @@ export class OrcaRuntimeRpcServer { ? (response: string): void => reply(injectDeviceScope(response, device.scope)) : reply - const connectionId = ws ? this.wsConnectionIds.get(ws) : undefined + const connectionId = ws ? this.mobileSocketWiring?.getConnectionId(ws) : undefined try { await this.dispatcher.dispatchStreaming(request, replyForRequest, { connectionId,