diff --git a/src/main/agent-hooks/wsl-hook-fs-adapter.ts b/src/main/agent-hooks/wsl-hook-fs-adapter.ts index d47dd400199..db6553b09de 100644 --- a/src/main/agent-hooks/wsl-hook-fs-adapter.ts +++ b/src/main/agent-hooks/wsl-hook-fs-adapter.ts @@ -5,10 +5,20 @@ // `installer-utils-remote.ts` touches are implemented. import type { SFTPWrapper } from 'ssh2' +import { assertFilesystemDirectoryWithinLimit } from '../../shared/filesystem-directory-listing-limit' +import { NodeFileReadTooLargeError } from '../../shared/memory-safety/node-bounded-file-reader' import type { installRemoteManagedAgentHooks } from './remote-managed-hook-installers' import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' import { wslCodexRuntimeHomeForGuestHome } from '../pty/codex-home-wsl-env' -import { WSL_HOOK_FS_METHODS, type WslFsResult } from '../../shared/wsl-hook-relay-contract' +import { + WSL_HOOK_FS_METHODS, + WSL_HOOK_FS_MAX_DIRECTORY_ENTRIES, + WSL_HOOK_FS_MAX_DIRECTORY_RETAINED_BYTES, + WSL_HOOK_FS_MAX_READ_BYTES, + type WslFsFailure, + type WslFsResult, + type WslHookFsDirectoryLimits +} from '../../shared/wsl-hook-relay-contract' /** Run the shared remote hook installers against a WSL guest over the relay's * fs bridge. Codex is the one agent whose home Orca redirects for WSL @@ -44,7 +54,23 @@ const ERRNO_TO_SFTP_CODE: Record = { EEXIST: 4 } -function toSftpError(failure: { errno?: string; message?: string }): Error { +function toSftpError(failure: { + errno?: string + message?: string + fileCapacity?: { observedBytes: number; maxBytes: number } +}): Error { + if ( + failure.errno === 'EFBIG' && + Number.isSafeInteger(failure.fileCapacity?.observedBytes) && + Number.isSafeInteger(failure.fileCapacity?.maxBytes) && + failure.fileCapacity!.observedBytes >= 0 && + failure.fileCapacity!.maxBytes >= 0 + ) { + return new NodeFileReadTooLargeError( + failure.fileCapacity!.observedBytes, + failure.fileCapacity!.maxBytes + ) + } const err = new Error(failure.message ?? 'wsl fs bridge failure') as Error & { code?: number } err.code = ERRNO_TO_SFTP_CODE[failure.errno ?? ''] ?? 5 return err @@ -62,7 +88,7 @@ export function createWslHookSftpAdapter(mux: SshChannelMultiplexer): SFTPWrappe .then((raw) => { const result = raw as WslFsResult if (!result || typeof result !== 'object' || result.ok !== true) { - callback(toSftpError((result ?? {}) as { errno?: string; message?: string })) + callback(toSftpError((result ?? {}) as WslFsFailure)) return } callback(null, pick(result)) @@ -81,7 +107,15 @@ export function createWslHookSftpAdapter(mux: SshChannelMultiplexer): SFTPWrappe readFile(path: string, _encoding: unknown, callback: SftpCallback): void { call<{ content: string }, string>( WSL_HOOK_FS_METHODS.readFile, - { path }, + { path, maxBytes: WSL_HOOK_FS_MAX_READ_BYTES }, + callback, + (r) => r.content + ) + }, + orcaReadFileWithinLimit(path: string, maxBytes: number, callback: SftpCallback): void { + call<{ content: string }, string>( + WSL_HOOK_FS_METHODS.readFile, + { path, maxBytes }, callback, (r) => r.content ) @@ -122,9 +156,45 @@ export function createWslHookSftpAdapter(mux: SshChannelMultiplexer): SFTPWrappe readdir(path: string, callback: SftpCallback<{ filename: string }[]>): void { call<{ entries: { filename: string }[] }, { filename: string }[]>( WSL_HOOK_FS_METHODS.readdir, - { path }, + { + path, + maxEntries: WSL_HOOK_FS_MAX_DIRECTORY_ENTRIES, + maxRetainedBytes: WSL_HOOK_FS_MAX_DIRECTORY_RETAINED_BYTES + }, callback, - (r) => r.entries + (r) => { + const limits = { + maxEntries: WSL_HOOK_FS_MAX_DIRECTORY_ENTRIES, + maxRetainedBytes: WSL_HOOK_FS_MAX_DIRECTORY_RETAINED_BYTES + } + assertFilesystemDirectoryWithinLimit( + r.entries.map((entry) => ({ name: entry.filename })), + limits + ) + return r.entries + } + ) + }, + orcaReaddirWithinLimit( + path: string, + limits: WslHookFsDirectoryLimits, + callback: SftpCallback<{ filename: string }[]> + ): void { + call<{ entries: { filename: string }[] }, { filename: string }[]>( + WSL_HOOK_FS_METHODS.readdir, + { + path, + maxEntries: limits.maxEntries, + maxRetainedBytes: limits.maxRetainedBytes + }, + callback, + (r) => { + assertFilesystemDirectoryWithinLimit( + r.entries.map((entry) => ({ name: entry.filename })), + limits + ) + return r.entries + } ) }, mkdir(path: string, callback: SftpCallback): void { diff --git a/src/main/providers/ssh-filesystem-provider-stream.test.ts b/src/main/providers/ssh-filesystem-provider-stream.test.ts index b66b5633503..2563fbeaf7b 100644 --- a/src/main/providers/ssh-filesystem-provider-stream.test.ts +++ b/src/main/providers/ssh-filesystem-provider-stream.test.ts @@ -105,6 +105,54 @@ describe('SshFilesystemProvider readFile streaming', () => { expect(result).toEqual(legacyResult) }) + it('propagates validated raster dimensions from stream metadata', async () => { + const content = Buffer.from('png-header') + mux.request.mockImplementation(async () => { + setImmediate(() => { + mux._emitMethod('fs.streamChunk', { + streamId: 3, + seq: 0, + data: content.toString('base64') + }) + mux._emitMethod('fs.streamEnd', { streamId: 3 }) + }) + return { + streamId: 3, + totalSize: content.length, + isBinary: true, + isImage: true, + mimeType: 'image/png', + imageDimensions: { width: 640, height: 480 }, + resultEncoding: 'base64' + } + }) + + await expect(provider.readFile('/home/image.png')).resolves.toEqual({ + content: content.toString('base64'), + isBinary: true, + isImage: true, + mimeType: 'image/png', + imageDimensions: { width: 640, height: 480 } + }) + }) + + it('rejects unsafe raster dimensions in stream metadata', async () => { + mux.request.mockResolvedValue({ + streamId: 4, + totalSize: 1, + isBinary: true, + isImage: true, + mimeType: 'image/png', + imageDimensions: { width: 32_769, height: 1 }, + resultEncoding: 'base64' + }) + + await expect(provider.readFile('/home/bomb.png')).rejects.toThrow( + 'Malformed file stream metadata' + ) + expect(mux.notify).toHaveBeenCalledWith('fs.cancelStream', { streamId: 4 }) + }) + it('rejects when chunk arrives out of order', async () => { const totalSize = 256 * 1024 * 2 mux.request.mockImplementation(async () => { @@ -178,7 +226,7 @@ describe('SshFilesystemProvider readFile streaming', () => { resultEncoding: 'base64' } }) - await expect(provider.readFile('/home/x.bin')).rejects.toThrow(/count mismatch/i) + await expect(provider.readFile('/home/x.bin')).rejects.toThrow(/incomplete/i) }) it('rejects a short final chunk instead of zero-filling the buffer', async () => { diff --git a/src/main/providers/types.ts b/src/main/providers/types.ts index 7ebe3de9f17..6fe63f70cc1 100644 --- a/src/main/providers/types.ts +++ b/src/main/providers/types.ts @@ -225,6 +225,7 @@ export type FileReadResult = { isBinary: boolean isImage?: boolean mimeType?: string + imageDimensions?: { width: number; height: number } } export type IFilesystemProvider = { diff --git a/src/main/ssh/relay-protocol.test.ts b/src/main/ssh/relay-protocol.test.ts index 459eaef88df..3e2ada59323 100644 --- a/src/main/ssh/relay-protocol.test.ts +++ b/src/main/ssh/relay-protocol.test.ts @@ -1,6 +1,7 @@ import { describe, expect, it, vi } from 'vitest' import { HEADER_LENGTH, + MAX_BUFFERED_FRAME_CHUNKS, MessageType, encodeFrame, encodeJsonRpcFrame, @@ -12,6 +13,7 @@ import { type JsonRpcRequest, type DecodedFrame } from './relay-protocol' +import { RELAY_JSON_MAX_STRUCTURAL_TOKENS } from '../../shared/relay-json-admission' describe('git response stream marker', () => { it('accepts only complete non-negative integer metadata', () => { @@ -26,6 +28,15 @@ describe('git response stream marker', () => { __orcaGitResponseStream: { streamId: -1, totalBytes: 1024, chunkCount: 2 } }) ).toBe(false) + expect( + isGitResponseStreamMarker({ + __orcaGitResponseStream: { + streamId: 1, + totalBytes: Number.MAX_SAFE_INTEGER + 1, + chunkCount: 2 + } + }) + ).toBe(false) }) }) @@ -79,11 +90,37 @@ describe('frame encoding', () => { method: 'x', params: { data: 'a'.repeat(17 * 1024 * 1024) } } - expect(() => encodeJsonRpcFrame(bigPayload, 1, 0)).toThrow('Message too large') + const fromSpy = vi.spyOn(Buffer, 'from') + try { + expect(() => encodeJsonRpcFrame(bigPayload, 1, 0)).toThrow('Message too large') + expect(fromSpy).not.toHaveBeenCalled() + } finally { + fromSpy.mockRestore() + } }) }) describe('FrameDecoder', () => { + it('bounds retained chunk wrappers for a byte-fragmented frame', () => { + const frames: DecodedFrame[] = [] + const decoder = new FrameDecoder((frame) => frames.push(frame)) + const payloadLength = MAX_BUFFERED_FRAME_CHUNKS + 512 + const header = Buffer.alloc(HEADER_LENGTH) + header[0] = MessageType.Regular + header.writeUInt32BE(payloadLength, 9) + decoder.feed(header) + + const byte = Buffer.from('x') + const state = decoder as unknown as { chunks: Buffer[] } + for (let index = 0; index < payloadLength; index += 1) { + decoder.feed(byte) + expect(state.chunks.length).toBeLessThanOrEqual(MAX_BUFFERED_FRAME_CHUNKS) + } + + expect(frames).toHaveLength(1) + expect(frames[0].payload.equals(Buffer.alloc(payloadLength, 0x78))).toBe(true) + }) + it('decodes a complete frame', () => { const frames: DecodedFrame[] = [] const decoder = new FrameDecoder((f) => frames.push(f)) @@ -163,6 +200,33 @@ describe('FrameDecoder', () => { expect(errors[0].message).toContain('discarded') }) + it('rejects an oversized header immediately without retaining its payload', () => { + const errors: Error[] = [] + const decoder = new FrameDecoder( + () => {}, + (err) => errors.push(err) + ) + const header = Buffer.alloc(HEADER_LENGTH) + header[0] = MessageType.Regular + header.writeUInt32BE(0xffffffff, 9) + + decoder.feed(header) + const state = decoder as unknown as { + bufferedLength: number + oversizedPayloadBytesRemaining: number + } + expect(errors).toHaveLength(1) + expect(state.bufferedLength).toBe(0) + expect(state.oversizedPayloadBytesRemaining).toBe(0xffffffff) + + const payloadChunk = Buffer.alloc(64 * 1024) + for (let index = 0; index < 128; index += 1) { + decoder.feed(payloadChunk) + expect(state.bufferedLength).toBe(0) + } + expect(errors).toHaveLength(1) + }) + it('reset clears internal buffer', () => { const frames: DecodedFrame[] = [] const decoder = new FrameDecoder((f) => frames.push(f)) @@ -224,26 +288,31 @@ describe('FrameDecoder', () => { expect(frames[0].payload.toString()).toBe('after') }) - it('never rebuilds the buffered stream per feed while assembling a large frame', () => { - // Regression: feed() used Buffer.concat([buffered, chunk]) per data event, - // re-copying the whole backlog for every TCP chunk — O(n²) memcpy on the - // Electron main thread while fs.streamChunk frames arrive (SSH typing lag). + it('retains ordinary transport chunks without copying them while a frame is incomplete', () => { const frames: DecodedFrame[] = [] const decoder = new FrameDecoder((f) => frames.push(f)) - const frame = encodeFrame(MessageType.Regular, 1, 0, Buffer.alloc(512 * 1024, 0x61)) + const payloadLength = 1024 * 1024 + const header = Buffer.alloc(HEADER_LENGTH) + header[0] = MessageType.Regular + header.writeUInt32BE(1, 1) + header.writeUInt32BE(payloadLength, 9) + decoder.feed(header) + const ordinaryChunks = Array.from({ length: 16 }, () => Buffer.alloc(32 * 1024, 0x61)) - const concatSpy = vi.spyOn(Buffer, 'concat') - try { - for (let i = 0; i < frame.length; i += 32 * 1024) { - decoder.feed(frame.subarray(i, i + 32 * 1024)) - } - } finally { - concatSpy.mockRestore() + for (const chunk of ordinaryChunks) { + decoder.feed(chunk) } + const state = decoder as unknown as { chunks: Buffer[] } + expect(state.chunks).toHaveLength(ordinaryChunks.length + 1) + for (const [index, chunk] of ordinaryChunks.entries()) { + expect(state.chunks[index + 1]).toBe(chunk) + } + + decoder.feed(Buffer.alloc(payloadLength / 2, 0x61)) + expect(frames).toHaveLength(1) - expect(frames[0].payload.length).toBe(512 * 1024) - expect(concatSpy).not.toHaveBeenCalled() + expect(frames[0].payload.length).toBe(payloadLength) }) }) @@ -270,6 +339,21 @@ describe('parseJsonRpcMessage', () => { const payload = Buffer.from('not json') expect(() => parseJsonRpcMessage(payload)).toThrow() }) + + it('rejects structurally amplified JSON before parsing', () => { + const payload = Buffer.from( + `{"jsonrpc":"2.0","id":1,"method":"x","params":{"values":[${'0,'.repeat( + RELAY_JSON_MAX_STRUCTURAL_TOKENS + )}0]}}` + ) + const parseSpy = vi.spyOn(JSON, 'parse') + try { + expect(() => parseJsonRpcMessage(payload)).toThrow(/JSON structure exceeds/) + expect(parseSpy).not.toHaveBeenCalled() + } finally { + parseSpy.mockRestore() + } + }) }) describe('parseUnameToRelayPlatform', () => { diff --git a/src/main/ssh/relay-protocol.ts b/src/main/ssh/relay-protocol.ts index c05b6c33bfe..b8f86cee34e 100644 --- a/src/main/ssh/relay-protocol.ts +++ b/src/main/ssh/relay-protocol.ts @@ -3,6 +3,11 @@ // See design-ssh-support.md § JSON-RPC Protocol Specification. import { DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS } from '../../shared/ssh-types' +import { + JsonStringifyByteLimitError, + stringifyJsonWithinByteLimit +} from '../../shared/memory-safety/node-bounded-json-stringify' +import { parseRelayJsonText } from '../../shared/relay-json-admission' export const RELAY_VERSION = '0.1.0' export const RELAY_SENTINEL = `ORCA-RELAY v${RELAY_VERSION} READY\n` @@ -13,6 +18,8 @@ export const RELAY_REMOTE_DIR = '.orca-remote' export const HEADER_LENGTH = 13 export const MAX_MESSAGE_SIZE = 16 * 1024 * 1024 // 16 MB +export const MAX_BUFFERED_FRAME_CHUNKS = 4_096 +const MAX_COALESCED_TINY_FRAME_CHUNK_BYTES = 8 * 1024 /** Message type byte. */ export const MessageType = { @@ -68,6 +75,11 @@ export const GIT_RESPONSE_STREAM_THRESHOLD = 256 * 1024 * so it stays cross-version safe. */ export const GIT_RESPONSE_CHUNK_SIZE = 128 * 1024 +/** Bounds retained serialized Git results while preserving responses far above + * the per-file render limit and git.exec's output ceiling. */ +export const MAX_GIT_RESPONSE_STREAM_BYTES = 64 * 1024 * 1024 +export const MAX_GIT_RESPONSE_STREAM_CHUNKS = 4_096 + /** Sentinel the relay returns as the RPC result when the real payload streams * as git.responseChunk frames. Absent from old relays, so a new client falls * back to the plain result they return. */ @@ -85,11 +97,11 @@ export function isGitResponseStreamMarker(value: unknown): value is GitResponseS } const fields = marker as Record return ( - Number.isInteger(fields.streamId) && + Number.isSafeInteger(fields.streamId) && (fields.streamId as number) > 0 && - Number.isInteger(fields.totalBytes) && + Number.isSafeInteger(fields.totalBytes) && (fields.totalBytes as number) >= 0 && - Number.isInteger(fields.chunkCount) && + Number.isSafeInteger(fields.chunkCount) && (fields.chunkCount as number) >= 0 ) } @@ -144,10 +156,16 @@ export function encodeFrame( } export function encodeJsonRpcFrame(msg: JsonRpcMessage, id: number, ack: number): Buffer { - const payload = Buffer.from(JSON.stringify(msg), 'utf-8') - if (payload.length > MAX_MESSAGE_SIZE) { - throw new Error(`Message too large: ${payload.length} bytes (max ${MAX_MESSAGE_SIZE})`) + let serialized: string + try { + serialized = stringifyJsonWithinByteLimit(msg, MAX_MESSAGE_SIZE).serialized + } catch (error) { + if (error instanceof JsonStringifyByteLimitError) { + throw new Error(`Message too large: more than ${MAX_MESSAGE_SIZE} bytes`) + } + throw error } + const payload = Buffer.from(serialized, 'utf-8') return encodeFrame(MessageType.Regular, id, ack, payload) } @@ -173,6 +191,7 @@ export class FrameDecoder { // list assembles each frame exactly once instead. private chunks: Buffer[] = [] private bufferedLength = 0 + private oversizedPayloadBytesRemaining = 0 private onFrame: (frame: DecodedFrame) => void private onError: ((err: Error) => void) | null @@ -182,12 +201,18 @@ export class FrameDecoder { } feed(chunk: Buffer | Uint8Array): void { - const buf = Buffer.isBuffer(chunk) + let buf = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk.buffer, chunk.byteOffset, chunk.byteLength) + if (this.oversizedPayloadBytesRemaining > 0 && buf.length > 0) { + const discarded = Math.min(this.oversizedPayloadBytesRemaining, buf.length) + this.oversizedPayloadBytesRemaining -= discarded + buf = buf.subarray(discarded) + } if (buf.length > 0) { this.chunks.push(buf) this.bufferedLength += buf.length + this.coalesceBufferedChunks() } while (this.bufferedLength >= HEADER_LENGTH) { @@ -195,25 +220,26 @@ export class FrameDecoder { 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) + // Why: retain only the header, then stream-discard the advertised + // payload so a 13-byte malicious header cannot pin up to 4 GiB. + this.discardBytes(HEADER_LENGTH) + this.oversizedPayloadBytesRemaining = length + this.discardBufferedOversizedPayload() const err = new Error(`Frame payload too large: ${length} bytes — discarded`) if (this.onError) { this.onError(err) } + if (this.oversizedPayloadBytesRemaining > 0) { + break + } continue } + if (this.bufferedLength < totalLength) { + break + } + const framed = this.takeBytes(totalLength) const frame: DecodedFrame = { type: framed[0], @@ -228,6 +254,7 @@ export class FrameDecoder { reset(): void { this.chunks = [] this.bufferedLength = 0 + this.oversizedPayloadBytesRemaining = 0 } /** View of the first `count` buffered bytes without consuming them. */ @@ -292,14 +319,40 @@ export class FrameDecoder { } this.bufferedLength -= count } + + private discardBufferedOversizedPayload(): void { + const discarded = Math.min(this.oversizedPayloadBytesRemaining, this.bufferedLength) + if (discarded === 0) { + return + } + this.discardBytes(discarded) + this.oversizedPayloadBytesRemaining -= discarded + } + + private coalesceBufferedChunks(): void { + // Why: merge adversarial tiny fragments without copying ordinary transport chunks. + while (this.chunks.length >= 2) { + const right = this.chunks.at(-1)! + const left = this.chunks.at(-2)! + if ( + left.length > right.length || + left.length + right.length > MAX_COALESCED_TINY_FRAME_CHUNK_BYTES + ) { + break + } + this.chunks.splice(-2, 2, Buffer.concat([left, right], left.length + right.length)) + } + if (this.chunks.length > MAX_BUFFERED_FRAME_CHUNKS) { + this.chunks = [Buffer.concat(this.chunks, this.bufferedLength)] + } + } } /** * 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 + const msg = parseRelayJsonText(payload.toString('utf-8')) if (msg.jsonrpc !== '2.0') { throw new Error(`Invalid JSON-RPC version: ${(msg as Record).jsonrpc}`) } diff --git a/src/main/ssh/ssh-channel-multiplexer.ts b/src/main/ssh/ssh-channel-multiplexer.ts index 9886b724755..1550eb5e210 100644 --- a/src/main/ssh/ssh-channel-multiplexer.ts +++ b/src/main/ssh/ssh-channel-multiplexer.ts @@ -305,14 +305,14 @@ export class SshChannelMultiplexer { this.decoder.reset() this.transport.close?.() - for (const handler of this.disposeHandlers) { + // Why: stream handlers unsubscribe while failing; detach first so one cleanup cannot skip the next. + for (const handler of this.disposeHandlers.splice(0)) { try { handler(reason) } catch { // Don't let a handler error prevent other handlers from running } } - this.disposeHandlers.length = 0 } isDisposed(): boolean { diff --git a/src/main/ssh/ssh-file-stream-assembler.ts b/src/main/ssh/ssh-file-stream-assembler.ts new file mode 100644 index 00000000000..4a40c6adf21 --- /dev/null +++ b/src/main/ssh/ssh-file-stream-assembler.ts @@ -0,0 +1,196 @@ +import type { FileReadResult } from '../providers/types' +import { RelayErrorCode, STREAM_CHUNK_SIZE } from './relay-protocol' +import { + defaultSshStreamAssemblyBudget, + type SshStreamAssemblyBudget, + base64EncodedLength +} from './ssh-stream-reader-memory' +import type { RasterImageDimensions } from '../../shared/raster-image-dimensions' +import { isRasterImagePreviewDimensions } from '../../shared/raster-image-preview-limits' + +const MAX_PREVIEWABLE_BINARY_SIZE = 50 * 1024 * 1024 +const MAX_TEXT_FILE_SIZE = 10 * 1024 * 1024 + +type FileStreamMetadata = { + streamId?: number + totalSize: number + isBinary: boolean + isImage?: boolean + mimeType?: string + imageDimensions?: RasterImageDimensions + resultEncoding?: 'base64' | 'utf-8' + empty?: boolean +} + +export class StreamProtocolError extends Error { + readonly code = RelayErrorCode.StreamProtocolError + constructor(message: string) { + super(message) + } +} + +export type FileStreamSetup = + | { kind: 'empty'; result: FileReadResult } + | { kind: 'stream'; streamId: number; assembler: FileStreamAssembler } + +export function createFileStreamSetup( + rawMetadata: unknown, + assemblyBudget: SshStreamAssemblyBudget = defaultSshStreamAssemblyBudget +): FileStreamSetup { + const metadata = rawMetadata as FileStreamMetadata + if ( + typeof metadata !== 'object' || + metadata === null || + typeof metadata.isBinary !== 'boolean' || + !Number.isSafeInteger(metadata.totalSize) || + (metadata.empty !== undefined && typeof metadata.empty !== 'boolean') || + (metadata.isImage !== undefined && typeof metadata.isImage !== 'boolean') || + (metadata.mimeType !== undefined && typeof metadata.mimeType !== 'string') || + (metadata.imageDimensions !== undefined && + !isRasterImagePreviewDimensions(metadata.imageDimensions)) || + (metadata.resultEncoding !== undefined && + metadata.resultEncoding !== 'base64' && + metadata.resultEncoding !== 'utf-8') + ) { + throw new StreamProtocolError('Malformed file stream metadata') + } + if (metadata.empty) { + return { + kind: 'empty', + result: { + content: '', + isBinary: metadata.isBinary, + ...(metadata.isImage !== undefined ? { isImage: metadata.isImage } : {}), + ...(metadata.mimeType !== undefined ? { mimeType: metadata.mimeType } : {}), + ...(metadata.imageDimensions !== undefined + ? { imageDimensions: metadata.imageDimensions } + : {}) + } + } + } + if ( + typeof metadata.streamId !== 'number' || + !Number.isSafeInteger(metadata.streamId) || + metadata.streamId <= 0 + ) { + throw new StreamProtocolError('Metadata missing streamId for non-empty stream') + } + const cap = metadata.isBinary ? MAX_PREVIEWABLE_BINARY_SIZE : MAX_TEXT_FILE_SIZE + if (metadata.totalSize < 0 || metadata.totalSize > cap) { + throw new StreamProtocolError( + `Reported totalSize ${metadata.totalSize} exceeds client cap ${cap}` + ) + } + return { + kind: 'stream', + streamId: metadata.streamId, + assembler: new FileStreamAssembler(metadata, assemblyBudget) + } +} + +export class FileStreamAssembler { + private buffer: Buffer | null + private expectedSeq = 0 + private bytesReceived = 0 + private readonly totalChunks: number + private releaseBudget: (() => void) | null + + constructor( + private readonly metadata: FileStreamMetadata, + assemblyBudget: SshStreamAssemblyBudget + ) { + this.releaseBudget = assemblyBudget.reserve(metadata.totalSize) + if (!this.releaseBudget) { + throw new StreamProtocolError( + `Active SSH stream assembly would exceed ${assemblyBudget.maxBytes} bytes` + ) + } + try { + this.buffer = Buffer.alloc(metadata.totalSize) + } catch (error) { + this.releaseMemory() + throw new StreamProtocolError( + `Failed to allocate ${metadata.totalSize} bytes: ${(error as Error).message}` + ) + } + this.totalChunks = + metadata.totalSize === 0 ? 0 : Math.ceil(metadata.totalSize / STREAM_CHUNK_SIZE) + } + + acceptChunk(params: Record, streamId: number): number { + const seq = params.seq + const data = params.data + if (typeof seq !== 'number' || !Number.isSafeInteger(seq) || typeof data !== 'string') { + throw new StreamProtocolError(`Malformed chunk for stream ${streamId}`) + } + if (seq !== this.expectedSeq) { + throw new StreamProtocolError( + `Out-of-order chunk for stream ${streamId}: expected ${this.expectedSeq}, got ${seq}` + ) + } + if (seq >= this.totalChunks) { + throw new StreamProtocolError( + `File stream ${streamId} exceeded declared chunk count ${this.totalChunks}` + ) + } + const offset = seq * STREAM_CHUNK_SIZE + const expectedLength = Math.min(STREAM_CHUNK_SIZE, this.metadata.totalSize - offset) + if (expectedLength < 0 || data.length !== base64EncodedLength(expectedLength)) { + throw new StreamProtocolError( + `Encoded chunk length mismatch for stream ${streamId}: seq=${seq}` + ) + } + const decoded = Buffer.from(data, 'base64') + if (decoded.length !== expectedLength) { + throw new StreamProtocolError( + `Chunk length mismatch for stream ${streamId}: seq=${seq} expected=${expectedLength} got=${decoded.length}` + ) + } + if (!this.buffer) { + throw new StreamProtocolError(`Chunk arrived after completion for stream ${streamId}`) + } + decoded.copy(this.buffer, offset) + this.expectedSeq += 1 + this.bytesReceived += decoded.length + return seq + } + + finish(streamId: number): FileReadResult { + if (this.expectedSeq !== this.totalChunks || this.bytesReceived !== this.metadata.totalSize) { + throw new StreamProtocolError( + `File stream ${streamId} incomplete: chunks ${this.expectedSeq}/${this.totalChunks}, bytes ${this.bytesReceived}/${this.metadata.totalSize}` + ) + } + if (!this.buffer) { + throw new StreamProtocolError(`Stream ${streamId} already completed`) + } + const buffer = this.buffer + this.buffer = null + try { + return { + content: + (this.metadata.resultEncoding ?? 'base64') === 'base64' + ? buffer.toString('base64') + : buffer.toString('utf-8'), + isBinary: this.metadata.isBinary, + ...(this.metadata.isImage !== undefined ? { isImage: this.metadata.isImage } : {}), + ...(this.metadata.mimeType !== undefined ? { mimeType: this.metadata.mimeType } : {}), + ...(this.metadata.imageDimensions !== undefined + ? { imageDimensions: this.metadata.imageDimensions } + : {}) + } + } finally { + this.releaseMemory() + } + } + + release(): void { + this.buffer = null + this.releaseMemory() + } + + private releaseMemory(): void { + this.releaseBudget?.() + this.releaseBudget = null + } +} diff --git a/src/main/ssh/ssh-filesystem-stream-reader.ts b/src/main/ssh/ssh-filesystem-stream-reader.ts index 2147ed0ef4f..d05089abf80 100644 --- a/src/main/ssh/ssh-filesystem-stream-reader.ts +++ b/src/main/ssh/ssh-filesystem-stream-reader.ts @@ -1,23 +1,25 @@ import type { SshChannelMultiplexer } from './ssh-channel-multiplexer' -import { STREAM_CHUNK_SIZE, JsonRpcErrorCode, RelayErrorCode } from './relay-protocol' +import { JsonRpcErrorCode } from './relay-protocol' import type { FileReadResult } from '../providers/types' +import { + PreMetadataStreamFrameBuffer, + STREAM_READER_INACTIVITY_TIMEOUT_MS, + defaultSshPreMetadataStreamBudget, + defaultSshStreamAssemblyBudget, + type SshPreMetadataStreamBudget, + type SshStreamAssemblyBudget, + createStreamInactivityDeadline +} from './ssh-stream-reader-memory' +import { + createFileStreamSetup, + type FileStreamAssembler, + StreamProtocolError +} from './ssh-file-stream-assembler' + +export { StreamProtocolError } from './ssh-file-stream-assembler' -const RESULT_ENCODING_BASE64 = 'base64' const SENTINEL_STREAM_ID = -1 -const MAX_PREVIEWABLE_BINARY_SIZE = 50 * 1024 * 1024 -const MAX_TEXT_FILE_SIZE = 10 * 1024 * 1024 - -type StreamMetadataResponse = { - streamId?: number - totalSize: number - isBinary: boolean - isImage?: boolean - mimeType?: string - resultEncoding?: 'base64' | 'utf-8' - empty?: boolean -} - export function isMethodNotFoundError(err: unknown): boolean { if (!err || typeof err !== 'object') { return false @@ -26,16 +28,12 @@ export function isMethodNotFoundError(err: unknown): boolean { return code === JsonRpcErrorCode.MethodNotFound } -export class StreamProtocolError extends Error { - readonly code = RelayErrorCode.StreamProtocolError - constructor(message: string) { - super(message) - } -} - export async function readFileViaStream( mux: SshChannelMultiplexer, - filePath: string + filePath: string, + options?: { inactivityTimeoutMs?: number }, + assemblyBudget: SshStreamAssemblyBudget = defaultSshStreamAssemblyBudget, + preMetadataBudget: SshPreMetadataStreamBudget = defaultSshPreMetadataStreamBudget ): Promise { // Why: subscribe BEFORE awaiting the metadata response so a chunk arriving // immediately after the response cannot beat the listener registration. @@ -55,42 +53,44 @@ export async function readFileViaStream( } return new Promise((resolve, reject) => { - let buffer: Buffer | null = null - let resultEncoding: 'base64' | 'utf-8' = RESULT_ENCODING_BASE64 - let isBinary = false - let isImage: boolean | undefined - let mimeType: string | undefined - let totalSize = 0 - let expectedSeq = 0 - let receivedChunks = 0 - let totalChunks = 0 - let bytesReceived = 0 + let assembler: FileStreamAssembler | null = null let settled = false // Why: chunk/end/error frames may arrive in the same dispatch tick as the // metadata response. Queue them until streamIdRef is set, then drain. - type PendingFrame = - | { kind: 'chunk'; params: Record } - | { kind: 'end'; params: Record } - | { kind: 'error'; params: Record } - const pending: PendingFrame[] = [] + const pending = new PreMetadataStreamFrameBuffer(preMetadataBudget) let metadataReady = false - const cancel = (): void => { - if (streamIdRef.current !== SENTINEL_STREAM_ID && !mux.isDisposed()) { + const inactivityMs = options?.inactivityTimeoutMs ?? STREAM_READER_INACTIVITY_TIMEOUT_MS + const inactivity = createStreamInactivityDeadline(inactivityMs, () => { + fail(new StreamProtocolError(`File stream stalled (>${inactivityMs}ms without data)`)) + }) + + const cancelStreamId = (streamId: number): void => { + if (!mux.isDisposed()) { try { - mux.notify('fs.cancelStream', { streamId: streamIdRef.current }) + mux.notify('fs.cancelStream', { streamId }) } catch { // Best-effort } } } + const cancel = (): void => { + if (streamIdRef.current !== SENTINEL_STREAM_ID) { + cancelStreamId(streamIdRef.current) + } + } + const fail = (err: Error): void => { if (settled) { return } settled = true + inactivity.clear() + pending.clear() + assembler?.release() + assembler = null cancel() cleanup() reject(err) @@ -101,6 +101,8 @@ export async function readFileViaStream( return } settled = true + inactivity.clear() + pending.clear() cleanup() resolve(value) } @@ -113,45 +115,26 @@ export async function readFileViaStream( if (id !== streamIdRef.current) { return } - const seq = params.seq as number - const data = params.data as string - if (typeof seq !== 'number' || typeof data !== 'string') { - fail(new StreamProtocolError(`Malformed chunk for stream ${id}`)) - return - } - if (seq !== expectedSeq) { - fail( - new StreamProtocolError( - `Out-of-order chunk for stream ${id}: expected ${expectedSeq}, got ${seq}` - ) - ) - return - } - const offset = seq * STREAM_CHUNK_SIZE - const decoded = Buffer.from(data, 'base64') - // Why: a short chunk would leave the pre-allocated buffer zero-filled and - // resolve as silently-corrupt data; validate each chunk's exact length. - const expectedLength = Math.min(STREAM_CHUNK_SIZE, totalSize - offset) - if (decoded.length !== expectedLength) { - fail( - new StreamProtocolError( - `Chunk length mismatch for stream ${id}: seq=${seq} expected=${expectedLength} got=${decoded.length}` - ) - ) - return - } - if (!buffer) { + if (!assembler) { fail(new StreamProtocolError(`Chunk arrived before metadata for stream ${id}`)) return } - decoded.copy(buffer, offset) - expectedSeq += 1 - receivedChunks += 1 - bytesReceived += decoded.length + let seq: number + try { + seq = assembler.acceptChunk(params, id) + } catch (error) { + fail(error as Error) + return + } + inactivity.reset() // Why: credit-based flow control — the relay caps unacked chunks so bulk // stream frames cannot queue unbounded ahead of interactive pty.data // frames on the shared SSH channel. Old relays ignore this notification. - mux.notify('fs.streamAck', { streamId: id, seq }) + try { + mux.notify('fs.streamAck', { streamId: id, seq }) + } catch { + // Disposal can race the write; teardown will settle the reader. + } } const handleEnd = (params: Record): void => { @@ -162,38 +145,15 @@ export async function readFileViaStream( if (id !== streamIdRef.current) { return } - if (receivedChunks !== totalChunks) { - fail( - new StreamProtocolError( - `Chunk count mismatch for stream ${id}: expected ${totalChunks}, received ${receivedChunks}` - ) - ) - return - } - // Why: redundant given the per-chunk length + count checks, but kept as a - // last-line invariant guard; never resolve with fewer bytes than declared. - if (bytesReceived !== totalSize) { - fail( - new StreamProtocolError( - `Byte count mismatch for stream ${id}: expected ${totalSize}, received ${bytesReceived}` - ) - ) - return - } - if (!buffer) { + if (!assembler) { fail(new StreamProtocolError(`Stream end before metadata for stream ${id}`)) return } - const content = - resultEncoding === RESULT_ENCODING_BASE64 - ? buffer.toString('base64') - : buffer.toString('utf-8') - succeed({ - content, - isBinary, - ...(isImage !== undefined ? { isImage } : {}), - ...(mimeType !== undefined ? { mimeType } : {}) - }) + try { + succeed(assembler.finish(id)) + } catch (error) { + fail(error as Error) + } } const handleStreamError = (params: Record): void => { @@ -213,7 +173,10 @@ export async function readFileViaStream( const drainPending = (): void => { while (!settled && pending.length > 0) { - const frame = pending.shift()! + const frame = pending.shift() + if (!frame) { + break + } if (frame.kind === 'chunk') { handleChunk(frame.params) } else if (frame.kind === 'end') { @@ -224,10 +187,18 @@ export async function readFileViaStream( } } + const pushPending = ( + kind: 'chunk' | 'end' | 'error', + params: Record + ): void => { + // Why: the stream id is unknown here; overload drops must not let one foreign frame fail every reader. + pending.push({ kind, params }) + } + unsubscribers.push( mux.onNotificationByMethod('fs.streamChunk', (params) => { if (!metadataReady) { - pending.push({ kind: 'chunk', params }) + pushPending('chunk', params) return } handleChunk(params) @@ -236,7 +207,7 @@ export async function readFileViaStream( unsubscribers.push( mux.onNotificationByMethod('fs.streamEnd', (params) => { if (!metadataReady) { - pending.push({ kind: 'end', params }) + pushPending('end', params) return } handleEnd(params) @@ -245,7 +216,7 @@ export async function readFileViaStream( unsubscribers.push( mux.onNotificationByMethod('fs.streamError', (params) => { if (!metadataReady) { - pending.push({ kind: 'error', params }) + pushPending('error', params) return } handleStreamError(params) @@ -269,51 +240,31 @@ export async function readFileViaStream( .request('fs.readFileStream', { filePath, flowControl: 'ack' }) .then((rawMetadata) => { if (settled) { + const streamId = (rawMetadata as { streamId?: unknown } | null)?.streamId + if (typeof streamId === 'number') { + cancelStreamId(streamId) + } return } - const metadata = rawMetadata as StreamMetadataResponse - isBinary = metadata.isBinary - isImage = metadata.isImage - mimeType = metadata.mimeType - resultEncoding = metadata.resultEncoding ?? RESULT_ENCODING_BASE64 - - if (metadata.empty) { - succeed({ - content: '', - isBinary: metadata.isBinary, - ...(metadata.isImage !== undefined ? { isImage: metadata.isImage } : {}), - ...(metadata.mimeType !== undefined ? { mimeType: metadata.mimeType } : {}) - }) - return - } - - if (typeof metadata.streamId !== 'number') { - fail(new StreamProtocolError('Metadata missing streamId for non-empty stream')) - return - } - - const cap = metadata.isBinary ? MAX_PREVIEWABLE_BINARY_SIZE : MAX_TEXT_FILE_SIZE - if (metadata.totalSize < 0 || metadata.totalSize > cap) { - streamIdRef.current = metadata.streamId - fail( - new StreamProtocolError( - `Reported totalSize ${metadata.totalSize} exceeds client cap ${cap}` - ) - ) - return - } - - totalSize = metadata.totalSize - totalChunks = totalSize === 0 ? 0 : Math.ceil(totalSize / STREAM_CHUNK_SIZE) + let setup try { - buffer = Buffer.alloc(totalSize) - } catch (err) { - streamIdRef.current = metadata.streamId - fail(new Error(`Failed to allocate ${totalSize} bytes: ${(err as Error).message}`)) + setup = createFileStreamSetup(rawMetadata, assemblyBudget) + } catch (error) { + const streamId = (rawMetadata as { streamId?: unknown } | null)?.streamId + if (typeof streamId === 'number') { + streamIdRef.current = streamId + } + fail(error as Error) return } - streamIdRef.current = metadata.streamId + if (setup.kind === 'empty') { + succeed(setup.result) + return + } + streamIdRef.current = setup.streamId + assembler = setup.assembler metadataReady = true + inactivity.reset() drainPending() }) .catch((err) => { diff --git a/src/main/ssh/ssh-git-response-stream-assembler.ts b/src/main/ssh/ssh-git-response-stream-assembler.ts new file mode 100644 index 00000000000..7153826ca80 --- /dev/null +++ b/src/main/ssh/ssh-git-response-stream-assembler.ts @@ -0,0 +1,108 @@ +import { RelayErrorCode } from './relay-protocol' +import { type SshStreamAssemblyBudget, base64EncodedLength } from './ssh-stream-reader-memory' + +export class GitResponseStreamError extends Error { + readonly code = RelayErrorCode.StreamProtocolError + constructor(message: string) { + super(message) + } +} + +export class GitResponseStreamAssembler { + private buffer: Buffer | null + private expectedSeq = 0 + private receivedBytes = 0 + private releaseBudget: (() => void) | null + + constructor( + private readonly streamId: number, + private readonly totalBytes: number, + private readonly chunkCount: number, + assemblyBudget: SshStreamAssemblyBudget + ) { + this.releaseBudget = assemblyBudget.reserve(totalBytes) + if (!this.releaseBudget) { + throw new GitResponseStreamError( + `Active SSH stream assembly would exceed ${assemblyBudget.maxBytes} bytes` + ) + } + try { + this.buffer = Buffer.alloc(totalBytes) + } catch (error) { + this.release() + throw new GitResponseStreamError( + `Failed to allocate ${totalBytes} bytes: ${(error as Error).message}` + ) + } + } + + acceptChunk(params: Record): number { + const seq = params.seq + const data = params.data + if ( + typeof seq !== 'number' || + !Number.isSafeInteger(seq) || + seq < 0 || + typeof data !== 'string' + ) { + throw new GitResponseStreamError(`Malformed chunk for git stream ${this.streamId}`) + } + if (seq !== this.expectedSeq) { + throw new GitResponseStreamError( + `Out-of-order chunk for git stream ${this.streamId}: expected ${this.expectedSeq}, got ${seq}` + ) + } + if (this.expectedSeq >= this.chunkCount) { + throw new GitResponseStreamError( + `Git stream ${this.streamId} exceeded declared chunk count ${this.chunkCount}` + ) + } + const remainingBytes = this.totalBytes - this.receivedBytes + if (data.length > base64EncodedLength(remainingBytes)) { + throw new GitResponseStreamError( + `Git stream ${this.streamId} chunk exceeds ${remainingBytes} remaining bytes` + ) + } + const decoded = Buffer.from(data, 'base64') + if (decoded.length > remainingBytes) { + throw new GitResponseStreamError( + `Git stream ${this.streamId} decoded beyond its declared size` + ) + } + if (!this.buffer) { + throw new GitResponseStreamError( + `Chunk arrived after completion for git stream ${this.streamId}` + ) + } + decoded.copy(this.buffer, this.receivedBytes) + this.receivedBytes += decoded.length + this.expectedSeq += 1 + return seq + } + + finish(): unknown { + if (this.expectedSeq !== this.chunkCount || this.receivedBytes !== this.totalBytes) { + throw new GitResponseStreamError( + `Git stream ${this.streamId} incomplete: chunks ${this.expectedSeq}/${this.chunkCount}, bytes ${this.receivedBytes}/${this.totalBytes}` + ) + } + if (!this.buffer) { + throw new GitResponseStreamError(`Git stream ${this.streamId} already completed`) + } + const serialized = this.buffer.toString('utf-8') + this.release() + try { + return JSON.parse(serialized) + } catch (error) { + throw new GitResponseStreamError( + `Git stream ${this.streamId} JSON parse failed: ${String(error)}` + ) + } + } + + release(): void { + this.buffer = null + this.releaseBudget?.() + this.releaseBudget = null + } +} diff --git a/src/main/ssh/ssh-git-response-stream-reader.ts b/src/main/ssh/ssh-git-response-stream-reader.ts index ef63c7a6dc0..94f9d5b8228 100644 --- a/src/main/ssh/ssh-git-response-stream-reader.ts +++ b/src/main/ssh/ssh-git-response-stream-reader.ts @@ -1,31 +1,26 @@ import type { SshChannelMultiplexer } from './ssh-channel-multiplexer' -import { RelayErrorCode, isGitResponseStreamMarker } from './relay-protocol' +import { + MAX_GIT_RESPONSE_STREAM_BYTES, + MAX_GIT_RESPONSE_STREAM_CHUNKS, + isGitResponseStreamMarker +} from './relay-protocol' +import { + PreMetadataStreamFrameBuffer, + STREAM_READER_INACTIVITY_TIMEOUT_MS, + defaultSshPreMetadataStreamBudget, + defaultSshStreamAssemblyBudget, + type SshPreMetadataStreamBudget, + type SshStreamAssemblyBudget, + createStreamInactivityDeadline +} from './ssh-stream-reader-memory' +import { + GitResponseStreamAssembler, + GitResponseStreamError +} from './ssh-git-response-stream-assembler' const SENTINEL_STREAM_ID = -1 -/** Reject if no stream frame (chunk/end/error) arrives within this window, - * reset on each frame. mux.request's own timeout only bounds the fast sentinel - * response; without this, a relay pump that breaks on staleness (which sends no - * responseEnd) while the SSH channel stays up would hang the client forever. */ -const STREAM_INACTIVITY_TIMEOUT_MS = 30_000 - -/** Bound transient buffering of other concurrent streams' chunks while this - * reader awaits its sentinel: every reader sees all git.responseChunk frames - * and can't filter by streamId until its own sentinel resolves. Foreign frames - * are dropped on drain anyway; this just caps the pre-sentinel backlog. */ -const MAX_PENDING_FRAMES = 64 - -export class GitResponseStreamError extends Error { - readonly code = RelayErrorCode.StreamProtocolError - constructor(message: string) { - super(message) - } -} - -type PendingFrame = - | { kind: 'chunk'; params: Record } - | { kind: 'end'; params: Record } - | { kind: 'error'; params: Record } +export { GitResponseStreamError } from './ssh-git-response-stream-assembler' /** * Request a git method that may return a large payload, opting into response @@ -49,7 +44,9 @@ export function requestGitStreamable( timeoutMs?: number /** Bounds the post-sentinel reassembly stall; resets on each chunk. */ inactivityTimeoutMs?: number - } + }, + assemblyBudget: SshStreamAssemblyBudget = defaultSshStreamAssemblyBudget, + preMetadataBudget: SshPreMetadataStreamBudget = defaultSshPreMetadataStreamBudget ): Promise { // Why: subscribe to chunk/end/error BEFORE awaiting the sentinel response so a // chunk that lands in the same dispatch tick as the response is not dropped @@ -68,42 +65,22 @@ export function requestGitStreamable( } return new Promise((resolve, reject) => { - const parts: Buffer[] = [] - let expectedSeq = 0 - let receivedBytes = 0 - let totalBytes = 0 - let chunkCount = 0 + let assembler: GitResponseStreamAssembler | null = null let settled = false let metadataReady = false - const pending: PendingFrame[] = [] + const pending = new PreMetadataStreamFrameBuffer(preMetadataBudget) - const inactivityMs = options?.inactivityTimeoutMs ?? STREAM_INACTIVITY_TIMEOUT_MS - let inactivityTimer: ReturnType | null = null - const clearInactivity = (): void => { - if (inactivityTimer) { - clearTimeout(inactivityTimer) - inactivityTimer = null - } - } - // Why: reset on every stream frame so a legitimately long stream is not - // killed, but a wedged stream (no frames arriving) rejects instead of - // hanging the caller forever. - const armInactivity = (): void => { - clearInactivity() - inactivityTimer = setTimeout(() => { - fail( - new GitResponseStreamError( - `Git response stream stalled (>${inactivityMs}ms without data)` - ) - ) - }, inactivityMs) - inactivityTimer.unref?.() - } + const inactivityMs = options?.inactivityTimeoutMs ?? STREAM_READER_INACTIVITY_TIMEOUT_MS + const inactivity = createStreamInactivityDeadline(inactivityMs, () => { + fail( + new GitResponseStreamError(`Git response stream stalled (>${inactivityMs}ms without data)`) + ) + }) - const cancel = (): void => { - if (streamIdRef.current !== SENTINEL_STREAM_ID && !mux.isDisposed()) { + const cancelStreamId = (streamId: number): void => { + if (streamId !== SENTINEL_STREAM_ID && !mux.isDisposed()) { try { - mux.notify('git.cancelResponseStream', { streamId: streamIdRef.current }) + mux.notify('git.cancelResponseStream', { streamId }) } catch { // best-effort } @@ -114,8 +91,11 @@ export function requestGitStreamable( return } settled = true - clearInactivity() - cancel() + inactivity.clear() + pending.clear() + assembler?.release() + assembler = null + cancelStreamId(streamIdRef.current) cleanup() reject(err) } @@ -124,7 +104,8 @@ export function requestGitStreamable( return } settled = true - clearInactivity() + inactivity.clear() + pending.clear() cleanup() resolve(value) } @@ -133,25 +114,22 @@ export function requestGitStreamable( if (settled || p.streamId !== streamIdRef.current) { return } - const seq = p.seq as number - const data = p.data as string - if (typeof seq !== 'number' || typeof data !== 'string') { - fail(new GitResponseStreamError(`Malformed chunk for git stream ${streamIdRef.current}`)) - return - } - if (seq !== expectedSeq) { + if (!assembler) { fail( new GitResponseStreamError( - `Out-of-order chunk for git stream ${streamIdRef.current}: expected ${expectedSeq}, got ${seq}` + `Chunk arrived before metadata for git stream ${streamIdRef.current}` ) ) return } - const decoded = Buffer.from(data, 'base64') - parts.push(decoded) - receivedBytes += decoded.length - expectedSeq += 1 - armInactivity() + let seq: number + try { + seq = assembler.acceptChunk(p) + } catch (error) { + fail(error as Error) + return + } + inactivity.reset() // Why: credit-based flow control — the relay caps unacked chunks so a big // response cannot queue unbounded ahead of interactive pty.data frames. if (!mux.isDisposed()) { @@ -167,22 +145,18 @@ export function requestGitStreamable( if (settled || p.streamId !== streamIdRef.current) { return } - if (expectedSeq !== chunkCount || receivedBytes !== totalBytes) { + if (!assembler) { fail( new GitResponseStreamError( - `Git stream ${streamIdRef.current} incomplete: chunks ${expectedSeq}/${chunkCount}, bytes ${receivedBytes}/${totalBytes}` + `Stream end before metadata for git stream ${streamIdRef.current}` ) ) return } try { - succeed(JSON.parse(Buffer.concat(parts).toString('utf-8'))) - } catch (err) { - fail( - new GitResponseStreamError( - `Git stream ${streamIdRef.current} JSON parse failed: ${String(err)}` - ) - ) + succeed(assembler.finish()) + } catch (error) { + fail(error as Error) } } @@ -195,7 +169,10 @@ export function requestGitStreamable( const drainPending = (): void => { while (!settled && pending.length > 0) { - const frame = pending.shift()! + const frame = pending.shift() + if (!frame) { + break + } if (frame.kind === 'chunk') { handleChunk(frame.params) } else if (frame.kind === 'end') { @@ -206,22 +183,18 @@ export function requestGitStreamable( } } - // Why: pre-sentinel we cannot filter by streamId (our id is unknown yet), so - // every concurrent reader transiently buffers all readers' chunks. Cap the - // backlog by dropping the oldest; foreign frames are dropped on drain anyway, - // and if our own seq-0 were ever dropped the seq check fails loudly rather - // than corrupting. The sentinel normally resolves long before this cap. - const pushPending = (frame: PendingFrame): void => { - pending.push(frame) - if (pending.length > MAX_PENDING_FRAMES) { - pending.shift() - } + const pushPending = ( + kind: 'chunk' | 'end' | 'error', + params: Record + ): void => { + // Why: the stream id is unknown here; overload drops must not let one foreign frame fail every reader. + pending.push({ kind, params }) } unsubscribers.push( mux.onNotificationByMethod('git.responseChunk', (p) => { if (!metadataReady) { - pushPending({ kind: 'chunk', params: p }) + pushPending('chunk', p) return } handleChunk(p) @@ -230,7 +203,7 @@ export function requestGitStreamable( unsubscribers.push( mux.onNotificationByMethod('git.responseEnd', (p) => { if (!metadataReady) { - pushPending({ kind: 'end', params: p }) + pushPending('end', p) return } handleEnd(p) @@ -239,7 +212,7 @@ export function requestGitStreamable( unsubscribers.push( mux.onNotificationByMethod('git.responseError', (p) => { if (!metadataReady) { - pushPending({ kind: 'error', params: p }) + pushPending('error', p) return } handleStreamError(p) @@ -289,21 +262,52 @@ export function requestGitStreamable( void requestPromise .then((result) => { if (settled) { + if (isGitResponseStreamMarker(result)) { + cancelStreamId(result.__orcaGitResponseStream.streamId) + } return } // Old relay / small result: plain single-frame value, no stream follows. if (!isGitResponseStreamMarker(result)) { + if ( + typeof result === 'object' && + result !== null && + '__orcaGitResponseStream' in result + ) { + fail(new GitResponseStreamError('Malformed Git response stream metadata')) + return + } succeed(result) return } const marker = result.__orcaGitResponseStream - totalBytes = marker.totalBytes - chunkCount = marker.chunkCount streamIdRef.current = marker.streamId + if ( + marker.totalBytes > MAX_GIT_RESPONSE_STREAM_BYTES || + marker.chunkCount > MAX_GIT_RESPONSE_STREAM_CHUNKS + ) { + fail( + new GitResponseStreamError( + `Git response stream exceeds client limit (${marker.totalBytes} bytes, ${marker.chunkCount} chunks)` + ) + ) + return + } + try { + assembler = new GitResponseStreamAssembler( + marker.streamId, + marker.totalBytes, + marker.chunkCount, + assemblyBudget + ) + } catch (error) { + fail(error as Error) + return + } metadataReady = true // Why: start the inactivity deadline now — mux.request's timeout only // covered the sentinel; the reassembly phase needs its own guard. - armInactivity() + inactivity.reset() drainPending() }) .catch((err) => fail(err as Error)) diff --git a/src/main/ssh/ssh-stream-reader-memory.test.ts b/src/main/ssh/ssh-stream-reader-memory.test.ts new file mode 100644 index 00000000000..8cbe291dc92 --- /dev/null +++ b/src/main/ssh/ssh-stream-reader-memory.test.ts @@ -0,0 +1,432 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + MAX_GIT_RESPONSE_STREAM_BYTES, + MAX_GIT_RESPONSE_STREAM_CHUNKS, + STREAM_CHUNK_SIZE, + MessageType, + encodeFrame +} from './relay-protocol' +import { SshChannelMultiplexer, type MultiplexerTransport } from './ssh-channel-multiplexer' +import { readFileViaStream } from './ssh-filesystem-stream-reader' +import { requestGitStreamable } from './ssh-git-response-stream-reader' +import { + MAX_PRE_METADATA_STREAM_ENCODED_BYTES, + PreMetadataStreamFrameBuffer, + SshPreMetadataStreamBudget, + SshStreamAssemblyBudget +} from './ssh-stream-reader-memory' + +type Notification = { method: string; params: Record | undefined } + +function createStreamMux(): { + mux: Record + emit: (method: string, params: Record) => void + resolveRequest: (value: unknown) => void + rejectRequest: (error: Error) => void + dispose: (reason: 'shutdown' | 'connection_lost') => void + notifications: Notification[] +} { + const handlers = new Map) => void>>() + const disposeHandlers = new Set<(reason: 'shutdown' | 'connection_lost') => void>() + const notifications: Notification[] = [] + let resolveRequest!: (value: unknown) => void + let rejectRequest!: (error: Error) => void + const requestPromise = new Promise((resolve, reject) => { + resolveRequest = resolve + rejectRequest = reject + }) + const mux = { + request: vi.fn(() => requestPromise), + notify: vi.fn((method: string, params?: Record) => { + notifications.push({ method, params }) + }), + onNotificationByMethod: vi.fn( + (method: string, handler: (params: Record) => void) => { + const current = handlers.get(method) ?? new Set() + current.add(handler) + handlers.set(method, current) + return () => current.delete(handler) + } + ), + onDispose: vi.fn((handler: (reason: 'shutdown' | 'connection_lost') => void) => { + disposeHandlers.add(handler) + return () => disposeHandlers.delete(handler) + }), + isDisposed: vi.fn(() => false) + } + return { + mux, + emit: (method, params) => { + for (const handler of handlers.get(method) ?? []) { + handler(params) + } + }, + resolveRequest, + rejectRequest, + dispose: (reason) => { + for (const handler of disposeHandlers) { + handler(reason) + } + }, + notifications + } +} + +async function flushRequestResolution(): Promise { + await Promise.resolve() + await Promise.resolve() +} + +function encodeMuxResponse(requestId: number, result: unknown, sequence: number): Buffer { + return encodeFrame( + MessageType.Regular, + sequence, + 0, + Buffer.from(JSON.stringify({ jsonrpc: '2.0', id: requestId, result })) + ) +} + +afterEach(() => { + vi.useRealTimers() +}) + +describe('SSH stream assembly aggregate budget', () => { + it('bounds concurrent reservations and releases them idempotently', () => { + const budget = new SshStreamAssemblyBudget(10) + const releaseSix = budget.reserve(6) + const releaseFour = budget.reserve(4) + + expect(releaseSix).not.toBeNull() + expect(releaseFour).not.toBeNull() + expect(budget.reserve(1)).toBeNull() + expect(budget.retainedBytes).toBe(10) + + releaseSix?.() + releaseSix?.() + expect(budget.retainedBytes).toBe(4) + expect(budget.reserve(6)).not.toBeNull() + }) +}) + +describe('SSH pre-metadata stream aggregate budget', () => { + it('bounds concurrent frame buffers by count and bytes', () => { + const budget = new SshPreMetadataStreamBudget(2, 5) + const first = new PreMetadataStreamFrameBuffer(budget) + const second = new PreMetadataStreamFrameBuffer(budget) + const third = new PreMetadataStreamFrameBuffer(budget) + + expect(first.push({ kind: 'chunk', params: { streamId: 1, seq: 0, data: 'abc' } })).toBe(true) + expect(second.push({ kind: 'chunk', params: { streamId: 2, seq: 0, data: 'de' } })).toBe(true) + expect(third.push({ kind: 'end', params: { streamId: 3 } })).toBe(false) + expect(budget.retainedFrames).toBe(2) + expect(budget.retainedBytes).toBe(5) + + first.clear() + expect(third.push({ kind: 'end', params: { streamId: 3 } })).toBe(true) + expect(budget.retainedFrames).toBe(2) + second.clear() + third.clear() + expect(budget.retainedFrames).toBe(0) + expect(budget.retainedBytes).toBe(0) + }) + + it('releases shared reservations on success, stream error, request timeout, and disconnect', async () => { + const budget = new SshPreMetadataStreamBudget(4, 64) + const fileHarness = createStreamMux() + const errorHarness = createStreamMux() + const gitHarness = createStreamMux() + const disconnectHarness = createStreamMux() + const fileResult = readFileViaStream( + fileHarness.mux as never, + 'empty.txt', + undefined, + new SshStreamAssemblyBudget(1), + budget + ) + const errorResult = readFileViaStream( + errorHarness.mux as never, + 'error.txt', + undefined, + new SshStreamAssemblyBudget(1), + budget + ) + const gitResult = requestGitStreamable( + gitHarness.mux as never, + 'git.diff', + {}, + undefined, + new SshStreamAssemblyBudget(1), + budget + ) + const disconnectResult = readFileViaStream( + disconnectHarness.mux as never, + 'disconnected.txt', + undefined, + new SshStreamAssemblyBudget(1), + budget + ) + fileHarness.emit('fs.streamEnd', { streamId: 1 }) + errorHarness.emit('fs.streamError', { streamId: 2, message: 'remote read failed' }) + gitHarness.emit('git.responseEnd', { streamId: 3 }) + disconnectHarness.emit('fs.streamEnd', { streamId: 4 }) + expect(budget.retainedFrames).toBe(4) + + const streamRejection = expect(errorResult).rejects.toThrow('remote read failed') + const gitRejection = expect(gitResult).rejects.toThrow(/timed out/) + const disconnectRejection = expect(disconnectResult).rejects.toMatchObject({ + code: 'CONNECTION_LOST' + }) + fileHarness.resolveRequest({ streamId: 1, totalSize: 0, isBinary: false }) + errorHarness.resolveRequest({ streamId: 2, totalSize: 1, isBinary: false }) + gitHarness.rejectRequest(new Error('Request "git.diff" timed out after 30ms')) + disconnectHarness.dispose('connection_lost') + + await expect(fileResult).resolves.toEqual({ content: '', isBinary: false }) + await streamRejection + await gitRejection + await disconnectRejection + expect(budget.retainedFrames).toBe(0) + expect(budget.retainedBytes).toBe(0) + }) +}) + +describe('SSH Git response stream memory bounds', () => { + it('rejects and cancels an oversized stream marker', async () => { + const harness = createStreamMux() + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}) + harness.resolveRequest({ + __orcaGitResponseStream: { + streamId: 7, + totalBytes: MAX_GIT_RESPONSE_STREAM_BYTES + 1, + chunkCount: 1 + } + }) + + await expect(result).rejects.toThrow(/exceeds client limit/) + expect(harness.notifications).toContainEqual({ + method: 'git.cancelResponseStream', + params: { streamId: 7 } + }) + }) + + it('rejects an excessive declared chunk count', async () => { + const harness = createStreamMux() + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}) + harness.resolveRequest({ + __orcaGitResponseStream: { + streamId: 8, + totalBytes: 1, + chunkCount: MAX_GIT_RESPONSE_STREAM_CHUNKS + 1 + } + }) + + await expect(result).rejects.toThrow(/exceeds client limit/) + }) + + it('drops an oversized foreign frame received before the stream marker', async () => { + const harness = createStreamMux() + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}) + harness.emit('git.responseChunk', { + streamId: 99, + seq: 0, + data: 'A'.repeat(MAX_PRE_METADATA_STREAM_ENCODED_BYTES + 1) + }) + const payload = Buffer.from('{"ok":true}') + harness.resolveRequest({ + __orcaGitResponseStream: { streamId: 9, totalBytes: payload.length, chunkCount: 1 } + }) + await flushRequestResolution() + harness.emit('git.responseChunk', { streamId: 9, seq: 0, data: payload.toString('base64') }) + harness.emit('git.responseEnd', { streamId: 9 }) + + await expect(result).resolves.toEqual({ ok: true }) + }) + + it('rejects an encoded chunk larger than the remaining declared bytes', async () => { + const harness = createStreamMux() + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}) + harness.resolveRequest({ + __orcaGitResponseStream: { streamId: 10, totalBytes: 1, chunkCount: 1 } + }) + await flushRequestResolution() + harness.emit('git.responseChunk', { streamId: 10, seq: 0, data: 'A'.repeat(1_000_000) }) + + await expect(result).rejects.toThrow(/exceeds 1 remaining bytes/) + }) + + it('reassembles a valid result without changing its value', async () => { + const harness = createStreamMux() + const expected = { answer: 'unchanged', count: 2 } + const payload = Buffer.from(JSON.stringify(expected)) + const budget = new SshStreamAssemblyBudget(payload.length) + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}, undefined, budget) + const first = payload.subarray(0, 7) + const second = payload.subarray(7) + harness.resolveRequest({ + __orcaGitResponseStream: { streamId: 11, totalBytes: payload.length, chunkCount: 2 } + }) + await flushRequestResolution() + expect(budget.retainedBytes).toBe(payload.length) + harness.emit('git.responseChunk', { streamId: 11, seq: 0, data: first.toString('base64') }) + harness.emit('git.responseChunk', { streamId: 11, seq: 1, data: second.toString('base64') }) + harness.emit('git.responseEnd', { streamId: 11 }) + + await expect(result).resolves.toEqual(expected) + expect(budget.retainedBytes).toBe(0) + }) + + it('rejects and cancels when the aggregate assembly budget is unavailable', async () => { + const harness = createStreamMux() + const budget = new SshStreamAssemblyBudget(1) + const held = budget.reserve(1) + const result = requestGitStreamable(harness.mux as never, 'git.diff', {}, undefined, budget) + harness.resolveRequest({ + __orcaGitResponseStream: { streamId: 16, totalBytes: 1, chunkCount: 1 } + }) + + await expect(result).rejects.toThrow(/Active SSH stream assembly would exceed/) + expect(harness.notifications).toContainEqual({ + method: 'git.cancelResponseStream', + params: { streamId: 16 } + }) + held?.() + }) +}) + +describe('SSH filesystem stream memory bounds', () => { + it('releases every active assembly immediately when the shared mux disconnects', async () => { + const dataHandlers: ((data: Buffer) => void)[] = [] + const transport: MultiplexerTransport = { + write: vi.fn(), + onData: (handler) => dataHandlers.push(handler), + onClose: vi.fn(), + close: vi.fn() + } + const mux = new SshChannelMultiplexer(transport) + const budget = new SshStreamAssemblyBudget(2) + const first = readFileViaStream(mux, 'first.txt', undefined, budget) + const second = readFileViaStream(mux, 'second.txt', undefined, budget) + const firstRejection = expect(first).rejects.toMatchObject({ code: 'CONNECTION_LOST' }) + const secondRejection = expect(second).rejects.toMatchObject({ code: 'CONNECTION_LOST' }) + + dataHandlers[0]( + Buffer.concat([ + encodeMuxResponse(1, { streamId: 21, totalSize: 1, isBinary: false }, 1), + encodeMuxResponse(2, { streamId: 22, totalSize: 1, isBinary: false }, 2) + ]) + ) + await flushRequestResolution() + expect(budget.retainedBytes).toBe(2) + + mux.dispose('connection_lost') + + expect(budget.retainedBytes).toBe(0) + await firstRejection + await secondRejection + }) + + it('drops an oversized foreign frame received before file metadata', async () => { + const harness = createStreamMux() + const result = readFileViaStream(harness.mux as never, 'tiny.txt') + harness.emit('fs.streamChunk', { + streamId: 90, + seq: 0, + data: 'A'.repeat(MAX_PRE_METADATA_STREAM_ENCODED_BYTES + 1) + }) + harness.resolveRequest({ + streamId: 12, + totalSize: 1, + isBinary: false, + resultEncoding: 'utf-8' + }) + await flushRequestResolution() + harness.emit('fs.streamChunk', { + streamId: 12, + seq: 0, + data: Buffer.from('a').toString('base64') + }) + harness.emit('fs.streamEnd', { streamId: 12 }) + + await expect(result).resolves.toEqual({ content: 'a', isBinary: false }) + }) + + it('validates encoded size before decoding a file chunk', async () => { + const harness = createStreamMux() + const result = readFileViaStream(harness.mux as never, 'tiny.txt') + harness.resolveRequest({ streamId: 13, totalSize: 1, isBinary: false }) + await flushRequestResolution() + harness.emit('fs.streamChunk', { streamId: 13, seq: 0, data: 'A'.repeat(1_000_000) }) + + await expect(result).rejects.toThrow(/Encoded chunk length mismatch/) + expect(harness.notifications).toContainEqual({ + method: 'fs.cancelStream', + params: { streamId: 13 } + }) + }) + + it('cancels a file stream that stops making progress', async () => { + vi.useFakeTimers() + const harness = createStreamMux() + const result = readFileViaStream(harness.mux as never, 'stalled.txt', { + inactivityTimeoutMs: 50 + }) + harness.resolveRequest({ streamId: 14, totalSize: 1, isBinary: false }) + await flushRequestResolution() + + const rejection = expect(result).rejects.toThrow(/stalled/) + await vi.advanceTimersByTimeAsync(51) + await rejection + expect(harness.notifications).toContainEqual({ + method: 'fs.cancelStream', + params: { streamId: 14 } + }) + }) + + it('resets the inactivity deadline on progress and preserves valid content', async () => { + vi.useFakeTimers() + const harness = createStreamMux() + const first = Buffer.alloc(STREAM_CHUNK_SIZE, 0x61) + const second = Buffer.from('z') + const budget = new SshStreamAssemblyBudget(first.length + second.length) + const result = readFileViaStream( + harness.mux as never, + 'progress.txt', + { inactivityTimeoutMs: 50 }, + budget + ) + harness.resolveRequest({ + streamId: 15, + totalSize: first.length + second.length, + isBinary: false, + resultEncoding: 'utf-8' + }) + await flushRequestResolution() + expect(budget.retainedBytes).toBe(first.length + second.length) + await vi.advanceTimersByTimeAsync(40) + harness.emit('fs.streamChunk', { streamId: 15, seq: 0, data: first.toString('base64') }) + await vi.advanceTimersByTimeAsync(40) + harness.emit('fs.streamChunk', { streamId: 15, seq: 1, data: second.toString('base64') }) + harness.emit('fs.streamEnd', { streamId: 15 }) + + await expect(result).resolves.toEqual({ + content: `${'a'.repeat(STREAM_CHUNK_SIZE)}z`, + isBinary: false + }) + expect(budget.retainedBytes).toBe(0) + }) + + it('rejects and cancels when file assembly would exceed the aggregate budget', async () => { + const harness = createStreamMux() + const budget = new SshStreamAssemblyBudget(1) + const held = budget.reserve(1) + const result = readFileViaStream(harness.mux as never, 'tiny.txt', undefined, budget) + harness.resolveRequest({ streamId: 17, totalSize: 1, isBinary: false }) + + await expect(result).rejects.toThrow(/Active SSH stream assembly would exceed/) + expect(harness.notifications).toContainEqual({ + method: 'fs.cancelStream', + params: { streamId: 17 } + }) + held?.() + }) +}) diff --git a/src/main/ssh/ssh-stream-reader-memory.ts b/src/main/ssh/ssh-stream-reader-memory.ts new file mode 100644 index 00000000000..c599a11a58b --- /dev/null +++ b/src/main/ssh/ssh-stream-reader-memory.ts @@ -0,0 +1,205 @@ +export const STREAM_READER_INACTIVITY_TIMEOUT_MS = 30_000 +export const MAX_PRE_METADATA_STREAM_FRAMES = 64 +export const MAX_PRE_METADATA_STREAM_ENCODED_BYTES = 8 * 1024 * 1024 +export const MAX_SSH_PRE_METADATA_STREAM_FRAMES = 256 +export const MAX_SSH_PRE_METADATA_STREAM_RETAINED_BYTES = 32 * 1024 * 1024 +export const MAX_SSH_STREAM_ASSEMBLY_BYTES = 128 * 1024 * 1024 + +export type BufferedStreamFrame = { + kind: 'chunk' | 'end' | 'error' + params: Record +} + +type RetainedStreamFrame = { + frame: BufferedStreamFrame + encodedBytes: number + release: () => void +} + +export class SshPreMetadataStreamBudget { + private frames = 0 + private bytes = 0 + + constructor( + readonly maxFrames: number, + readonly maxBytes: number + ) {} + + reserve(bytes: number): (() => void) | null { + if ( + !Number.isSafeInteger(bytes) || + bytes < 0 || + this.frames >= this.maxFrames || + bytes > this.maxBytes - this.bytes + ) { + return null + } + this.frames += 1 + this.bytes += bytes + let released = false + return () => { + if (released) { + return + } + released = true + this.frames -= 1 + this.bytes -= bytes + } + } + + get retainedFrames(): number { + return this.frames + } + + get retainedBytes(): number { + return this.bytes + } +} + +export const defaultSshPreMetadataStreamBudget = new SshPreMetadataStreamBudget( + MAX_SSH_PRE_METADATA_STREAM_FRAMES, + MAX_SSH_PRE_METADATA_STREAM_RETAINED_BYTES +) + +export class PreMetadataStreamFrameBuffer { + private frames: RetainedStreamFrame[] = [] + private encodedBytes = 0 + + constructor(private readonly budget: SshPreMetadataStreamBudget) {} + + push(frame: BufferedStreamFrame): boolean { + const retainedFrame = retainStreamFrameFields(frame) + const frameBytes = retainedStreamFrameBytes(retainedFrame) + if (frameBytes > MAX_PRE_METADATA_STREAM_ENCODED_BYTES) { + return false + } + while ( + this.frames.length >= MAX_PRE_METADATA_STREAM_FRAMES || + this.encodedBytes + frameBytes > MAX_PRE_METADATA_STREAM_ENCODED_BYTES + ) { + this.shift() + } + let release = this.budget.reserve(frameBytes) + while (!release && this.frames.length > 0) { + this.shift() + release = this.budget.reserve(frameBytes) + } + if (!release) { + return false + } + this.frames.push({ frame: retainedFrame, encodedBytes: frameBytes, release }) + this.encodedBytes += frameBytes + return true + } + + shift(): BufferedStreamFrame | undefined { + const retained = this.frames.shift() + if (retained) { + this.encodedBytes -= retained.encodedBytes + retained.release() + } + return retained?.frame + } + + get length(): number { + return this.frames.length + } + + clear(): void { + while (this.frames.length > 0) { + this.shift() + } + } +} + +function retainStreamFrameFields(frame: BufferedStreamFrame): BufferedStreamFrame { + const { params } = frame + if (frame.kind === 'chunk') { + return { + kind: frame.kind, + params: { + ...(typeof params.streamId === 'number' ? { streamId: params.streamId } : {}), + ...(typeof params.seq === 'number' ? { seq: params.seq } : {}), + ...(typeof params.data === 'string' ? { data: params.data } : {}) + } + } + } + if (frame.kind === 'error') { + return { + kind: frame.kind, + params: { + ...(typeof params.streamId === 'number' ? { streamId: params.streamId } : {}), + ...(typeof params.message === 'string' ? { message: params.message } : {}), + ...(typeof params.code === 'string' ? { code: params.code } : {}) + } + } + } + return { + kind: frame.kind, + params: typeof params.streamId === 'number' ? { streamId: params.streamId } : {} + } +} + +function retainedStreamFrameBytes(frame: BufferedStreamFrame): number { + let bytes = 0 + for (const value of Object.values(frame.params)) { + if (typeof value === 'string') { + bytes += Buffer.byteLength(value, 'utf-8') + } + } + return bytes +} + +export class SshStreamAssemblyBudget { + private retained = 0 + + constructor(readonly maxBytes: number) {} + + reserve(bytes: number): (() => void) | null { + if (!Number.isSafeInteger(bytes) || bytes < 0 || bytes > this.maxBytes - this.retained) { + return null + } + this.retained += bytes + let released = false + return () => { + if (released) { + return + } + released = true + this.retained -= bytes + } + } + + get retainedBytes(): number { + return this.retained + } +} + +export const defaultSshStreamAssemblyBudget = new SshStreamAssemblyBudget( + MAX_SSH_STREAM_ASSEMBLY_BYTES +) + +export function base64EncodedLength(rawBytes: number): number { + return Math.ceil(rawBytes / 3) * 4 +} + +export function createStreamInactivityDeadline( + timeoutMs: number, + onTimeout: () => void +): { reset: () => void; clear: () => void } { + let timer: ReturnType | null = null + const clear = (): void => { + if (timer) { + clearTimeout(timer) + timer = null + } + } + return { + reset: () => { + clear() + timer = setTimeout(onTimeout, timeoutMs) + timer.unref?.() + }, + clear + } +} diff --git a/src/relay/wsl-hook-fs-bridge.ts b/src/relay/wsl-hook-fs-bridge.ts index 3ebb6047b1b..d0eaff4f7c5 100644 --- a/src/relay/wsl-hook-fs-bridge.ts +++ b/src/relay/wsl-hook-fs-bridge.ts @@ -7,17 +7,54 @@ import { promises as fs } from 'node:fs' import { posix } from 'node:path' import type { RelayDispatcher } from './dispatcher' +import { + createFilesystemDirectoryLimitState, + trackFilesystemDirectoryEntry +} from '../shared/filesystem-directory-listing-limit' +import { + NodeFileReadTooLargeError, + readNodeFileWithinLimit +} from '../shared/memory-safety/node-bounded-file-reader' import { WSL_HOOK_FS_METHODS, + WSL_HOOK_FS_MAX_DIRECTORY_ENTRIES, + WSL_HOOK_FS_MAX_DIRECTORY_RETAINED_BYTES, + WSL_HOOK_FS_MAX_READ_BYTES, type WslFsFailure, type WslFsResult } from '../shared/wsl-hook-relay-contract' function failure(err: unknown): WslFsFailure { + if (err instanceof NodeFileReadTooLargeError) { + return { + ok: false, + errno: 'EFBIG', + message: err.message, + fileCapacity: { observedBytes: err.observedBytes, maxBytes: err.maxBytes } + } + } const e = err as NodeJS.ErrnoException return { ok: false, errno: e?.code ?? 'EUNKNOWN', message: e?.message ?? String(err) } } +function requestedLimit(value: unknown, maximum: number): number { + if (value === undefined) { + return maximum + } + if (typeof value !== 'number' || !Number.isSafeInteger(value) || value < 0) { + throw Object.assign(new Error('invalid capacity limit'), { code: 'EINVAL' }) + } + return Math.min(value, maximum) +} + +function requestedDirectoryLimit(value: unknown, maximum: number): number { + const limit = requestedLimit(value, maximum) + if (limit === 0) { + throw Object.assign(new Error('directory capacity limit must be positive'), { code: 'EINVAL' }) + } + return limit +} + export function registerWslHookFsHandlers( dispatcher: RelayDispatcher, home: string, @@ -46,7 +83,7 @@ export function registerWslHookFsHandlers( return resolved } // Why: the installers' mkdir-p walks top-down from `/`, probing every - // ancestor of home with readdir before it ever creates a dir. Allow + // ancestor of home before it ever creates a dir. Allow // read-only existence probes on those ancestors; everything else stays // home-scoped. const scopedProbe = (rawPath: unknown): string => { @@ -68,7 +105,10 @@ export function registerWslHookFsHandlers( WSL_HOOK_FS_METHODS.readFile, async (params): Promise> => { try { - const content = await fs.readFile(scoped(params.path), 'utf8') + const maxBytes = requestedLimit(params.maxBytes, WSL_HOOK_FS_MAX_READ_BYTES) + const content = ( + await readNodeFileWithinLimit(scoped(params.path), maxBytes) + ).buffer.toString('utf8') return { ok: true, content } } catch (err) { return failure(err) @@ -93,7 +133,7 @@ export function registerWslHookFsHandlers( WSL_HOOK_FS_METHODS.stat, async (params): Promise> => { try { - const stats = await fs.stat(scoped(params.path)) + const stats = await fs.stat(scopedProbe(params.path)) return { ok: true, mode: stats.mode } } catch (err) { return failure(err) @@ -133,11 +173,39 @@ export function registerWslHookFsHandlers( dispatcher.onRequest( WSL_HOOK_FS_METHODS.readdir, async (params): Promise> => { + let directory: Awaited> | undefined try { - const names = await fs.readdir(scopedProbe(params.path)) - return { ok: true, entries: names.map((filename) => ({ filename })) } + const limits = { + maxEntries: requestedDirectoryLimit(params.maxEntries, WSL_HOOK_FS_MAX_DIRECTORY_ENTRIES), + maxRetainedBytes: requestedDirectoryLimit( + params.maxRetainedBytes, + WSL_HOOK_FS_MAX_DIRECTORY_RETAINED_BYTES + ) + } + const limit = createFilesystemDirectoryLimitState(limits) + const entries: { filename: string }[] = [] + const directoryPath = scopedProbe(params.path) + directory = await fs.opendir(directoryPath, { bufferSize: 32 }) + // Ancestors outside home are authorized only as existence probes. + if ( + directoryPath !== homeRoot && + (directoryPath === '/' || homeRoot.startsWith(`${directoryPath}/`)) + ) { + return { ok: true, entries } + } + for (;;) { + const entry = await directory.read() + if (entry === null) { + break + } + trackFilesystemDirectoryEntry(limit, { name: entry.name }) + entries.push({ filename: entry.name }) + } + return { ok: true, entries } } catch (err) { return failure(err) + } finally { + await directory?.close().catch(() => undefined) } } )