From 5ed7a279718cd80f1be829c30176e70e128bbac2 Mon Sep 17 00:00:00 2001 From: OrcaWin <293788423+OrcaWin@users.noreply.github.com> Date: Mon, 27 Jul 2026 20:11:45 -0700 Subject: [PATCH] feat(ssh): negotiate bounded PTY consumer sessions --- .../wsl-hook-relay-sentinel.test.ts | 14 +- .../agent-hooks/wsl-hook-relay-sentinel.ts | 2 + .../providers/ssh-pty-provider-contract.ts | 1 + src/main/providers/ssh-pty-spawn-env.ts | 3 + src/main/ssh/relay-frame-decoder.ts | 9 +- .../ssh/relay-protocol-backpressure.test.ts | 13 + src/main/ssh/relay-protocol.ts | 6 +- src/main/ssh/ssh-pty-consumer-session.test.ts | 103 ++++++ src/main/ssh/ssh-pty-consumer-session.ts | 104 ++++++ src/main/ssh/ssh-relay-deploy-helpers.ts | 2 + src/main/ssh/ssh-relay-deploy.test.ts | 1 + src/main/ssh/ssh-relay-deploy.ts | 72 +++- ...ay-session-agent-hooks.integration.test.ts | 31 +- .../ssh-relay-session-data-delivery.test.ts | 8 + .../ssh/ssh-relay-session-incarnation.test.ts | 8 + ...elay-session-reconnect-incarnation.test.ts | 8 + .../ssh-relay-session-terminal-error.test.ts | 9 + src/main/ssh/ssh-relay-session.test.ts | 9 + src/main/ssh/ssh-relay-session.ts | 86 ++++- src/main/ssh/ssh-remote-cli-launcher.ts | 11 +- src/relay/dispatcher.test.ts | 29 ++ src/relay/dispatcher.ts | 110 +++++- src/relay/protocol-backpressure.test.ts | 13 + src/relay/protocol.ts | 8 +- src/relay/pty-handler.ts | 15 + src/relay/relay-frame-decoder.ts | 12 +- src/relay/relay-handshake-roundtrip.test.ts | 23 +- src/relay/relay-handshake.ts | 30 +- src/relay/relay.ts | 203 +++++++---- .../ssh-pty-consumer-session-adapter.test.ts | 210 ++++++++++++ src/relay/ssh-pty-consumer-session-adapter.ts | 120 +++++++ src/shared/pty-consumer-session-contract.ts | 58 ++++ src/shared/pty-consumer-session.test.ts | 242 ++++++++++++++ src/shared/pty-consumer-session.ts | 315 ++++++++++++++++++ 34 files changed, 1770 insertions(+), 118 deletions(-) create mode 100644 src/main/ssh/ssh-pty-consumer-session.test.ts create mode 100644 src/main/ssh/ssh-pty-consumer-session.ts create mode 100644 src/relay/ssh-pty-consumer-session-adapter.test.ts create mode 100644 src/relay/ssh-pty-consumer-session-adapter.ts create mode 100644 src/shared/pty-consumer-session-contract.ts create mode 100644 src/shared/pty-consumer-session.test.ts create mode 100644 src/shared/pty-consumer-session.ts diff --git a/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts b/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts index a8085485553..b806f43da74 100644 --- a/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts +++ b/src/main/agent-hooks/wsl-hook-relay-sentinel.test.ts @@ -16,12 +16,18 @@ type FakeChild = ChildProcessWithoutNullStreams & { kill: ReturnType + resume: ReturnType + } stderr: EventEmitter stdin: { write: ReturnType } kill: ReturnType } - child.stdout = new EventEmitter() + child.stdout = Object.assign(new EventEmitter(), { + pause: vi.fn(), + resume: vi.fn() + }) child.stderr = new EventEmitter() child.stdin = { write: vi.fn(() => true) } child.kill = vi.fn() @@ -62,6 +68,10 @@ describe('waitForWslRelaySentinel', () => { const transport = await promise expect(typeof transport.write).toBe('function') expect(typeof transport.onData).toBe('function') + transport.pauseReads?.() + transport.resumeReads?.() + expect(child.stdout.pause).toHaveBeenCalledOnce() + expect(child.stdout.resume).toHaveBeenCalledOnce() }) it('resolves past leading garbage and hands trailing bytes to onData', async () => { diff --git a/src/main/agent-hooks/wsl-hook-relay-sentinel.ts b/src/main/agent-hooks/wsl-hook-relay-sentinel.ts index fdabf70fe15..30dd59241ca 100644 --- a/src/main/agent-hooks/wsl-hook-relay-sentinel.ts +++ b/src/main/agent-hooks/wsl-hook-relay-sentinel.ts @@ -143,6 +143,8 @@ export function waitForWslRelaySentinel( } }, onClose: (cb) => closeCallbacks.push(cb), + pauseReads: () => child.stdout.pause(), + resumeReads: () => child.stdout.resume(), close: () => child.kill() } resolve(transport) diff --git a/src/main/providers/ssh-pty-provider-contract.ts b/src/main/providers/ssh-pty-provider-contract.ts index 9b709d9e08d..57ace057dcc 100644 --- a/src/main/providers/ssh-pty-provider-contract.ts +++ b/src/main/providers/ssh-pty-provider-contract.ts @@ -5,6 +5,7 @@ export type RemoteCliBridgeEnv = { relayDir: string nodePath: string sockPath: string + credentialFile?: string pathDelimiter?: ':' | ';' } diff --git a/src/main/providers/ssh-pty-spawn-env.ts b/src/main/providers/ssh-pty-spawn-env.ts index ad9cc7377a4..595e750c90d 100644 --- a/src/main/providers/ssh-pty-spawn-env.ts +++ b/src/main/providers/ssh-pty-spawn-env.ts @@ -22,6 +22,9 @@ export function buildSshPtySpawnEnv(args: { merged.ORCA_RELAY_DIR = args.remoteCliBridgeEnv.relayDir merged.ORCA_RELAY_NODE_PATH = args.remoteCliBridgeEnv.nodePath merged.ORCA_RELAY_SOCKET_PATH = args.remoteCliBridgeEnv.sockPath + if (args.remoteCliBridgeEnv.credentialFile) { + merged.ORCA_RELAY_CREDENTIAL_FILE = args.remoteCliBridgeEnv.credentialFile + } } // Why: match local/daemon precedence—managed defaults cannot restore explicitly removed values. for (const key of args.envToDelete ?? []) { diff --git a/src/main/ssh/relay-frame-decoder.ts b/src/main/ssh/relay-frame-decoder.ts index 2f4683a89fc..009f799c4fb 100644 --- a/src/main/ssh/relay-frame-decoder.ts +++ b/src/main/ssh/relay-frame-decoder.ts @@ -2,7 +2,8 @@ export const HEADER_LENGTH = 13 export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 export const FRAME_DECODER_MAX_FRAMES_PER_TURN = 64 export const FRAME_DECODER_MAX_BYTES_PER_TURN = MAX_MESSAGE_SIZE + HEADER_LENGTH -export const FRAME_DECODER_MAX_TURN_MS = 4 +export const FRAME_DECODER_MAX_TURN_MS = 4, + FRAME_DECODER_MAX_RETAINED_BYTES = 2 * (MAX_MESSAGE_SIZE + HEADER_LENGTH) export type DecodedFrame = { type: number @@ -68,6 +69,12 @@ export class FrameDecoder { const buf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + const retained = this.bufferedLength + buf.length + if (retained > FRAME_DECODER_MAX_RETAINED_BYTES) { + this.reset() + this.onError?.(new Error(`Frame decoder retained-input limit exceeded: ${retained}`)) + return + } if (buf.length > 0) { this.chunks.push(buf) this.bufferedLength += buf.length diff --git a/src/main/ssh/relay-protocol-backpressure.test.ts b/src/main/ssh/relay-protocol-backpressure.test.ts index d78b64ff1fe..3cd42375821 100644 --- a/src/main/ssh/relay-protocol-backpressure.test.ts +++ b/src/main/ssh/relay-protocol-backpressure.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from 'vitest' import { FrameDecoder, + FRAME_DECODER_MAX_RETAINED_BYTES, HEADER_LENGTH, MAX_MESSAGE_SIZE, MessageType, @@ -40,6 +41,18 @@ function frame(id: number, payload = `${id}`): Buffer { } describe('FrameDecoder bounded turns', () => { + it('fails closed when one delivered chunk exceeds retained-input capacity', () => { + const onError = vi.fn() + const decoder = new FrameDecoder(vi.fn(), onError) + + decoder.feed(Buffer.alloc(FRAME_DECODER_MAX_RETAINED_BYTES + 1)) + + expect(onError).toHaveBeenCalledWith( + expect.objectContaining({ message: expect.stringContaining('retained-input') }) + ) + expect(decoder.drain()).toHaveLength(0) + }) + it('emits the first frame synchronously and preserves order through self-pause', () => { const scheduler = createScheduler() const seen: number[] = [] diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index afc7a9ea6a4..00f6c1605d6 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -9,7 +9,8 @@ import { MAX_MESSAGE_SIZE, FRAME_DECODER_MAX_FRAMES_PER_TURN, FRAME_DECODER_MAX_BYTES_PER_TURN, - FRAME_DECODER_MAX_TURN_MS + FRAME_DECODER_MAX_TURN_MS, + FRAME_DECODER_MAX_RETAINED_BYTES } from './relay-frame-decoder' export { @@ -18,7 +19,8 @@ export { MAX_MESSAGE_SIZE, FRAME_DECODER_MAX_FRAMES_PER_TURN, FRAME_DECODER_MAX_BYTES_PER_TURN, - FRAME_DECODER_MAX_TURN_MS + FRAME_DECODER_MAX_TURN_MS, + FRAME_DECODER_MAX_RETAINED_BYTES } export type { DecodedFrame, FrameDecoderOptions } from './relay-frame-decoder' diff --git a/src/main/ssh/ssh-pty-consumer-session.test.ts b/src/main/ssh/ssh-pty-consumer-session.test.ts new file mode 100644 index 00000000000..46936d09b55 --- /dev/null +++ b/src/main/ssh/ssh-pty-consumer-session.test.ts @@ -0,0 +1,103 @@ +import { describe, expect, it, vi } from 'vitest' +import type { SshChannelMultiplexer } from './ssh-channel-multiplexer' +import { openSshPtyConsumerSession } from './ssh-pty-consumer-session' + +function muxReturning(result: unknown): { + mux: SshChannelMultiplexer + request: ReturnType +} { + const request = vi.fn().mockResolvedValue(result) + return { mux: { request } as unknown as SshChannelMultiplexer, request } +} + +function legacyOwnerGrant(overrides: Record = {}): Record { + return { + protocolVersion: 1, + serverBuildId: 'build-a', + clientGeneration: 3, + role: 'session-owner', + ownerGeneration: 7, + ownerLease: 'lease-a', + ...overrides + } +} + +describe('openSshPtyConsumerSession', () => { + it('makes openClient the one request needed for token-free legacy readiness', async () => { + const { mux, request } = muxReturning(legacyOwnerGrant()) + + await expect( + openSshPtyConsumerSession(mux, { + clientInstanceId: 'client-a', + expectedServerBuildId: 'build-a' + }) + ).resolves.toEqual({ + clientInstanceId: 'client-a', + clientGeneration: 3, + ownerGeneration: 7, + ownerLease: 'lease-a' + }) + expect(request).toHaveBeenCalledWith( + 'pty.openClient', + { + protocolVersion: 1, + clientInstanceId: 'client-a', + requestedRole: 'session-owner' + }, + { timeoutMs: 10_000 } + ) + }) + + it('carries recovery generation and lease on reconnect', async () => { + const { mux, request } = muxReturning( + legacyOwnerGrant({ ownerGeneration: 8, ownerLease: 'lease-b' }) + ) + await openSshPtyConsumerSession(mux, { + clientInstanceId: 'client-a', + expectedServerBuildId: 'build-a', + resume: { ownerGeneration: 7, ownerLease: 'lease-a' } + }) + + expect(request.mock.calls[0][1]).toMatchObject({ + resume: { ownerGeneration: 7, ownerLease: 'lease-a' } + }) + }) + + it('rejects a prior or mismatched relay build', async () => { + const { mux } = muxReturning(legacyOwnerGrant({ serverBuildId: 'old-build' })) + + await expect( + openSshPtyConsumerSession(mux, { + clientInstanceId: 'client-a', + expectedServerBuildId: 'build-a' + }) + ).rejects.toThrow('session contract mismatch') + }) + + it('does not silently downgrade when V1 was offered', async () => { + const { mux } = muxReturning(legacyOwnerGrant()) + + await expect( + openSshPtyConsumerSession(mux, { + clientInstanceId: 'client-a', + expectedServerBuildId: 'build-a', + outputFlowControl: { requestedWindowSu: 64 } + }) + ).rejects.toThrow('did not grant') + }) + + it('rejects an unoffered V1 capability in a legacy session', async () => { + const { mux } = muxReturning( + legacyOwnerGrant({ + capabilities: { outputFlowControl: { version: 1, windowSu: 64 } } + }) + ) + + await expect( + openSshPtyConsumerSession(mux, { + clientInstanceId: 'client-a', + expectedServerBuildId: 'build-a' + }) + ).rejects.toThrow('unoffered') + }) +}) diff --git a/src/main/ssh/ssh-pty-consumer-session.ts b/src/main/ssh/ssh-pty-consumer-session.ts new file mode 100644 index 00000000000..f106ede4187 --- /dev/null +++ b/src/main/ssh/ssh-pty-consumer-session.ts @@ -0,0 +1,104 @@ +import { + PTY_CONSUMER_SESSION_PROTOCOL_VERSION, + type PtyConsumerSessionGrant +} from '../../shared/pty-consumer-session' +import type { SshChannelMultiplexer } from './ssh-channel-multiplexer' + +export const SSH_PTY_OPEN_CLIENT_METHOD = 'pty.openClient' +export const SSH_PTY_OPEN_CLIENT_TIMEOUT_MS = 10_000 + +export type SshPtyConsumerOwnerState = { + clientInstanceId: string + clientGeneration: number + ownerGeneration: number + ownerLease: string +} + +export type OpenSshPtyConsumerSessionOptions = { + clientInstanceId: string + expectedServerBuildId: string | undefined + resume?: Pick + outputFlowControl?: { + requestedWindowSu: number + } +} + +function validateGrant( + value: unknown, + options: OpenSshPtyConsumerSessionOptions +): PtyConsumerSessionGrant { + if (typeof value !== 'object' || value === null) { + throw new Error('Remote relay returned an invalid pty.openClient grant') + } + if (!options.expectedServerBuildId) { + throw new Error('Local relay build identity is unavailable') + } + const grant = value as Partial + if ( + grant.protocolVersion !== PTY_CONSUMER_SESSION_PROTOCOL_VERSION || + grant.serverBuildId !== options.expectedServerBuildId + ) { + throw new Error( + `Remote relay session contract mismatch — expected build ${options.expectedServerBuildId}, got ${grant.serverBuildId ?? 'unknown'}` + ) + } + if ( + !Number.isSafeInteger(grant.clientGeneration) || + grant.clientGeneration! <= 0 || + grant.role !== 'session-owner' || + !Number.isSafeInteger(grant.ownerGeneration) || + grant.ownerGeneration! <= 0 || + typeof grant.ownerLease !== 'string' || + grant.ownerLease.length === 0 + ) { + throw new Error('Remote relay did not grant an authenticated PTY session owner') + } + const requestedFlow = options.outputFlowControl + const grantedFlow = grant.capabilities?.outputFlowControl + if (requestedFlow) { + if ( + grantedFlow?.version !== 1 || + !Number.isSafeInteger(grantedFlow.windowSu) || + grantedFlow.windowSu <= 0 || + grantedFlow.windowSu > requestedFlow.requestedWindowSu + ) { + throw new Error('Remote relay did not grant the offered PTY output-flow-control capability') + } + } else if (grantedFlow) { + throw new Error('Remote relay granted an unoffered PTY output-flow-control capability') + } + return grant as PtyConsumerSessionGrant +} + +export async function openSshPtyConsumerSession( + mux: SshChannelMultiplexer, + options: OpenSshPtyConsumerSessionOptions +): Promise { + const result = await mux.request( + SSH_PTY_OPEN_CLIENT_METHOD, + { + protocolVersion: PTY_CONSUMER_SESSION_PROTOCOL_VERSION, + clientInstanceId: options.clientInstanceId, + requestedRole: 'session-owner', + ...(options.resume ? { resume: options.resume } : {}), + ...(options.outputFlowControl + ? { + capabilities: { + outputFlowControl: { + versions: [1], + requestedWindowSu: options.outputFlowControl.requestedWindowSu + } + } + } + : {}) + }, + { timeoutMs: SSH_PTY_OPEN_CLIENT_TIMEOUT_MS } + ) + const grant = validateGrant(result, options) + return { + clientInstanceId: options.clientInstanceId, + clientGeneration: grant.clientGeneration, + ownerGeneration: grant.ownerGeneration!, + ownerLease: grant.ownerLease! + } +} diff --git a/src/main/ssh/ssh-relay-deploy-helpers.ts b/src/main/ssh/ssh-relay-deploy-helpers.ts index 80d2b78ced2..6cd5a652430 100644 --- a/src/main/ssh/ssh-relay-deploy-helpers.ts +++ b/src/main/ssh/ssh-relay-deploy-helpers.ts @@ -233,6 +233,8 @@ export function waitForSentinel( cb() } }, + pauseReads: () => channel.pause(), + resumeReads: () => channel.resume(), close: () => { channel.close() } diff --git a/src/main/ssh/ssh-relay-deploy.test.ts b/src/main/ssh/ssh-relay-deploy.test.ts index 723cbbbf413..244d0c43f09 100644 --- a/src/main/ssh/ssh-relay-deploy.test.ts +++ b/src/main/ssh/ssh-relay-deploy.test.ts @@ -115,6 +115,7 @@ function makeMockConnection(): SshConnection { stdout: { on: vi.fn() }, close: vi.fn() }), + writeFile: vi.fn().mockResolvedValue(undefined), sftp: vi.fn().mockResolvedValue({ mkdir: vi.fn((_p: string, cb: (err: Error | null) => void) => cb(null)), createWriteStream: vi.fn().mockReturnValue({ diff --git a/src/main/ssh/ssh-relay-deploy.ts b/src/main/ssh/ssh-relay-deploy.ts index 81165efc4d0..ac9bdcaffc9 100644 --- a/src/main/ssh/ssh-relay-deploy.ts +++ b/src/main/ssh/ssh-relay-deploy.ts @@ -1,4 +1,5 @@ import { join } from 'node:path' +import { randomBytes } from 'node:crypto' /* eslint-disable max-lines -- Why: one cohesive contract (version detect, install-locked deploy, native-deps probe, launch, GC); splitting risks install/GC drift. */ import { existsSync } from 'node:fs' import { app } from 'electron' @@ -75,12 +76,14 @@ import { export type RelayDeployResult = { transport: MultiplexerTransport + serverBuildId?: string platform: RelayPlatform hostPlatform?: RemoteHostPlatform remoteHome?: string remoteRelayDir?: string nodePath?: string sockPath?: string + credentialFile?: string } class RelayDirectoryGcConflictError extends Error { @@ -434,12 +437,14 @@ async function deployAndLaunchRelayAttempt( return { transport: launched.transport, + serverBuildId: fullVersion, platform, hostPlatform, remoteHome, remoteRelayDir, nodePath: launched.nodePath, - sockPath: launched.sockPath + sockPath: launched.sockPath, + credentialFile: launched.credentialFile } } @@ -1157,7 +1162,12 @@ async function launchRelay( graceTimeSeconds?: number, relayInstanceId?: string, signal?: AbortSignal -): Promise<{ transport: MultiplexerTransport; nodePath: string; sockPath: string }> { +): Promise<{ + transport: MultiplexerTransport + nodePath: string + sockPath: string + credentialFile: string +}> { // Why: graceTimeSeconds comes from user-editable SshTarget config; floor+clamp to an integer prevents shell injection if the type ever loosened. const requestedGraceTime = Math.floor(graceTimeSeconds ?? DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS) const graceTime = @@ -1173,6 +1183,7 @@ async function launchRelay( const sockName = relaySocketNameForInstanceId(relayInstanceId) const sockFile = relayEndpointForHost(hostPlatform, remoteDir, sockName) const endpointDir = relayHookEndpointDirForHost(hostPlatform, remoteDir, sockFile) + const credentialFile = joinRemotePath(hostPlatform, remoteDir, `${sockName}.credential`) if (isWindowsRemoteHost(hostPlatform)) { const activePipeMarkerPath = windowsActivePipeMarkerPath(hostPlatform, remoteDir, sockName) @@ -1187,7 +1198,7 @@ async function launchRelay( endpointDir } const fallbackEndpoint = buildWindowsRelayFallbackEndpoint(hostPlatform, remoteDir, sockName) - return launchWindowsRelay( + const launched = await launchWindowsRelay( conn, hostPlatform, { @@ -1197,10 +1208,12 @@ async function launchRelay( endpointDir: activeEndpoint.endpointDir, graceTime, activePipeMarkerPath, - reconnectFallback: fallbackEndpoint + reconnectFallback: fallbackEndpoint, + credentialFile }, signal ) + return { ...launched, credentialFile } } // Why: after a restart the relay may still be alive in its grace period; --connect to its socket preserves PTY state and scrollback. @@ -1215,12 +1228,12 @@ async function launchRelay( console.log('[ssh-relay] Existing relay socket found, attempting reconnect...') try { const channel = await conn.exec( - `cd ${escapedDir} && ${escapedNode} relay.js --connect --sock-path ${shellEscape(sockFile)}`, + `cd ${escapedDir} && ${escapedNode} relay.js --connect --sock-path ${shellEscape(sockFile)} --credential-file ${shellEscape(credentialFile)}`, { signal } ) const transport = await waitForSentinel(channel, signal) console.log('[ssh-relay] Reconnected to existing relay via socket') - return { transport, nodePath, sockPath: sockFile } + return { transport, nodePath, sockPath: sockFile, credentialFile } } catch (err) { signal?.throwIfAborted() console.warn( @@ -1249,8 +1262,9 @@ async function launchRelay( // Why: relay must outlive the SSH connection so PTY sessions survive app restarts — nohup + ${shellEscape(logFile)} 2>&1 ${shellEscape(logFile)} 2>&1 {}) launchChannel.on('error', () => {}) @@ -1300,10 +1314,25 @@ async function launchRelay( // Why: backgrounded relay's stdout goes to a log file, not the exec channel; --connect bridges this channel to its Unix socket. const channel = await conn.exec( - `cd ${escapedDir} && ${escapedNode} relay.js --connect --sock-path ${shellEscape(sockFile)}`, + `cd ${escapedDir} && ${escapedNode} relay.js --connect --sock-path ${shellEscape(sockFile)} --credential-file ${shellEscape(credentialFile)}`, { signal } ) - return { transport: await waitForSentinel(channel, signal), nodePath, sockPath: sockFile } + return { + transport: await waitForSentinel(channel, signal), + nodePath, + sockPath: sockFile, + credentialFile + } +} + +async function writeRelayEndpointCredential( + conn: SshConnection, + hostPlatform: RemoteHostPlatform, + credentialFile: string, + signal?: AbortSignal +): Promise { + const credential = randomBytes(32).toString('base64url') + await writeRelayFile(conn, hostPlatform, credentialFile, credential, { signal }) } function waitForRelayPoll(delayMs: number, signal?: AbortSignal): Promise { @@ -1398,6 +1427,7 @@ type WindowsRelayLaunchOptions = { nodePath: string graceTime: number activePipeMarkerPath: string + credentialFile: string } & WindowsRelayEndpoint & { reconnectFallback?: WindowsRelayEndpoint } @@ -1467,6 +1497,7 @@ async function launchWindowsRelay( const logFile = joinRemotePath(hostPlatform, launchOpts.remoteDir, 'relay.log') const errFile = joinRemotePath(hostPlatform, launchOpts.remoteDir, 'relay.err.log') + await writeRelayEndpointCredential(conn, hostPlatform, launchOpts.credentialFile, signal) await execHostCommand( conn, hostPlatform, @@ -1478,7 +1509,8 @@ async function launchWindowsRelay( launchOpts.endpointDir, launchOpts.graceTime, logFile, - errFile + errFile, + launchOpts.credentialFile ), { signal } ) @@ -1529,11 +1561,18 @@ async function connectWindowsRelay( remoteDir: string nodePath: string sockPath: string + credentialFile: string }, signal?: AbortSignal ): Promise { const channel = await conn.exec( - windowsRelayConnectCommand(hostPlatform, opts.nodePath, opts.remoteDir, opts.sockPath), + windowsRelayConnectCommand( + hostPlatform, + opts.nodePath, + opts.remoteDir, + opts.sockPath, + opts.credentialFile + ), { wrapCommand: false, signal } ) return waitForSentinel(channel, signal) @@ -1543,13 +1582,14 @@ function windowsRelayConnectCommand( hostPlatform: RemoteHostPlatform, nodePath: string, remoteDir: string, - sockPath: string + sockPath: string, + credentialFile: string ): string { return commandWithNodePath( hostPlatform, nodePath, remoteDir, - `& ${powerShellLiteral(nodePath)} relay.js --connect --sock-path ${powerShellLiteral(sockPath)}` + `& ${powerShellLiteral(nodePath)} relay.js --connect --sock-path ${powerShellLiteral(sockPath)} --credential-file ${powerShellLiteral(credentialFile)}` ) } @@ -1561,7 +1601,8 @@ function windowsRelayLaunchCommand( endpointDir: string, graceTime: number, logFile: string, - errFile: string + errFile: string, + credentialFile: string ): string { const relayScript = joinRemotePath(hostPlatform, remoteDir, 'relay.js') // Why: Windows sshd kills the exec channel's process tree on close; WMI re-parents the detached relay to survive. @@ -1574,6 +1615,8 @@ function windowsRelayLaunchCommand( String(graceTime), '--sock-path', quoted(sockPath), + '--credential-file', + quoted(credentialFile), '--endpoint-dir', quoted(endpointDir), // Why: --log-file owns rotation; shell redirects still capture pre-JS boot/crash output. @@ -1588,6 +1631,7 @@ function windowsRelayLaunchCommand( nodePath, remoteDir, [ + `& icacls.exe ${powerShellLiteral(credentialFile)} /inheritance:r /grant:r "$($env:USERNAME):(R,W)" | Out-Null`, `$result = Invoke-CimMethod -ClassName Win32_Process -MethodName Create -Arguments @{ CommandLine = ${powerShellLiteral(wmiCommandLine)}; CurrentDirectory = ${powerShellLiteral(remoteDir)} }`, `if ($result.ReturnValue -ne 0) { throw "Win32_Process.Create failed with $($result.ReturnValue)" }` ].join('; ') diff --git a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts index d6c98bf03a0..cd96cb8f3af 100644 --- a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts +++ b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts @@ -97,6 +97,18 @@ function createFakeRelay(): FakeRelay { }) relayFeed = (data) => dispatcher.feed(data) + dispatcher.onRequest('pty.openClient', async (params) => ({ + protocolVersion: 1, + serverBuildId: 'test-relay-build', + clientGeneration: 1, + role: 'session-owner', + ownerGeneration: + typeof (params.resume as { ownerGeneration?: unknown } | undefined)?.ownerGeneration === + 'number' + ? (params.resume as { ownerGeneration: number }).ownerGeneration + 1 + : 1, + ownerLease: 'test-owner-lease' + })) dispatcher.onRequest('session.resolveHome', async (params) => ({ resolvedPath: params.path === '~' ? '/home/orca' : params.path })) @@ -224,6 +236,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) const events: CapturedStatus[] = [] @@ -277,8 +290,16 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { const initialRelay = createFakeRelay() relay = createFakeRelay() vi.mocked(deployAndLaunchRelay) - .mockResolvedValueOnce({ transport: initialRelay.transport, platform: 'linux-x64' }) - .mockResolvedValueOnce({ transport: relay.transport, platform: 'linux-x64' }) + .mockResolvedValueOnce({ + transport: initialRelay.transport, + serverBuildId: 'test-relay-build', + platform: 'linux-x64' + }) + .mockResolvedValueOnce({ + transport: relay.transport, + serverBuildId: 'test-relay-build', + platform: 'linux-x64' + }) const clearListener = vi.fn() agentHookServer.setPaneStatusClearListener(clearListener) session = createSession('conn-clear') @@ -318,6 +339,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { ) vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) const events: CapturedStatus[] = [] @@ -345,6 +367,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) const events: CapturedStatus[] = [] @@ -375,6 +398,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) @@ -434,6 +458,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) const ingestSpy = vi.spyOn(agentHookServer, 'ingestRemote') @@ -488,6 +513,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) @@ -535,6 +561,7 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ transport: relay.transport, + serverBuildId: 'test-relay-build', platform: 'linux-x64' }) diff --git a/src/main/ssh/ssh-relay-session-data-delivery.test.ts b/src/main/ssh/ssh-relay-session-data-delivery.test.ts index 06c09d1b082..7f7e95d963a 100644 --- a/src/main/ssh/ssh-relay-session-data-delivery.test.ts +++ b/src/main/ssh/ssh-relay-session-data-delivery.test.ts @@ -8,6 +8,14 @@ const { acceptOutputDataMock, muxRequestMock } = vi.hoisted(() => ({ })) vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputData: acceptOutputDataMock, acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined), diff --git a/src/main/ssh/ssh-relay-session-incarnation.test.ts b/src/main/ssh/ssh-relay-session-incarnation.test.ts index b06a78bfca1..517201f684b 100644 --- a/src/main/ssh/ssh-relay-session-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-incarnation.test.ts @@ -8,6 +8,14 @@ const { acceptOutputExitMock, muxRequestMock } = vi.hoisted(() => ({ })) vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined), acceptSshPtyOutputExit: acceptOutputExitMock, diff --git a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts index 0fd63d68d6a..1d2ae2a1a19 100644 --- a/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts +++ b/src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts @@ -8,6 +8,14 @@ const { acceptOutputExitMock, muxRequestMock } = vi.hoisted(() => ({ })) vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined), acceptSshPtyOutputExit: acceptOutputExitMock, diff --git a/src/main/ssh/ssh-relay-session-terminal-error.test.ts b/src/main/ssh/ssh-relay-session-terminal-error.test.ts index ce9317c6660..6b982fd1397 100644 --- a/src/main/ssh/ssh-relay-session-terminal-error.test.ts +++ b/src/main/ssh/ssh-relay-session-terminal-error.test.ts @@ -15,6 +15,15 @@ vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) + vi.mock('./ssh-channel-multiplexer', () => { return { SshChannelMultiplexer: class MockSshChannelMultiplexer { diff --git a/src/main/ssh/ssh-relay-session.test.ts b/src/main/ssh/ssh-relay-session.test.ts index 653bb28237d..16196e95b34 100644 --- a/src/main/ssh/ssh-relay-session.test.ts +++ b/src/main/ssh/ssh-relay-session.test.ts @@ -17,6 +17,15 @@ vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() })) +vi.mock('./ssh-pty-consumer-session', () => ({ + openSshPtyConsumerSession: vi.fn(async (_mux, options) => ({ + clientInstanceId: options.clientInstanceId, + clientGeneration: 1, + ownerGeneration: 1, + ownerLease: 'test-owner-lease' + })) +})) + vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({ acceptSshPtyOutputData: acceptOutputDataMock, acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined), diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index af35f6416a9..dd6e7792a47 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1,6 +1,7 @@ /* oxlint-disable max-lines */ // Why: single authority for all relay lifecycle state per SSH target (previously scattered across module Maps/Sets with duplicated paths). +import { randomUUID } from 'node:crypto' import type { BrowserWindow } from 'electron' import { deployAndLaunchRelay } from './ssh-relay-deploy' import { execCommand } from './ssh-relay-deploy-helpers' @@ -69,6 +70,10 @@ import { runRemoteOrcaCli } from './ssh-remote-orca-cli' import { toSshExecutionHostId, type ExecutionHostId } from '../../shared/execution-host' import { isTerminalLeafId, makePaneKey } from '../../shared/stable-pane-id' import { isValidTerminalTabId } from '../../shared/terminal-tab-id' +import { + openSshPtyConsumerSession, + type SshPtyConsumerOwnerState +} from './ssh-pty-consumer-session' export type RelaySessionState = 'idle' | 'deploying' | 'ready' | 'reconnecting' | 'disposed' @@ -81,6 +86,7 @@ type RemoteCliBridgeEnv = { relayDir: string nodePath: string sockPath: string + credentialFile?: string hostPlatform: RemoteHostPlatform pathDelimiter?: ':' | ';' } @@ -153,6 +159,8 @@ export class SshRelaySession { private forwardedReattachReplayByPty = new Map() private pendingPtyReattaches = new Map() private activePtyProviderGeneration: number | null = null + private readonly ptyConsumerClientInstanceId = randomUUID() + private ptyConsumerOwnerState: SshPtyConsumerOwnerState | null = null constructor( readonly targetId: string, @@ -251,8 +259,16 @@ export class SshRelaySession { this.currentConnection = conn try { - const { transport, remoteHome, remoteRelayDir, nodePath, sockPath, hostPlatform } = - await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) + const { + transport, + serverBuildId, + remoteHome, + remoteRelayDir, + nodePath, + sockPath, + credentialFile, + hostPlatform + } = await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) this.hostPlatform = hostPlatform ?? null this.remoteCliBridgeEnv = remoteHome && remoteRelayDir && nodePath && sockPath && hostPlatform @@ -262,6 +278,7 @@ export class SshRelaySession { relayDir: remoteRelayDir, nodePath, sockPath, + ...(credentialFile ? { credentialFile } : {}), hostPlatform, pathDelimiter: hostPlatform.pathDelimiter } @@ -278,7 +295,19 @@ export class SshRelaySession { this.mux = mux const ownsAttempt = (): boolean => this.mux === mux && !this.isDisposed() - // Why: round-trip the relay before registering providers so a closed --connect bridge fails fast instead of leaving a 'ready' session on a dead mux. + this.ptyConsumerOwnerState = await openSshPtyConsumerSession(mux, { + clientInstanceId: this.ptyConsumerClientInstanceId, + expectedServerBuildId: serverBuildId, + ...(this.ptyConsumerOwnerState + ? { + resume: { + ownerGeneration: this.ptyConsumerOwnerState.ownerGeneration, + ownerLease: this.ptyConsumerOwnerState.ownerLease + } + } + : {}) + }) + await mux.request('session.resolveHome', { path: '~' }) const registered = await this.registerProviders(mux, ownsAttempt) @@ -350,8 +379,16 @@ export class SshRelaySession { this.teardownProviders('connection_lost') try { - const { transport, remoteHome, remoteRelayDir, nodePath, sockPath, hostPlatform } = - await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) + const { + transport, + serverBuildId, + remoteHome, + remoteRelayDir, + nodePath, + sockPath, + credentialFile, + hostPlatform + } = await deployAndLaunchRelay(conn, undefined, graceTimeSeconds, this.targetId) this.hostPlatform = hostPlatform ?? null this.remoteCliBridgeEnv = remoteHome && remoteRelayDir && nodePath && sockPath && hostPlatform @@ -361,6 +398,7 @@ export class SshRelaySession { relayDir: remoteRelayDir, nodePath, sockPath, + ...(credentialFile ? { credentialFile } : {}), hostPlatform, pathDelimiter: hostPlatform.pathDelimiter } @@ -381,7 +419,25 @@ export class SshRelaySession { !abortController.signal.aborted && !this.isDisposed() - // Why: same health check as establish() — round-trip the relay before registering providers so a dead --connect bridge fails fast. + this.ptyConsumerOwnerState = await openSshPtyConsumerSession(mux, { + clientInstanceId: this.ptyConsumerClientInstanceId, + expectedServerBuildId: serverBuildId, + ...(this.ptyConsumerOwnerState + ? { + resume: { + ownerGeneration: this.ptyConsumerOwnerState.ownerGeneration, + ownerLease: this.ptyConsumerOwnerState.ownerLease + } + } + : {}) + }) + if (!ownsAttempt()) { + if (!mux.isDisposed()) { + mux.dispose() + } + return + } + await mux.request('session.resolveHome', { path: '~' }) if (!ownsAttempt()) { if (!mux.isDisposed()) { @@ -541,6 +597,24 @@ export class SshRelaySession { this.remoteCliBridgeEnv ?? undefined, providerGeneration ) + const consumerOwnerState = this.ptyConsumerOwnerState + if (consumerOwnerState) { + ptyProvider.setPtyDeliveryPauseAdapter?.(({ id, providerGeneration: generation, paused }) => { + if ( + generation !== providerGeneration || + this.activePtyProviderGeneration !== providerGeneration || + this.mux !== mux + ) { + return + } + mux.notify('pty.setDeliveryPaused', { + id, + paused, + clientGeneration: consumerOwnerState.clientGeneration, + ownerGeneration: consumerOwnerState.ownerGeneration + }) + }) + } this.activePtyProviderGeneration = providerGeneration registerSshPtyProvider(this.targetId, ptyProvider) diff --git a/src/main/ssh/ssh-remote-cli-launcher.ts b/src/main/ssh/ssh-remote-cli-launcher.ts index 4122e0dc09d..ca206d7312a 100644 --- a/src/main/ssh/ssh-remote-cli-launcher.ts +++ b/src/main/ssh/ssh-remote-cli-launcher.ts @@ -7,6 +7,7 @@ type RemoteCliInstallEnv = { relayDir: string nodePath: string sockPath: string + credentialFile?: string hostPlatform: RemoteHostPlatform } @@ -35,6 +36,7 @@ internal static class OrcaRemoteCliLauncher string nodePath = RequireEnvironmentVariable("ORCA_RELAY_NODE_PATH"); string relayDirectory = RequireEnvironmentVariable("ORCA_RELAY_DIR"); string socketPath = RequireEnvironmentVariable("ORCA_RELAY_SOCKET_PATH"); + string credentialFile = RequireEnvironmentVariable("ORCA_RELAY_CREDENTIAL_FILE"); string relayPath = Path.Combine(relayDirectory, "relay.js"); if (!File.Exists(nodePath)) @@ -51,7 +53,7 @@ internal static class OrcaRemoteCliLauncher ProcessStartInfo startInfo = new ProcessStartInfo { FileName = nodePath, - Arguments = BuildArguments(relayPath, socketPath, args), + Arguments = BuildArguments(relayPath, socketPath, credentialFile, args), UseShellExecute = false }; @@ -78,12 +80,14 @@ internal static class OrcaRemoteCliLauncher return value; } - private static string BuildArguments(string relayPath, string socketPath, string[] args) + private static string BuildArguments(string relayPath, string socketPath, string credentialFile, string[] args) { StringBuilder commandLine = new StringBuilder(); AppendArgument(commandLine, relayPath); AppendArgument(commandLine, "--sock-path"); AppendArgument(commandLine, socketPath); + AppendArgument(commandLine, "--credential-file"); + AppendArgument(commandLine, credentialFile); AppendArgument(commandLine, "--orca-cli"); foreach (string arg in args) { @@ -223,11 +227,12 @@ export function createRemoteCliInstallPlan(env: RemoteCliInstallEnv): RemoteCliI `ORCA_RELAY_NODE_PATH=\${ORCA_RELAY_NODE_PATH:-${quoteSh(env.nodePath)}}`, `ORCA_RELAY_DIR=\${ORCA_RELAY_DIR:-${quoteSh(env.relayDir)}}`, `ORCA_RELAY_SOCKET_PATH=\${ORCA_RELAY_SOCKET_PATH:-${quoteSh(env.sockPath)}}`, + `ORCA_RELAY_CREDENTIAL_FILE=\${ORCA_RELAY_CREDENTIAL_FILE:-${quoteSh(env.credentialFile ?? `${env.sockPath}.credential`)}}`, 'if [ ! -S "$ORCA_RELAY_SOCKET_PATH" ]; then', ' echo "Orca SSH CLI bridge cannot find the relay socket: $ORCA_RELAY_SOCKET_PATH" >&2', ' exit 1', 'fi', - 'exec "$ORCA_RELAY_NODE_PATH" "$ORCA_RELAY_DIR/relay.js" --sock-path "$ORCA_RELAY_SOCKET_PATH" --orca-cli "$@"', + 'exec "$ORCA_RELAY_NODE_PATH" "$ORCA_RELAY_DIR/relay.js" --sock-path "$ORCA_RELAY_SOCKET_PATH" --credential-file "$ORCA_RELAY_CREDENTIAL_FILE" --orca-cli "$@"', '' ].join('\n') } diff --git a/src/relay/dispatcher.test.ts b/src/relay/dispatcher.test.ts index f76e4ae809f..60bf9a79e92 100644 --- a/src/relay/dispatcher.test.ts +++ b/src/relay/dispatcher.test.ts @@ -513,6 +513,35 @@ describe('RelayDispatcher', () => { } }) + it('does not retry accepted clients when a later broadcast member closes reentrantly', () => { + const primary: string[] = [] + const secondary: string[] = [] + let secondaryId = 0 + const readData = (frame: Buffer): string => { + const message = JSON.parse(decodeFirstFrame(frame).payload.toString()) as JsonRpcNotification + return String(message.params?.data) + } + const dispatcher = new RelayDispatcher((frame) => { + primary.push(readData(frame)) + if (secondaryId !== 0) { + dispatcher.detachClient(secondaryId) + } + return true + }) + secondaryId = dispatcher.attachClient((frame) => { + secondary.push(readData(frame)) + return true + }) + + try { + expect(dispatcher.tryNotifyPtyData({ id: 'pty-1', data: 'once' })).toBe(true) + expect(primary).toEqual(['once']) + expect(secondary).toEqual([]) + } finally { + dispatcher.dispose() + } + }) + describe('notifyBulk (bulk lane backpressure)', () => { it('resolves immediately when the sink accepts the frame', async () => { const frames: Buffer[] = [] diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts index 9d1752565f9..7649141227a 100644 --- a/src/relay/dispatcher.ts +++ b/src/relay/dispatcher.ts @@ -37,6 +37,20 @@ export type RequestContext = { clientId: number isStale: () => boolean signal?: AbortSignal + sessionIdentity?: RelayClientSessionIdentity + onResponseSettled?: (handler: (result: SinkWriteSettlement) => void) => void +} + +export type RelayClientSessionIdentity = { + principal: string + authenticated: boolean + allowSessionOwner: boolean + authenticationKind: 'unproved' | 'launch-nonce' | 'endpoint-credential' +} + +export type RelayClientSourceOptions = { + pauseReads?: () => void + resumeReads?: () => void } export type MethodHandler = ( @@ -55,6 +69,7 @@ type RelayClient = { highestReceivedSeq: number generation: number closed: boolean + sessionIdentity: RelayClientSessionIdentity } type PendingRelayRequest = { @@ -83,8 +98,13 @@ export class RelayDispatcher { private nextClientId = 1 private nextRequestId = 1 - constructor(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions) { - this.primaryClient = this.createClient(write, sinkOptions) + constructor( + write: RelayClientWrite, + sinkOptions?: RelayClientSinkOptions, + sessionIdentity?: RelayClientSessionIdentity, + sourceOptions?: RelayClientSourceOptions + ) { + this.primaryClient = this.createClient(write, sinkOptions, sessionIdentity, sourceOptions) this.clients.set(this.primaryClient.id, this.primaryClient) this.startKeepalive() } @@ -105,8 +125,13 @@ export class RelayDispatcher { } // Why: seq numbers and request ids are per SSH channel, so each attached client needs independent protocol state. - attachClient(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): number { - const client = this.createClient(write, sinkOptions) + attachClient( + write: RelayClientWrite, + sinkOptions?: RelayClientSinkOptions, + sessionIdentity?: RelayClientSessionIdentity, + sourceOptions?: RelayClientSourceOptions + ): number { + const client = this.createClient(write, sinkOptions, sessionIdentity, sourceOptions) this.clients.set(client.id, client) return client.id } @@ -427,7 +452,12 @@ export class RelayDispatcher { this.legacyCapacityListeners.clear() } - private createClient(write: RelayClientWrite, sinkOptions?: RelayClientSinkOptions): RelayClient { + private createClient( + write: RelayClientWrite, + sinkOptions?: RelayClientSinkOptions, + sessionIdentity?: RelayClientSessionIdentity, + sourceOptions?: RelayClientSourceOptions + ): RelayClient { const id = this.nextClientId++ const client = { id, @@ -437,9 +467,19 @@ export class RelayDispatcher { nextOutgoingSeq: 1, highestReceivedSeq: 0, generation: 0, - closed: false + closed: false, + sessionIdentity: sessionIdentity ?? { + principal: `unproved:${id}`, + authenticated: false, + allowSessionOwner: false, + authenticationKind: 'unproved' + } } satisfies RelayClient - client.decoder = new FrameDecoder((frame) => this.handleFrame(client, frame)) + client.decoder = new FrameDecoder( + (frame) => this.handleFrame(client, frame), + (error) => this.closeClient(client, error, client !== this.primaryClient), + { pause: sourceOptions?.pauseReads, resume: sourceOptions?.resumeReads } + ) client.writer = this.createWriter(client, write, sinkOptions) return client } @@ -519,25 +559,57 @@ export class RelayDispatcher { client.id, req.id ) + const responseSettledHandlers = new Set<(result: SinkWriteSettlement) => void>() + let responseSettled = false + const settleResponse = (result: SinkWriteSettlement): void => { + if (responseSettled) { + return + } + responseSettled = true + for (const callback of responseSettledHandlers) { + try { + callback(result) + } catch (err) { + process.stderr.write( + `[relay] Response settlement callback failed: ${err instanceof Error ? err.message : String(err)}\n` + ) + } + } + responseSettledHandlers.clear() + } const context: RequestContext = { clientId: client.id, isStale: () => client.generation !== gen || !this.clients.has(client.id) || abortController.signal.aborted, - signal: abortController.signal + signal: abortController.signal, + sessionIdentity: client.sessionIdentity, + onResponseSettled: (handler) => { + if (responseSettled) { + throw new Error('Response settlement callback registered after settlement') + } + responseSettledHandlers.add(handler) + } } try { const result = await handler(req.params ?? {}, context) if (context.isStale()) { + settleResponse({ ok: false, error: new Error('Relay request became stale') }) return } - this.sendResponse(client, req.id, result) + this.sendResponse(client, req.id, result, undefined, settleResponse) } catch (err) { if (context.isStale()) { + settleResponse({ ok: false, error: new Error('Relay request became stale') }) return } const message = err instanceof Error ? err.message : String(err) const code = (err as { code?: number }).code ?? -32000 - this.sendResponse(client, req.id, undefined, { code, message }) + this.sendResponse(client, req.id, undefined, { code, message }, (result) => { + settleResponse({ + ok: false, + error: result.ok ? new Error(message) : result.error + }) + }) } finally { this.requestAborts.delete(abortKey) } @@ -555,7 +627,11 @@ export class RelayDispatcher { const gen = client.generation handler(notif.params ?? {}, { clientId: client.id, - isStale: () => client.generation !== gen || !this.clients.has(client.id) + isStale: () => client.generation !== gen || !this.clients.has(client.id), + sessionIdentity: client.sessionIdentity, + onResponseSettled: () => { + throw new Error('Notifications do not have response publication fences') + } }) } } @@ -564,8 +640,9 @@ export class RelayDispatcher { client: RelayClient, id: number, result?: unknown, - error?: { code: number; message: string; data?: unknown } - ): void { + error?: { code: number; message: string; data?: unknown }, + onSettled: (result: SinkWriteSettlement) => void = () => {} + ): boolean { const msg: JsonRpcResponse = { jsonrpc: '2.0', id, @@ -573,7 +650,8 @@ export class RelayDispatcher { } const estimatedBytes = this.estimateFrameBytes(msg) const lane = estimatedBytes > DISPATCHER_CONTROL_QUEUE_MAX_BYTES ? 'legacy-response' : 'control' - if (!this.enqueueFrame(client, msg, lane)) { + const accepted = this.enqueueFrame(client, msg, lane, onSettled) + if (!accepted) { this.closeClient( client, new Error( @@ -582,6 +660,7 @@ export class RelayDispatcher { client !== this.primaryClient ) } + return accepted } private enqueueFrame( @@ -665,6 +744,9 @@ export class RelayDispatcher { } for (let index = 0; index < clients.length; index++) { if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index])) { + if (this.disposed || clients[index].closed) { + continue + } for (let remaining = index; remaining < leases.length; remaining++) { leases[remaining].release() } diff --git a/src/relay/protocol-backpressure.test.ts b/src/relay/protocol-backpressure.test.ts index 18ed2e35495..a6569ec8ce9 100644 --- a/src/relay/protocol-backpressure.test.ts +++ b/src/relay/protocol-backpressure.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from 'vitest' import { FrameDecoder, + FRAME_DECODER_MAX_RETAINED_BYTES, HEADER_LENGTH, MAX_MESSAGE_SIZE, MessageType, @@ -40,6 +41,18 @@ function frame(id: number, payload = `${id}`): Buffer { } describe('relay FrameDecoder bounded turns', () => { + it('fails closed when one delivered chunk exceeds retained-input capacity', () => { + const onError = vi.fn() + const decoder = new FrameDecoder(vi.fn(), onError) + + decoder.feed(Buffer.alloc(FRAME_DECODER_MAX_RETAINED_BYTES + 1)) + + expect(onError).toHaveBeenCalledWith( + expect.objectContaining({ message: expect.stringContaining('retained-input') }) + ) + expect(decoder.drain()).toHaveLength(0) + }) + it('emits the first frame synchronously and preserves order through self-pause', () => { const scheduler = createScheduler() const seen: number[] = [] diff --git a/src/relay/protocol.ts b/src/relay/protocol.ts index 356655d64e4..f946b3e8985 100644 --- a/src/relay/protocol.ts +++ b/src/relay/protocol.ts @@ -7,7 +7,8 @@ import { MAX_MESSAGE_SIZE, FRAME_DECODER_MAX_FRAMES_PER_TURN, FRAME_DECODER_MAX_BYTES_PER_TURN, - FRAME_DECODER_MAX_TURN_MS + FRAME_DECODER_MAX_TURN_MS, + FRAME_DECODER_MAX_RETAINED_BYTES } from './relay-frame-decoder' export { @@ -16,7 +17,8 @@ export { MAX_MESSAGE_SIZE, FRAME_DECODER_MAX_FRAMES_PER_TURN, FRAME_DECODER_MAX_BYTES_PER_TURN, - FRAME_DECODER_MAX_TURN_MS + FRAME_DECODER_MAX_TURN_MS, + FRAME_DECODER_MAX_RETAINED_BYTES } export type { DecodedFrame, FrameDecoderOptions } from './relay-frame-decoder' @@ -34,7 +36,7 @@ export const MessageType = { // to refuse mismatched-version --connect bridges that would otherwise drive a // stale daemon. export type HandshakeMessage = - | { type: 'orca-relay-handshake'; version: string } + | { type: 'orca-relay-handshake'; version: string; endpointCredential?: string } | { type: 'orca-relay-handshake-ok'; version: string } | { type: 'orca-relay-handshake-mismatch'; expected: string; got: string } diff --git a/src/relay/pty-handler.ts b/src/relay/pty-handler.ts index c0a07ac22ee..0117907eda2 100644 --- a/src/relay/pty-handler.ts +++ b/src/relay/pty-handler.ts @@ -315,6 +315,7 @@ export class PtyHandler { private pendingOutputByPty = new Map() private pendingExitByPty = new Map() private pausedOutputPtys = new Set() + private consumerPausedOutputPtys = new Set() private removeLegacyCapacityListener: (() => void) | null = null private lastInputAtByPty = new Map() private interactiveOutputCharsByPty = new Map() @@ -345,6 +346,16 @@ export class PtyHandler { this.dispatcher.onLegacyPtyCapacity?.(() => this.handleLegacyCapacity()) ?? null } + setConsumerDeliveryPaused(id: string, paused: boolean): void { + if (paused) { + this.consumerPausedOutputPtys.add(id) + this.pausePtyOutput(id) + return + } + this.consumerPausedOutputPtys.delete(id) + this.maybeResumePtyOutput(id) + } + private async loadPty(): Promise { if (this.ptyModule) { return this.ptyModule @@ -587,6 +598,7 @@ export class PtyHandler { this.clearStartupCommandTimer(managed) this.releaseRelayIngress(managed) this.pausedOutputPtys.delete(managed.id) + this.consumerPausedOutputPtys.delete(managed.id) this.flushPtyOutput(managed.id) this.pendingExitByPty.set(managed.id, { id: managed.id, @@ -874,6 +886,7 @@ export class PtyHandler { this.pendingOutputByPty.delete(id) this.pendingExitByPty.delete(id) this.pausedOutputPtys.delete(id) + this.consumerPausedOutputPtys.delete(id) this.clearPtyInputState(id) this.clearOutputFlushTimerIfIdle() } @@ -946,6 +959,7 @@ export class PtyHandler { private maybeResumePtyOutput(id: string): void { if ( !this.pausedOutputPtys.has(id) || + this.consumerPausedOutputPtys.has(id) || this.pendingProducerBytes(id) > PTY_OUTPUT_PRODUCER_LOW_BYTES || this.dispatcher.legacyRetentionBelowLowWater === false ) { @@ -1798,6 +1812,7 @@ export class PtyHandler { this.pendingOutputByPty.clear() this.pendingExitByPty.clear() this.pausedOutputPtys.clear() + this.consumerPausedOutputPtys.clear() this.lastInputAtByPty.clear() this.interactiveOutputCharsByPty.clear() const results = await Promise.allSettled( diff --git a/src/relay/relay-frame-decoder.ts b/src/relay/relay-frame-decoder.ts index f98b3647646..2751d984bfd 100644 --- a/src/relay/relay-frame-decoder.ts +++ b/src/relay/relay-frame-decoder.ts @@ -2,7 +2,8 @@ export const HEADER_LENGTH = 13 export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 export const FRAME_DECODER_MAX_FRAMES_PER_TURN = 64 export const FRAME_DECODER_MAX_BYTES_PER_TURN = MAX_MESSAGE_SIZE + HEADER_LENGTH -export const FRAME_DECODER_MAX_TURN_MS = 4 +export const FRAME_DECODER_MAX_TURN_MS = 4, + FRAME_DECODER_MAX_RETAINED_BYTES = 2 * (MAX_MESSAGE_SIZE + HEADER_LENGTH) export type DecodedFrame = { type: number @@ -68,6 +69,12 @@ export class FrameDecoder { const buf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + const retained = this.bufferedLength + buf.length + if (retained > FRAME_DECODER_MAX_RETAINED_BYTES) { + this.reset() + this.onError(new Error(`Frame decoder retained-input limit exceeded: ${retained}`)) + return + } if (buf.length > 0) { this.chunks.push(buf) this.bufferedLength += buf.length @@ -126,8 +133,7 @@ export class FrameDecoder { this.discardBytes(HEADER_LENGTH) this.oversizedPayloadBytesRemaining = length bytes += HEADER_LENGTH - const error = new Error(`Frame payload too large: ${length} bytes — discarded`) - this.onError(error) + this.onError(new Error(`Frame payload too large: ${length} bytes — discarded`)) continue } const totalLength = HEADER_LENGTH + length diff --git a/src/relay/relay-handshake-roundtrip.test.ts b/src/relay/relay-handshake-roundtrip.test.ts index bc7ba71c19d..0713d62295e 100644 --- a/src/relay/relay-handshake-roundtrip.test.ts +++ b/src/relay/relay-handshake-roundtrip.test.ts @@ -76,7 +76,10 @@ describe('handshake round-trip over a real Socket pair', () => { return s } - function startDaemon(version: string): Promise<{ + function startDaemon( + version: string, + endpointCredential?: string + ): Promise<{ accepted: Promise<{ sock: Socket; leftover: Buffer }> }> { return new Promise((resolve) => { @@ -95,6 +98,7 @@ describe('handshake round-trip over a real Socket pair', () => { trackServerSocket(sock) setupDaemonHandshake(sock, { launchVersion: version, + endpointCredential, onAccepted: (s, leftover) => acceptedDeferred.resolve({ sock: s, leftover }) }) }) @@ -120,6 +124,23 @@ describe('handshake round-trip over a real Socket pair', () => { bridgeSock.destroy() }) + it('rejects a same-build socket that lacks the detached endpoint credential', async () => { + const { accepted } = await startDaemon('0.1.0+match', 'secret-credential') + const bridgeSock = connect(sockPath) + await new Promise((resolve) => bridgeSock.once('connect', resolve)) + const closed = new Promise((resolve) => bridgeSock.once('close', () => resolve())) + + runConnectHandshake(bridgeSock, '0.1.0+match', { onAccepted: vi.fn() }, 'wrong-credential') + + await closed + await expect( + Promise.race([ + accepted.then(() => 'accepted'), + new Promise((resolve) => setTimeout(() => resolve('closed'), 20)) + ]) + ).resolves.toBe('closed') + }) + it('preserves leftover bytes on the daemon side when an extra frame is coalesced after the handshake', async () => { // Why: simulate an aggressive client that pipelines a frame immediately // after the handshake. We bypass runConnectHandshake here and write the diff --git a/src/relay/relay-handshake.ts b/src/relay/relay-handshake.ts index e021c392561..bb61e8a08af 100644 --- a/src/relay/relay-handshake.ts +++ b/src/relay/relay-handshake.ts @@ -51,6 +51,7 @@ export type DaemonHandshakeCallbacks = { // leftover: bytes buffered after the handshake frame; caller must feed the dispatcher before attaching the data listener or they're lost. onAccepted: (sock: Socket, leftover: Buffer) => void launchVersion: string + endpointCredential?: string } // Why: read one handshake frame before attaching the dispatcher; version mismatch closes the socket so the bridge exits 42. @@ -61,7 +62,12 @@ export function setupDaemonHandshake(sock: Socket, cb: DaemonHandshakeCallbacks) if (handshakeResolved) { return } - const accepted = handleDaemonHandshakeFrame(sock, frame, cb.launchVersion) + const accepted = handleDaemonHandshakeFrame( + sock, + frame, + cb.launchVersion, + cb.endpointCredential + ) if (accepted) { handshakeResolved = true const leftover = decoder.drain() @@ -94,7 +100,8 @@ export function detachHandshakeListener(sock: Socket): void { function handleDaemonHandshakeFrame( sock: Socket, frame: DecodedFrame, - launchVersion: string + launchVersion: string, + endpointCredential?: string ): boolean { if (frame.type !== MessageType.Handshake) { process.stderr.write( @@ -134,6 +141,14 @@ function handleDaemonHandshakeFrame( sock.end() return false } + if ( + endpointCredential !== undefined && + ('endpointCredential' in msg ? msg.endpointCredential : undefined) !== endpointCredential + ) { + relayLogLine('[relay] Endpoint credential mismatch; closing socket') + sock.destroy() + return false + } process.stderr.write(`[relay] Handshake OK from version=${msg.version}\n`) sock.write(encodeHandshakeFrame({ type: 'orca-relay-handshake-ok', version: launchVersion })) return true @@ -150,7 +165,8 @@ export type ConnectHandshakeCallbacks = { export function runConnectHandshake( sock: Socket, myVersion: string, - cb: ConnectHandshakeCallbacks + cb: ConnectHandshakeCallbacks, + endpointCredential?: string ): void { let handshakeDone = false @@ -212,5 +228,11 @@ export function runConnectHandshake( } }) - sock.write(encodeHandshakeFrame({ type: 'orca-relay-handshake', version: myVersion })) + sock.write( + encodeHandshakeFrame({ + type: 'orca-relay-handshake', + version: myVersion, + ...(endpointCredential ? { endpointCredential } : {}) + }) + ) } diff --git a/src/relay/relay.ts b/src/relay/relay.ts index c45345ef71d..31e7fab109a 100644 --- a/src/relay/relay.ts +++ b/src/relay/relay.ts @@ -10,7 +10,7 @@ import { createServer, createConnection, type Socket, type Server } from 'node:net' import { join } from 'node:path' -import { unlinkSync, existsSync, statSync } from 'node:fs' +import { unlinkSync, existsSync, statSync, readFileSync, chmodSync } from 'node:fs' import { RELAY_SENTINEL, FrameDecoder, @@ -54,6 +54,7 @@ import { shouldReadRemoteCliStdin } from './remote-cli-stdin' import { registerManagedHookInstaller } from './managed-hook-installer' import { registerRelayPluginHostCallHandlers } from './plugin-host-call-handler' import { DispatcherClientWriter } from './dispatcher-client-writer' +import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' const DEFAULT_GRACE_MS = DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 const SOCK_NAME = 'relay.sock' @@ -107,6 +108,7 @@ function parseArgs(argv: string[]): { sockPath: string endpointDir?: string logFile?: string + credentialFile?: string } { let graceTimeMs = DEFAULT_GRACE_MS let connectMode = false @@ -115,6 +117,7 @@ function parseArgs(argv: string[]): { let sockPath = '' let endpointDir: string | undefined let logFile: string | undefined + let credentialFile: string | undefined for (let i = 2; i < argv.length; i++) { if (argv[i] === '--grace-time' && argv[i + 1]) { const parsed = Number.parseInt(argv[i + 1], 10) @@ -138,18 +141,44 @@ function parseArgs(argv: string[]): { } else if (argv[i] === '--log-file' && argv[i + 1]) { logFile = argv[i + 1] i++ + } else if (argv[i] === '--credential-file' && argv[i + 1]) { + credentialFile = argv[i + 1] + i++ } } if (!sockPath) { sockPath = join(process.cwd(), SOCK_NAME) } - return { graceTimeMs, connectMode, detached, cliMode, sockPath, endpointDir, logFile } + return { + graceTimeMs, + connectMode, + detached, + cliMode, + sockPath, + endpointDir, + logFile, + credentialFile + } +} + +function readEndpointCredential(credentialFile: string | undefined): string | undefined { + if (!credentialFile) { + return undefined + } + const credential = readFileSync(credentialFile, 'utf8').trim() + if (!/^[A-Za-z0-9_-]{32,256}$/.test(credential)) { + throw new Error('Relay endpoint credential is missing or invalid') + } + if (process.platform !== 'win32') { + chmodSync(credentialFile, 0o600) + } + return credential } // ── Connect mode ───────────────────────────────────────────────────── // Why: --connect bridges a new SSH channel's stdin/stdout to the existing relay's socket so the client keeps talking to the process that owns the live PTYs. -function runConnectMode(sockPath: string): void { +function runConnectMode(sockPath: string, endpointCredential?: string): void { const myVersion = readLaunchVersion() const sock = createConnection({ path: sockPath }) const stdoutWriter = new DispatcherClientWriter( @@ -180,48 +209,53 @@ function runConnectMode(sockPath: string): void { sock.on('connect', () => { clearTimeout(connectTimeout) - runConnectHandshake(sock, myVersion, { - onAccepted: (leftover: Buffer) => { - stdoutWriter.enqueue('control', () => Buffer.from(RELAY_SENTINEL), RELAY_SENTINEL.length) - if (leftover.length > 0) { - stdoutWriter.enqueue('control', () => leftover, leftover.length) - } - process.stdin.pipe(sock) - sock.on('data', (data: Buffer) => { - sock.pause() - let offset = 0 - const writeNext = (): void => { - if (offset >= data.length) { - sock.resume() - return - } - const bytes = Math.min(stdoutWriter.producerFrameCapacity, data.length - offset) - if (bytes <= 0) { - stdoutWriter.close(new Error('Relay stdout has no producer capacity')) - return - } - const chunk = data.subarray(offset, offset + bytes) - if ( - !stdoutWriter.enqueue( - 'ordinary', - () => chunk, - chunk.length, - (result) => { - if (!result.ok) { - return - } - offset += bytes - writeNext() - } - ) - ) { - stdoutWriter.close(new Error('Relay stdout bridge capacity exceeded')) - } + runConnectHandshake( + sock, + myVersion, + { + onAccepted: (leftover: Buffer) => { + stdoutWriter.enqueue('control', () => Buffer.from(RELAY_SENTINEL), RELAY_SENTINEL.length) + if (leftover.length > 0) { + stdoutWriter.enqueue('control', () => leftover, leftover.length) } - writeNext() - }) - } - }) + process.stdin.pipe(sock) + sock.on('data', (data: Buffer) => { + sock.pause() + let offset = 0 + const writeNext = (): void => { + if (offset >= data.length) { + sock.resume() + return + } + const bytes = Math.min(stdoutWriter.producerFrameCapacity, data.length - offset) + if (bytes <= 0) { + stdoutWriter.close(new Error('Relay stdout has no producer capacity')) + return + } + const chunk = data.subarray(offset, offset + bytes) + if ( + !stdoutWriter.enqueue( + 'ordinary', + () => chunk, + chunk.length, + (result) => { + if (!result.ok) { + return + } + offset += bytes + writeNext() + } + ) + ) { + stdoutWriter.close(new Error('Relay stdout bridge capacity exceeded')) + } + } + writeNext() + }) + } + }, + endpointCredential + ) }) // Why: Node swallows EPIPE on stdout, so the bridge would zombie and drop frames; exit on stdout error so the relay enters grace promptly. @@ -241,7 +275,11 @@ function runConnectMode(sockPath: string): void { }) } -async function runOrcaCliMode(sockPath: string, argv: string[]): Promise { +async function runOrcaCliMode( + sockPath: string, + argv: string[], + endpointCredential?: string +): Promise { const myVersion = readLaunchVersion() const stdin = shouldReadRemoteCliStdin(argv) ? await readOrcaCliStdin() : undefined const sock = createConnection({ path: sockPath }) @@ -335,17 +373,22 @@ async function runOrcaCliMode(sockPath: string, argv: string[]): Promise { sock.on('connect', () => { clearTimeout(connectTimeout) - runConnectHandshake(sock, myVersion, { - onAccepted: (leftover) => { - if (leftover.length > 0) { - decoder.feed(leftover) + runConnectHandshake( + sock, + myVersion, + { + onAccepted: (leftover) => { + if (leftover.length > 0) { + decoder.feed(leftover) + } + sock.on('data', (chunk) => + decoder.feed(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)) + ) + sendRequest() } - sock.on('data', (chunk) => - decoder.feed(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk)) - ) - sendRequest() - } - }) + }, + endpointCredential + ) }) sock.on('error', (err) => { @@ -369,17 +412,29 @@ async function readOrcaCliStdin(): Promise { // ── Normal mode ────────────────────────────────────────────────────── async function main(): Promise { - const { graceTimeMs, connectMode, detached, cliMode, sockPath, endpointDir, logFile } = parseArgs( - process.argv - ) + const { + graceTimeMs, + connectMode, + detached, + cliMode, + sockPath, + endpointDir, + logFile, + credentialFile + } = parseArgs(process.argv) + const endpointCredential = readEndpointCredential(credentialFile) if (connectMode) { - runConnectMode(sockPath) + runConnectMode(sockPath, endpointCredential) return } if (cliMode) { const marker = process.argv.indexOf('--orca-cli') - await runOrcaCliMode(sockPath, marker >= 0 ? process.argv.slice(marker + 1) : []) + await runOrcaCliMode( + sockPath, + marker >= 0 ? process.argv.slice(marker + 1) : [], + endpointCredential + ) return } @@ -464,8 +519,11 @@ async function main(): Promise { stdoutDrainWaiters.add(cb) return () => stdoutDrainWaiters.delete(cb) } - } + }, + undefined, + { pauseReads: () => process.stdin.pause(), resumeReads: () => process.stdin.resume() } ) + const launchVersion = readLaunchVersion() const context = new RelayContext() @@ -495,6 +553,12 @@ async function main(): Promise { }) const ptyHandler = new PtyHandler(dispatcher, graceTimeMs) + const ptyConsumerSessionAdapter = new SshPtyConsumerSessionAdapter( + dispatcher, + launchVersion, + (id, paused) => ptyHandler.setConsumerDeliveryPaused(id, paused) + ) + void ptyConsumerSessionAdapter const fsHandler = new FsHandler(dispatcher, context) const watchRegistry = fsHandler.getWatchRegistry() ptyHandler.setWorktreeRemovalCoordinator(watchRegistry) @@ -663,7 +727,6 @@ async function main(): Promise { const socketClients = new Map() let socketServer: Server | null = null - const launchVersion = readLaunchVersion() const startedAt = Date.now() let acceptedSocketConnections = 0 let hasAcceptedSocketClient = false @@ -748,6 +811,16 @@ async function main(): Promise { sockDrainWaiters.add(cb) return () => sockDrainWaiters.delete(cb) } + }, + { + principal: `relay-endpoint:${launchVersion}`, + authenticated: endpointCredential !== undefined, + allowSessionOwner: endpointCredential !== undefined, + authenticationKind: endpointCredential ? 'endpoint-credential' : 'unproved' + }, + { + pauseReads: () => sock.pause(), + resumeReads: () => sock.resume() } ) socketClients.set(sock, clientId) @@ -766,7 +839,11 @@ async function main(): Promise { async function startSocketServer(): Promise { const server = createServer((sock) => { // Why: pre-dispatcher version handshake — see relay-handshake.ts. - setupDaemonHandshake(sock, { launchVersion, onAccepted: attachAcceptedSocket }) + setupDaemonHandshake(sock, { + launchVersion, + endpointCredential, + onAccepted: attachAcceptedSocket + }) // Why: destroy on 'end' (FIN from --connect's dying channel) so the 'close' handler fires promptly and the daemon enters grace. sock.on('end', () => { diff --git a/src/relay/ssh-pty-consumer-session-adapter.test.ts b/src/relay/ssh-pty-consumer-session-adapter.test.ts new file mode 100644 index 00000000000..93ec7d032a4 --- /dev/null +++ b/src/relay/ssh-pty-consumer-session-adapter.test.ts @@ -0,0 +1,210 @@ +import { afterEach, describe, expect, it } from 'vitest' +import { RelayDispatcher, type RelayClientSessionIdentity } from './dispatcher' +import { encodeJsonRpcFrame, MessageType } from './protocol' +import { SshPtyConsumerSessionAdapter } from './ssh-pty-consumer-session-adapter' + +const endpointIdentity: RelayClientSessionIdentity = { + principal: 'endpoint-principal', + authenticated: true, + allowSessionOwner: true, + authenticationKind: 'endpoint-credential' +} + +function openFrame(id: number, overrides: Record = {}): Buffer { + return encodeJsonRpcFrame( + { + jsonrpc: '2.0', + id, + method: 'pty.openClient', + params: { + protocolVersion: 1, + clientInstanceId: `client-${id}`, + requestedRole: 'session-owner', + ...overrides + } + }, + 1, + 0 + ) +} + +function responseResult(buffer: Buffer): Record { + expect(buffer[0]).toBe(MessageType.Regular) + const length = buffer.readUInt32BE(9) + return JSON.parse(buffer.subarray(13, 13 + length).toString('utf8')).result +} + +async function flushRequests(): Promise { + await new Promise((resolve) => setImmediate(resolve)) +} + +describe('SshPtyConsumerSessionAdapter', () => { + let dispatcher: RelayDispatcher | null = null + + afterEach(() => { + dispatcher?.dispose() + dispatcher = null + }) + + it('does not activate owner authority until the grant write settles', async () => { + const firstWrites: Buffer[] = [] + const firstSettlements: ((result: { ok: true } | { ok: false; error: Error }) => void)[] = [] + dispatcher = new RelayDispatcher( + (data, onSettled) => { + firstWrites.push(Buffer.from(data)) + firstSettlements.push(onSettled) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(openFrame(1)) + await flushRequests() + + const secondWrites: Buffer[] = [] + const secondId = dispatcher.attachClient( + (data, onSettled) => { + secondWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + { ...endpointIdentity, principal: 'competitor' } + ) + dispatcher.feedClient(secondId, openFrame(2)) + await flushRequests() + + expect(responseResult(firstWrites[0])).toMatchObject({ + role: 'session-owner', + ownerGeneration: 1 + }) + expect(responseResult(secondWrites[0])).toMatchObject({ role: 'subscriber' }) + firstSettlements[0]({ ok: true }) + }) + + it('rolls back owner election when the grant write fails', async () => { + dispatcher = new RelayDispatcher( + (_data, onSettled) => { + onSettled({ ok: false, error: new Error('send failed') }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + dispatcher.feed(openFrame(1)) + await flushRequests() + + const retryWrites: Buffer[] = [] + const retryId = dispatcher.attachClient( + (data, onSettled) => { + retryWrites.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + dispatcher.feedClient(retryId, openFrame(2, { clientInstanceId: 'client-1' })) + await flushRequests() + + expect(responseResult(retryWrites[0])).toMatchObject({ + role: 'session-owner', + ownerGeneration: 2 + }) + }) + + it('rejects an unproved constructor stream as an owner principal', async () => { + const writes: Buffer[] = [] + dispatcher = new RelayDispatcher((data, onSettled) => { + writes.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(openFrame(1)) + await flushRequests() + + const response = JSON.parse( + writes[0].subarray(13, 13 + writes[0].readUInt32BE(9)).toString('utf8') + ) + expect(response.error.message).toContain('authentication') + }) + + it('rejects an invalid requested role instead of promoting it to owner', async () => { + const writes: Buffer[] = [] + dispatcher = new RelayDispatcher( + (data, onSettled) => { + writes.push(Buffer.from(data)) + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a') + + dispatcher.feed(openFrame(1, { requestedRole: 'administrator' })) + await flushRequests() + + const response = JSON.parse( + writes[0].subarray(13, 13 + writes[0].readUInt32BE(9)).toString('utf8') + ) + expect(response.error.message).toContain('requestedRole') + }) + + it('generation-fences per-PTY delivery pause notifications', async () => { + const setPaused: { id: string; paused: boolean }[] = [] + dispatcher = new RelayDispatcher( + (_data, onSettled) => { + onSettled({ ok: true }) + return true + }, + { supportsWriteCallback: true }, + endpointIdentity + ) + new SshPtyConsumerSessionAdapter(dispatcher, 'build-a', (id, paused) => { + setPaused.push({ id, paused }) + }) + dispatcher.feed(openFrame(1)) + await flushRequests() + + dispatcher.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.setDeliveryPaused', + params: { + id: 'pty-1', + paused: true, + clientGeneration: 1, + ownerGeneration: 1 + } + }, + 2, + 0 + ) + ) + dispatcher.feed( + encodeJsonRpcFrame( + { + jsonrpc: '2.0', + method: 'pty.setDeliveryPaused', + params: { + id: 'pty-1', + paused: false, + clientGeneration: 99, + ownerGeneration: 1 + } + }, + 3, + 0 + ) + ) + + expect(setPaused).toEqual([{ id: 'pty-1', paused: true }]) + }) +}) diff --git a/src/relay/ssh-pty-consumer-session-adapter.ts b/src/relay/ssh-pty-consumer-session-adapter.ts new file mode 100644 index 00000000000..efbcbf9202e --- /dev/null +++ b/src/relay/ssh-pty-consumer-session-adapter.ts @@ -0,0 +1,120 @@ +import { + PTY_CONSUMER_SESSION_PROTOCOL_VERSION, + PtyConsumerSession, + type PtyConsumerSessionGrant, + type PtyConsumerSessionHello +} from '../shared/pty-consumer-session' +import type { RelayClientSessionIdentity, RelayDispatcher, RequestContext } from './dispatcher' + +export const SSH_PTY_OPEN_CLIENT_METHOD = 'pty.openClient' + +type OpenClientParams = PtyConsumerSessionHello & { + protocolVersion: number +} + +function parseOpenClientParams(params: Record): OpenClientParams { + const resume = + typeof params.resume === 'object' && params.resume !== null + ? (params.resume as Record) + : undefined + const capabilities = + typeof params.capabilities === 'object' && params.capabilities !== null + ? (params.capabilities as Record) + : undefined + const outputFlowControl = + typeof capabilities?.outputFlowControl === 'object' && capabilities.outputFlowControl !== null + ? (capabilities.outputFlowControl as Record) + : undefined + return { + protocolVersion: Number(params.protocolVersion), + clientInstanceId: String(params.clientInstanceId ?? ''), + requestedRole: String(params.requestedRole ?? '') as PtyConsumerSessionHello['requestedRole'], + ...(resume + ? { + resume: { + ownerGeneration: Number(resume.ownerGeneration), + ownerLease: String(resume.ownerLease ?? '') + } + } + : {}), + ...(outputFlowControl + ? { + capabilities: { + outputFlowControl: { + versions: Array.isArray(outputFlowControl.versions) + ? outputFlowControl.versions.map(Number) + : [], + requestedWindowSu: Number(outputFlowControl.requestedWindowSu) + } + } + } + : {}) + } +} + +function requireIdentity(context: RequestContext): RelayClientSessionIdentity { + if (!context.sessionIdentity) { + throw new Error('SSH PTY consumer transport identity is unavailable') + } + return context.sessionIdentity +} + +export class SshPtyConsumerSessionAdapter { + private readonly session: PtyConsumerSession + + constructor( + dispatcher: RelayDispatcher, + serverBuildId: string, + setDeliveryPaused?: (id: string, paused: boolean) => void + ) { + this.session = new PtyConsumerSession({ serverBuildId }) + dispatcher.onRequest(SSH_PTY_OPEN_CLIENT_METHOD, (params, context) => + this.openClient(params, context) + ) + dispatcher.onClientDetached((clientId) => this.session.close(String(clientId))) + dispatcher.onNotification('pty.setDeliveryPaused', (params, context) => { + const grant = this.session.activeGrant(String(context.clientId)) + if ( + !grant || + grant.clientGeneration !== params.clientGeneration || + grant.ownerGeneration !== params.ownerGeneration || + typeof params.id !== 'string' || + typeof params.paused !== 'boolean' + ) { + return + } + setDeliveryPaused?.(params.id, params.paused) + }) + } + + private async openClient( + rawParams: Record, + context: RequestContext + ): Promise { + const params = parseOpenClientParams(rawParams) + if (params.protocolVersion !== PTY_CONSUMER_SESSION_PROTOCOL_VERSION) { + throw new Error( + `Unsupported pty.openClient protocol version: ${params.protocolVersion || 'missing'}` + ) + } + const identity = requireIdentity(context) + const admission = this.session.admit(params, { + connectionId: String(context.clientId), + principal: identity.principal, + authenticated: identity.authenticated, + allowSessionOwner: identity.allowSessionOwner + }) + if (!context.onResponseSettled) { + admission.rollbackPublication() + throw new Error('SSH PTY consumer response publication fence is unavailable') + } + context.onResponseSettled((result) => { + if (result.ok) { + admission.commitPublication() + } else { + admission.rollbackPublication() + } + }) + return admission.grant + } +} diff --git a/src/shared/pty-consumer-session-contract.ts b/src/shared/pty-consumer-session-contract.ts new file mode 100644 index 00000000000..fb068ce0e51 --- /dev/null +++ b/src/shared/pty-consumer-session-contract.ts @@ -0,0 +1,58 @@ +export const PTY_CONSUMER_SESSION_PROTOCOL_VERSION = 1 +export const PTY_CONSUMER_OWNER_GRACE_MS = 30_000 + +export type PtyConsumerRole = 'session-owner' | 'subscriber' + +export type PtyConsumerSessionHello = { + clientInstanceId: string + requestedRole: PtyConsumerRole + resume?: { + ownerGeneration: number + ownerLease: string + } + capabilities?: { + outputFlowControl?: { + versions: number[] + requestedWindowSu: number + } + } +} + +export type PtyConsumerSessionGrant = { + protocolVersion: typeof PTY_CONSUMER_SESSION_PROTOCOL_VERSION + serverBuildId: string + clientGeneration: number + role: PtyConsumerRole + ownerGeneration?: number + ownerLease?: string + capabilities?: { + outputFlowControl?: { + version: 1 + windowSu: number + } + } +} + +export type PtyConsumerAuthentication = { + connectionId: string + principal: string + authenticated: boolean + allowSessionOwner: boolean +} + +export type PtyConsumerSessionAdmission = { + grant: Readonly + commitPublication: () => void + rollbackPublication: () => void +} + +export type PtyConsumerSessionOptions = { + serverBuildId: string + outputFlowControl?: { + versions: readonly number[] + maxWindowSu: number + } + ownerGraceMs?: number + now?: () => number + createLease?: () => string +} diff --git a/src/shared/pty-consumer-session.test.ts b/src/shared/pty-consumer-session.test.ts new file mode 100644 index 00000000000..3ec266f9902 --- /dev/null +++ b/src/shared/pty-consumer-session.test.ts @@ -0,0 +1,242 @@ +import { describe, expect, it } from 'vitest' +import { + PtyConsumerSession, + type PtyConsumerAuthentication, + type PtyConsumerSessionHello +} from './pty-consumer-session' + +function auth( + connectionId: string, + overrides: Partial = {} +): PtyConsumerAuthentication { + return { + connectionId, + principal: 'desktop', + authenticated: true, + allowSessionOwner: true, + ...overrides + } +} + +function ownerHello(overrides: Partial = {}): PtyConsumerSessionHello { + return { + clientInstanceId: 'client-a', + requestedRole: 'session-owner', + ...overrides + } +} + +function createSession(options: { now?: () => number } = {}): PtyConsumerSession { + let lease = 0 + return new PtyConsumerSession({ + serverBuildId: 'relay-build', + createLease: () => `lease-${++lease}`, + ownerGraceMs: 30_000, + ...options + }) +} + +describe('PtyConsumerSession', () => { + it('activates an authenticated owner only after its publication fence', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + const competitor = session.admit( + ownerHello({ clientInstanceId: 'client-b' }), + auth('connection-2', { principal: 'other' }) + ) + + expect(first.grant).toMatchObject({ + clientGeneration: 1, + role: 'session-owner', + ownerGeneration: 1, + ownerLease: 'lease-1' + }) + expect(competitor.grant).toMatchObject({ clientGeneration: 2, role: 'subscriber' }) + first.commitPublication() + }) + + it('rolls back an unpublished owner without consuming authority', () => { + const session = createSession() + session.admit(ownerHello(), auth('failed')).rollbackPublication() + + const retry = session.admit(ownerHello(), auth('retry')) + expect(retry.grant).toMatchObject({ + role: 'session-owner', + ownerGeneration: 2, + ownerLease: 'lease-2' + }) + }) + + it('returns the same generation and lease for duplicate opens on one connection', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + const duplicate = session.admit(ownerHello(), auth('connection-1')) + + expect(duplicate.grant).toBe(first.grant) + first.commitPublication() + duplicate.commitPublication() + }) + + it('rejects a second, different open on one connection', () => { + const session = createSession() + session.admit(ownerHello(), auth('connection-1')) + + expect(() => + session.admit(ownerHello({ requestedRole: 'subscriber' }), auth('connection-1')) + ).toThrow('only once') + }) + + it('cannot self-promote an authenticated but owner-ineligible principal', () => { + const session = createSession() + const admission = session.admit( + ownerHello(), + auth('connection-1', { allowSessionOwner: false }) + ) + + expect(admission.grant.role).toBe('subscriber') + expect(admission.grant.ownerLease).toBeUndefined() + }) + + it('rejects an unauthenticated transport', () => { + const session = createSession() + expect(() => + session.admit(ownerHello(), auth('connection-1', { authenticated: false })) + ).toThrow('authentication required') + }) + + it('rotates the lease and increments owner generation on valid recovery', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + session.close('connection-1') + + const recovered = session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-2') + ) + recovered.commitPublication() + + expect(recovered.grant).toMatchObject({ + role: 'session-owner', + ownerGeneration: 2, + ownerLease: 'lease-2' + }) + }) + + it('rejects recovery while the current owner is still active', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + + expect(() => + session.admit( + ownerHello({ + resume: { + ownerGeneration: first.grant.ownerGeneration!, + ownerLease: first.grant.ownerLease! + } + }), + auth('connection-2') + ) + ).toThrow('Active owner') + }) + + it('rejects stale recovery without disturbing the retained owner', () => { + const session = createSession() + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + session.close('connection-1') + + expect(() => + session.admit( + ownerHello({ resume: { ownerGeneration: 1, ownerLease: 'wrong' } }), + auth('connection-2') + ) + ).toThrow('stale') + + const recovered = session.admit( + ownerHello({ resume: { ownerGeneration: 1, ownerLease: 'lease-1' } }), + auth('connection-3') + ) + expect(recovered.grant.role).toBe('session-owner') + }) + + it('elects a new owner after disconnected-owner grace expires', () => { + let now = 10 + const session = createSession({ now: () => now }) + const first = session.admit(ownerHello(), auth('connection-1')) + first.commitPublication() + session.close('connection-1') + now += 30_000 + + const next = session.admit( + ownerHello({ clientInstanceId: 'client-b' }), + auth('connection-2', { principal: 'other' }) + ) + expect(next.grant).toMatchObject({ role: 'session-owner', ownerGeneration: 2 }) + }) + + it('intersects V1 capability and clamps its source-unit window', () => { + const session = new PtyConsumerSession({ + serverBuildId: 'relay-build', + outputFlowControl: { versions: [1], maxWindowSu: 64 }, + createLease: () => 'lease' + }) + const admission = session.admit( + ownerHello({ + capabilities: { + outputFlowControl: { versions: [1, 2], requestedWindowSu: 128 } + } + }), + auth('connection-1') + ) + + expect(admission.grant.capabilities?.outputFlowControl).toEqual({ + version: 1, + windowSu: 64 + }) + }) + + it('makes token-free bounded legacy an explicit capability omission', () => { + const session = createSession() + const admission = session.admit(ownerHello(), auth('connection-1')) + + expect(admission.grant.capabilities).toBeUndefined() + expect(admission.grant).not.toHaveProperty('deliveryToken') + }) + + it('bounds capability offers before fingerprinting them', () => { + const session = createSession() + expect(() => + session.admit( + ownerHello({ + capabilities: { + outputFlowControl: { + versions: Array.from({ length: 9 }, (_, index) => index + 1), + requestedWindowSu: 64 + } + } + }), + auth('connection-1') + ) + ).toThrow('versions') + }) + + it('rejects invalid server windows and owner grace', () => { + expect( + () => + new PtyConsumerSession({ + serverBuildId: 'build', + outputFlowControl: { versions: [1], maxWindowSu: 0 } + }) + ).toThrow('support') + expect( + () => new PtyConsumerSession({ serverBuildId: 'build', ownerGraceMs: Number.MAX_VALUE }) + ).toThrow('ownerGraceMs') + }) +}) diff --git a/src/shared/pty-consumer-session.ts b/src/shared/pty-consumer-session.ts new file mode 100644 index 00000000000..c6d83d0fb77 --- /dev/null +++ b/src/shared/pty-consumer-session.ts @@ -0,0 +1,315 @@ +import { randomUUID } from 'node:crypto' +import { + PTY_CONSUMER_OWNER_GRACE_MS, + PTY_CONSUMER_SESSION_PROTOCOL_VERSION, + type PtyConsumerAuthentication, + type PtyConsumerSessionAdmission, + type PtyConsumerSessionGrant, + type PtyConsumerSessionHello, + type PtyConsumerSessionOptions +} from './pty-consumer-session-contract' + +export * from './pty-consumer-session-contract' + +const MAX_CAPABILITY_VERSIONS = 8 + +type ClientRecord = { + fingerprint: string + principal: string + clientInstanceId: string + grant: Readonly + state: 'pending' | 'active' + publicationState: 'pending' | 'committed' | 'rolled-back' +} + +type OwnerRecord = { + connectionId: string + principal: string + clientInstanceId: string + generation: number + lease: string + state: 'pending' | 'active' | 'disconnected' + disconnectedAt?: number + replaces?: OwnerRecord +} + +function assertNonEmptyString(value: unknown, name: string): asserts value is string { + if (typeof value !== 'string' || value.length === 0 || value.length > 512) { + throw new Error(`${name} must be a non-empty string of at most 512 characters`) + } +} + +function validateHello(hello: PtyConsumerSessionHello): void { + assertNonEmptyString(hello.clientInstanceId, 'clientInstanceId') + if (hello.requestedRole !== 'session-owner' && hello.requestedRole !== 'subscriber') { + throw new Error('requestedRole must be session-owner or subscriber') + } + if (hello.resume) { + if (!Number.isSafeInteger(hello.resume.ownerGeneration) || hello.resume.ownerGeneration <= 0) { + throw new Error('resume.ownerGeneration must be a positive safe integer') + } + assertNonEmptyString(hello.resume.ownerLease, 'resume.ownerLease') + } + const flow = hello.capabilities?.outputFlowControl + if (!flow) { + return + } + if ( + !Array.isArray(flow.versions) || + flow.versions.length > MAX_CAPABILITY_VERSIONS || + flow.versions.some((version) => !Number.isSafeInteger(version) || version <= 0) + ) { + throw new Error('outputFlowControl.versions must contain positive safe integers') + } + if (!Number.isSafeInteger(flow.requestedWindowSu) || flow.requestedWindowSu <= 0) { + throw new Error('outputFlowControl.requestedWindowSu must be a positive safe integer') + } +} + +function helloFingerprint(hello: PtyConsumerSessionHello): string { + const flow = hello.capabilities?.outputFlowControl + return JSON.stringify({ + clientInstanceId: hello.clientInstanceId, + requestedRole: hello.requestedRole, + resume: hello.resume, + outputFlowControl: flow + ? { + versions: [...flow.versions].sort((a, b) => a - b), + requestedWindowSu: flow.requestedWindowSu + } + : undefined + }) +} + +export class PtyConsumerSession { + private readonly clients = new Map() + private readonly now: () => number + private readonly createLease: () => string + private readonly ownerGraceMs: number + private nextClientGeneration = 1 + private nextOwnerGeneration = 1 + private owner: OwnerRecord | null = null + + constructor(private readonly options: PtyConsumerSessionOptions) { + assertNonEmptyString(options.serverBuildId, 'serverBuildId') + if ( + options.outputFlowControl && + (!Number.isSafeInteger(options.outputFlowControl.maxWindowSu) || + options.outputFlowControl.maxWindowSu <= 0 || + options.outputFlowControl.versions.length > MAX_CAPABILITY_VERSIONS || + options.outputFlowControl.versions.some( + (version) => !Number.isSafeInteger(version) || version <= 0 + )) + ) { + throw new Error('outputFlowControl support is invalid') + } + if ( + options.ownerGraceMs !== undefined && + (!Number.isSafeInteger(options.ownerGraceMs) || options.ownerGraceMs < 0) + ) { + throw new Error('ownerGraceMs must be a non-negative safe integer') + } + this.now = options.now ?? Date.now + this.createLease = options.createLease ?? randomUUID + this.ownerGraceMs = options.ownerGraceMs ?? PTY_CONSUMER_OWNER_GRACE_MS + } + + admit( + hello: PtyConsumerSessionHello, + authentication: PtyConsumerAuthentication + ): PtyConsumerSessionAdmission { + validateHello(hello) + assertNonEmptyString(authentication.connectionId, 'connectionId') + assertNonEmptyString(authentication.principal, 'principal') + if (!authentication.authenticated) { + throw new Error('PTY consumer authentication required') + } + this.expireOwner() + + const fingerprint = helloFingerprint(hello) + const duplicate = this.clients.get(authentication.connectionId) + if (duplicate) { + if ( + duplicate.fingerprint !== fingerprint || + duplicate.principal !== authentication.principal + ) { + throw new Error('pty.openClient may be used only once per transport connection') + } + return this.admissionFor(duplicate) + } + + const owner = this.selectOwner(hello, authentication) + const grant = Object.freeze({ + protocolVersion: PTY_CONSUMER_SESSION_PROTOCOL_VERSION, + serverBuildId: this.options.serverBuildId, + clientGeneration: this.nextClientGeneration++, + role: owner ? ('session-owner' as const) : ('subscriber' as const), + ...(owner ? { ownerGeneration: owner.generation, ownerLease: owner.lease } : {}), + ...this.intersectCapabilities(hello) + }) + const client: ClientRecord = { + fingerprint, + principal: authentication.principal, + clientInstanceId: hello.clientInstanceId, + grant, + state: 'pending', + publicationState: 'pending' + } + this.clients.set(authentication.connectionId, client) + if (owner) { + this.owner = owner + } + return this.admissionFor(client) + } + + close(connectionId: string): void { + const client = this.clients.get(connectionId) + if (!client) { + return + } + this.clients.delete(connectionId) + if (this.owner?.connectionId !== connectionId) { + return + } + if (this.owner.state === 'pending') { + this.owner = this.owner.replaces ?? null + return + } + this.owner = { + ...this.owner, + state: 'disconnected', + disconnectedAt: this.now() + } + } + + sweepExpired(): void { + this.expireOwner() + } + + activeGrant(connectionId: string): Readonly | null { + const client = this.clients.get(connectionId) + return client?.state === 'active' ? client.grant : null + } + + private admissionFor(client: ClientRecord): PtyConsumerSessionAdmission { + return { + grant: client.grant, + commitPublication: () => { + if (client.publicationState !== 'pending') { + return + } + client.publicationState = 'committed' + if (client.state !== 'pending') { + return + } + client.state = 'active' + const owner = this.owner + if (owner?.connectionId === this.connectionIdFor(client) && owner.state === 'pending') { + this.owner = { ...owner, state: 'active', replaces: undefined } + } + }, + rollbackPublication: () => { + if (client.publicationState !== 'pending') { + return + } + client.publicationState = 'rolled-back' + if (client.state !== 'pending') { + return + } + const connectionId = this.connectionIdFor(client) + this.clients.delete(connectionId) + if (this.owner?.connectionId === connectionId && this.owner.state === 'pending') { + this.owner = this.owner.replaces ?? null + } + } + } + } + + private connectionIdFor(client: ClientRecord): string { + for (const [connectionId, candidate] of this.clients) { + if (candidate === client) { + return connectionId + } + } + return '' + } + + private selectOwner( + hello: PtyConsumerSessionHello, + authentication: PtyConsumerAuthentication + ): OwnerRecord | null { + if (hello.requestedRole !== 'session-owner' || !authentication.allowSessionOwner) { + return null + } + const current = this.owner + if (!current) { + if (hello.resume) { + throw new Error('Owner recovery lease is stale') + } + return this.newOwner(hello, authentication, null) + } + if (!hello.resume) { + return null + } + const recoveryMatches = + hello.resume.ownerGeneration === current.generation && + hello.resume.ownerLease === current.lease && + hello.clientInstanceId === current.clientInstanceId && + authentication.principal === current.principal + if (!recoveryMatches) { + throw new Error('Owner recovery lease is stale or belongs to another principal') + } + if (current.state === 'pending') { + throw new Error('Owner grant publication is still pending') + } + if (current.state !== 'disconnected') { + throw new Error('Active owner cannot be replaced by recovery') + } + return this.newOwner(hello, authentication, current) + } + + private newOwner( + hello: PtyConsumerSessionHello, + authentication: PtyConsumerAuthentication, + replaces: OwnerRecord | null + ): OwnerRecord { + const lease = this.createLease() + assertNonEmptyString(lease, 'ownerLease') + return { + connectionId: authentication.connectionId, + principal: authentication.principal, + clientInstanceId: hello.clientInstanceId, + generation: this.nextOwnerGeneration++, + lease, + state: 'pending', + ...(replaces ? { replaces } : {}) + } + } + + private intersectCapabilities( + hello: PtyConsumerSessionHello + ): Pick { + const offer = hello.capabilities?.outputFlowControl + const support = this.options.outputFlowControl + if (!offer || !support || !offer.versions.includes(1) || !support.versions.includes(1)) { + return {} + } + return { + capabilities: { + outputFlowControl: { + version: 1, + windowSu: Math.min(offer.requestedWindowSu, support.maxWindowSu) + } + } + } + } + + private expireOwner(): void { + if ( + this.owner?.state === 'disconnected' && + this.now() - (this.owner.disconnectedAt ?? this.now()) >= this.ownerGraceMs + ) { + this.owner = null + } + } +}