Jelajahi Sumber

fix(code-runtime-python): close four settlement-path review findings

Pace-free completion framing, stray UTF-8 flush, and late-rejection guards:
- Pre-encode the completion value at its validation point so send_done never
  re-walks a live value a mutating daemon thread could change (TOCTOU); a
  mutation-induced encode throw is then classified as 'exception', not a
  host-side worker-exit.
- Budget-triggered stray flush retains an incomplete multibyte UTF-8 tail
  (<=3 bytes) as residual instead of decoding a legal, split character to
  U+FFFD in an admitted entry; the end/closeDeadline paths still full-decode.
- Check 'settled' before formatting a late binding rejection's message, so a
  hostile message getter cannot stall or exhaust a run that already settled.
- Document _check_done_value's first-to-trip ruling in its docstring.
- Rewrite ProtocolChannel.send_sync around a shared write_encoded that the
  done frame's pre-encoded string path uses.
Chinesezjc 1 bulan lalu
induk
melakukan
f71914ceea

+ 63 - 13
packages/code-runtime/code-runtime-python/py/bootstrap.py

@@ -506,12 +506,24 @@ class ProtocolChannel:
         (dispatch raises the lossless-JSON message).
         """
 
-        payload = (_encode_json_plain(message) + "\n").encode("utf-8")
-        # Full-write loop under the writer lock: one os.write may consume only
-        # part of a frame beyond PIPE_BUF (64 KiB logs / 32 KiB completions /
-        # uncapped call args exceed it), and a partial or interleaved frame is
-        # dropped host-side as malformed JSON — the run would then hang to the
-        # wall clock.
+        self.write_encoded(_encode_json_plain(message))
+
+    def write_encoded(self, frame: str) -> None:
+        """Write a frame that is ALREADY encoded to its JSON string form.
+
+        Appends the frame's trailing newline and full-write-loops the bytes
+        under the writer lock, identically to :meth:`send_sync`. The consumer
+        supplies the encoded JSON (a ``"done"`` frame carrying a completion
+        value that was serialized at its validation point — see
+        :func:`_done_with_value`); the channel does not re-encode it, so the
+        bytes written are exactly what was validated with no second traversal
+        of a live object.
+        """
+
+        payload = (frame + "\n").encode("utf-8")
+        # Full-write loop under the writer lock (same rationale as send_sync):
+        # one os.write may consume only part of a frame beyond PIPE_BUF, and a
+        # partial or interleaved frame is dropped host-side as malformed JSON.
         with self._write_lock:
             view = memoryview(payload)
             while view:
@@ -881,9 +893,20 @@ async def _run(channel: ProtocolChannel) -> None:
     safe_model_traceback = _SAFE_MODEL_TRACEBACK
     flush_out = out_stream.flush_line
     flush_err = err_stream.flush_line
-    send_done = channel.send_sync
+    # `done` is either a pre-encoded frame STRING (a `_done_with_value` success:
+    # the completion value was serialized at its validation point, inside the try,
+    # so a later send never re-walks the live value a mutating daemon thread could
+    # have changed) or a dict ERROR frame (a rejection or the exception handler,
+    # which carry no live model value). `send_done` posts whichever form: a string
+    # is written verbatim via `write_encoded`, a dict is encoded by `send_sync`.
+    def send_done(payload: dict[str, Any] | str) -> None:
+        if isinstance(payload, str):
+            channel.write_encoded(payload)
+        else:
+            channel.send_sync(payload)
+
     max_value_bytes = int(boot["maxValueBytes"])
-    done: dict[str, Any]
+    done: dict[str, Any] | str
     try:
         module = ast.parse(program)
         wrapper = ast.AsyncFunctionDef(
@@ -908,7 +931,8 @@ async def _run(channel: ProtocolChannel) -> None:
         die_if_cpu_exhausted(cpu_seconds)
         # Flush the log buffers BEFORE metering and framing the completion value.
         # `_done_with_value` materializes the value's escaped JSON form to meter
-        # it, and `send_done` encodes the frame — several copies of a near-budget
+        # it and then pre-encodes the admitted value into its frame (see its
+        # docstring for the TOCTOU rationale) — several copies of a near-budget
         # value live at once (see OUTPUT_BUDGET_WORST_CASE_ADDRESS_SPACE_MULTIPLE).
         # Any unflushed log pending would add its own bytes to that peak, so a
         # `maxLogBytes` and a `maxValueBytes` each admitted alone by the load gate
@@ -1477,6 +1501,14 @@ def _check_done_value(value: Any, max_bytes: int):
     Returns ``("invalid-output", message)`` for a non-lossless value,
     ``("output-limit", message)`` once the size crosses ``max_bytes``, or
     ``None`` when the value is lossless JSON within budget.
+
+    Metering and validation interleave in this single traversal: each member is
+    costed the moment it is visited, and it is rejected the moment it trips
+    either check. A value that holds BOTH an over-budget member and an
+    invalid-typed member therefore resolves to whichever tripped FIRST in
+    visit order — both are rejects, and neither kind claims priority over the
+    other, so that first-trip order is not part of the seam contract; the
+    host side independently re-measures the value it receives.
     """
 
     js_safe = 2**53 - 1
