diff --git a/packages/code-runtime/code-runtime-python/src/index.ts b/packages/code-runtime/code-runtime-python/src/index.ts index 3068eef776..bc2ee80a4e 100644 --- a/packages/code-runtime/code-runtime-python/src/index.ts +++ b/packages/code-runtime/code-runtime-python/src/index.ts @@ -1320,13 +1320,17 @@ export class PythonCodeRuntime extends CodeRuntime { // defaults. // // The cap used HERE is FRAME_PARSE_CAP_BYTES, not the 256 MiB wire - // ceiling: a single frame between 64 MiB and the ceiling would otherwise - // be fully `Buffer.concat`-ed (a second copy of its bytes) and only then - // dropped in the line loop — the peak-memory doubling this pre-concat - // check exists to prevent, now for a frame the parser is guaranteed to - // discard. Dropping the oversized unframed buffer before the join keeps - // the peak at one copy of the wire bytes. - if (pendingBytes > FRAME_PARSE_CAP_BYTES) { + // ceiling, and ONLY when the held bytes are still a single unframed + // line (this chunk carries no newline, and earlier newline-bearing + // chunks were joined immediately): a frame between 64 MiB and the + // ceiling would otherwise be fully `Buffer.concat`-ed (a second copy + // of its bytes) and only then dropped in the line loop — the + // peak-memory doubling this pre-concat check exists to prevent. + // Dropping the oversized unframed buffer before the join keeps the + // peak at one copy of the wire bytes. When this chunk DOES carry a + // newline the buffer holds several frames, so the FIRST-FRAME check + // below (not this counter, which charges them all) decides. + if (pendingBytes > FRAME_PARSE_CAP_BYTES && !chunk.includes(0x0a)) { pendingChunks = [] sealedBlocks = [] pendingBytes = 0 @@ -1361,6 +1365,35 @@ export class PythonCodeRuntime extends CodeRuntime { pendingChunks = [] } if (chunk.includes(0x0a)) { + // First-FRAME check before the join: measure the bytes up to the + // first newline across the held chunks. The byte counter cannot + // serve here — it charges the whole buffer, which legitimately + // holds several frames each within the cap. A first frame past the + // cap is dropped before the join (one copy of its wire bytes); + // later frames in the same buffer are handled by the per-line check + // in the loop below. + let firstFrameLen = 0 + let sawNewline = false + // Sealed blocks hold newline-free prefixes only (a newline-bearing + // chunk is joined immediately), so they are entirely part of the + // first frame. + for (const b of sealedBlocks) firstFrameLen += b.length + for (const c of pendingChunks) { + const nl = c.indexOf(0x0a) + if (nl >= 0) { + firstFrameLen += nl + sawNewline = true + break + } + firstFrameLen += c.length + } + if (sawNewline && firstFrameLen > FRAME_PARSE_CAP_BYTES) { + pendingChunks = [] + sealedBlocks = [] + pendingBytes = 0 + finish({ error: { kind: 'worker-exit', message: `protocol frame exceeded ${FRAME_PARSE_CAP_BYTES} bytes on fd 3` } }) + return + } let buffered = Buffer.concat(sealedBlocks.length > 0 ? [...sealedBlocks, ...pendingChunks] : pendingChunks) sealedBlocks = [] let newline: number @@ -1369,10 +1402,10 @@ export class PythonCodeRuntime extends CodeRuntime { buffered = buffered.subarray(newline + 1) /* v8 ignore next -- an empty line comes only from a forged `\n\n` write. */ if (line.length === 0) continue - // No per-line cap check here: the unframed-buffer counter above - // already guarantees every line is within FRAME_PARSE_CAP_BYTES - // before this join runs, so a cap check on the line would be - // dead code. + // A later frame in this buffer may still exceed the cap; drop that + // single line like any junk frame (the first frame was already + // bounded by the check above). + if (line.length > FRAME_PARSE_CAP_BYTES) continue const text = line.toString('utf8') // JSON.parse would silently ROUND an integer token outside the // safe range before validation could see it, so a forged frame diff --git a/packages/code-runtime/code-runtime-python/tests/runtime.spec.ts b/packages/code-runtime/code-runtime-python/tests/runtime.spec.ts index 9e45577806..0accc86bfc 100644 --- a/packages/code-runtime/code-runtime-python/tests/runtime.spec.ts +++ b/packages/code-runtime/code-runtime-python/tests/runtime.spec.ts @@ -4847,35 +4847,25 @@ describe('PythonCodeRuntime — hostile peer', () => { expect(result.error?.message).toContain(`protocol frame exceeded ${ceiling} bytes`) }, 90_000) - it('rejects an over-ceiling fd-3 buffer without first joining it into one line', async () => { - // The ceiling has to be enforced on the byte COUNTER before Buffer.concat, + it('rejects an over-cap newline-free fd-3 buffer without first joining it into one line', async () => { + // The cap has to be enforced on the byte COUNTER before Buffer.concat, // not on the joined line afterwards: the join is a second copy of // everything held, so a program could force roughly twice the advertised - // 256 MiB of host memory before anything rejected it. + // 64 MiB of host memory before anything rejected it. // - // This program makes the two orders observably different rather than merely - // differently sized. It writes exactly the ceiling with no newline (at the - // limit, so nothing trips), then a newline followed by 8 MiB more. Checking - // the counter first sees more than the ceiling on the newline-bearing pipe - // chunk and rejects. Checking the joined line instead produced a FIRST LINE - // of exactly the ceiling — inside the per-line bound, so it passed as a junk - // frame — and left an 8 MiB residual well under the bound, so the breach was - // never reported: measured, the run settled as - // `python exited (code=0, signal=null) before completing` after the host had - // held the ceiling AND copied it, which is the doubling this check prevents. + // This program writes past the cap with no newline: the counter crosses on + // the 9th 8 MiB write (72 MiB) while the buffer is still a single unframed + // line, so the pre-join check rejects it without concat-ing a second copy. + // Checking the joined line instead would have produced a 72 MiB FIRST LINE + // that the per-line bound then dropped only after the doubling had happened. const ceiling = 64 * 1024 * 1024 const { runtime } = await setup({ maxWallMs: 60_000, addressSpaceMb: 2048 }) const result = await runtime.run({ program: [ 'import os', 'chunk = b"A" * (8 * 1024 * 1024)', - `for _ in range(${ceiling / (8 * 1024 * 1024)}):`, + `for _ in range(${ceiling / (8 * 1024 * 1024) + 1}):`, ' os.write(3, chunk)', - // One drain loop: a single os.write past the pipe buffer returns short, - // and a truncated tail would change which bytes cross the ceiling. - 'view = memoryview(b"\\n" + b"B" * (8 * 1024 * 1024))', - 'while view:', - ' view = view[os.write(3, view):]', 'return "never"', ].join('\n'), bindings: [], @@ -4885,4 +4875,30 @@ describe('PythonCodeRuntime — hostile peer', () => { expect(result.error?.message).toContain(`protocol frame exceeded ${ceiling} bytes`) }, 120_000) + it('keeps two within-cap frames whose combined buffer crosses the cap', async () => { + // The unframed byte counter charges the WHOLE buffer, which legitimately + // holds several frames each within FRAME_PARSE_CAP_BYTES. A first frame of + // exactly the cap followed by a second frame crosses the counter without + // either frame exceeding the cap; the first-frame check (not the counter) + // must let them through, or a legitimate near-cap frame plus a trailing + // frame would be misreported as a worker-exit. + const { runtime } = await setup({ maxWallMs: 60_000, addressSpaceMb: 2048 }) + const result = await runtime.run({ + program: [ + 'import os', + 'chunk = b"A" * (8 * 1024 * 1024)', + // Exactly the cap, no newline — at the limit, so nothing trips. + 'for _ in range(8):', + ' os.write(3, chunk)', + // A newline, then a small legitimate log frame. + 'os.write(3, b"\\n{\\"type\\":\\"log\\",\\"text\\":\\"after-cap-frames\\"}\\n")', + 'return "done"', + ].join('\n'), + bindings: [], + }) + expect(result.error).toBeUndefined() + expect(result.value).toBe('done') + expect(result.logs).toContain('after-cap-frames') + }, 120_000) + })