diff --git a/packages/code-runtime/code-runtime-python/src/protocol.ts b/packages/code-runtime/code-runtime-python/src/protocol.ts index 10d13a211c..a2aaecb5be 100644 --- a/packages/code-runtime/code-runtime-python/src/protocol.ts +++ b/packages/code-runtime/code-runtime-python/src/protocol.ts @@ -189,20 +189,23 @@ function scalarJson(current: unknown): string { } /** - * Meter a forged done value's compact-JSON byte length AND its number - * losslessness in one bounded traversal, stopping the instant `maxBytes` is - * crossed. A forged `done.value` arrives straight off fd 3 and can sit anywhere - * below the 256 MiB frame ceiling while `maxValueBytes` defaults to 32 KiB. The - * previous split — an unbounded `hasNonLosslessNumber` scan in - * {@link validateChildFrame} followed by a separate byte meter — pushed every - * member of a wide flat payload onto a scan stack before any cap check ran, so - * a below-ceiling forgery could still force a hundreds-of-megabytes host - * allocation. Folding both jobs here rejects over-budget BEFORE enqueuing an - * array's or object's children, keeping the traversal O(cap). A non-lossless - * number (non-finite, negative zero) is caught only when the value fits the - * budget — an over-budget value is rejected regardless, so the distinction is - * moot. Same JSON-plain precondition and traversal shape as - * {@link encodeJsonPlain}; per-scalar byte length is measured through + * Meter a `JSON.parse`-produced done value's compact-JSON byte length AND its + * number losslessness in one traversal, stopping the instant `maxBytes` is + * crossed. This bounds the INCREMENTAL allocation the check itself would add on + * top of the already-parsed value — the escaped-string copy, the enqueued + * children, the per-key `JSON.stringify` — not the parse that produced `value`. + * That upstream width is bounded separately: the host reads fd 3 into a fixed + * 256 MiB receive buffer (a later stack layer), so `value` cannot already be + * larger than that when it reaches here, while `maxValueBytes` defaults to + * 32 KiB. The traversal rejects over-budget BEFORE materializing a string's + * escaped form or enqueuing an array's/object's children, so a below-ceiling + * forgery cannot force those secondary allocations. Object key COUNTING is + * unavoidably O(keys) — JS has no lazy own-key iterator, and the parse already + * built the key set — but the check still refuses the per-entry work before the + * enqueue loop. A non-lossless number (non-finite, negative zero) is caught only + * when the value fits the budget — an over-budget value is rejected regardless, + * so the distinction is moot. Same JSON-plain precondition and traversal shape + * as {@link encodeJsonPlain}; per-scalar byte length is measured through * {@link scalarJson} (matching the encoder, so a beyond-safe-range integer * meters its exact BigInt digits, not `JSON.stringify`'s rounded spelling) and * `JSON.stringify` for strings. @@ -230,30 +233,23 @@ export function checkDoneValue(value: unknown, maxBytes: number): { ok: true; by } else if (Array.isArray(current)) { // Brackets plus one comma per gap; elements add themselves. Reject // BEFORE enqueuing children: every element serializes to at least one - // byte, so a forged flat array below the frame ceiling but far above - // the budget fails here without growing the host stack by millions of - // entries first. + // byte, so a forged flat array far above the budget fails here without + // pushing its elements onto the host stack. (The array itself is already + // materialized by the upstream parse; this only bounds the extra stack.) bytes += 2 + (current.length > 1 ? current.length - 1 : 0) if (bytes + current.length > maxBytes) return { ok: false, reason: 'over-budget' } for (const item of current) stack.push(item) } else if (typeof current === 'object' && current !== null) { const record = current as Record - // Count own keys WITHOUT Object.entries/Object.keys (either allocates one - // slot per member up front), AND bail mid-count the instant the minimum - // encoding exceeds the budget: braces (+2), each entry a quoted key - // (>= 2 bytes) + colon + >= 1-byte value (>= 4 bytes), and a comma per - // gap. A forged wide object with millions of keys and a small cap must - // fail in O(cap), not walk its whole breadth first. `bytes` still holds - // the pre-object total throughout this loop. + // Count own keys with for...in + hasOwn. This IS O(keys) — JS has no lazy + // own-key iterator and the parse already built the key set — so the count + // cannot be sublinear; what the bound below buys is refusing the per-entry + // work (key escaping, value enqueue) before it runs. Each entry costs at + // least a quoted key (>= 2 bytes) + colon + >= 1-byte value. let count = 0 - for (const key in record) { - if (!Object.hasOwn(record, key)) continue - count += 1 - if (bytes + 2 + count * 4 + (count - 1) > maxBytes) return { ok: false, reason: 'over-budget' } - } - // The loop's final iteration already proved the whole object's lower - // bound fits, so no separate post-count check is needed here. + for (const key in record) if (Object.hasOwn(record, key)) count += 1 bytes += 2 + (count > 1 ? count - 1 : 0) + if (bytes + count * 4 > maxBytes) return { ok: false, reason: 'over-budget' } for (const key in record) { if (!Object.hasOwn(record, key)) continue // The same string lower bound, before escaping the key. diff --git a/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts b/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts index f0e06af723..b57ae178c6 100644 --- a/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts +++ b/packages/code-runtime/code-runtime-python/tests/protocol.spec.ts @@ -194,34 +194,23 @@ describe('checkDoneValue', () => { } }) - it('stops early on a huge value instead of measuring it whole', () => { + it('rejects an over-budget value before its secondary allocations', () => { + // A huge string is refused on the cheap length lower bound, before its + // escaped copy is built. const huge = { data: 'x'.repeat(1_000_000), tail: 'y' } expect(checkDoneValue(huge, 1024)).toEqual({ ok: false, reason: 'over-budget' }) - // A forged flat array below the frame ceiling must fail BEFORE its - // elements are enqueued — the pre-enqueue bound keeps the walk O(cap). + // A flat array far above the budget fails on the brackets+length bound, + // before its elements are pushed onto the traversal stack. (The array is + // already materialized by the upstream parse; this only avoids the extra + // per-element stack growth.) const flat = new Array(10_000_000).fill(0) expect(checkDoneValue(flat, 1024)).toEqual({ ok: false, reason: 'over-budget' }) - // Same bound for a wide object: braces+commas fit the cap, but the - // per-entry lower bound (quoted key + colon + value) does not, so it fails - // before any key is metered or any value enqueued. + // A wide object: braces+commas fit the cap, but the per-entry lower bound + // (quoted key + colon + value = count*4) does not, so it fails before any + // key is escaped or any value enqueued. const wide: Record = {} for (let i = 0; i < 10; i++) wide[`k${i}`] = i expect(checkDoneValue(wide, 12)).toEqual({ ok: false, reason: 'over-budget' }) - // A forged object with millions of keys and a small cap must reject in - // O(cap): the key COUNT loop itself bails once the running minimum encoding - // (braces + 4 bytes/entry + commas) crosses the budget, rather than walking - // the whole breadth before checking. Observable as a bounded key subset: - // build a Proxy whose ownKeys would yield far more than the cap admits and - // assert the metered walk never enumerates past it. - let enumerated = 0 - const millionKeys = new Proxy({}, { - ownKeys() { return Array.from({ length: 2_000_000 }, (_unused, i) => `k${i}`) }, - getOwnPropertyDescriptor() { enumerated += 1; return { enumerable: true, configurable: true, value: 0 } }, - }) - expect(checkDoneValue(millionKeys, 64)).toEqual({ ok: false, reason: 'over-budget' }) - // With cap 64, at most ~16 entries (4 bytes each) can fit before the bound - // trips, so the walk enumerates far fewer than the 2,000,000 declared keys. - expect(enumerated).toBeLessThan(1000) }) it('rejects an over-budget string on its length before escaping it', () => {