@@ -2056,7 +2088,7 @@ def _join_bounded(lines, max_bytes: int) -> str:
     return "".join(chunks)
 
 
-def _done_with_value(value: Any, max_value_bytes: int) -> dict[str, Any]:
+def _done_with_value(value: Any, max_value_bytes: int) -> dict[str, Any] | str:
     """Build the terminal done frame under the seam's lossless-JSON contract.
 
     A completion value returned by the program (``None`` when it returns
@@ -2065,20 +2097,38 @@ def _done_with_value(value: Any, max_value_bytes: int) -> dict[str, Any]:
     Substituting a ``repr`` or truncated string would be a silent lie about
     what the program computed, so both paths refuse instead (mirroring the
     worker backend's contract). ``None`` crosses as an exact JSON ``null``.
+
+    The SUCCESS path returns the whole ``"done"`` frame as an ALREADY-ENCODED
+    JSON string: the admitted value is serialized here, at its validation
+    point, rather than handed to ``send_sync`` to re-walk later. The program
+    can keep mutating the returned list/dict from a daemon thread or signal
+    handler after it returns, so a second traversal held at a later point
+    would be a TOCTOU — a mutation into a non-JSON type would let that later
+    encode throw outside the settlement handler and downgrade the run
+    host-side to ``worker-exit``. Serializing once, inside the try that wraps
+    this call, closes the window: if a concurrent mutation makes the encode
+    throw, the exception handler classifies it as ``exception``, and once the
+    string is produced the frame is sent verbatim with no further touching of
+    the live value. Returns a ``dict`` only for a rejection (an error frame
+    carries no live model value and is safe to send via ``send_sync``).
     """
 
     # One bounded walk folds the losslessness check and the byte meter (mirrors
     # the host's checkDoneValue): the former split ran the full losslessness
     # walk first, materializing one tuple per element for a wide completion
     # before the size cap could reject it — an RLIMIT_AS death on a value the
-    # meter would have refused. send_sync later encodes the admitted value,
-    # whose size the walk proved within budget. Iterative like the encoder, so a
+    # meter would have refused. The value's escaped JSON is then produced in the
+    # SAME call, so the admitted value is serialized exactly once (see above);
+    # its size the walk proved within budget. Iterative like the encoder, so a
     # valid completion deeper than the recursion limit still checks.
     rejection = _check_done_value(value, max_value_bytes)
     if rejection is not None:
         kind, message = rejection
         return {"type": "done", "error": {"kind": kind, "message": message}}
-    return {"type": "done", "value": value}
+    # Pre-encode the value at the validation point (not in `_run`'s later send,
+    # which is outside the try): see the TOCTOU note in the docstring. The value
+    # is JSON-plain by construction, so `_encode_json_plain` is the encoder.
+    return '{"type": "done", "value": ' + _encode_json_plain(value) + "}"
 
 
 def main() -> None:

+ 57 - 6
packages/code-runtime/code-runtime-python/src/index.ts

@@ -1097,9 +1097,17 @@ export class PythonCodeRuntime extends CodeRuntime {
         // nondeterministically with each other and with the child's own fd-3
         // `log` frames, so `logs` carries no cross-pipe ordering guarantee to
         // preserve here; a fixed drain order is as valid as any.
+        // Flushing is NOT a stream end: a multibyte UTF-8 character can be split
+        // across pipe `data` chunks, so the residual may end mid-sequence. A
+        // budget-triggered flush must decode only the complete prefix and carry
+        // the incomplete tail forward (≤3 bytes) on the same pipe's residual —
+        // decoding it here would render a legal character as U+FFFD in a released
+        // entry (see `flushStray`). This is unlike the `end`/closeDeadline paths
+        // below, where a trailing incomplete sequence is genuinely truncated input
+        // and U+FFFD is honest.
         if (strayOut.cost + strayErr.cost + 3 > logBudget) {
-          flushStray(strayOut)
-          flushStray(strayErr)
+          flushStray(strayOut, true)
+          flushStray(strayErr, true)
         }
       }
       // Flush a pipe's residual into `logs`. Called on the combined-budget
