mirror of
https://github.com/stablyai/orca.git
synced 2026-10-17 08:03:14 +00:00
refactor(runtime): unify mobile socket wiring
Co-authored-by: Orca <help@stably.ai>
This commit is contained in:
Jinwoo-H
Orca
@@ -34,7 +34,10 @@ function setup(overrides?: Partial<E2EEChannelOptions>) {
|
||||
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
|
||||
|
||||
@@ -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<typeof setup>) {
|
||||
@@ -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()
|
||||
})
|
||||
|
||||
|
||||
@@ -27,7 +27,10 @@ function setup(overrides?: Partial<E2EEChannelOptions>) {
|
||||
|
||||
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')
|
||||
|
||||
|
||||
@@ -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<typeof setTimeout> | 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<string> | 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<ArrayBufferLike>): 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
|
||||
|
||||
@@ -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
|
||||
})
|
||||
|
||||
@@ -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<MobileSocketTransport['onMessage']>[0] | null = null
|
||||
private closeHandler: Parameters<MobileSocketTransport['onConnectionClose']>[0] | null = null
|
||||
readonly setClientId = vi.fn()
|
||||
readonly terminateClientConnections = vi.fn(() => 0)
|
||||
|
||||
onMessage(handler: Parameters<MobileSocketTransport['onMessage']>[0]): void {
|
||||
this.messageHandler = handler
|
||||
}
|
||||
|
||||
onConnectionClose(handler: Parameters<MobileSocketTransport['onConnectionClose']>[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')
|
||||
})
|
||||
})
|
||||
@@ -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<ArrayBufferLike>
|
||||
|
||||
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<ArrayBufferLike>) => boolean | void
|
||||
) => void
|
||||
onBinary: (socket: AuthenticatedMobileSocket, bytes: Uint8Array<ArrayBufferLike>) => 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<WebSocket, E2EEChannel>()
|
||||
private readonly connectionIds = new Map<WebSocket, string>()
|
||||
private readonly authenticatedSockets = new Map<WebSocket, AuthenticatedMobileSocket>()
|
||||
private readonly transports = new Set<MobileSocketTransport>()
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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<string, unknown>[] = []
|
||||
|
||||
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 {
|
||||
|
||||
@@ -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<WebSocket, E2EEChannel>()
|
||||
// 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<WebSocket, string>()
|
||||
private mobileSocketWiring: MobileSocketWiring | null = null
|
||||
private readonly binaryStreamHandlers = new Map<
|
||||
string,
|
||||
Map<number, (frame: TerminalStreamFrame) => 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<ArrayBufferLike>, 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,
|
||||
|
||||
Reference in New Issue
Block a user