kobako 0.2.1 → 0.3.0

data/ext/kobako/src/wasm/host_state.rs

@@ -4,13 +4,17 @@
 //! and threaded through every host import — the `__kobako_dispatch`
 //! dispatcher reads the server handle, while the run-path methods on
 //! [`crate::wasm::Instance`] install fresh WASI context + pipes before
-//! every `#run` (SPEC.md B-03 / B-04).
+//! every `#run` (docs/behavior.md B-03 / B-04).
 //!
-//! The state also carries the per-run wall-clock deadline (SPEC.md B-01,
-//! E-19) and the linear-memory cap [`KobakoLimiter`] (SPEC.md B-01,
-//! E-20). Both are read from the wasmtime `epoch_deadline_callback` /
-//! `ResourceLimiter` callbacks installed in
-//! [`crate::wasm::Instance::from_path`].
+//! The state also carries the per-invocation wall-clock deadline
+//! (docs/behavior.md B-01, E-19) and the per-invocation linear-memory
+//! delta cap [`KobakoLimiter`] (docs/behavior.md B-01, E-20). Both are
+//! read from the wasmtime `epoch_deadline_callback` / `ResourceLimiter`
+//! callbacks installed in [`crate::wasm::Instance::from_path`]. The
+//! memory cap measures only the `memory.grow` delta past the linear-
+//! memory size captured at invocation entry — the mruby image's
+//! initial allocation and prior invocations' watermark are outside the
+//! budget.
 
 use std::cell::{Ref, RefCell, RefMut};
 use std::time::Instant;