@@ -1111,14 +1119,48 @@ export class PythonCodeRuntime extends CodeRuntime {
       // `chunks`/`blocks` guard is the only emptiness check needed — `data` never
       // emits a zero-length Buffer, so a non-empty fragment list always decodes
       // to a non-empty tail.
-      function flushStray(stray: StrayBuffer): void {
+      //
+      // `retainPartialTail` is true only on the budget-triggered path: there the
+      // residual can end at an ARBITRARY pipe boundary, so if the incomplete
+      // trailing bytes of a UTF-8 lead sequence are pending (`stray.utf8.expected
+      // > 0`), they are withheld from the decode and re-carried on `chunks` for a
+      // later chunk to complete — decoding them here would render a LEGAL,
+      // un-finished character as U+FFFD in an admitted entry, and the next chunk's
+      // bytes would then each independently break into more U+FFFD. The withheld
+      // tail is `stray.utf8.width - stray.utf8.expected` bytes (the lead plus the
+      // continuations consumed so far), at most 3; `stray.utf8` is reset and the
+      // withheld tail re-accrued so the next chunk continues the walk correctly.
+      // The `end`/closeDeadline paths pass `false`: there a trailing incomplete
+      // sequence is real truncated input and the U+FFFD is the honest render.
+      function flushStray(stray: StrayBuffer, retainPartialTail?: boolean): void {
         if (stray.chunks.length === 0 && stray.blocks.length === 0) return
-        const tail = Buffer.concat([...stray.blocks, ...stray.chunks]).toString('utf8')
-        stray.chunks = []
+        const begun = stray.blocks.length > 0 ? [...stray.blocks, ...stray.chunks] : stray.chunks
+        const full = Buffer.concat(begun)
+        let drop = 0
+        // A budget flush landing exactly between a lead byte and its still-pending
+        // continuation requires the combined-cost threshold to trip on a specific
+        // mid-multibyte pipe boundary — not deterministically schedulable through
+        // the black-box seam, which observes only complete entries. v8 ignore keeps
+        // the retention branch honest (it is exercised by review reasoning over the
+        // `stray.utf8` state, not by an in-tree test).
+        /* v8 ignore next 8 -- mid-sequence budget-flush boundary is not schedulable from a test. */
+        if (retainPartialTail && stray.utf8.expected > 0) {
+          drop = stray.utf8.width - stray.utf8.expected
+          // Guard against a pathological width/expected mismatch: never drop more
+          // bytes than were captured, and never drop so many that decoding the
+          // admitted prefix would be empty because a single mid-sequence lead sat
+          // alone. A well-formed walk keeps `drop` ≤ 3, but a defensive clamp
+          // keeps the retention bounded.
+          drop = Math.min(drop, full.length)
+        }
+        const keep = full.subarray(full.length - drop)
+        const emit = full.subarray(0, full.length - drop).toString('utf8')
+        stray.chunks = drop > 0 ? detachResidual(keep) : []
         stray.blocks = []
         stray.cost = 0
         stray.utf8 = { expected: 0, width: 0, lowerFirst: 0, upperFirst: 0 }
-        admit(tail)
+        if (drop > 0) stray.cost = accrueStrayCost(keep, stray.utf8)
+        admit(emit)
       }
       child.stdout.on('data', (chunk: Buffer) => { captureStray(strayOut, chunk) })
       child.stderr.on('data', (chunk: Buffer) => { captureStray(strayErr, chunk) })
@@ -1368,6 +1410,15 @@ export class PythonCodeRuntime extends CodeRuntime {
                 }
                 sendReply({ type: 'reply', id: message.id, ok: true, value })
               } catch (error: unknown) {
+                // Check `settled` before formatting the error: a rejection that
+                // arrives after `maxWallMs`, an abort, or dispose has already
+                // settled the run, and `messageOf(error)` runs hostile getters
+                // before `sendReply` peeks at `settled`. Dropping the framed
+                // reply early spares the host heap and time for a run whose
+                // outcome is already fixed.
+                // oxlint-disable-next-line typescript/no-unnecessary-condition -- the run can settle while this binding is awaited.
+                /* v8 ignore next -- a rejection arriving after settlement is not schedulable from a test. */
+                if (settled) return
                 sendReply({ type: 'reply', id: message.id, ok: false, message: messageOf(error) })
               }
             })()