diff --git a/crates/bashkit/src/limits.rs b/crates/bashkit/src/limits.rs index 254b3d9a..07ce732a 100644 --- a/crates/bashkit/src/limits.rs +++ b/crates/bashkit/src/limits.rs @@ -409,8 +409,8 @@ impl ExecutionCounters { Ok(()) }); - self.commands += 1; - self.session_commands += 1; + self.commands = self.commands.saturating_add(1); + self.session_commands = self.session_commands.saturating_add(1); if self.commands > limits.max_commands { return Err(LimitExceeded::MaxCommands(limits.max_commands)); } @@ -437,7 +437,7 @@ impl ExecutionCounters { /// Increment exec call counter for session tracking. pub fn tick_exec_call(&mut self) { - self.session_exec_calls += 1; + self.session_exec_calls = self.session_exec_calls.saturating_add(1); } /// Increment loop iteration counter, returns error if limit exceeded @@ -924,6 +924,31 @@ mod tests { )); } + #[test] + fn test_command_counter_saturates_on_overflow() { + let limits = ExecutionLimits::new().max_commands(5); + let mut counters = ExecutionCounters::new(); + counters.commands = usize::MAX; + counters.session_commands = u64::MAX; + + assert!(matches!( + counters.tick_command(&limits), + Err(LimitExceeded::MaxCommands(5)) + )); + assert_eq!(counters.commands, usize::MAX); + assert_eq!(counters.session_commands, u64::MAX); + } + + #[test] + fn test_exec_counter_saturates_on_overflow() { + let mut counters = ExecutionCounters::new(); + counters.session_exec_calls = u64::MAX; + + counters.tick_exec_call(); + + assert_eq!(counters.session_exec_calls, u64::MAX); + } + #[test] fn test_loop_counter() { let limits = ExecutionLimits::new().max_loop_iterations(3); diff --git a/crates/bashkit/tests/integration/snapshot_tests.rs b/crates/bashkit/tests/integration/snapshot_tests.rs index 5eb776d6..b8980247 100644 --- a/crates/bashkit/tests/integration/snapshot_tests.rs +++ b/crates/bashkit/tests/integration/snapshot_tests.rs @@ -607,6 +607,32 @@ async fn snapshot_restore_does_not_reset_session_exec_limit_with_tampered_counte ); } +#[tokio::test] +async fn snapshot_restore_extreme_exec_counter_errors_without_overflow() { + let session_limits = SessionLimits::new().max_exec_calls(2); + let mut bash = Bash::builder().session_limits(session_limits).build(); + bash.exec("echo first").await.unwrap(); + let bytes = bash.snapshot().unwrap(); + + let mut tampered_json: serde_json::Value = serde_json::from_slice(&bytes[32..]).unwrap(); + tampered_json["session_exec_calls"] = serde_json::json!(u64::MAX); + let tampered_snapshot: Snapshot = serde_json::from_value(tampered_json).unwrap(); + let tampered_bytes = tampered_snapshot.to_bytes().unwrap(); + + bash.restore_snapshot(&tampered_bytes).unwrap(); + let result = bash.exec("echo must-not-wrap").await; + assert!( + result.is_err(), + "restored u64::MAX exec counter must not panic or wrap below the limit" + ); + let err_str = result.unwrap_err().to_string(); + assert!( + err_str.contains("session exec() call limit"), + "error must be session exec limit, got: {err_str}" + ); + assert_eq!(bash.session_counters().1, u64::MAX); +} + #[tokio::test] async fn keyed_snapshot_restore_carries_session_exec_budget_forward() { let key = b"session-budget-hmac-key";