From eb316f37d57bc2497270292ce4cc55ad895afbde Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Fri, 10 Jul 2026 04:53:36 -0400 Subject: [PATCH] fix(remote): gate ws sends with an order-preserving backpressure queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two remote-terminal ws paths wrote to the socket with no backpressure handling, so a fast producer over a slow/congested link ballooned ws.bufferedAmount and main-process RSS without bound: - Server text/JSON reply path (e2ee-channel encryptedReply): had NO bufferedAmount gate (only the binary path did). A legacy client without the terminalBinaryStream capability streams terminal output as JSON, and that path has no seq/resync — so silently dropping frames there would recreate the corruption bug. Instead we now hold replies in order and drain on recovery. - Client binary send path (remote-runtime-client sendBinary): carries keystrokes and must never drop. Same held-and-drained treatment. Adds a shared, generic createWsOutboundBackpressureQueue: while bufferedAmount is over a soft cap (8 MiB) it parks frames in FIFO order and drains as the socket clears — no drop, no reorder. Only when a hard byte cap (64 MiB) is exceeded (link effectively dead) does it fire onOverflow, which closes the socket so the client reconnects and replays a fresh authoritative snapshot instead of growing memory unbounded. The existing binary drop+resync path (#8059) is untouched. Repro-first: deterministic queue unit tests (injected bufferedAmount + timer) and an e2ee text-path repro that fails on main (replies shoved onto a congested socket) and passes after the fix. Co-authored-by: Orca --- .../e2ee-channel-text-backpressure.test.ts | 98 ++++++++++++ src/main/runtime/rpc/e2ee-channel.ts | 29 +++- src/shared/remote-runtime-client.ts | 31 +++- .../ws-outbound-backpressure-queue.test.ts | 118 +++++++++++++++ src/shared/ws-outbound-backpressure-queue.ts | 141 ++++++++++++++++++ 5 files changed, 415 insertions(+), 2 deletions(-) create mode 100644 src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts create mode 100644 src/shared/ws-outbound-backpressure-queue.test.ts create mode 100644 src/shared/ws-outbound-backpressure-queue.ts diff --git a/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts b/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts new file mode 100644 index 00000000000..5d0b2d3da92 --- /dev/null +++ b/src/main/runtime/rpc/e2ee-channel-text-backpressure.test.ts @@ -0,0 +1,98 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type { WebSocket } from 'ws' +import { E2EEChannel, type E2EEChannelOptions } from './e2ee-channel' +import { deriveSharedKey, decrypt, encrypt, generateKeyPair } from './e2ee-crypto' + +// Repro for gap (a): the streaming JSON reply path (encryptedReply) had no +// bufferedAmount gate, so a fast producer over a slow link (legacy +// terminal.subscribe, which has NO seq/resync) ballooned ws.bufferedAmount +// without bound. The fix holds replies in order and drains on recovery — never +// dropping a frame (which would recreate the corruption bug on the legacy path). + +function publicKeyToBase64(key: Uint8Array): string { + return Buffer.from(key).toString('base64') +} + +function createMockWs() { + const sent: string[] = [] + return { + OPEN: 1 as const, + readyState: 1, + bufferedAmount: 0, + send: vi.fn((data: string) => { + sent.push(data) + }), + close: vi.fn(), + sent + } +} + +function setup(overrides?: Partial) { + const serverKeys = generateKeyPair() + const clientKeys = generateKeyPair() + const ws = createMockWs() + const onError = vi.fn() + const channel = new E2EEChannel(ws as unknown as WebSocket, { + serverSecretKey: serverKeys.secretKey, + validateToken: (token) => token === 'valid-token', + onReady: vi.fn(), + onError, + ...overrides + }) + const sharedKey = deriveSharedKey(clientKeys.secretKey, serverKeys.publicKey) + channel.handleRawMessage( + JSON.stringify({ type: 'e2ee_hello', publicKeyB64: publicKeyToBase64(clientKeys.publicKey) }) + ) + channel.handleRawMessage( + encrypt(JSON.stringify({ type: 'e2ee_auth', deviceToken: 'valid-token' }), sharedKey) + ) + return { channel, ws, sharedKey, onError } +} + +/** Fire a streaming reply through the real channel's encryptedReply closure. */ +function emitReply(ctx: ReturnType, payload: string): void { + ctx.channel.onMessage((_plaintext, encryptedReply) => { + encryptedReply(payload) + }) + ctx.channel.handleRawMessage(encrypt('{"id":"x","method":"status.get"}', ctx.sharedKey)) +} + +describe('E2EE text reply backpressure', () => { + beforeEach(() => { + vi.useFakeTimers() + }) + afterEach(() => { + vi.useRealTimers() + }) + + it('holds text replies while over the buffer cap and drains them in order', () => { + const ctx = setup() + const baseline = ctx.ws.sent.length // ready + authenticated control frames + + // Simulate a congested socket: bufferedAmount pinned over the 8 MiB cap. + ctx.ws.bufferedAmount = 9 * 1024 * 1024 + + emitReply(ctx, '{"seq":1}') + emitReply(ctx, '{"seq":2}') + emitReply(ctx, '{"seq":3}') + + // Not dropped, not sent yet — parked in order on the queue. + expect(ctx.ws.sent.length).toBe(baseline) + + // Link drains; the queue flushes every reply, in order, none lost. + ctx.ws.bufferedAmount = 0 + vi.runOnlyPendingTimers() + + const replies = ctx.ws.sent.slice(baseline).map((frame) => decrypt(frame, ctx.sharedKey)) + expect(replies).toEqual(['{"seq":1}', '{"seq":2}', '{"seq":3}']) + expect(ctx.onError).not.toHaveBeenCalled() + }) + + it('sends straight through when the socket is not congested', () => { + const ctx = setup() + const baseline = ctx.ws.sent.length + emitReply(ctx, '{"ok":true}') + expect(ctx.ws.sent.length).toBe(baseline + 1) + expect(decrypt(ctx.ws.sent[baseline]!, ctx.sharedKey)).toBe('{"ok":true}') + }) +}) diff --git a/src/main/runtime/rpc/e2ee-channel.ts b/src/main/runtime/rpc/e2ee-channel.ts index 97ce9c5c813..63cc2c0dfcd 100644 --- a/src/main/runtime/rpc/e2ee-channel.ts +++ b/src/main/runtime/rpc/e2ee-channel.ts @@ -3,6 +3,10 @@ // handler only sees plaintext JSON, identical to the Unix socket path. import type { WebSocket } from 'ws' import { deriveSharedKey, encrypt, decrypt, encryptBytes, decryptBytes } from './e2ee-crypto' +import { + createWsOutboundBackpressureQueue, + type WsOutboundBackpressureQueue +} from '../../../shared/ws-outbound-backpressure-queue' type ChannelState = 'awaiting_hello' | 'awaiting_auth' | 'ready' @@ -48,6 +52,11 @@ export class E2EEChannel { ) => void) | null = null private binaryMessageHandler: ((plaintext: Uint8Array) => void) | null = null + // Why: the streaming JSON reply path (e.g. legacy terminal.subscribe) has no + // seq/resync, so it must never drop frames under backpressure. Hold text + // replies in order while bufferedAmount is over the cap and drain as it + // clears; only a wedged link (hard cap) closes the socket for a clean resync. + private textReplyQueue: WsOutboundBackpressureQueue | null = null deviceToken: string | null = null @@ -128,7 +137,7 @@ export class E2EEChannel { if (!this.sharedKey || this.ws.readyState !== this.ws.OPEN) { return } - this.ws.send(encrypt(response, this.sharedKey)) + this.ensureTextReplyQueue().enqueue(encrypt(response, this.sharedKey)) } const encryptedBinaryReply = (response: Uint8Array): boolean => { if (!this.sharedKey || this.ws.readyState !== this.ws.OPEN) { @@ -215,6 +224,22 @@ export class E2EEChannel { this.onReady(this) } + private ensureTextReplyQueue(): WsOutboundBackpressureQueue { + if (!this.textReplyQueue) { + this.textReplyQueue = createWsOutboundBackpressureQueue({ + send: (frame) => this.ws.send(frame), + // Encrypted replies are base64 ASCII strings, so length === byte count. + byteLengthOf: (frame) => frame.length, + getBufferedAmount: () => this.ws.bufferedAmount, + isWritable: () => Boolean(this.sharedKey) && this.ws.readyState === this.ws.OPEN, + // 1013 (Try Again Later): the link is wedged; drop the channel so the + // client reconnects and replays a full snapshot instead of unbounded RSS. + onOverflow: () => this.onError(1013, 'Outbound reply buffer overflow') + }) + } + return this.textReplyQueue + } + private sendEncryptedControl(message: unknown): void { if (this.ws.readyState === this.ws.OPEN && this.sharedKey) { this.ws.send(encrypt(JSON.stringify(message), this.sharedKey)) @@ -229,5 +254,7 @@ export class E2EEChannel { this.sharedKey = null this.messageHandler = null this.binaryMessageHandler = null + this.textReplyQueue?.dispose() + this.textReplyQueue = null } } diff --git a/src/shared/remote-runtime-client.ts b/src/shared/remote-runtime-client.ts index ecd9f1df790..15ec38e4005 100644 --- a/src/shared/remote-runtime-client.ts +++ b/src/shared/remote-runtime-client.ts @@ -29,6 +29,7 @@ import { type RemoteRuntimeSocketLivenessMonitor, type RemoteRuntimeSocketLivenessOptions } from './remote-runtime-socket-liveness' +import { createWsOutboundBackpressureQueue } from './ws-outbound-backpressure-queue' export { RemoteRuntimeClientError } from './remote-runtime-client-error' @@ -375,6 +376,8 @@ export async function subscribeRemoteRuntimeRequest( const cleanupSocketListeners = (): WebSocket | null => { liveness?.stop() liveness = null + sendQueue?.dispose() + sendQueue = null const socket = ws if (!socket) { return null @@ -420,11 +423,37 @@ export async function subscribeRemoteRuntimeRequest( } } + // Why: client input (keystrokes) must never be dropped under backpressure. + // Hold encrypted frames in order while bufferedAmount is over the cap and + // drain as it clears; a wedged link (hard cap) fails the socket so the + // renderer resubscribes and replays a fresh snapshot. + let sendQueue: ReturnType> | null = null + const ensureSendQueue = ( + socket: WebSocket + ): ReturnType> => { + if (!sendQueue) { + sendQueue = createWsOutboundBackpressureQueue({ + send: (frame) => socket.send(frame, { binary: true }), + byteLengthOf: (frame) => frame.byteLength, + getBufferedAmount: () => socket.bufferedAmount, + isWritable: () => socket.readyState === WebSocket.OPEN, + onOverflow: () => + fail( + new RemoteRuntimeClientError( + 'remote_runtime_unavailable', + 'Remote Orca runtime send buffer overflow; reconnecting.' + ) + ) + }) + } + return sendQueue + } + const sendBinary = (bytes: Uint8Array): boolean => { if (state !== 'ready' || !ws || ws.readyState !== WebSocket.OPEN) { return false } - ws.send(Buffer.from(encryptBytes(bytes, sharedKey)), { binary: true }) + ensureSendQueue(ws).enqueue(Buffer.from(encryptBytes(bytes, sharedKey))) return true } diff --git a/src/shared/ws-outbound-backpressure-queue.test.ts b/src/shared/ws-outbound-backpressure-queue.test.ts new file mode 100644 index 00000000000..674987ca073 --- /dev/null +++ b/src/shared/ws-outbound-backpressure-queue.test.ts @@ -0,0 +1,118 @@ +import { describe, expect, it, vi } from 'vitest' +import { createWsOutboundBackpressureQueue } from './ws-outbound-backpressure-queue' + +// Deterministic harness: bufferedAmount and the drain timer are both injected, +// so no wall-clock races. `runTimers` fires the single parked drain callback. + +function createHarness(overrides?: { + softCapBytes?: number + maxQueuedBytes?: number + writable?: boolean +}) { + const sent: string[] = [] + let bufferedAmount = 0 + let writable = overrides?.writable ?? true + const overflow = vi.fn() + let pendingTimer: (() => void) | null = null + + const queue = createWsOutboundBackpressureQueue({ + send: (frame) => sent.push(frame), + byteLengthOf: (frame) => frame.length, + getBufferedAmount: () => bufferedAmount, + isWritable: () => writable, + onOverflow: overflow, + softCapBytes: overrides?.softCapBytes ?? 100, + maxQueuedBytes: overrides?.maxQueuedBytes ?? 1000, + drainPollMs: 10, + setTimer: (cb) => { + pendingTimer = cb + return 1 as unknown as ReturnType + }, + clearTimer: () => { + pendingTimer = null + } + }) + + return { + queue, + sent, + overflow, + setBuffered: (value: number) => { + bufferedAmount = value + }, + setWritable: (value: boolean) => { + writable = value + }, + runTimer: () => { + const cb = pendingTimer + pendingTimer = null + cb?.() + }, + hasTimer: () => pendingTimer !== null + } +} + +describe('ws outbound backpressure queue', () => { + it('sends straight through while under the soft cap', () => { + const h = createHarness() + h.queue.enqueue('a') + h.queue.enqueue('b') + expect(h.sent).toEqual(['a', 'b']) + expect(h.hasTimer()).toBe(false) + }) + + it('parks frames in order while over the cap and drains on recovery without loss', () => { + const h = createHarness({ softCapBytes: 100 }) + h.setBuffered(200) // over cap + h.queue.enqueue('one') + h.queue.enqueue('two') + h.queue.enqueue('three') + // Nothing sent yet; all held in order. + expect(h.sent).toEqual([]) + expect(h.queue.queuedBytes()).toBe('one'.length + 'two'.length + 'three'.length) + + // Link recovers; the drain timer flushes everything in FIFO order. + h.setBuffered(0) + h.runTimer() + expect(h.sent).toEqual(['one', 'two', 'three']) + expect(h.queue.queuedBytes()).toBe(0) + }) + + it('keeps ordering when a frame arrives while a backlog is parked', () => { + const h = createHarness({ softCapBytes: 100 }) + h.setBuffered(200) + h.queue.enqueue('first') + h.setBuffered(0) + // Even though the wire is now clear, an existing backlog means the new + // frame must queue behind it, not jump the line. + h.queue.enqueue('second') + expect(h.sent).toEqual([]) + h.runTimer() + expect(h.sent).toEqual(['first', 'second']) + }) + + it('signals overflow (and drops backlog) when the hard cap is exceeded', () => { + const h = createHarness({ softCapBytes: 10, maxQueuedBytes: 8 }) + h.setBuffered(100) // over soft cap: everything queues + h.queue.enqueue('12345') // 5 bytes queued + expect(h.overflow).not.toHaveBeenCalled() + h.queue.enqueue('67890') // 10 bytes total > 8 -> overflow + expect(h.overflow).toHaveBeenCalledTimes(1) + // After overflow the queue is inert: no sends, no further overflow calls. + h.setBuffered(0) + h.runTimer() + h.queue.enqueue('later') + expect(h.sent).toEqual([]) + expect(h.overflow).toHaveBeenCalledTimes(1) + }) + + it('drops the backlog if the socket becomes unwritable mid-park', () => { + const h = createHarness({ softCapBytes: 10 }) + h.setBuffered(100) + h.queue.enqueue('data') + h.setWritable(false) + h.runTimer() + expect(h.sent).toEqual([]) + expect(h.queue.queuedBytes()).toBe(0) + }) +}) diff --git a/src/shared/ws-outbound-backpressure-queue.ts b/src/shared/ws-outbound-backpressure-queue.ts new file mode 100644 index 00000000000..7efc7a4dfdf --- /dev/null +++ b/src/shared/ws-outbound-backpressure-queue.ts @@ -0,0 +1,141 @@ +// Why: both the server reply path (e2ee-channel) and the client send path +// (remote-runtime-client) write to a ws with no backpressure handling. A fast +// producer over a slow link balloons ws.bufferedAmount / RSS without bound, or +// (binary path) silently drops frames. This queue holds outbound frames in +// order while bufferedAmount is over a soft cap and flushes as it drains, so no +// frame is dropped or reordered. It only signals overflow when a hard byte +// bound is exceeded (the link is effectively dead), letting the caller force a +// clean reconnect/resync instead of growing memory without limit. +// +// Generic over the frame type so it serves both the text reply path (encrypted +// base64 strings) and the binary send path (Uint8Array frames). + +export type WsOutboundBackpressureQueueOptions = { + /** Send a frame on the wire. Called only when under the soft cap. */ + send: (frame: TFrame) => void + /** Serialized byte length of a frame, for cap accounting. */ + byteLengthOf: (frame: TFrame) => number + /** Current ws.bufferedAmount in bytes. */ + getBufferedAmount: () => number + /** True when the socket can still accept sends (OPEN and keyed). */ + isWritable: () => boolean + /** + * Called once when queued bytes exceed maxQueuedBytes — the link is wedged. + * The caller should tear the connection down so a fresh subscription can + * replay an authoritative snapshot. The queue drops its backlog afterward. + */ + onOverflow: () => void + /** Soft cap: stop draining onto the wire while bufferedAmount is above this. */ + softCapBytes?: number + /** Hard cap on bytes held in this queue before onOverflow fires. */ + maxQueuedBytes?: number + /** Poll interval used to re-check bufferedAmount while parked. */ + drainPollMs?: number + /** Injectable scheduler for deterministic tests. */ + setTimer?: (cb: () => void, ms: number) => ReturnType + clearTimer?: (timer: ReturnType) => void +} + +export type WsOutboundBackpressureQueue = { + /** Queue-or-send a frame. Preserves order across all prior frames. */ + enqueue: (frame: TFrame) => void + /** Bytes currently held (not yet handed to the wire). */ + queuedBytes: () => number + /** Drop the backlog and stop the drain timer (call on close). */ + dispose: () => void +} + +const DEFAULT_SOFT_CAP_BYTES = 8 * 1024 * 1024 +// Why: tolerate a large transient burst (e.g. a build log spike) before +// declaring the link dead; 64 MiB is ~8x the soft cap yet still bounds RSS. +const DEFAULT_MAX_QUEUED_BYTES = 64 * 1024 * 1024 +const DEFAULT_DRAIN_POLL_MS = 25 + +export function createWsOutboundBackpressureQueue( + options: WsOutboundBackpressureQueueOptions +): WsOutboundBackpressureQueue { + const softCapBytes = options.softCapBytes ?? DEFAULT_SOFT_CAP_BYTES + const maxQueuedBytes = options.maxQueuedBytes ?? DEFAULT_MAX_QUEUED_BYTES + const drainPollMs = options.drainPollMs ?? DEFAULT_DRAIN_POLL_MS + const setTimer = options.setTimer ?? ((cb, ms) => setTimeout(cb, ms)) + const clearTimer = options.clearTimer ?? ((timer) => clearTimeout(timer)) + + // Why: a ws without a numeric bufferedAmount (some mocks/transports) must not + // strand frames in the queue forever; treat unknown backpressure as "clear". + const bufferedAmount = (): number => { + const value = options.getBufferedAmount() + return Number.isFinite(value) ? value : 0 + } + + const queue: { frame: TFrame; bytes: number }[] = [] + let queued = 0 + let timer: ReturnType | null = null + let overflowed = false + let disposed = false + + const stopTimer = (): void => { + if (timer !== null) { + clearTimer(timer) + timer = null + } + } + + const dropBacklog = (): void => { + queue.length = 0 + queued = 0 + stopTimer() + } + + // Drain as many queued frames as the wire will take without crossing the + // soft cap; re-arm the poll timer if frames remain. + const drain = (): void => { + if (disposed || overflowed) { + return + } + if (!options.isWritable()) { + // Socket went away mid-park; let the transport's own close path clean up. + dropBacklog() + return + } + while (queue.length > 0 && bufferedAmount() <= softCapBytes) { + const entry = queue.shift()! + queued -= entry.bytes + options.send(entry.frame) + } + if (queue.length > 0) { + timer = setTimer(drain, drainPollMs) + } else { + stopTimer() + } + } + + return { + enqueue(frame: TFrame): void { + if (disposed || overflowed) { + return + } + // Fast path: nothing parked and the wire is under the cap — send directly. + if (queue.length === 0 && bufferedAmount() <= softCapBytes) { + options.send(frame) + return + } + const bytes = options.byteLengthOf(frame) + queue.push({ frame, bytes }) + queued += bytes + if (queued > maxQueuedBytes) { + overflowed = true + dropBacklog() + options.onOverflow() + return + } + if (timer === null) { + timer = setTimer(drain, drainPollMs) + } + }, + queuedBytes: () => queued, + dispose(): void { + disposed = true + dropBacklog() + } + } +}