From 22f77635bfa5489a24f4f450950679d678f0a31c Mon Sep 17 00:00:00 2001 From: OrcaWin <293788423+OrcaWin@users.noreply.github.com> Date: Mon, 27 Jul 2026 19:04:01 -0700 Subject: [PATCH] fix(ssh): bound relay frame decoding --- src/main/ssh/relay-frame-decoder.ts | 316 +++++++++++++++++ .../ssh/relay-protocol-backpressure.test.ts | 187 +++++++++++ src/main/ssh/relay-protocol.ts | 165 +-------- ...h-channel-multiplexer-backpressure.test.ts | 138 ++++++++ src/main/ssh/ssh-channel-multiplexer.ts | 78 ++++- src/relay/protocol-backpressure.test.ts | 211 ++++++++++++ src/relay/protocol.ts | 171 ++-------- src/relay/relay-frame-decoder.ts | 317 ++++++++++++++++++ 8 files changed, 1271 insertions(+), 312 deletions(-) create mode 100644 src/main/ssh/relay-frame-decoder.ts create mode 100644 src/main/ssh/relay-protocol-backpressure.test.ts create mode 100644 src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts create mode 100644 src/relay/protocol-backpressure.test.ts create mode 100644 src/relay/relay-frame-decoder.ts diff --git a/src/main/ssh/relay-frame-decoder.ts b/src/main/ssh/relay-frame-decoder.ts new file mode 100644 index 00000000000..2f4683a89fc --- /dev/null +++ b/src/main/ssh/relay-frame-decoder.ts @@ -0,0 +1,316 @@ +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 type DecodedFrame = { + type: number + id: number + ack: number + payload: Buffer +} + +export type FrameDecoderOptions = { + maxFramesPerTurn?: number + maxBytesPerTurn?: number + maxTurnMs?: number + now?: () => number + schedule?: (callback: () => void) => unknown + cancelScheduled?: (handle: unknown) => void + pause?: () => void + resume?: () => void +} + +export class FrameDecoder { + private chunks: Buffer[] = [] + private bufferedLength = 0 + private oversizedPayloadBytesRemaining = 0 + private onFrame: (frame: DecodedFrame) => void + private onError: ((err: Error) => void) | null + private maxFramesPerTurn: number + private maxBytesPerTurn: number + private maxTurnMs: number + private now: () => number + private schedule: (callback: () => void) => unknown + private cancelScheduled: (handle: unknown) => void + private pause: (() => void) | null + private resume: (() => void) | null + private continuationHandle: unknown + private continuationHandleAssigned = false + private continuationScheduled = false + private paused = false + private draining = false + private generation = 0 + + constructor( + onFrame: (frame: DecodedFrame) => void, + onError?: (err: Error) => void, + options: FrameDecoderOptions = {} + ) { + this.onFrame = onFrame + this.onError = onError ?? null + this.maxFramesPerTurn = positiveLimit( + options.maxFramesPerTurn, + FRAME_DECODER_MAX_FRAMES_PER_TURN + ) + this.maxBytesPerTurn = positiveLimit(options.maxBytesPerTurn, FRAME_DECODER_MAX_BYTES_PER_TURN) + this.maxTurnMs = positiveLimit(options.maxTurnMs, FRAME_DECODER_MAX_TURN_MS) + this.now = options.now ?? Date.now + this.schedule = options.schedule ?? ((callback) => setImmediate(callback)) + this.cancelScheduled = + options.cancelScheduled ?? ((handle) => clearImmediate(handle as NodeJS.Immediate)) + this.pause = options.pause ?? null + this.resume = options.resume ?? null + } + + feed(chunk: Buffer | Uint8Array): void { + const buf = Buffer.isBuffer(chunk) + ? chunk + : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + if (buf.length > 0) { + this.chunks.push(buf) + this.bufferedLength += buf.length + } + if (!this.draining && !this.continuationScheduled) { + this.drainTurn() + } + } + + reset(): void { + this.generation += 1 + this.cancelContinuation() + this.chunks = [] + this.bufferedLength = 0 + this.oversizedPayloadBytesRemaining = 0 + this.releasePause() + } + + drain(): Buffer { + const out = + this.chunks.length === 1 ? this.chunks[0] : Buffer.concat(this.chunks, this.bufferedLength) + this.reset() + return out + } + + private drainTurn(): void { + if (this.draining) { + return + } + this.draining = true + const generation = this.generation + const startedAt = this.now() + let frames = 0 + let bytes = 0 + + try { + while (generation === this.generation) { + if ( + frames >= this.maxFramesPerTurn || + bytes >= this.maxBytesPerTurn || + (frames > 0 && this.now() - startedAt >= this.maxTurnMs) + ) { + break + } + const discarded = this.discardOversizedPayload(bytes) + if (discarded > 0) { + bytes += discarded + continue + } + if (this.bufferedLength < HEADER_LENGTH) { + break + } + const header = this.peekBytes(HEADER_LENGTH) + const length = header.readUInt32BE(9) + if (length > MAX_MESSAGE_SIZE) { + this.discardBytes(HEADER_LENGTH) + this.oversizedPayloadBytesRemaining = length + bytes += HEADER_LENGTH + this.onError?.(new Error(`Frame payload too large: ${length} bytes — discarded`)) + continue + } + const totalLength = HEADER_LENGTH + length + if (this.bufferedLength < totalLength) { + break + } + if (frames > 0 && bytes + totalLength > this.maxBytesPerTurn) { + break + } + const framed = this.takeBytes(totalLength) + frames += 1 + bytes += totalLength + this.onFrame({ + type: framed[0], + id: framed.readUInt32BE(1), + ack: framed.readUInt32BE(5), + payload: framed.subarray(HEADER_LENGTH, totalLength) + }) + } + } finally { + this.draining = false + } + + if (generation !== this.generation) { + return + } + if (this.hasRunnableWork()) { + this.scheduleContinuation() + } else { + this.releasePause() + } + } + + private discardOversizedPayload(bytes: number): number { + if (this.oversizedPayloadBytesRemaining === 0 || this.bufferedLength === 0) { + return 0 + } + const discarded = Math.min( + this.oversizedPayloadBytesRemaining, + this.bufferedLength, + Math.max(1, this.maxBytesPerTurn - bytes) + ) + this.discardBytes(discarded) + this.oversizedPayloadBytesRemaining -= discarded + return discarded + } + + private hasRunnableWork(): boolean { + if (this.oversizedPayloadBytesRemaining > 0) { + return this.bufferedLength > 0 + } + if (this.bufferedLength < HEADER_LENGTH) { + return false + } + const length = this.peekBytes(HEADER_LENGTH).readUInt32BE(9) + return length > MAX_MESSAGE_SIZE || this.bufferedLength >= HEADER_LENGTH + length + } + + private scheduleContinuation(): void { + if (this.continuationScheduled) { + return + } + const generation = this.generation + this.continuationScheduled = true + try { + this.acquirePause() + } catch (error) { + this.continuationScheduled = false + throw error + } + if (generation !== this.generation) { + this.continuationScheduled = false + return + } + try { + this.continuationHandle = this.schedule(() => { + if (!this.continuationScheduled || generation !== this.generation) { + return + } + this.continuationScheduled = false + this.continuationHandleAssigned = false + this.continuationHandle = undefined + this.drainTurn() + }) + this.continuationHandleAssigned = true + } catch (error) { + this.continuationScheduled = false + this.continuationHandle = undefined + this.releasePause() + throw error + } + } + + private cancelContinuation(): void { + if (!this.continuationScheduled) { + return + } + this.continuationScheduled = false + if (this.continuationHandleAssigned) { + this.cancelScheduled(this.continuationHandle) + } + this.continuationHandleAssigned = false + this.continuationHandle = undefined + } + + private acquirePause(): void { + if (!this.paused) { + this.paused = true + try { + this.pause?.() + } catch (error) { + this.paused = false + throw error + } + } + } + + private releasePause(): void { + if (this.paused) { + this.paused = false + this.resume?.() + } + } + + private peekBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length >= count) { + return first + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + for (const part of this.chunks) { + copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) + if (copied >= count) { + break + } + } + return out + } + + private takeBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length === count) { + this.chunks.shift() + this.bufferedLength -= count + return first + } + if (first.length > count) { + this.chunks[0] = first.subarray(count) + this.bufferedLength -= count + return first.subarray(0, count) + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + while (copied < count) { + const part = this.chunks[0] + const take = Math.min(part.length, count - copied) + part.copy(out, copied, 0, take) + copied += take + if (take === part.length) { + this.chunks.shift() + } else { + this.chunks[0] = part.subarray(take) + } + } + this.bufferedLength -= count + return out + } + + private discardBytes(count: number): void { + let remaining = count + while (remaining > 0) { + const part = this.chunks[0] + if (part.length <= remaining) { + this.chunks.shift() + remaining -= part.length + } else { + this.chunks[0] = part.subarray(remaining) + remaining = 0 + } + } + this.bufferedLength -= count + } +} + +const positiveLimit = (value: number | undefined, fallback: number): number => + value !== undefined && Number.isFinite(value) && value > 0 ? value : fallback diff --git a/src/main/ssh/relay-protocol-backpressure.test.ts b/src/main/ssh/relay-protocol-backpressure.test.ts new file mode 100644 index 00000000000..d78b64ff1fe --- /dev/null +++ b/src/main/ssh/relay-protocol-backpressure.test.ts @@ -0,0 +1,187 @@ +import { describe, expect, it, vi } from 'vitest' +import { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + MessageType, + encodeFrame, + type DecodedFrame +} from './relay-protocol' + +function createScheduler(): { + schedule: (callback: () => void) => number + cancel: (handle: unknown) => void + runNext: () => void + pending: () => number +} { + let nextHandle = 1 + const callbacks = new Map void>() + return { + schedule: (callback) => { + const handle = nextHandle++ + callbacks.set(handle, callback) + return handle + }, + cancel: (handle) => callbacks.delete(handle as number), + runNext: () => { + const entry = callbacks.entries().next().value as [number, () => void] | undefined + if (!entry) { + throw new Error('No decoder continuation scheduled') + } + callbacks.delete(entry[0]) + entry[1]() + }, + pending: () => callbacks.size + } +} + +function frame(id: number, payload = `${id}`): Buffer { + return encodeFrame(MessageType.Regular, id, 0, Buffer.from(payload)) +} + +describe('FrameDecoder bounded turns', () => { + it('emits the first frame synchronously and preserves order through self-pause', () => { + const scheduler = createScheduler() + const seen: number[] = [] + let decoder: FrameDecoder + const pause = vi.fn(() => decoder.feed(frame(4))) + const resume = vi.fn() + decoder = new FrameDecoder((decoded) => seen.push(decoded.id), undefined, { + maxFramesPerTurn: 1, + schedule: scheduler.schedule, + cancelScheduled: scheduler.cancel, + pause, + resume + }) + + decoder.feed(Buffer.concat([frame(1), frame(2), frame(3)])) + + expect(seen).toEqual([1]) + expect(pause).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(1) + + scheduler.runNext() + scheduler.runNext() + scheduler.runNext() + + expect(seen).toEqual([1, 2, 3, 4]) + expect(pause).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(0) + }) + + it('bounds decoded bytes and time independently from the frame count', () => { + const byteScheduler = createScheduler() + const byteSeen: number[] = [] + const first = frame(1, 'one') + const second = frame(2, 'two') + const byteDecoder = new FrameDecoder((decoded) => byteSeen.push(decoded.id), undefined, { + maxFramesPerTurn: 64, + maxBytesPerTurn: first.length, + schedule: byteScheduler.schedule, + cancelScheduled: byteScheduler.cancel + }) + + byteDecoder.feed(Buffer.concat([first, second])) + expect(byteSeen).toEqual([1]) + byteScheduler.runNext() + expect(byteSeen).toEqual([1, 2]) + + const timeScheduler = createScheduler() + const timeSeen: number[] = [] + let nowCalls = 0 + const timeDecoder = new FrameDecoder((decoded) => timeSeen.push(decoded.id), undefined, { + maxFramesPerTurn: 64, + maxBytesPerTurn: MAX_MESSAGE_SIZE + HEADER_LENGTH, + maxTurnMs: 4, + now: () => (nowCalls++ === 0 ? 0 : 5), + schedule: timeScheduler.schedule, + cancelScheduled: timeScheduler.cancel + }) + + timeDecoder.feed(Buffer.concat([frame(3), frame(4)])) + expect(timeSeen).toEqual([3]) + timeScheduler.runNext() + expect(timeSeen).toEqual([3, 4]) + }) + + it('releases its pause epoch when continuation scheduling throws', () => { + const pause = vi.fn() + const resume = vi.fn() + const decoder = new FrameDecoder(() => {}, undefined, { + maxFramesPerTurn: 1, + pause, + resume, + schedule: () => { + throw new Error('scheduler unavailable') + } + }) + + expect(() => decoder.feed(Buffer.concat([frame(1), frame(2)]))).toThrow('scheduler unavailable') + expect(pause).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + decoder.reset() + expect(resume).toHaveBeenCalledTimes(1) + }) + + it('keeps reads active for partial frames and incrementally discards oversized payloads', () => { + const errors: Error[] = [] + const seen: DecodedFrame[] = [] + const pause = vi.fn() + const decoder = new FrameDecoder( + (decoded) => seen.push(decoded), + (error) => errors.push(error), + { pause } + ) + const valid = frame(2, 'complete') + + decoder.feed(valid.subarray(0, HEADER_LENGTH + 2)) + expect(seen).toHaveLength(0) + expect(pause).not.toHaveBeenCalled() + decoder.feed(valid.subarray(HEADER_LENGTH + 2)) + expect(seen.map(({ id }) => id)).toEqual([2]) + + const oversizedHeader = Buffer.alloc(HEADER_LENGTH) + oversizedHeader[0] = MessageType.Regular + oversizedHeader.writeUInt32BE(3, 1) + oversizedHeader.writeUInt32BE(MAX_MESSAGE_SIZE + 1, 9) + decoder.feed(Buffer.concat([oversizedHeader, Buffer.alloc(32)])) + + expect(errors).toHaveLength(1) + expect(pause).not.toHaveBeenCalled() + decoder.reset() + decoder.feed(frame(4, 'after-reset')) + expect(seen.map(({ id }) => id)).toEqual([2, 4]) + }) + + it('drain and reset cancel continuation ownership without replaying residue', () => { + const scheduler = createScheduler() + const seen: number[] = [] + const resume = vi.fn() + const cancel = vi.fn(scheduler.cancel) + const decoder = new FrameDecoder((decoded) => seen.push(decoded.id), undefined, { + maxFramesPerTurn: 1, + schedule: scheduler.schedule, + cancelScheduled: cancel, + resume + }) + const second = frame(2, 'residue') + + decoder.feed(Buffer.concat([frame(1), second])) + const residue = decoder.drain() + + expect(seen).toEqual([1]) + expect(residue.equals(second)).toBe(true) + expect(cancel).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(0) + + decoder.feed(Buffer.concat([frame(3), frame(4)])) + decoder.reset() + expect(cancel).toHaveBeenCalledTimes(2) + expect(scheduler.pending()).toBe(0) + + decoder.feed(frame(5)) + expect(seen).toEqual([1, 3, 5]) + }) +}) diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index 0412c97994d..afc7a9ea6a4 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -3,17 +3,30 @@ // See design-ssh-support.md § JSON-RPC Protocol Specification. import { DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../../shared/ssh-types' +import { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + FRAME_DECODER_MAX_FRAMES_PER_TURN, + FRAME_DECODER_MAX_BYTES_PER_TURN, + FRAME_DECODER_MAX_TURN_MS +} from './relay-frame-decoder' + +export { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + FRAME_DECODER_MAX_FRAMES_PER_TURN, + FRAME_DECODER_MAX_BYTES_PER_TURN, + FRAME_DECODER_MAX_TURN_MS +} +export type { DecodedFrame, FrameDecoderOptions } from './relay-frame-decoder' export const RELAY_VERSION = '0.1.0' export const RELAY_SENTINEL = `ORCA-RELAY v${RELAY_VERSION} READY\n` export const RELAY_SENTINEL_TIMEOUT_MS = 10_000 export const RELAY_REMOTE_DIR = '.orca-remote' -// ── Framing constants (VS Code ProtocolConstants) ─────────────────── - -export const HEADER_LENGTH = 13 -export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 // 16 MB - /** Message type byte. */ export const MessageType = { Regular: 1, @@ -155,148 +168,6 @@ export function encodeKeepAliveFrame(id: number, ack: number): Buffer { return encodeFrame(MessageType.KeepAlive, id, ack, Buffer.alloc(0)) } -export type DecodedFrame = { - type: number - id: number - ack: number - payload: Buffer -} - -/** - * Incremental frame parser. Feed it chunks of data; it emits complete frames. - */ -export class FrameDecoder { - // Why: feed() runs on the Electron main thread for every SSH channel data - // event. Rebuilding one contiguous buffer per feed (Buffer.concat) re-copies - // every already-buffered byte for each incoming ~32KB TCP chunk — O(n²) per - // large frame (a 340KB fs.streamChunk frame cost ~2MB of memcpy). A chunk - // list assembles each frame exactly once instead. - private chunks: Buffer[] = [] - private bufferedLength = 0 - private onFrame: (frame: DecodedFrame) => void - private onError: ((err: Error) => void) | null - - constructor(onFrame: (frame: DecodedFrame) => void, onError?: (err: Error) => void) { - this.onFrame = onFrame - this.onError = onError ?? null - } - - feed(chunk: Buffer | Uint8Array): void { - const buf = Buffer.isBuffer(chunk) - ? chunk - : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) - if (buf.length > 0) { - this.chunks.push(buf) - this.bufferedLength += buf.length - } - - while (this.bufferedLength >= HEADER_LENGTH) { - const header = this.peekBytes(HEADER_LENGTH) - const length = header.readUInt32BE(9) - const totalLength = HEADER_LENGTH + length - - if (this.bufferedLength < totalLength) { - // Not fully received yet (also holds oversized frames until they can - // be skipped whole, keeping the decoder synchronized). - break - } - - // Why: throwing here would leave the buffer in a partially consumed - // state — subsequent feed() calls would try to parse leftover payload - // bytes as a new header, corrupting every future frame. Instead we - // skip the entire oversized frame so the decoder stays synchronized. - if (length > MAX_MESSAGE_SIZE) { - this.discardBytes(totalLength) - const err = new Error(`Frame payload too large: ${length} bytes — discarded`) - if (this.onError) { - this.onError(err) - } - continue - } - - const framed = this.takeBytes(totalLength) - const frame: DecodedFrame = { - type: framed[0], - id: framed.readUInt32BE(1), - ack: framed.readUInt32BE(5), - payload: framed.subarray(HEADER_LENGTH, totalLength) - } - this.onFrame(frame) - } - } - - reset(): void { - this.chunks = [] - this.bufferedLength = 0 - } - - /** View of the first `count` buffered bytes without consuming them. */ - private peekBytes(count: number): Buffer { - const first = this.chunks[0] - if (first.length >= count) { - return first - } - const out = Buffer.allocUnsafe(count) - let copied = 0 - for (const part of this.chunks) { - copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) - if (copied >= count) { - break - } - } - return out - } - - /** Consume and return the first `count` buffered bytes (single copy). */ - private takeBytes(count: number): Buffer { - const first = this.chunks[0] - if (first.length === count) { - this.chunks.shift() - this.bufferedLength -= count - return first - } - if (first.length > count) { - this.chunks[0] = first.subarray(count) - this.bufferedLength -= count - return first.subarray(0, count) - } - const out = Buffer.allocUnsafe(count) - let copied = 0 - while (copied < count) { - const part = this.chunks[0] - const take = Math.min(part.length, count - copied) - part.copy(out, copied, 0, take) - copied += take - if (take === part.length) { - this.chunks.shift() - } else { - this.chunks[0] = part.subarray(take) - } - } - this.bufferedLength -= count - return out - } - - /** Consume the first `count` buffered bytes without assembling them. */ - private discardBytes(count: number): void { - let remaining = count - while (remaining > 0) { - const part = this.chunks[0] - if (part.length <= remaining) { - this.chunks.shift() - remaining -= part.length - } else { - this.chunks[0] = part.subarray(remaining) - remaining = 0 - } - } - this.bufferedLength -= count - } -} - -/** - * Parse a JSON-RPC message from a frame payload. - */ export function parseJsonRpcMessage(payload: Buffer): JsonRpcMessage { const text = payload.toString('utf-8') const msg = JSON.parse(text) as JsonRpcMessage diff --git a/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts b/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts new file mode 100644 index 00000000000..1e3d415dbfe --- /dev/null +++ b/src/main/ssh/ssh-channel-multiplexer-backpressure.test.ts @@ -0,0 +1,138 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { SshChannelMultiplexer, type MultiplexerTransport } from './ssh-channel-multiplexer' +import { encodeFrame, encodeKeepAliveFrame, MessageType } from './relay-protocol' + +type MockTransport = MultiplexerTransport & { + data: (chunk: Buffer) => void + written: Buffer[] + pauseReads: ReturnType + resumeReads: ReturnType +} + +type MuxInternals = { + nextOutgoingSeq: number + highestAckedBySelf: number + lastReceivedAt: number + decoderReadPaused: boolean + unackedTimestamps: Map +} + +function createTransport(): MockTransport { + let onData: (chunk: Buffer) => void = () => {} + const written: Buffer[] = [] + const pauseReads = vi.fn() + const resumeReads = vi.fn() + return { + write: (data) => { + written.push(data) + }, + onData: (callback) => { + onData = callback + }, + onClose: () => {}, + pauseReads, + resumeReads, + data: (chunk) => onData(chunk), + written + } +} + +function internals(mux: SshChannelMultiplexer): MuxInternals { + return mux as unknown as MuxInternals +} + +describe('SshChannelMultiplexer backpressure hardening', () => { + let transport: MockTransport + let mux: SshChannelMultiplexer + + beforeEach(() => { + vi.useFakeTimers() + vi.setSystemTime(0) + transport = createTransport() + mux = new SshChannelMultiplexer(transport) + }) + + afterEach(() => { + mux.dispose() + vi.restoreAllMocks() + vi.useRealTimers() + }) + + it('clamps ack=0xffffffff and deletes only retained sequence keys', () => { + mux.notify('one') + mux.notify('two') + const state = internals(mux) + const deleteSpy = vi.spyOn(state.unackedTimestamps, 'delete') + + transport.data(encodeKeepAliveFrame(1, 0xffffffff)) + + expect(deleteSpy).toHaveBeenCalledTimes(2) + expect(state.unackedTimestamps.size).toBe(0) + expect(state.highestAckedBySelf).toBe(state.nextOutgoingSeq - 1) + }) + + it('caps timestamps while preserving a reserved liveness entry', () => { + for (let i = 0; i < 5000; i += 1) { + mux.notify('bounded') + } + const state = internals(mux) + + expect(state.unackedTimestamps.size).toBe(4095) + expect(state.unackedTimestamps.has(1)).toBe(true) + void mux.probeLiveness(60_000) + expect(state.unackedTimestamps.size).toBe(4096) + void mux.probeLiveness(60_000) + expect(state.unackedTimestamps.size).toBe(4096) + }) + + it('suppresses timeout while self-paused and rebases both clocks on resume', () => { + mux.dispose() + const continuations: (() => void)[] = [] + vi.spyOn(globalThis, 'setImmediate').mockImplementation( + (callback: (...args: never[]) => void) => { + continuations.push(callback) + return {} as NodeJS.Immediate + } + ) + transport = createTransport() + mux = new SshChannelMultiplexer(transport) + mux.notify('tracked') + + const incoming = Buffer.concat( + Array.from({ length: 65 }, (_, index) => encodeKeepAliveFrame(index + 1, 0)) + ) + transport.data(incoming) + const state = internals(mux) + + expect(state.decoderReadPaused).toBe(true) + expect(transport.pauseReads).toHaveBeenCalledTimes(1) + expect(continuations).toHaveLength(1) + + vi.advanceTimersByTime(25_000) + expect(mux.isDisposed()).toBe(false) + + continuations.shift()!() + const resumedAt = Date.now() + expect(state.decoderReadPaused).toBe(false) + expect(transport.resumeReads).toHaveBeenCalledTimes(1) + expect(state.lastReceivedAt).toBe(resumedAt) + expect(new Set(state.unackedTimestamps.values())).toEqual(new Set([resumedAt])) + }) + + it('keeps frame dispatch ordered across a decoder continuation', () => { + const seen: number[] = [] + mux.onNotification((_method, params) => seen.push(params.index as number)) + const frames = Array.from({ length: 65 }, (_, index) => { + const payload = Buffer.from( + JSON.stringify({ jsonrpc: '2.0', method: 'ordered', params: { index } }) + ) + return encodeFrame(MessageType.Regular, index + 1, 0, payload) + }) + + transport.data(Buffer.concat(frames)) + expect(seen).toEqual(Array.from({ length: 64 }, (_, index) => index)) + vi.runAllTicks() + vi.advanceTimersByTime(0) + expect(seen).toEqual(Array.from({ length: 65 }, (_, index) => index)) + }) +}) diff --git a/src/main/ssh/ssh-channel-multiplexer.ts b/src/main/ssh/ssh-channel-multiplexer.ts index 9886b724755..490e5be74fa 100644 --- a/src/main/ssh/ssh-channel-multiplexer.ts +++ b/src/main/ssh/ssh-channel-multiplexer.ts @@ -19,6 +19,8 @@ export type MultiplexerTransport = { write: (data: Buffer) => void onData: (cb: (data: Buffer) => void) => void onClose: (cb: () => void) => void + pauseReads?: () => void + resumeReads?: () => void close?: () => void } @@ -34,6 +36,8 @@ export type MethodNotificationHandler = (params: Record) => voi export type RequestHandler = (params: Record) => Promise | unknown const REQUEST_TIMEOUT_MS = 30_000 +const MAX_ORDINARY_UNACKED_TIMESTAMPS = 4095 +const MAX_UNACKED_TIMESTAMPS = MAX_ORDINARY_UNACKED_TIMESTAMPS + 1 // Why: a tick gap far beyond the interval means the process was paused // (system sleep, App Nap timer throttling) — not that the link is dead (#7773). const WAKE_GAP_MS = KEEPALIVE_SEND_MS * 3 @@ -56,6 +60,7 @@ export class SshChannelMultiplexer { private disposeHandlers: ((reason: 'shutdown' | 'connection_lost') => void)[] = [] private connectionHealthTimer: ReturnType | null = null private disposed = false + private decoderReadPaused = false // Track the oldest unacked outgoing message timestamp private unackedTimestamps = new Map() @@ -69,7 +74,11 @@ export class SshChannelMultiplexer { this.decoder = new FrameDecoder( (frame) => this.handleFrame(frame), - (err) => this.handleProtocolError(err) + (err) => this.handleProtocolError(err), + { + pause: () => this.pauseDecoderReads(), + resume: () => this.resumeDecoderReads() + } ) transport.onData((data) => { @@ -324,7 +333,7 @@ export class SshChannelMultiplexer { private sendMessage(msg: JsonRpcMessage): void { const seq = this.nextOutgoingSeq++ const frame = encodeJsonRpcFrame(msg, seq, this.highestReceivedSeq) - this.unackedTimestamps.set(seq, Date.now()) + this.trackOutgoingTimestamp(seq, false) try { this.transport.write(frame) } catch (err) { @@ -341,7 +350,7 @@ export class SshChannelMultiplexer { } const seq = this.nextOutgoingSeq++ const frame = encodeKeepAliveFrame(seq, this.highestReceivedSeq) - this.unackedTimestamps.set(seq, Date.now()) + this.trackOutgoingTimestamp(seq, true) try { this.transport.write(frame) } catch (err) { @@ -363,12 +372,16 @@ export class SshChannelMultiplexer { this.highestReceivedSeq = frame.id } - // Process ack from remote: discard timestamps for acked messages - if (frame.ack > this.highestAckedBySelf) { - for (let i = this.highestAckedBySelf + 1; i <= frame.ack; i++) { - this.unackedTimestamps.delete(i) + // Header ACKs are untrusted uint32 values; work stays proportional to the + // bounded set of sequence keys we actually retained. + const acknowledgedSeq = Math.min(frame.ack, this.nextOutgoingSeq - 1) + if (acknowledgedSeq > this.highestAckedBySelf) { + for (const seq of this.unackedTimestamps.keys()) { + if (seq <= acknowledgedSeq) { + this.unackedTimestamps.delete(seq) + } } - this.highestAckedBySelf = frame.ack + this.highestAckedBySelf = acknowledgedSeq } if (frame.type === MessageType.KeepAlive) { @@ -495,15 +508,12 @@ export class SshChannelMultiplexer { // before this tick's fresh probe, then allow the next full window. const resumedAfterWake = sinceLastTick > WAKE_GAP_MS if (resumedAfterWake) { - this.lastReceivedAt = now - for (const seq of this.unackedTimestamps.keys()) { - this.unackedTimestamps.set(seq, now) - } + this.rebaseHealthClocks(now) } this.sendKeepAlive() - if (this.disposed || resumedAfterWake) { + if (this.disposed || resumedAfterWake || this.decoderReadPaused) { return } @@ -529,4 +539,46 @@ export class SshChannelMultiplexer { console.warn(`[ssh-mux] Protocol error: ${err instanceof Error ? err.message : String(err)}`) this.dispose('connection_lost') } + + private trackOutgoingTimestamp(seq: number, liveness: boolean): void { + const limit = liveness ? MAX_UNACKED_TIMESTAMPS : MAX_ORDINARY_UNACKED_TIMESTAMPS + if (this.unackedTimestamps.size < limit) { + this.unackedTimestamps.set(seq, Date.now()) + } + } + + private pauseDecoderReads(): void { + if (this.disposed || this.decoderReadPaused) { + return + } + this.decoderReadPaused = true + try { + this.transport.pauseReads?.() + } catch (error) { + this.handleProtocolError(error) + } + } + + private resumeDecoderReads(): void { + if (!this.decoderReadPaused) { + return + } + this.decoderReadPaused = false + if (this.disposed) { + return + } + this.rebaseHealthClocks(Date.now()) + try { + this.transport.resumeReads?.() + } catch (error) { + this.handleProtocolError(error) + } + } + + private rebaseHealthClocks(now: number): void { + this.lastReceivedAt = now + for (const seq of this.unackedTimestamps.keys()) { + this.unackedTimestamps.set(seq, now) + } + } } diff --git a/src/relay/protocol-backpressure.test.ts b/src/relay/protocol-backpressure.test.ts new file mode 100644 index 00000000000..18ed2e35495 --- /dev/null +++ b/src/relay/protocol-backpressure.test.ts @@ -0,0 +1,211 @@ +import { describe, expect, it, vi } from 'vitest' +import { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + MessageType, + encodeFrame, + type DecodedFrame +} from './protocol' + +function createScheduler(): { + schedule: (callback: () => void) => number + cancel: (handle: unknown) => void + runNext: () => void + pending: () => number +} { + let nextHandle = 1 + const callbacks = new Map void>() + return { + schedule: (callback) => { + const handle = nextHandle++ + callbacks.set(handle, callback) + return handle + }, + cancel: (handle) => callbacks.delete(handle as number), + runNext: () => { + const entry = callbacks.entries().next().value as [number, () => void] | undefined + if (!entry) { + throw new Error('No decoder continuation scheduled') + } + callbacks.delete(entry[0]) + entry[1]() + }, + pending: () => callbacks.size + } +} + +function frame(id: number, payload = `${id}`): Buffer { + return encodeFrame(MessageType.Regular, id, 0, Buffer.from(payload)) +} + +describe('relay FrameDecoder bounded turns', () => { + it('emits the first frame synchronously and preserves order through self-pause', () => { + const scheduler = createScheduler() + const seen: number[] = [] + let decoder: FrameDecoder + const pause = vi.fn(() => decoder.feed(frame(4))) + const resume = vi.fn() + decoder = new FrameDecoder((decoded) => seen.push(decoded.id), undefined, { + maxFramesPerTurn: 1, + schedule: scheduler.schedule, + cancelScheduled: scheduler.cancel, + pause, + resume + }) + + decoder.feed(Buffer.concat([frame(1), frame(2), frame(3)])) + + expect(seen).toEqual([1]) + expect(pause).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(1) + + scheduler.runNext() + scheduler.runNext() + scheduler.runNext() + + expect(seen).toEqual([1, 2, 3, 4]) + expect(pause).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(0) + }) + + it('bounds decoded bytes and time independently from the frame count', () => { + const byteScheduler = createScheduler() + const byteSeen: number[] = [] + const first = frame(1, 'one') + const second = frame(2, 'two') + const byteDecoder = new FrameDecoder((decoded) => byteSeen.push(decoded.id), undefined, { + maxFramesPerTurn: 64, + maxBytesPerTurn: first.length, + schedule: byteScheduler.schedule, + cancelScheduled: byteScheduler.cancel + }) + + byteDecoder.feed(Buffer.concat([first, second])) + expect(byteSeen).toEqual([1]) + byteScheduler.runNext() + expect(byteSeen).toEqual([1, 2]) + + const timeScheduler = createScheduler() + const timeSeen: number[] = [] + let nowCalls = 0 + const timeDecoder = new FrameDecoder((decoded) => timeSeen.push(decoded.id), undefined, { + maxFramesPerTurn: 64, + maxBytesPerTurn: MAX_MESSAGE_SIZE + HEADER_LENGTH, + maxTurnMs: 4, + now: () => (nowCalls++ === 0 ? 0 : 5), + schedule: timeScheduler.schedule, + cancelScheduled: timeScheduler.cancel + }) + + timeDecoder.feed(Buffer.concat([frame(3), frame(4)])) + expect(timeSeen).toEqual([3]) + timeScheduler.runNext() + expect(timeSeen).toEqual([3, 4]) + }) + + it('releases its pause epoch when continuation scheduling throws', () => { + const pause = vi.fn() + const resume = vi.fn() + const decoder = new FrameDecoder(() => {}, undefined, { + maxFramesPerTurn: 1, + pause, + resume, + schedule: () => { + throw new Error('scheduler unavailable') + } + }) + + expect(() => decoder.feed(Buffer.concat([frame(1), frame(2)]))).toThrow('scheduler unavailable') + expect(pause).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + decoder.reset() + expect(resume).toHaveBeenCalledTimes(1) + }) + + it('keeps reads active for partial frames and incrementally discards oversized payloads', () => { + const errors: Error[] = [] + const seen: DecodedFrame[] = [] + const pause = vi.fn() + const decoder = new FrameDecoder( + (decoded) => seen.push(decoded), + (error) => errors.push(error), + { pause } + ) + const valid = frame(2, 'complete') + + decoder.feed(valid.subarray(0, HEADER_LENGTH + 2)) + expect(seen).toHaveLength(0) + expect(pause).not.toHaveBeenCalled() + decoder.feed(valid.subarray(HEADER_LENGTH + 2)) + expect(seen.map(({ id }) => id)).toEqual([2]) + + const oversizedHeader = Buffer.alloc(HEADER_LENGTH) + oversizedHeader[0] = MessageType.Regular + oversizedHeader.writeUInt32BE(3, 1) + oversizedHeader.writeUInt32BE(MAX_MESSAGE_SIZE + 1, 9) + decoder.feed(Buffer.concat([oversizedHeader, Buffer.alloc(32)])) + + expect(errors).toHaveLength(1) + expect(pause).not.toHaveBeenCalled() + decoder.reset() + decoder.feed(frame(4, 'after-reset')) + expect(seen.map(({ id }) => id)).toEqual([2, 4]) + }) + + it('resynchronizes after an oversized payload across bounded discard turns', () => { + const scheduler = createScheduler() + const errors: Error[] = [] + const seen: number[] = [] + const decoder = new FrameDecoder( + (decoded) => seen.push(decoded.id), + (error) => errors.push(error), + { schedule: scheduler.schedule, cancelScheduled: scheduler.cancel } + ) + const oversizedLength = MAX_MESSAGE_SIZE + 1 + const header = Buffer.alloc(HEADER_LENGTH) + header[0] = MessageType.Regular + header.writeUInt32BE(1, 1) + header.writeUInt32BE(oversizedLength, 9) + + decoder.feed(Buffer.concat([header, Buffer.alloc(oversizedLength), frame(2, 'after')])) + + expect(errors).toHaveLength(1) + expect(seen).toHaveLength(0) + expect(scheduler.pending()).toBe(1) + scheduler.runNext() + expect(seen).toEqual([2]) + }) + + it('drain and reset cancel continuation ownership without replaying residue', () => { + const scheduler = createScheduler() + const seen: number[] = [] + const resume = vi.fn() + const cancel = vi.fn(scheduler.cancel) + const decoder = new FrameDecoder((decoded) => seen.push(decoded.id), undefined, { + maxFramesPerTurn: 1, + schedule: scheduler.schedule, + cancelScheduled: cancel, + resume + }) + const second = frame(2, 'residue') + + decoder.feed(Buffer.concat([frame(1), second])) + const residue = decoder.drain() + + expect(seen).toEqual([1]) + expect(residue.equals(second)).toBe(true) + expect(cancel).toHaveBeenCalledTimes(1) + expect(resume).toHaveBeenCalledTimes(1) + expect(scheduler.pending()).toBe(0) + + decoder.feed(Buffer.concat([frame(3), frame(4)])) + decoder.reset() + expect(cancel).toHaveBeenCalledTimes(2) + expect(scheduler.pending()).toBe(0) + + decoder.feed(frame(5)) + expect(seen).toEqual([1, 3, 5]) + }) +}) diff --git a/src/relay/protocol.ts b/src/relay/protocol.ts index 617ab84e679..356655d64e4 100644 --- a/src/relay/protocol.ts +++ b/src/relay/protocol.ts @@ -1,12 +1,28 @@ // Self-contained relay protocol — mirrors src/main/ssh/relay-protocol.ts // but has no Electron dependencies. Deployed standalone to remote hosts. +import { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + FRAME_DECODER_MAX_FRAMES_PER_TURN, + FRAME_DECODER_MAX_BYTES_PER_TURN, + FRAME_DECODER_MAX_TURN_MS +} from './relay-frame-decoder' + +export { + FrameDecoder, + HEADER_LENGTH, + MAX_MESSAGE_SIZE, + FRAME_DECODER_MAX_FRAMES_PER_TURN, + FRAME_DECODER_MAX_BYTES_PER_TURN, + FRAME_DECODER_MAX_TURN_MS +} +export type { DecodedFrame, FrameDecoderOptions } from './relay-frame-decoder' + export const RELAY_VERSION = '0.1.0' export const RELAY_SENTINEL = `ORCA-RELAY v${RELAY_VERSION} READY\n` -export const HEADER_LENGTH = 13 -export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 - export const MessageType = { Regular: 1, Handshake: 2, @@ -107,13 +123,6 @@ export type JsonRpcNotification = { export type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse | JsonRpcNotification -export type DecodedFrame = { - type: number - id: number - ack: number - payload: Buffer -} - export function encodeFrame( type: number, id: number, @@ -140,148 +149,6 @@ export function encodeKeepAliveFrame(id: number, ack: number): Buffer { return encodeFrame(MessageType.KeepAlive, id, ack, Buffer.alloc(0)) } -export class FrameDecoder { - // Why: feed() sits on the hot receive path. Rebuilding one contiguous - // buffer per feed (Buffer.concat) re-copies every already-buffered byte for - // each incoming chunk — O(n²) per large frame. A chunk list assembles each - // frame exactly once instead. - private chunks: Buffer[] = [] - private bufferedLength = 0 - private onFrame: (frame: DecodedFrame) => void - private onError: ((err: Error) => void) | null - - constructor(onFrame: (frame: DecodedFrame) => void, onError?: (err: Error) => void) { - this.onFrame = onFrame - this.onError = onError ?? null - } - - feed(chunk: Buffer | Uint8Array): void { - const buf = Buffer.isBuffer(chunk) - ? chunk - : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) - if (buf.length > 0) { - this.chunks.push(buf) - this.bufferedLength += buf.length - } - - while (this.bufferedLength >= HEADER_LENGTH) { - const header = this.peekBytes(HEADER_LENGTH) - const length = header.readUInt32BE(9) - const totalLength = HEADER_LENGTH + length - - if (this.bufferedLength < totalLength) { - // Not fully received yet (also holds oversized frames until they can - // be skipped whole, keeping the decoder synchronized). - break - } - - if (length > MAX_MESSAGE_SIZE) { - // Why: Throwing here would leave the buffer in a partially consumed - // state — subsequent feed() calls would try to parse the leftover - // payload bytes as a new header, corrupting every future frame. - // Instead we skip the entire oversized frame so the decoder stays - // synchronized with the stream. - this.discardBytes(totalLength) - const err = new Error(`Frame payload too large: ${length} bytes — discarded`) - if (this.onError) { - this.onError(err) - } else { - process.stderr.write(`[relay] ${err.message}\n`) - } - continue - } - - const framed = this.takeBytes(totalLength) - const frame: DecodedFrame = { - type: framed[0], - id: framed.readUInt32BE(1), - ack: framed.readUInt32BE(5), - payload: framed.subarray(HEADER_LENGTH, totalLength) - } - this.onFrame(frame) - } - } - - reset(): void { - this.chunks = [] - this.bufferedLength = 0 - } - - // Why: at the handshake → dispatcher transition, the next consumer must - // pick up any bytes that arrived in the same TCP chunk as the handshake - // frame. This returns and clears the decoder's internal residue so the - // caller can hand it to the dispatcher (or stdout pipe) without loss. - drain(): Buffer { - const out = - this.chunks.length === 1 ? this.chunks[0] : Buffer.concat(this.chunks, this.bufferedLength) - this.reset() - return out - } - - /** View of the first `count` buffered bytes without consuming them. */ - private peekBytes(count: number): Buffer { - const first = this.chunks[0] - if (first.length >= count) { - return first - } - const out = Buffer.allocUnsafe(count) - let copied = 0 - for (const part of this.chunks) { - copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) - if (copied >= count) { - break - } - } - return out - } - - /** Consume and return the first `count` buffered bytes (single copy). */ - private takeBytes(count: number): Buffer { - const first = this.chunks[0] - if (first.length === count) { - this.chunks.shift() - this.bufferedLength -= count - return first - } - if (first.length > count) { - this.chunks[0] = first.subarray(count) - this.bufferedLength -= count - return first.subarray(0, count) - } - const out = Buffer.allocUnsafe(count) - let copied = 0 - while (copied < count) { - const part = this.chunks[0] - const take = Math.min(part.length, count - copied) - part.copy(out, copied, 0, take) - copied += take - if (take === part.length) { - this.chunks.shift() - } else { - this.chunks[0] = part.subarray(take) - } - } - this.bufferedLength -= count - return out - } - - /** Consume the first `count` buffered bytes without assembling them. */ - private discardBytes(count: number): void { - let remaining = count - while (remaining > 0) { - const part = this.chunks[0] - if (part.length <= remaining) { - this.chunks.shift() - remaining -= part.length - } else { - this.chunks[0] = part.subarray(remaining) - remaining = 0 - } - } - this.bufferedLength -= count - } -} - export function parseJsonRpcMessage(payload: Buffer): JsonRpcMessage { const text = payload.toString('utf-8') const msg = JSON.parse(text) as JsonRpcMessage diff --git a/src/relay/relay-frame-decoder.ts b/src/relay/relay-frame-decoder.ts new file mode 100644 index 00000000000..f98b3647646 --- /dev/null +++ b/src/relay/relay-frame-decoder.ts @@ -0,0 +1,317 @@ +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 type DecodedFrame = { + type: number + id: number + ack: number + payload: Buffer +} + +export type FrameDecoderOptions = { + maxFramesPerTurn?: number + maxBytesPerTurn?: number + maxTurnMs?: number + now?: () => number + schedule?: (callback: () => void) => unknown + cancelScheduled?: (handle: unknown) => void + pause?: () => void + resume?: () => void +} + +export class FrameDecoder { + private chunks: Buffer[] = [] + private bufferedLength = 0 + private oversizedPayloadBytesRemaining = 0 + private onFrame: (frame: DecodedFrame) => void + private onError: (err: Error) => void + private maxFramesPerTurn: number + private maxBytesPerTurn: number + private maxTurnMs: number + private now: () => number + private schedule: (callback: () => void) => unknown + private cancelScheduled: (handle: unknown) => void + private pause: (() => void) | null + private resume: (() => void) | null + private continuationHandle: unknown + private continuationHandleAssigned = false + private continuationScheduled = false + private paused = false + private draining = false + private generation = 0 + + constructor( + onFrame: (frame: DecodedFrame) => void, + onError?: (err: Error) => void, + options: FrameDecoderOptions = {} + ) { + this.onFrame = onFrame + this.onError = onError ?? ((error) => process.stderr.write(`[relay] ${error.message}\n`)) + this.maxFramesPerTurn = positiveLimit( + options.maxFramesPerTurn, + FRAME_DECODER_MAX_FRAMES_PER_TURN + ) + this.maxBytesPerTurn = positiveLimit(options.maxBytesPerTurn, FRAME_DECODER_MAX_BYTES_PER_TURN) + this.maxTurnMs = positiveLimit(options.maxTurnMs, FRAME_DECODER_MAX_TURN_MS) + this.now = options.now ?? Date.now + this.schedule = options.schedule ?? ((callback) => setImmediate(callback)) + this.cancelScheduled = + options.cancelScheduled ?? ((handle) => clearImmediate(handle as NodeJS.Immediate)) + this.pause = options.pause ?? null + this.resume = options.resume ?? null + } + + feed(chunk: Buffer | Uint8Array): void { + const buf = Buffer.isBuffer(chunk) + ? chunk + : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + if (buf.length > 0) { + this.chunks.push(buf) + this.bufferedLength += buf.length + } + if (!this.draining && !this.continuationScheduled) { + this.drainTurn() + } + } + + reset(): void { + this.generation += 1 + this.cancelContinuation() + this.chunks = [] + this.bufferedLength = 0 + this.oversizedPayloadBytesRemaining = 0 + this.releasePause() + } + + drain(): Buffer { + const out = + this.chunks.length === 1 ? this.chunks[0] : Buffer.concat(this.chunks, this.bufferedLength) + this.reset() + return out + } + + private drainTurn(): void { + if (this.draining) { + return + } + this.draining = true + const generation = this.generation + const startedAt = this.now() + let frames = 0 + let bytes = 0 + + try { + while (generation === this.generation) { + if ( + frames >= this.maxFramesPerTurn || + bytes >= this.maxBytesPerTurn || + (frames > 0 && this.now() - startedAt >= this.maxTurnMs) + ) { + break + } + const discarded = this.discardOversizedPayload(bytes) + if (discarded > 0) { + bytes += discarded + continue + } + if (this.bufferedLength < HEADER_LENGTH) { + break + } + const header = this.peekBytes(HEADER_LENGTH) + const length = header.readUInt32BE(9) + if (length > MAX_MESSAGE_SIZE) { + this.discardBytes(HEADER_LENGTH) + this.oversizedPayloadBytesRemaining = length + bytes += HEADER_LENGTH + const error = new Error(`Frame payload too large: ${length} bytes — discarded`) + this.onError(error) + continue + } + const totalLength = HEADER_LENGTH + length + if (this.bufferedLength < totalLength) { + break + } + if (frames > 0 && bytes + totalLength > this.maxBytesPerTurn) { + break + } + const framed = this.takeBytes(totalLength) + frames += 1 + bytes += totalLength + this.onFrame({ + type: framed[0], + id: framed.readUInt32BE(1), + ack: framed.readUInt32BE(5), + payload: framed.subarray(HEADER_LENGTH, totalLength) + }) + } + } finally { + this.draining = false + } + + if (generation !== this.generation) { + return + } + if (this.hasRunnableWork()) { + this.scheduleContinuation() + } else { + this.releasePause() + } + } + + private discardOversizedPayload(bytes: number): number { + if (this.oversizedPayloadBytesRemaining === 0 || this.bufferedLength === 0) { + return 0 + } + const discarded = Math.min( + this.oversizedPayloadBytesRemaining, + this.bufferedLength, + Math.max(1, this.maxBytesPerTurn - bytes) + ) + this.discardBytes(discarded) + this.oversizedPayloadBytesRemaining -= discarded + return discarded + } + + private hasRunnableWork(): boolean { + if (this.oversizedPayloadBytesRemaining > 0) { + return this.bufferedLength > 0 + } + if (this.bufferedLength < HEADER_LENGTH) { + return false + } + const length = this.peekBytes(HEADER_LENGTH).readUInt32BE(9) + return length > MAX_MESSAGE_SIZE || this.bufferedLength >= HEADER_LENGTH + length + } + + private scheduleContinuation(): void { + if (this.continuationScheduled) { + return + } + const generation = this.generation + this.continuationScheduled = true + try { + this.acquirePause() + } catch (error) { + this.continuationScheduled = false + throw error + } + if (generation !== this.generation) { + this.continuationScheduled = false + return + } + try { + this.continuationHandle = this.schedule(() => { + if (!this.continuationScheduled || generation !== this.generation) { + return + } + this.continuationScheduled = false + this.continuationHandleAssigned = false + this.continuationHandle = undefined + this.drainTurn() + }) + this.continuationHandleAssigned = true + } catch (error) { + this.continuationScheduled = false + this.continuationHandle = undefined + this.releasePause() + throw error + } + } + + private cancelContinuation(): void { + if (!this.continuationScheduled) { + return + } + this.continuationScheduled = false + if (this.continuationHandleAssigned) { + this.cancelScheduled(this.continuationHandle) + } + this.continuationHandleAssigned = false + this.continuationHandle = undefined + } + + private acquirePause(): void { + if (!this.paused) { + this.paused = true + try { + this.pause?.() + } catch (error) { + this.paused = false + throw error + } + } + } + + private releasePause(): void { + if (this.paused) { + this.paused = false + this.resume?.() + } + } + + private peekBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length >= count) { + return first + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + for (const part of this.chunks) { + copied += part.copy(out, copied, 0, Math.min(part.length, count - copied)) + if (copied >= count) { + break + } + } + return out + } + + private takeBytes(count: number): Buffer { + const first = this.chunks[0] + if (first.length === count) { + this.chunks.shift() + this.bufferedLength -= count + return first + } + if (first.length > count) { + this.chunks[0] = first.subarray(count) + this.bufferedLength -= count + return first.subarray(0, count) + } + const out = Buffer.allocUnsafe(count) + let copied = 0 + while (copied < count) { + const part = this.chunks[0] + const take = Math.min(part.length, count - copied) + part.copy(out, copied, 0, take) + copied += take + if (take === part.length) { + this.chunks.shift() + } else { + this.chunks[0] = part.subarray(take) + } + } + this.bufferedLength -= count + return out + } + + private discardBytes(count: number): void { + let remaining = count + while (remaining > 0) { + const part = this.chunks[0] + if (part.length <= remaining) { + this.chunks.shift() + remaining -= part.length + } else { + this.chunks[0] = part.subarray(remaining) + remaining = 0 + } + } + this.bufferedLength -= count + } +} + +const positiveLimit = (value: number | undefined, fallback: number): number => + value !== undefined && Number.isFinite(value) && value > 0 ? value : fallback