diff --git a/crates/bashkit/src/builtins/python.rs b/crates/bashkit/src/builtins/python.rs index 3c0af82f..e2361283 100644 --- a/crates/bashkit/src/builtins/python.rs +++ b/crates/bashkit/src/builtins/python.rs @@ -30,7 +30,7 @@ use std::future::Future; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use super::{Builtin, Context, resolve_path}; use crate::error::Result; @@ -151,6 +151,10 @@ impl PythonLimits { /// /// Receives `(function_name, positional_args, keyword_args)` directly from monty. /// Return `ExtFunctionResult::Return(value)` for success or `ExtFunctionResult::Error(exc)` for failure. +/// +/// Important security decision: the Python builtin wraps each awaited handler +/// in the remaining [`PythonLimits::max_duration`] budget. Host handlers are +/// trusted code, but should still enforce their own I/O and service limits. pub type PythonExternalFnHandler = Arc< dyn Fn( String, @@ -493,6 +497,9 @@ async fn run_python( .max_recursion_depth(Some(py_limits.max_recursion)); let tracker = LimitedTracker::new(limits); + // Important security decision: cap awaited host callbacks with the same wall-clock + // budget as Monty so external functions cannot pin execution between VM steps. + let python_deadline = Instant::now().checked_add(py_limits.max_duration); // Run the synchronous start() phase, then extract collected output. // PrintWriter::CollectString is not Send, so we scope it to avoid holding across .await. @@ -521,10 +528,12 @@ async fn run_python( } RunProgress::FunctionCall(call) => { let result = if let Some(ef) = external_fns { - (ef.handler)( + call_external_with_deadline( + ef, call.function_name.clone(), call.args.clone(), call.kwargs.clone(), + python_deadline, ) .await } else { @@ -594,12 +603,44 @@ async fn run_python( } } +fn python_external_timeout_error() -> ExtFunctionResult { + ExtFunctionResult::Error(MontyException::new( + ExcType::RuntimeError, + Some("Python external function exceeded max_duration".into()), + )) +} + +fn remaining_python_budget(deadline: Option) -> Option { + let deadline = deadline?; + deadline.checked_duration_since(Instant::now()) +} + +async fn call_external_with_deadline( + external_fns: &PythonExternalFns, + function_name: String, + args: Vec, + kwargs: Vec<(MontyObject, MontyObject)>, + deadline: Option, +) -> ExtFunctionResult { + let Some(remaining) = remaining_python_budget(deadline) else { + return python_external_timeout_error(); + }; + + match tokio::time::timeout( + remaining, + (external_fns.handler)(function_name, args, kwargs), + ) + .await + { + Ok(result) => result, + Err(_) => python_external_timeout_error(), + } +} + // --------------------------------------------------------------------------- // VFS bridging: Monty OsCall → Bashkit FileSystem // --------------------------------------------------------------------------- -/// Dispatch a Monty OsCall to the appropriate VFS operation. -/// /// Monty's `OsFunctionCall` is a tagged enum carrying typed args. We project it /// to the generic `(positional, keyword)` `MontyObject` view via the public /// `to_args()` and dispatch on the stable `name()` string (kept byte-identical @@ -1505,13 +1546,22 @@ mod tests { code: &str, fn_names: &[&str], handler: PythonExternalFnHandler, + ) -> ExecResult { + run_with_external_and_limits(code, fn_names, handler, PythonLimits::default()).await + } + + async fn run_with_external_and_limits( + code: &str, + fn_names: &[&str], + handler: PythonExternalFnHandler, + limits: PythonLimits, ) -> ExecResult { let args = vec!["-c".to_string(), code.to_string()]; let env = opt_in_env(); let mut variables = HashMap::new(); let mut cwd = PathBuf::from("/home/user"); let fs = Arc::new(InMemoryFs::new()); - let py = Python::with_limits(PythonLimits::default()) + let py = Python::with_limits(limits) .with_external_handler(fn_names.iter().map(|s| s.to_string()).collect(), handler); let ctx = Context::new_for_test(&args, &env, &mut variables, &mut cwd, fs, None); py.execute(ctx).await.unwrap() @@ -1527,6 +1577,33 @@ mod tests { assert_eq!(r.stdout, "42\n"); } + #[tokio::test] + async fn test_external_fn_handler_timeout_uses_python_deadline() { + let handler: PythonExternalFnHandler = Arc::new(|_name, _args, _kwargs| { + Box::pin(async { + tokio::time::sleep(Duration::from_secs(1)).await; + ExtFunctionResult::Return(MontyObject::Int(1)) + }) + }); + let started = Instant::now(); + let r = tokio::time::timeout( + Duration::from_secs(1), + run_with_external_and_limits( + "print(slow())", + &["slow"], + handler, + PythonLimits::default().max_duration(Duration::from_millis(25)), + ), + ) + .await + .expect("Python external call should observe max_duration"); + + assert_eq!(r.exit_code, 1); + assert!(r.stderr.contains("RuntimeError")); + assert!(r.stderr.contains("exceeded max_duration")); + assert!(started.elapsed() < Duration::from_millis(500)); + } + #[tokio::test] async fn test_external_fn_with_args() { let handler: PythonExternalFnHandler = Arc::new(|_name, args, _kwargs| { diff --git a/specs/python-builtin.md b/specs/python-builtin.md index 29f8c134..0f2ed59d 100644 --- a/specs/python-builtin.md +++ b/specs/python-builtin.md @@ -257,9 +257,16 @@ let bash = Bash::builder() **Dispatch:** A single handler receives all registered function names; dispatch by `function_name` inside the handler. +**Timeouts:** Each awaited handler call is wrapped in the remaining +`PythonLimits::max_duration` wall-clock budget for the current Python invocation. +If the budget expires while a handler is pending, Bashkit resumes Python with a +`RuntimeError` instead of waiting for the handler indefinitely. + **Trust model:** External function handlers follow the same trust model as `BashBuilder::builtin()` and `ScriptedTool` callbacks — the host application -registers trusted Rust code, untrusted scripts invoke it by name. +registers trusted Rust code, untrusted scripts invoke it by name. Handlers are +trusted host code and should still enforce independent limits for outbound I/O, +remote services, and other resources they consume. **Unstable re-exports:** `MontyObject`, `ExtFunctionResult`, `MontyException`, and `ExcType` are re-exported from the `monty` crate (git-pinned, not on crates.io).