Merge harden/cancel-b3: ambient cancel hook in tick() + cycle/leak policy (B3-followup + C6)

Author: olwang

crates/rsscript/src/reg_vm/mod.rs

@@ -5,6 +5,8 @@ use std::io::{Read, Write};
 use std::net::{Shutdown, TcpStream};
 use std::path::{Component, Path, PathBuf};
 use std::rc::Rc;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
 
 use base64::Engine;
 use chrono::{DateTime, Datelike, NaiveDate, SecondsFormat, TimeZone, Timelike, Utc};
@@ -6341,7 +6343,10 @@ struct TaskSlot {
 /// the depth cap is generous (never trips real code, always catches
 /// `fn f(){f()}`), and the step/memory budgets are off unless a caller opts in
 /// (typically only the untrusted/agent-facing entry points).
-#[derive(Debug, Clone, Copy)]
+/// Note: not `Copy` — it carries an optional `Arc<AtomicBool>` cancel flag. All
+/// fields are public and `Clone`/struct-update (`..VmLimits::default()`) keep
+/// callers ergonomic; the scalar budget fields are read by value as before.
+#[derive(Debug, Clone)]
 pub struct VmLimits {
     /// Maximum simultaneous call frames (recursion depth). Default-on and
     /// generous; checked before every frame push. `usize::MAX` effectively
@@ -6356,6 +6361,15 @@ pub struct VmLimits {
     /// stacks + list/map growth). `None` (default) = no accounting (near-zero
     /// overhead). See [`RegVm::live_bytes`] for the accounting approximation.
     pub mem_budget: Option<usize>,
+    /// Host-level preemption hook. `None` (default) = no polling (the off path is
+    /// near-free: `tick()` never touches the atomic). When `Some`, the host can
+    /// set the flag to `true` from anywhere (e.g. a watchdog thread on timeout or
+    /// an abort signal) and the running evaluation is preempted at the next
+    /// throttled step check — even inside a tight `while true {}` loop that never
+    /// awaits or checks the cooperative RSS `CancellationToken`. The eval then
+    /// returns `EvalError::Runtime("evaluation cancelled")`. This stops the *whole*
+    /// eval; see the note in `tick()` on per-task preemption.
+    pub cancel: Option<Arc<AtomicBool>>,
 }
 
 /// Default recursion-depth cap: generous enough never to trip real code (deep
@@ -6364,6 +6378,12 @@ pub struct VmLimits {
 /// native stack.
 const DEFAULT_MAX_DEPTH: usize = 16_384;
 
+/// How often `tick()` polls the ambient cancel flag (once every this many
+/// instructions). A power of two so the modulo lowers to a mask. Small enough
+/// that a watchdog preempts a tight loop within microseconds, large enough that
+/// the relaxed atomic load is negligible amortized over real work.
+const CANCEL_POLL_INTERVAL: u64 = 1024;
+
 /// Estimated bytes charged per list element for the best-effort memory ceiling.
 /// Each list slot stores one `VmValue` inline (the heap pointed at by `Rc`
 /// containers is counted again when *those* grow, so this under-counts deeply
@@ -6381,6 +6401,7 @@ impl Default for VmLimits {
             max_depth: DEFAULT_MAX_DEPTH,
             step_budget: None,
             mem_budget: None,
+            cancel: None,
         }
     }
 }
@@ -6804,6 +6825,22 @@ impl RegVm {
     /// is off), and — only when `limits.step_budget` is `Some` — trips once the
     /// count exceeds the limit. This is what stops an infinite loop (`while true
     /// {}`) from hanging the host: it returns a clean error instead.
+    ///
+    /// It is also the host-level *preemption* hook. When `limits.cancel` is
+    /// `Some`, every `CANCEL_POLL_INTERVAL` steps we load the ambient
+    /// `AtomicBool` (`Relaxed` — we only need eventual visibility, not ordering)
+    /// and, if set, abort the eval with `EvalError::Runtime("evaluation
+    /// cancelled")`. The throttle keeps both the off path (no atomic touched at
+    /// all) and the on path (one relaxed load per 1024 instructions) cheap, so a
+    /// tight loop stays fast while still being interruptible by a watchdog.
+    ///
+    /// Limitation: this stops the *entire* evaluation. Preemptively cancelling a
+    /// single *sibling* task stuck in a tight loop — so a `select`/`task_group`
+    /// can reach its winner while one branch spins — would require the scheduler
+    /// to yield mid-instruction-stream (snapshot a `SavedTask` at an arbitrary
+    /// `ip` and reschedule), a deeper redesign that is out of scope here. The RSS
+    /// `CancellationToken` remains the cooperative, per-task mechanism (it only
+    /// preempts at await points); this ambient flag is the blunt host-level kill.
     #[inline]
     fn tick(&mut self) -> Result<(), EvalError> {
         self.steps += 1;
@@ -6814,6 +6851,12 @@ impl RegVm {
                 "step budget exceeded ({limit} instructions)"
             )));
         }