@@ -41,7 +45,7 @@ impl HostState {
     /// Build a fresh per-Store host state. `memory_limit` carries the
     /// `Sandbox#memory_limit` cap in bytes (or `None` to disable the cap);
     /// it is read from the wasmtime [`ResourceLimiter`] callback every
-    /// time the guest grows linear memory (SPEC.md B-01, E-20).
+    /// time the guest grows linear memory (docs/behavior.md B-01, E-20).
     pub(super) fn new(memory_limit: Option<usize>) -> Self {
         Self {
             wasi: None,
@@ -54,8 +58,9 @@ impl HostState {
     }
 
     /// Install a freshly-built WASI context plus the matching stdout/stderr
-    /// pipe clones. Called from [`crate::wasm::Instance::run`] at the top
-    /// of every guest invocation (SPEC.md B-03 / B-04).
+    /// pipe clones. Called from [`crate::wasm::Instance::eval`] /
+    /// [`crate::wasm::Instance::run`] at the top of every guest
+    /// invocation (docs/behavior.md B-03 / B-04).
     pub(super) fn install_wasi(
         &mut self,
         wasi: WasiP1Ctx,
@@ -100,17 +105,18 @@ impl HostState {
 
     /// Mutable handle to the live WASI context. Panics if no context has
     /// been installed yet — every call site is downstream of
-    /// [`HostState::install_wasi`] running at the top of `Instance::run`,
-    /// so reaching this branch with `None` signals a host-side wiring bug.
+    /// [`HostState::install_wasi`] running at the top of
+    /// `Instance::eval` / `Instance::run`, so reaching this branch with
+    /// `None` signals a host-side wiring bug.
     pub(super) fn wasi_mut(&mut self) -> &mut WasiP1Ctx {
-        self.wasi
-            .as_mut()
-            .expect("WASI context not initialised — call Instance#run before any WASI use")
+        self.wasi.as_mut().expect(
+            "WASI context not initialised — call Instance#eval / Instance#run before any WASI use",
+        )
     }
 
     /// Replace the per-run wall-clock deadline. `Some(at)` makes the
     /// epoch-deadline callback trap once `Instant::now() >= at`; `None`
-    /// disables the cap. Called at the top of every `#run` (SPEC.md B-01).
+    /// disables the cap. Called at the top of every `#run` (docs/behavior.md B-01).
     pub(super) fn set_deadline(&mut self, deadline: Option<Instant>) {
         self.deadline = deadline;
     }
@@ -121,36 +127,61 @@ impl HostState {
         self.deadline
     }
 
-    /// Mutable handle to the embedded [`KobakoLimiter`]. Shared by the
-    /// wasmtime [`ResourceLimiter`] callback (set once at Store build
-    /// time) and by [`crate::wasm::Instance`] for arming / disarming the
-    /// memory cap around each guest run. Same shape as
-    /// [`HostState::wasi_mut`] — callers operate on the inner type
-    /// directly instead of going through a per-action passthrough.
+    /// Mutable handle to the embedded [`KobakoLimiter`]. Required by
+    /// the wasmtime [`ResourceLimiter`] callback wiring in
+    /// [`crate::wasm::Instance::from_path`]
+    /// (`store.limiter(|state| state.limiter_mut())`); kept private to
+    /// the wasm submodule so the only public surface for arming the
+    /// cap goes through [`HostState::arm_memory_cap`] /
+    /// [`HostState::disarm_memory_cap`].
     pub(super) fn limiter_mut(&mut self) -> &mut KobakoLimiter {
         &mut self.limiter
     }
+
+    /// Arm the docs/behavior.md E-20 memory cap for one guest run with
+    /// the current linear-memory size as the baseline. The limiter
+    /// charges only the `memory.grow` delta past `baseline` against
+    /// the cap, so the mruby image's initial allocation and the
+    /// high-water mark left by prior invocations do not consume the
+    /// budget. Paired with [`HostState::disarm_memory_cap`] around the
+    /// call to the corresponding `__kobako_*` export so post-run host
+    /// bookkeeping (e.g. fetching the OUTCOME_BUFFER) is not
+    /// attributed to the user script.
+    pub(super) fn arm_memory_cap(&mut self, baseline: usize) {
+        self.limiter.activate(baseline);
+    }
+
+    /// Disarm the docs/behavior.md E-20 memory cap. See
+    /// [`HostState::arm_memory_cap`].
+    pub(super) fn disarm_memory_cap(&mut self) {
+        self.limiter.deactivate();
+    }
 }
 
-/// Resource limiter that enforces the `memory_limit` cap from SPEC.md
-/// B-01 / E-20 on every guest `memory.grow`.
+/// Resource limiter that enforces the per-invocation `memory_limit`
+/// cap from docs/behavior.md B-01 / E-20.
 ///
-/// `max_memory` is the byte cap (`None` disables the cap). `cap_active`
-/// gates whether the cap is enforced — wasmtime's `ResourceLimiter`
-/// fires for both the module's declared initial allocation and every
-/// subsequent `memory.grow`, but SPEC.md E-20 scopes the trap to
-/// `memory.grow` specifically. [`KobakoLimiter::activate`] /
-/// [`KobakoLimiter::deactivate`] flip the flag for the lifetime of an
-/// `Instance::run` call. When `cap_active` is `false`, the limiter
-/// always allows growth.
+/// `max_memory` is the byte cap on per-invocation growth (`None` disables
+/// the cap). `baseline` is the linear-memory size captured at invocation
+/// entry by [`KobakoLimiter::activate`]; the limiter charges only the
+/// `memory.grow` delta past `baseline` against `max_memory`, so the
+/// mruby image's initial allocation and any high-water mark left by
+/// prior invocations on the same Sandbox do not consume the budget.
+/// `cap_active` gates whether the cap is enforced — wasmtime's
+/// `ResourceLimiter` also fires for the module's declared initial
+/// allocation at instantiation time, but the cap stays dormant until
+/// [`KobakoLimiter::activate`] flips the flag for one
+/// `Instance::eval` / `Instance::run` call. When `cap_active` is
+/// `false`, the limiter always allows growth.
 ///
-/// When `memory.grow` would push linear memory past the cap, the
-/// limiter returns [`MemoryLimitTrap`] from `memory_growing`; wasmtime
-/// turns that into the trap surfaced to the host as `__kobako_run`
-/// failure.
+/// When `memory.grow` would push the per-invocation delta past
+/// `max_memory`, the limiter returns [`MemoryLimitTrap`] from
+/// `memory_growing`; wasmtime turns that into the trap surfaced to the
+/// host as a guest invocation failure.
 #[derive(Debug, Clone, Copy)]
 pub(super) struct KobakoLimiter {
     max_memory: Option<usize>,
+    baseline: usize,
     cap_active: bool,
 }
 
@@ -158,18 +189,20 @@ impl KobakoLimiter {
     fn new(max_memory: Option<usize>) -> Self {
         Self {
             max_memory,
+            baseline: 0,
             cap_active: false,
         }
     }
 
-    /// Arm the cap so subsequent `memory.grow` calls are checked
-    /// against `memory_limit`. The cap is dormant by default — the
-    /// module's declared initial memory is allocated during
-    /// `Linker::instantiate` and SPEC.md E-20 scopes the trap to
-    /// `memory.grow` (not the instantiation-time initial allocation).
-    /// [`crate::wasm::Instance::run`] calls this right before
-    /// `__kobako_run`.
-    pub(super) fn activate(&mut self) {
+    /// Arm the cap so subsequent `memory.grow` calls are charged
+    /// against `max_memory` starting from `baseline` bytes. Called via
+    /// [`HostState::arm_memory_cap`] at the top of every invocation;
+    /// the cap is dormant by default — the module's declared initial
+    /// memory is allocated during `Linker::instantiate` and the
+    /// per-invocation budget excludes anything that existed before
+    /// arming (docs/behavior.md B-01 Notes, E-20).
+    fn activate(&mut self, baseline: usize) {
+        self.baseline = baseline;
         self.cap_active = true;
     }
 
@@ -177,7 +210,7 @@ impl KobakoLimiter {
     /// OUTCOME_BUFFER, which can grow guest memory transiently) is
     /// not attributed to the user script. Paired with
     /// [`KobakoLimiter::activate`].
-    pub(super) fn deactivate(&mut self) {
+    fn deactivate(&mut self) {
         self.cap_active = false;
     }
 }
@@ -193,7 +226,8 @@ impl ResourceLimiter for KobakoLimiter {
             return Ok(true);
         }
         if let Some(limit) = self.max_memory {
-            if desired > limit {
+            let delta = desired.saturating_sub(self.baseline);
+            if delta > limit {
                 return Err(wasmtime::Error::new(MemoryLimitTrap { desired, limit }));
             }
         }
@@ -211,7 +245,7 @@ impl ResourceLimiter for KobakoLimiter {
 }
 
 /// Marker error returned from [`KobakoLimiter::memory_growing`] on
-/// SPEC.md E-20. Downcast from the wasmtime trap error to surface as
+/// docs/behavior.md E-20. Downcast from the wasmtime trap error to surface as
 /// `Kobako::Wasm::MemoryLimitError` on the Ruby side. Callers use the
 /// `Display` impl below — no field is read directly — so the inner
 /// state stays private.
@@ -221,6 +255,17 @@ pub(crate) struct MemoryLimitTrap {
     limit: usize,
 }
 
+impl MemoryLimitTrap {
+    /// Construct a trap with the given +desired+ / +limit+ pair. Used
+    /// internally by [`KobakoLimiter::memory_growing`] in production and
+    /// by the sibling-module +classify_trap+ unit tests to materialise
+    /// a representative error for downcast routing.
+    #[cfg(test)]
+    pub(super) fn new(desired: usize, limit: usize) -> Self {
+        Self { desired, limit }
+    }
+}
+
 impl std::fmt::Display for MemoryLimitTrap {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
@@ -233,9 +278,9 @@ impl std::fmt::Display for MemoryLimitTrap {
 
 impl std::error::Error for MemoryLimitTrap {}
 
-/// Marker error returned from the epoch-deadline callback on SPEC.md
-/// E-19. Downcast from the wasmtime trap error to surface as
-/// `Kobako::Wasm::TimeoutError` on the Ruby side.
+/// Marker error returned from the epoch-deadline callback on
+/// docs/behavior.md E-19. Downcast from the wasmtime trap error to
+/// surface as `Kobako::Wasm::TimeoutError` on the Ruby side.
 #[derive(Debug)]
 pub(crate) struct TimeoutTrap;
 
@@ -279,8 +324,100 @@ impl StoreCell {
     }
 }
 
-// SAFETY: Ruby's GVL serialises access to magnus-wrapped objects on a single
-// OS thread at a time. `wasmtime::Store` is `Send` (verified upstream); the
-// `RefCell`-mediated mutation is therefore safe under the GVL invariant.
+// SAFETY: magnus requires `Send + Sync` on `#[magnus::wrap]` types. Both
+// claims hold under the GVL invariant:
+//
+//   * Send — `wasmtime::Store<HostState>` is itself `Send` (verified
+//     upstream by wasmtime; see `wasmtime::Store`'s trait impls).
+//     `RefCell<T>: Send` whenever `T: Send`. The remaining stored state
+//     (`HostState`) holds `Opaque<Value>` for the Ruby Server handle —
+//     `Opaque<Value>` is documented as `Send` by magnus precisely so
+//     wrapped objects can satisfy this bound.
+//
+//   * Sync — `RefCell` is *not* `Sync` in the general Rust sense
+//     (concurrent `borrow_mut` is UB). We assert `Sync` here because the
+//     GVL serialises every call into Ruby C and every entry into magnus-
+//     wrapped methods onto a single OS thread at a time: by the time the
+//     `Sync` bound matters, magnus has already established that only one
+//     thread can be inside the wrapper. Cross-thread mutation cannot
+//     occur. If a future magnus release adopts a thread model that
+//     permits concurrent access to wrapped objects, this assertion would
+//     have to revert and `StoreCell` would need to switch to
+//     `Mutex<wasmtime::Store<…>>` — but as of magnus 0.8 the contract
+//     holds.
 unsafe impl Send for StoreCell {}
 unsafe impl Sync for StoreCell {}
+
+#[cfg(test)]
+mod tests {
+    //! Unit tests for [`KobakoLimiter`] — the per-invocation memory
+    //! delta cap. The Ruby-facing E2E suite exercises the full path
+    //! through wasmtime; these tests pin the pure delta arithmetic so
+    //! a regression that breaks the baseline accounting (e.g. dropping
+    //! the `baseline` subtraction, or letting `activate` carry stale
+    //! state across invocations) is caught without spinning up a
+    //! Store.
+    use super::{KobakoLimiter, MemoryLimitTrap};
+    use wasmtime::ResourceLimiter;
+
+    fn assert_growing(limiter: &mut KobakoLimiter, desired: usize) {
+        assert!(
+            limiter.memory_growing(0, desired, None).unwrap(),
+            "expected memory_growing({desired}) to allow growth"
+        );
+    }
+
+    fn assert_trapping(limiter: &mut KobakoLimiter, desired: usize) {
+        let err = limiter
+            .memory_growing(0, desired, None)
+            .expect_err("expected memory_growing to trap");
+        assert!(
+            err.downcast_ref::<MemoryLimitTrap>().is_some(),
+            "expected MemoryLimitTrap, got {err:?}"
+        );
+    }
+
+    #[test]
+    fn dormant_limiter_allows_any_growth() {
+        let mut limiter = KobakoLimiter::new(Some(1 << 20));
+        // Without `activate`, the cap is dormant — the module's
+        // declared initial allocation must pass through unconditionally.
+        assert_growing(&mut limiter, 100 << 20);
+    }
+
+    #[test]
+    fn delta_below_cap_passes_after_activate() {
+        let mut limiter = KobakoLimiter::new(Some(1 << 20));
+        limiter.activate(2 << 20);
+        // baseline=2 MiB, desired=2.5 MiB → delta=0.5 MiB ≤ 1 MiB cap.
+        assert_growing(&mut limiter, (2 << 20) + (1 << 19));
+    }
+
+    #[test]
+    fn delta_past_cap_traps_with_memory_limit_trap() {
+        let mut limiter = KobakoLimiter::new(Some(1 << 20));
+        limiter.activate(2 << 20);
+        // baseline=2 MiB, desired=4 MiB → delta=2 MiB > 1 MiB cap.
+        assert_trapping(&mut limiter, 4 << 20);
+    }
+
+    #[test]
+    fn activate_resets_baseline_on_each_invocation() {
+        let mut limiter = KobakoLimiter::new(Some(1 << 20));
+        limiter.activate(2 << 20);
+        assert_growing(&mut limiter, (2 << 20) + (1 << 20));
+        // Second invocation: linear memory has grown to 3 MiB. Re-arming
+        // must re-anchor the baseline so the next 1 MiB of growth fits
+        // the per-invocation budget rather than being charged against
+        // the prior invocation's residue.
+        limiter.activate(3 << 20);
+        assert_growing(&mut limiter, (3 << 20) + (1 << 20));
+    }
+
+    #[test]
+    fn disabled_cap_ignores_delta_size() {
+        let mut limiter = KobakoLimiter::new(None);
+        limiter.activate(0);
+        assert_growing(&mut limiter, 100 << 20);
+    }
+}