From 8156feab527d48f159994d7eb9e789b8b29ec82d Mon Sep 17 00:00:00 2001 From: Alexandre MAI Date: Sat, 29 Aug 2026 19:40:13 +0200 Subject: [PATCH 1/2] fix: account for generator handler allocations Charge try-handler save and restore operations against resource limits and make sandbox limit errors terminal to avoid continuing from partial generator state. --- crates/zapcode-core/src/sandbox.rs | 16 +++- crates/zapcode-core/src/vm/mod.rs | 32 +++++++ crates/zapcode-core/tests/error_handling.rs | 93 ++++++++++++++++++++- 3 files changed, 137 insertions(+), 4 deletions(-) diff --git a/crates/zapcode-core/src/sandbox.rs b/crates/zapcode-core/src/sandbox.rs index c6e3322..dafda12 100644 --- a/crates/zapcode-core/src/sandbox.rs +++ b/crates/zapcode-core/src/sandbox.rs @@ -50,10 +50,22 @@ impl ResourceTracker { } pub fn track_allocation(&mut self, limits: &ResourceLimits) -> crate::error::Result<()> { - self.allocations += 1; - if self.allocations > limits.max_allocations { + self.track_allocations(limits, 1) + } + + pub fn track_allocations( + &mut self, + limits: &ResourceLimits, + count: usize, + ) -> crate::error::Result<()> { + let allocations = self + .allocations + .checked_add(count) + .ok_or(crate::ZapcodeError::AllocationLimitExceeded)?; + if allocations > limits.max_allocations { return Err(crate::ZapcodeError::AllocationLimitExceeded); } + self.allocations = allocations; Ok(()) } diff --git a/crates/zapcode-core/src/vm/mod.rs b/crates/zapcode-core/src/vm/mod.rs index 8f45e05..25f238b 100644 --- a/crates/zapcode-core/src/vm/mod.rs +++ b/crates/zapcode-core/src/vm/mod.rs @@ -235,6 +235,16 @@ impl Vm { } fn try_handle_error(&mut self, err: &ZapcodeError, min_frame_depth: usize) -> Result { + if matches!( + err, + ZapcodeError::MemoryLimitExceeded(_) + | ZapcodeError::TimeLimitExceeded + | ZapcodeError::StackOverflow(_) + | ZapcodeError::AllocationLimitExceeded + ) { + return Ok(false); + } + let Some(try_info) = self.take_try_handler(min_frame_depth) else { return Ok(false); }; @@ -1053,6 +1063,12 @@ impl Vm { self.run_generator_until_yield_or_return(gen_obj) } Some(suspended) => { + let handler_count = self + .generator_try_handlers + .get(&gen_obj.id) + .map_or(0, Vec::len); + self.tracker + .track_allocations(&self.limits, handler_count)?; self.tracker.push_frame(); let stack_base = self.stack.len(); for val in &suspended.stack { @@ -1134,6 +1150,21 @@ impl Vm { } let instr = instructions[frame.ip].clone(); if matches!(instr, Instruction::Yield) { + let frame_depth = self.frames.len(); + let frame_stack_base = self.current_frame().stack_base; + let handler_count = self + .try_stack + .iter() + .filter(|try_info| try_info.frame_depth == frame_depth) + .count(); + let frame_stack_len = self.stack.len().saturating_sub(frame_stack_base + 1); + let map_entry_count = usize::from(handler_count > 0); + self.tracker.track_allocations( + &self.limits, + handler_count + .saturating_add(frame_stack_len) + .saturating_add(map_entry_count), + )?; self.current_frame_mut().ip += 1; let yielded_value = self.pop()?; let frame_depth = self.frames.len(); @@ -2087,6 +2118,7 @@ impl Vm { // Error handling Instruction::SetupTry(catch_ip, _) => { + self.tracker.track_allocation(&self.limits)?; self.try_stack.push(TryInfo { catch_ip, frame_depth: self.frames.len(), diff --git a/crates/zapcode-core/tests/error_handling.rs b/crates/zapcode-core/tests/error_handling.rs index e0e4547..598e888 100644 --- a/crates/zapcode-core/tests/error_handling.rs +++ b/crates/zapcode-core/tests/error_handling.rs @@ -1,5 +1,94 @@ -use zapcode_core::vm::eval_ts; -use zapcode_core::Value; +use zapcode_core::vm::{eval_ts, VmState}; +use zapcode_core::{ResourceLimits, Value, ZapcodeError, ZapcodeRun}; + +#[test] +fn test_try_handler_respects_allocation_limit() { + let runner = ZapcodeRun::new( + "try {} catch (e) {}".to_string(), + Vec::new(), + Vec::new(), + ResourceLimits { + max_allocations: 0, + ..ResourceLimits::default() + }, + ) + .unwrap(); + + assert!(matches!( + runner.start(Vec::new()), + Err(ZapcodeError::AllocationLimitExceeded) + )); +} + +#[test] +fn test_generator_handler_allocation_limit_is_not_catchable() { + let code = r#" + function* values() { + try { + try { + try { + try { + try { yield 1; } catch (e) {} + } catch (e) {} + } catch (e) {} + } catch (e) {} + } catch (e) {} + } + const iterator = values(); + iterator.next(); + let outcome = 1; + try { iterator.next(); } catch (e) { outcome = 999; } + outcome + "#; + + let run_with_limit = |max_allocations| { + let runner = ZapcodeRun::new( + code.to_string(), + Vec::new(), + Vec::new(), + ResourceLimits { + max_allocations, + ..ResourceLimits::default() + }, + ) + .unwrap(); + runner.start(Vec::new()) + }; + + assert!(matches!( + run_with_limit(27), + Err(ZapcodeError::AllocationLimitExceeded) + )); + assert!(matches!( + run_with_limit(512), + Ok(VmState::Complete(Value::Int(1))) + )); +} + +#[test] +fn test_stack_overflow_is_not_catchable() { + let runner = ZapcodeRun::new( + r#" + function recurse() { recurse(); } + let caught = false; + try { recurse(); } catch (e) { caught = true; } + caught + "# + .to_string(), + Vec::new(), + Vec::new(), + ResourceLimits { + max_stack_depth: 8, + ..ResourceLimits::default() + }, + ) + .unwrap(); + + assert!(matches!( + runner.start(Vec::new()), + Err(ZapcodeError::StackOverflow(_)) + )); +} #[test] fn test_try_catch() { From 7b29b9fb82579e2424c9f7da5a1a7bc3e524117d Mon Sep 17 00:00:00 2001 From: Alexandre MAI Date: Sat, 29 Aug 2026 20:14:55 +0200 Subject: [PATCH 2/2] fix: enforce generator stack limits Check max_stack_depth before allocating fresh or resumed generator frames and cover both paths with bounded regressions. --- crates/zapcode-core/src/vm/mod.rs | 2 + crates/zapcode-core/tests/generators.rs | 51 ++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/crates/zapcode-core/src/vm/mod.rs b/crates/zapcode-core/src/vm/mod.rs index 25f238b..a1b4917 100644 --- a/crates/zapcode-core/src/vm/mod.rs +++ b/crates/zapcode-core/src/vm/mod.rs @@ -1025,6 +1025,7 @@ impl Vm { None => { let func = &self.program.functions[func_idx]; self.tracker.push_frame(); + self.tracker.check_stack(&self.limits)?; let mut locals = Vec::with_capacity(func.local_count); for param in func.params.iter() { match param { @@ -1070,6 +1071,7 @@ impl Vm { self.tracker .track_allocations(&self.limits, handler_count)?; self.tracker.push_frame(); + self.tracker.check_stack(&self.limits)?; let stack_base = self.stack.len(); for val in &suspended.stack { self.push(val.clone())?; diff --git a/crates/zapcode-core/tests/generators.rs b/crates/zapcode-core/tests/generators.rs index 46eaac5..dcb9f2a 100644 --- a/crates/zapcode-core/tests/generators.rs +++ b/crates/zapcode-core/tests/generators.rs @@ -1,10 +1,59 @@ -use zapcode_core::Value; +use zapcode_core::{ResourceLimits, Value, ZapcodeError, ZapcodeRun}; /// Helper to run TS code and get stdout + value fn eval_with_output(code: &str) -> (Value, String) { zapcode_core::vm::eval_ts_with_output(code).unwrap() } +#[test] +fn test_generator_next_respects_stack_limit() { + let runner = ZapcodeRun::new( + r#" + function* values() { yield 1; } + values().next() + "# + .to_string(), + Vec::new(), + Vec::new(), + ResourceLimits { + max_stack_depth: 0, + ..ResourceLimits::default() + }, + ) + .unwrap(); + + assert!(matches!( + runner.start(Vec::new()), + Err(ZapcodeError::StackOverflow(_)) + )); +} + +#[test] +fn test_resumed_generator_respects_stack_limit() { + let runner = ZapcodeRun::new( + r#" + function* values() { yield 1; yield 2; } + const iterator = values(); + iterator.next(); + function resume() { return iterator.next(); } + resume() + "# + .to_string(), + Vec::new(), + Vec::new(), + ResourceLimits { + max_stack_depth: 1, + ..ResourceLimits::default() + }, + ) + .unwrap(); + + assert!(matches!( + runner.start(Vec::new()), + Err(ZapcodeError::StackOverflow(_)) + )); +} + #[test] fn test_basic_generator_yield() { let (val, _) = eval_with_output(