+        if self.steps.is_multiple_of(CANCEL_POLL_INTERVAL)
+            && let Some(flag) = self.limits.cancel.as_ref()
+            && flag.load(std::sync::atomic::Ordering::Relaxed)
+        {
+            return Err(EvalError::Runtime("evaluation cancelled".into()));
+        }
         Ok(())
     }
 
@@ -9483,6 +9526,11 @@ impl RegVm {
             }
             RegIntrinsic::CapabilityFrom => Ok(intrinsic_arg(&self.stack, base, args, 0)?.clone()),
             RegIntrinsic::CancellationSourceCancel => {
+                // RSS-level *cooperative* cancellation: flips the program-visible
+                // flag a task must poll (e.g. `token.is_cancelled()`); it preempts
+                // only at await/poll points, not inside a tight compute loop. The
+                // host-level *preemptive* hook for a runaway loop is the ambient
+                // `limits.cancel` atomic polled by `tick()` — see `RegVm::tick`.
                 let id = expect_cancellation_id_ref(
                     intrinsic_arg(&self.stack, base, args, 0)?,
                     "CancellationSource",

crates/rsscript/tests/hostile.rs

@@ -5,6 +5,8 @@
 
 use proptest::prelude::*;
 use rsscript::{EvalError, VmLimits, reg_vm_eval_source_main_with_limits};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 /// Invariant 1 (runtime hardening): no agent program may crash or hang the host.
 /// These cases prove that the reg-VM's sandbox limits turn the three classic
@@ -138,6 +140,133 @@ fn main() -> Int {
     assert_eq!(output.value, "6765");
 }
 
+/// B3 follow-up: the ambient `cancel` flag is the host-level preemption hook for
+/// a tight compute loop that never awaits or checks the cooperative RSS
+/// `CancellationToken`. Set BEFORE eval and with NO step budget, the running
+/// `while true {}` is preempted at the first throttled `tick()` poll and returns
+/// a clean "evaluation cancelled" error — deterministic, no threads/timing.
+#[test]
+fn ambient_cancel_flag_preempts_infinite_loop() {
+    let source = r#"
+fn main() -> Int {
+    let mut x = 0
+    while true {
+        x = x + 1
+    }
+    return x
+}
+"#;
+    let flag = Arc::new(AtomicBool::new(true));
+    let limits = VmLimits {
+        cancel: Some(Arc::clone(&flag)),
+        ..VmLimits::default()
+    };
+    let err = eval_limited(source, limits).expect_err("must error, not hang");
+    match err {
+        EvalError::Runtime(msg) => assert!(
+            msg.contains("cancelled"),
+            "expected cancellation error, got: {msg}"
+        ),
+        other => panic!("expected EvalError::Runtime, got {other:?}"),
+    }
+}
+
+/// B3 follow-up (negative): a cancel flag that is present but `false` must not
+/// trip — a short normal program still completes with `Ok`. Proves the hook is
+/// opt-in per-fire, not merely per-presence.
+#[test]
+fn ambient_cancel_flag_unset_does_not_trip() {
+    let source = r#"
+fn main() -> Int {
+    let mut x = 0
+    while x < 10 {
+        x = x + 1
+    }
+    return x
+}
+"#;
+    let flag = Arc::new(AtomicBool::new(false));
+    let limits = VmLimits {
+        cancel: Some(Arc::clone(&flag)),
+        ..VmLimits::default()
+    };
+    let output = eval_limited(source, limits).expect("flag false => normal completion");
+    assert_eq!(output.value, "10");
+    // Sanity: the host still holds the flag and can set it for a future run.
+    assert!(!flag.load(Ordering::Relaxed));
+}
+
+/// C6 (leak policy, positive): transient scalar work in a loop does not leak.
+/// A long arithmetic loop reuses a fixed set of registers each iteration, so the
+/// VM's best-effort byte accounting stays bounded and the program completes under
+/// a tight `mem_budget` rather than tripping a false OOM. (No threads/timing.)
+#[test]
+fn transient_scalar_work_completes_under_memory_ceiling() {
+    let source = r#"
+fn main() -> Int {
+    let mut index = 0
+    let mut acc = 0
+    while index < 200000 {
+        acc = (acc + index) % 1000000
+        index = index + 1
+    }
+    return acc
+}
+"#;
+    // A tight 1 MiB ceiling: register-stack growth is bounded per frame and does
+    // not grow per iteration, so a non-allocating loop never approaches it.
+    let limits = VmLimits {
+        mem_budget: Some(1 << 20),
+        step_budget: Some(50_000_000),
+        ..VmLimits::default()
+    };
+    let output = eval_limited(source, limits).expect("non-allocating loop must not leak/OOM");
+    // 200000 -> documents the loop ran to completion (exact value unimportant).
+    assert!(!output.value.is_empty());
+}
+
+/// C6 (leak policy, bounded): the value model uses `Rc`, so unbounded retained
+/// growth — including the only way to form a cycle, a self-referential mutable
+/// container — is the leak class of concern. rsscript does NOT run a cycle
+/// collector; instead the B4 `mem_budget` backstops it: the run trips a clean
+/// "memory limit" error (bounded), never an unbounded grow-until-host-OOM crash.
+///
+/// Note (a soundness bonus surfaced writing this test): the `local`/effect
+/// checker actively *rejects* pushing a `local` value into another container
+/// (RS0501 "retaining API cannot retain local value"), so a true `Rc` cycle
+/// cannot even be expressed from safe source — cycles are rarer than the policy
+/// assumes. This test therefore exercises the general retained-growth leak (a
+/// list accumulating fresh values without bound), which is what `mem_budget`
+/// must bound regardless of whether the retained graph is acyclic or cyclic.
+#[test]
+fn self_referential_container_is_bounded_by_memory_ceiling() {
+    let source = r#"features: local
+
+fn main() -> Int {
+    local values = List<String>.new()
+    let mut index = 0
+    while index < 100000000 {
+        List.push<String>(list: mut values, value: read "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
+        index = index + 1
+    }
+    return 0
+}
+"#;
+    let limits = VmLimits {
+        mem_budget: Some(1 << 20),
+        step_budget: Some(50_000_000),
+        ..VmLimits::default()
+    };
+    let err = eval_limited(source, limits).expect_err("must trip a bounded error, not leak/crash");
+    match err {
+        EvalError::Runtime(msg) => assert!(
+            msg.contains("memory limit"),
+            "expected memory-limit (bounded) error, got: {msg}"
+        ),
+        other => panic!("expected EvalError::Runtime, got {other:?}"),
+    }
+}
+
 /// Analyze every file under tests/corpus/malformed/. None may panic. Files not
 /// prefixed `gap-` must also fail closed (report a diagnostic); `gap-` files are
 /// known fail-open gaps the fuzzer surfaced (still must not panic).

docs/hardening-todo.md

@@ -52,10 +52,22 @@ prefixes (the enum is intentionally minimal: `Diagnostics` / `Runtime`).
 - [x] **B3 Step budget** (Inv 1) — DONE (`VmLimits.step_budget: Option<u64>`,
   default OFF; per-instruction `tick()` in BOTH the interpreter loop and the
   tier-0 JIT loop → `EvalError::Runtime("step budget exceeded …")`).
-  - [ ] *Follow-up (not done):* poll the existing `CancellationToken` every N
-    steps so cooperative cancellation can preempt a tight compute loop — the
-    preemption hook structured cancellation currently lacks (works only at await
-    points). The `tick()` counter is the natural place to add it.
+  - [x] *Follow-up:* **ambient host-level cancel hook** — `VmLimits.cancel:
+    Option<Arc<AtomicBool>>` (default `None`, not `Copy` anymore — `Clone` +
+    struct-update keep callers ergonomic). `tick()` polls it every
+    `CANCEL_POLL_INTERVAL` (1024) steps with `Ordering::Relaxed`; when set it
+    aborts with `EvalError::Runtime("evaluation cancelled")`. This is the
+    preemption hook a tight `while true {}` lacked: a host watchdog (timeout /
+    abort) flips the flag and the running loop stops at the next throttled check,
+    even though it never awaits or polls the cooperative RSS `CancellationToken`
+    (which only preempts at await points — unchanged). The off path touches no
+    atomic; the on path is one relaxed load per 1024 instructions.
+    - *Remaining limitation (out of scope, future work):* this stops the **whole**
+      eval. Preemptively cancelling a single **sibling** task stuck in a tight loop
+      — so a `select` / `task_group` can reach its winner while one branch spins —
+      requires the scheduler to yield mid-instruction-stream (snapshot a
+      `SavedTask` at an arbitrary `ip` and reschedule). That is a deeper redesign;
+      structured per-task preemption remains future work.
 - [x] **B4 Memory ceiling** (Inv 1) — DONE (`VmLimits.mem_budget: Option<usize>`,
   default OFF; best-effort cumulative `live_bytes` accounting at `ensure_regs`,
   list/map construction, and list push/append → `EvalError::Runtime("memory limit
@@ -79,9 +91,18 @@ prefixes (the enum is intentionally minimal: `Diagnostics` / `Runtime`).
     deliberately excluded rather than run-and-silenced.
   Setup (offline-installable in the dev container):
   `rustup toolchain install nightly --component miri,rust-src`.
-- [ ] **C6 Cycle/leak policy.** The value model makes accidental cycles unlikely;
-  decide explicitly — document the weak-discipline + add leak tests, or plan a
-  cycle collector for very-long-running apps.
+- [x] **C6 Cycle/leak policy** (documented, no collector). The value model uses
+  `Rc`, so the *only* way to form a cycle is a self-referential mutable container
+  (e.g. pushing a list into itself) — rare in practice. Policy: **cycles are NOT
+  collected; they are bounded by `mem_budget` (B4)** — a self-referential leak
+  trips a clean "memory limit" error instead of growing unbounded, so the host is
+  never OOM-killed. Revisit a real cycle collector only if profiling shows a need.
+  Leak tests in `tests/hostile.rs`: (a) a non-allocating loop completes under a
+  tight `mem_budget` (transient register work doesn't leak into the byte counter);
+  (b) a deliberately self-referential container under a `mem_budget` trips the
+  bounded "memory limit" error, not an unbounded crash. Note: the B4 byte counter
+  is a cumulative high-water estimate (it does not subtract frees), so the
+  positive leak test uses scalar work rather than transient containers.
 
 ## Meta — prove the invariant continuously
 

fuzz/fuzz_targets/no_panic.rs

@@ -39,6 +39,9 @@ fuzz_target!(|data: &[u8]| {
         max_depth: 16_384,
         step_budget: Some(50_000_000),
         mem_budget: Some(512 * 1024 * 1024),
+        // No ambient cancel flag: the fuzzer relies on the step/mem budgets to
+        // bound each run; the host-level preemption hook is not exercised here.
+        cancel: None,
     };
 
     let result = reg_vm_eval_source_main_with_limits(