improv_extfn/
lib.rs

1//! External-language function runtime for Improv (Phase 6).
2//!
3//! This crate is the analogue of the engine's in-process scalar registry
4//! (`improv_engine`'s `scalar_arity`/`scalar_func`) for NON-builtin functions:
5//! it evaluates a registered external function on typed
6//! [`improv_core_model::Value`] arguments, deterministically and in a
7//! subprocess sandbox. The engine calls it host-side, outside the dataflow
8//! (`improv_engine::external::refresh_external_measure`), so external calls
9//! never run on the incremental hot path.
10//!
11//! # Invariant (see `AGENT_MASTER_STEERING.md` §7)
12//!
13//! External functions must be **pure**, return a **typed** value, and declare
14//! **dimensionality**, so they behave as ordinary operators and keep the engine
15//! deterministic. This runtime is deterministic *given the same inputs + body*;
16//! purity of the user's function body is the user's contract (the engine treats
17//! external calls as pure). Dimensionality is declared on the descriptor
18//! (`arg_types` / `return_type`); higher-arity dimension broadcasting is the
19//! engine's job at the call site, not this runtime's.
20//!
21//! # Language runtimes
22//!
23//! `eval` dispatches on `f.language` to a per-language runner:
24//! * **Python / R / Julia / Pure** — subprocess runners (shared plumbing in
25//!   the `runner` module): the interpreter reads a generated program on stdin
26//!   and prints one `{"ok":..}` / `{"error":..}` JSON envelope on stdout. A
27//!   wall-clock timeout kills a runaway child. If the interpreter binary is
28//!   absent, `eval` returns [`ExtFnError::LanguageUnavailable`] — it never
29//!   panics.
30//! * **Wasm** — an in-process `wasmi` interpreter (no subprocess); numeric f64
31//!   ABI (see the `wasm` module).
32//!
33//! # Sandboxing and its limits
34//!
35//! Python is invoked as `python3 -I -S -` (isolated, no site, program on stdin);
36//! R as `Rscript --vanilla -`; Julia as `julia --startup-file=no -`; Pure as
37//! `pure -q`. In every case the program arrives on **stdin**, so the function
38//! body is never interpolated into a shell command line (no shell is spawned).
39//! A wall-clock **timeout** (default 5s) kills a runaway child.
40//!
41//! On top of that, every subprocess runner is spawned through the `sandbox`
42//! module's OS sandbox (see its docs): on Linux, a `bwrap` (bubblewrap)
43//! filesystem/network sandbox plus `setrlimit` ceilings (CPU time, address
44//! space, open files, process count). It **fails closed**: without `bwrap`
45//! (always the case on macOS and Windows) a `Restricted` body is refused with
46//! [`ExtFnError::Sandbox`] unless the operator sets
47//! `IMPROV_EXTFN_ALLOW_UNSANDBOXED=1`, which falls back to rlimits only. The
48//! Wasm runtime runs in-process via `wasmi`, which has no syscall access at
49//! all by construction, independent of any OS mechanism.
50//!
51//! The host, never the model, picks the sandbox strictness: every body runs
52//! [`SandboxPolicy::Restricted`] unless the operator sets
53//! `IMPROV_EXTFN_TRUSTED=1` (see `sandbox::policy_for`). [`ExternalFn::pure`]
54//! is model data and only affects caching.
55#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
56// The attribute is only used inside #[cfg(test)] modules, so non-test
57// coverage builds would otherwise warn that the feature is unused.
58#![cfg_attr(coverage_nightly, allow(unused_features))]
59
60mod julia;
61mod marshal;
62mod pure;
63mod python;
64mod r;
65mod runner;
66mod sandbox;
67mod wasm;
68
69pub use marshal::{error_value, value_to_json};
70pub use sandbox::SandboxPolicy;
71
72use improv_core_model::{Value, ValueType};
73use std::collections::HashMap;
74use std::time::Duration;
75
76// `ExternalFn`/`Language` are the model-level *definitions* (plain, serializable
77// data); this crate is their *runtime*. Re-exported so callers can use
78// `improv_extfn::{ExternalFn, Language}` interchangeably.
79pub use improv_core_model::{ExternalFn, Language};
80
81/// Name -> descriptor. The engine will resolve a `FuncId`/name to a descriptor
82/// here for non-builtin calls.
83#[derive(Debug, Default, Clone)]
84pub struct Registry {
85    fns: HashMap<String, ExternalFn>,
86    /// Wall-clock kill deadline for a single evaluation.
87    timeout: Duration,
88}
89
90impl Registry {
91    pub fn new() -> Self {
92        Registry {
93            fns: HashMap::new(),
94            timeout: Duration::from_secs(5),
95        }
96    }
97
98    /// Override the per-call timeout (default 5s).
99    pub fn with_timeout(mut self, timeout: Duration) -> Self {
100        self.timeout = timeout;
101        self
102    }
103
104    pub fn timeout(&self) -> Duration {
105        self.timeout
106    }
107
108    /// Register (or replace) a function. Returns the previous descriptor if the
109    /// name was already registered.
110    pub fn register(&mut self, f: ExternalFn) -> Option<ExternalFn> {
111        self.fns.insert(f.name.clone(), f)
112    }
113
114    pub fn get(&self, name: &str) -> Option<&ExternalFn> {
115        self.fns.get(name)
116    }
117
118    pub fn len(&self) -> usize {
119        self.fns.len()
120    }
121
122    pub fn is_empty(&self) -> bool {
123        self.fns.is_empty()
124    }
125
126    /// Evaluate a registered function by name on typed arguments.
127    pub fn eval(&self, name: &str, args: &[Value]) -> Result<Value, ExtFnError> {
128        let f = self
129            .fns
130            .get(name)
131            .ok_or_else(|| ExtFnError::NotFound(name.to_string()))?;
132        eval(f, args, self.timeout)
133    }
134}
135
136/// Evaluate one descriptor on typed arguments with an explicit timeout.
137///
138/// Validates arity and argument types (no subprocess needed for those), then
139/// dispatches to the language runtime and type-checks the returned value
140/// against `return_type`.
141pub fn eval(f: &ExternalFn, args: &[Value], timeout: Duration) -> Result<Value, ExtFnError> {
142    if args.len() != f.arity() {
143        return Err(ExtFnError::ArityMismatch {
144            name: f.name.clone(),
145            expected: f.arity(),
146            got: args.len(),
147        });
148    }
149    for (i, (arg, &ty)) in args.iter().zip(&f.arg_types).enumerate() {
150        if !marshal::arg_matches(arg, ty) {
151            return Err(ExtFnError::TypeMismatch {
152                name: f.name.clone(),
153                position: i,
154                expected: ty,
155                got: arg.type_of(),
156            });
157        }
158    }
159    match f.language {
160        Language::Python => python::eval(f, args, timeout),
161        Language::R => r::eval(f, args, timeout),
162        Language::Julia => julia::eval(f, args, timeout),
163        Language::Wasm => wasm::eval(f, args, timeout),
164        Language::Pure => pure::eval(f, args, timeout),
165    }
166}
167
168/// Errors from resolving or evaluating an external function. The engine will
169/// convert these to a `Value::Error` at the call site (see [`error_value`]);
170/// this API returns `Result` so the caller decides.
171#[derive(Debug, thiserror::Error)]
172pub enum ExtFnError {
173    #[error("external function not found: {0}")]
174    NotFound(String),
175
176    #[error("function {name}: expected {expected} argument(s), got {got}")]
177    ArityMismatch {
178        name: String,
179        expected: usize,
180        got: usize,
181    },
182
183    #[error(
184        "function {name}: argument {position} type mismatch, expected {expected:?}, got {got:?}"
185    )]
186    TypeMismatch {
187        name: String,
188        position: usize,
189        expected: ValueType,
190        got: Option<ValueType>,
191    },
192
193    #[error("language runtime unavailable: {0}")]
194    LanguageUnavailable(String),
195
196    #[error("runtime error: {message}\n--- stderr ---\n{stderr}")]
197    Runtime { message: String, stderr: String },
198
199    #[error("function timed out after {0:?}")]
200    Timeout(Duration),
201
202    /// The OS sandbox the host policy requires is unavailable (fail closed).
203    #[error("sandbox unavailable: {0}")]
204    Sandbox(String),
205}
206
207#[cfg(test)]
208#[cfg_attr(coverage_nightly, coverage(off))]
209mod tests {
210    use super::*;
211    use improv_core_model::Value;
212
213    fn add_fn() -> ExternalFn {
214        ExternalFn {
215            name: "add".into(),
216            language: Language::Python,
217            body: "result = args[0] + args[1]".into(),
218            arg_types: vec![ValueType::Number, ValueType::Number],
219            return_type: ValueType::Number,
220            pure: true,
221        }
222    }
223
224    #[test]
225    fn registry_roundtrip() {
226        let mut reg = Registry::new();
227        assert_eq!(reg.timeout(), Duration::from_secs(5));
228        assert_eq!(
229            Registry::new()
230                .with_timeout(Duration::from_millis(7))
231                .timeout(),
232            Duration::from_millis(7)
233        );
234        assert!(reg.is_empty());
235        assert!(reg.register(add_fn()).is_none());
236        assert_eq!(reg.len(), 1);
237        let got = reg.get("add").expect("registered");
238        assert_eq!(got.arity(), 2);
239        assert_eq!(got.return_type, ValueType::Number);
240        assert!(got.pure);
241        // Re-register returns the previous descriptor.
242        assert!(reg.register(add_fn()).is_some());
243        assert!(reg.get("missing").is_none());
244    }
245
246    #[test]
247    fn not_found() {
248        let reg = Registry::new();
249        assert!(matches!(
250            reg.eval("nope", &[]),
251            Err(ExtFnError::NotFound(_))
252        ));
253    }
254
255    #[test]
256    fn arity_mismatch_without_python() {
257        let f = add_fn();
258        let err = eval(&f, &[Value::Number(1.0)], Duration::from_secs(2)).unwrap_err();
259        match err {
260            ExtFnError::ArityMismatch { expected, got, .. } => {
261                assert_eq!(expected, 2);
262                assert_eq!(got, 1);
263            }
264            other => panic!("expected ArityMismatch, got {other:?}"),
265        }
266    }
267
268    #[test]
269    fn type_mismatch_without_python() {
270        let f = add_fn();
271        let err = eval(
272            &f,
273            &[Value::Number(1.0), Value::Text("nope".into())],
274            Duration::from_secs(2),
275        )
276        .unwrap_err();
277        match err {
278            ExtFnError::TypeMismatch {
279                position, expected, ..
280            } => {
281                assert_eq!(position, 1);
282                assert_eq!(expected, ValueType::Number);
283            }
284            other => panic!("expected TypeMismatch, got {other:?}"),
285        }
286    }
287
288    // ---- Python-dependent tests (guarded) ----
289
290    /// True if `python3` runs under the host sandbox policy. Python tests
291    /// early-return (logging a skip) when false, so CI without Python, or
292    /// without `bwrap` (where `Restricted` fails closed), still passes.
293    fn python_available() -> bool {
294        crate::runner::run_interpreter(
295            "python3",
296            &["-I", "-S", "-"],
297            "print(1)",
298            Duration::from_secs(5),
299            "python",
300            crate::sandbox::policy_for(),
301        )
302        .is_ok_and(|o| o.stdout.trim() == "1")
303    }
304
305    #[test]
306    fn python_add_numbers() {
307        if !python_available() {
308            println!("skipped: python3 not found");
309            return;
310        }
311        let mut reg = Registry::new().with_timeout(Duration::from_secs(2));
312        reg.register(add_fn());
313        let out = reg
314            .eval("add", &[Value::Number(2.0), Value::Number(3.0)])
315            .expect("eval");
316        assert_eq!(out, Value::Number(5.0));
317    }
318
319    #[test]
320    fn python_text_function() {
321        if !python_available() {
322            println!("skipped: python3 not found");
323            return;
324        }
325        let f = ExternalFn {
326            name: "shout".into(),
327            language: Language::Python,
328            body: "result = args[0].upper() + '!'".into(),
329            arg_types: vec![ValueType::Text],
330            return_type: ValueType::Text,
331            pure: true,
332        };
333        let out = eval(&f, &[Value::Text("hi".into())], Duration::from_secs(2)).expect("eval");
334        assert_eq!(out, Value::Text("HI!".into()));
335    }
336
337    #[test]
338    fn python_raises_is_runtime_error() {
339        if !python_available() {
340            println!("skipped: python3 not found");
341            return;
342        }
343        let f = ExternalFn {
344            name: "boom".into(),
345            language: Language::Python,
346            body: "raise ValueError('nope')".into(),
347            arg_types: vec![],
348            return_type: ValueType::Number,
349            pure: true,
350        };
351        let err = eval(&f, &[], Duration::from_secs(2)).unwrap_err();
352        match err {
353            ExtFnError::Runtime { message, .. } => assert!(message.contains("ValueError")),
354            other => panic!("expected Runtime, got {other:?}"),
355        }
356    }
357
358    #[test]
359    fn python_return_type_mismatch() {
360        if !python_available() {
361            println!("skipped: python3 not found");
362            return;
363        }
364        // Declares Number but returns a string.
365        let f = ExternalFn {
366            name: "liar".into(),
367            language: Language::Python,
368            body: "result = 'not a number'".into(),
369            arg_types: vec![],
370            return_type: ValueType::Number,
371            pure: true,
372        };
373        let err = eval(&f, &[], Duration::from_secs(2)).unwrap_err();
374        assert!(matches!(err, ExtFnError::Runtime { .. }));
375    }
376
377    #[test]
378    fn python_infinite_loop_times_out() {
379        if !python_available() {
380            println!("skipped: python3 not found");
381            return;
382        }
383        let f = ExternalFn {
384            name: "spin".into(),
385            language: Language::Python,
386            body: "while True:\n    pass".into(),
387            arg_types: vec![],
388            return_type: ValueType::Number,
389            pure: true,
390        };
391        let start = std::time::Instant::now();
392        let err = eval(&f, &[], Duration::from_secs(2)).unwrap_err();
393        assert!(matches!(err, ExtFnError::Timeout(_)), "got {err:?}");
394        // Killed near the deadline, not hung forever.
395        assert!(start.elapsed() < Duration::from_secs(8));
396    }
397
398    // ---- Interpreter availability probe (shared) ----
399
400    /// True if a trivial `language` body (`body` sets `result` to 1) runs
401    /// under the host sandbox policy, like [`python_available`]. R/Julia/Pure
402    /// tests early-return (logging a skip) when false: the interpreter is
403    /// missing, or `Restricted` fails closed (no working `bwrap`), or it is
404    /// too slow to start within `timeout`.
405    fn sandboxed_available(language: Language, body: &str, timeout: Duration) -> bool {
406        let probe = ExternalFn {
407            name: "probe".into(),
408            language,
409            body: body.into(),
410            arg_types: vec![],
411            return_type: ValueType::Number,
412            pure: true,
413        };
414        eval(&probe, &[], timeout).is_ok_and(|v| v == Value::Number(1.0))
415    }
416
417    // ---- R (guarded) ----
418
419    #[test]
420    fn r_hypot() {
421        if !sandboxed_available(Language::R, "result <- 1", Duration::from_secs(10)) {
422            println!("skipped: Rscript not found (or no sandbox)");
423            return;
424        }
425        let f = ExternalFn {
426            name: "hyp".into(),
427            language: Language::R,
428            body: "result <- sqrt(args[[1]]^2 + args[[2]]^2)".into(),
429            arg_types: vec![ValueType::Number, ValueType::Number],
430            return_type: ValueType::Number,
431            pure: true,
432        };
433        let out = eval(
434            &f,
435            &[Value::Number(3.0), Value::Number(4.0)],
436            Duration::from_secs(10),
437        )
438        .expect("eval");
439        assert_eq!(out, Value::Number(5.0));
440    }
441
442    // ---- Julia (guarded) ----
443
444    #[test]
445    fn julia_sum() {
446        if !sandboxed_available(Language::Julia, "    result = 1", Duration::from_secs(30)) {
447            println!("skipped: julia not found (or no sandbox)");
448            return;
449        }
450        let f = ExternalFn {
451            name: "add".into(),
452            language: Language::Julia,
453            body: "    result = args[1] + args[2]".into(),
454            arg_types: vec![ValueType::Number, ValueType::Number],
455            return_type: ValueType::Number,
456            pure: true,
457        };
458        let out = eval(
459            &f,
460            &[Value::Number(2.0), Value::Number(40.0)],
461            Duration::from_secs(30),
462        )
463        .expect("eval");
464        assert_eq!(out, Value::Number(42.0));
465    }
466
467    // ---- Pure-lang (guarded) ----
468
469    #[test]
470    fn pure_sum() {
471        if !sandboxed_available(Language::Pure, "let result = 1;", Duration::from_secs(10)) {
472            println!("skipped: pure not found (or no sandbox)");
473            return;
474        }
475        let f = ExternalFn {
476            name: "add".into(),
477            language: Language::Pure,
478            body: "let result = args!0 + args!1;".into(),
479            arg_types: vec![ValueType::Number, ValueType::Number],
480            return_type: ValueType::Number,
481            pure: true,
482        };
483        let out = eval(
484            &f,
485            &[Value::Number(19.0), Value::Number(23.0)],
486            Duration::from_secs(10),
487        )
488        .expect("eval");
489        assert_eq!(out, Value::Number(42.0));
490    }
491
492    // ---- WASM (no external tool; wasmi + wat assembled at test time) ----
493
494    #[test]
495    fn wasm_doubles_input() {
496        // A 3-line module exporting `improv_call(f64) -> f64` = x + x.
497        let wat = "(module (func (export \"improv_call\") (param f64) (result f64) \
498             local.get 0 local.get 0 f64.add))";
499        let f = ExternalFn {
500            name: "double".into(),
501            language: Language::Wasm,
502            body: format!("wat:{wat}"),
503            arg_types: vec![ValueType::Number],
504            return_type: ValueType::Number,
505            pure: true,
506        };
507        let out = eval(&f, &[Value::Number(21.0)], Duration::from_secs(5)).expect("eval");
508        assert_eq!(out, Value::Number(42.0));
509    }
510
511    #[test]
512    fn wasm_two_args() {
513        // a*a + b*b, exercising a 2-arg numeric ABI.
514        let wat = "(module (func (export \"improv_call\") (param f64 f64) (result f64) \
515             local.get 0 local.get 0 f64.mul local.get 1 local.get 1 f64.mul f64.add))";
516        let f = ExternalFn {
517            name: "sqsum".into(),
518            language: Language::Wasm,
519            body: format!("wat:{wat}"),
520            arg_types: vec![ValueType::Number, ValueType::Number],
521            return_type: ValueType::Number,
522            pure: true,
523        };
524        let out = eval(
525            &f,
526            &[Value::Number(3.0), Value::Number(4.0)],
527            Duration::from_secs(5),
528        )
529        .expect("eval");
530        assert_eq!(out, Value::Number(25.0));
531    }
532
533    #[test]
534    fn wasm_missing_export_is_runtime_error() {
535        let f = ExternalFn {
536            name: "nope".into(),
537            language: Language::Wasm,
538            body: "wat:(module)".into(),
539            arg_types: vec![],
540            return_type: ValueType::Number,
541            pure: true,
542        };
543        let err = eval(&f, &[], Duration::from_secs(5)).unwrap_err();
544        match err {
545            ExtFnError::Runtime { message, .. } => assert!(message.contains("improv_call")),
546            other => panic!("expected Runtime, got {other:?}"),
547        }
548    }
549
550    #[test]
551    fn wasm_rejects_non_numeric_abi() {
552        let f = ExternalFn {
553            name: "bad".into(),
554            language: Language::Wasm,
555            body: "wat:(module)".into(),
556            arg_types: vec![ValueType::Text],
557            return_type: ValueType::Text,
558            pure: true,
559        };
560        let err = eval(&f, &[Value::Text("x".into())], Duration::from_secs(5)).unwrap_err();
561        assert!(matches!(err, ExtFnError::Runtime { .. }), "got {err:?}");
562    }
563
564    // ---- Unavailable runtime does not panic ----
565
566    #[test]
567    fn absent_interpreter_is_error_not_panic() {
568        // A language whose interpreter may not be installed yields an error, not
569        // a panic. If the runtime happens to be present, a benign body still
570        // succeeds — either way, no panic.
571        for lang in [Language::R, Language::Julia, Language::Pure] {
572            let body = match lang {
573                Language::R => "result <- 1",
574                Language::Julia => "    result = 1",
575                Language::Pure => "let result = 1;",
576                _ => unreachable!(),
577            };
578            let f = ExternalFn {
579                name: "probe".into(),
580                language: lang,
581                body: body.into(),
582                arg_types: vec![],
583                return_type: ValueType::Number,
584                pure: true,
585            };
586            // Must return a Result (no panic); Ok or Err both acceptable.
587            let _ = eval(&f, &[], Duration::from_secs(5));
588        }
589    }
590
591    // ---- OS sandbox (Linux: bwrap + rlimits; guarded like the other
592    // interpreter-dependent tests, skip cleanly when tools are absent) ----
593
594    fn bwrap_available() -> bool {
595        std::process::Command::new("bwrap")
596            .arg("--version")
597            .stdout(std::process::Stdio::null())
598            .stderr(std::process::Stdio::null())
599            .status()
600            .map(|s| s.success())
601            .unwrap_or(false)
602    }
603
604    /// A `Restricted` (the host default) body that tries to allocate
605    /// well past the sandbox's `RLIMIT_AS` ceiling must fail (Python raises
606    /// `MemoryError`, which the driver catches and reports as an error
607    /// envelope) rather than succeed. This is the rlimit half of the sandbox,
608    /// independent of whether `bwrap` is present. Linux only: macOS accepts
609    /// `RLIMIT_AS` but does not enforce it.
610    #[cfg(target_os = "linux")]
611    #[test]
612    fn restricted_python_hits_memory_limit() {
613        if !python_available() {
614            println!("skipped: python3 not found");
615            return;
616        }
617        let f = ExternalFn {
618            name: "blow_memory".into(),
619            language: Language::Python,
620            body: "x = bytearray(2 * 10**9)\nresult = len(x)".into(),
621            arg_types: vec![],
622            return_type: ValueType::Number,
623            pure: true,
624        };
625        let err = eval(&f, &[], Duration::from_secs(10)).unwrap_err();
626        // Either the driver's own try/except reports MemoryError as a Runtime
627        // error, or (if the allocator kills the process outright under the
628        // rlimit) the subprocess exits nonzero/empty-output, also a Runtime
629        // error. A Timeout would indicate the limit did nothing and the body
630        // ran to completion within the wall clock only by luck; assert it's
631        // specifically a Runtime error, not success.
632        assert!(matches!(err, ExtFnError::Runtime { .. }), "got {err:?}");
633    }
634
635    /// Sanity check that the SAME allocation succeeds under `Trusted` (no
636    /// sandbox applied), proving the failure above is the rlimit and not some
637    /// unrelated Python/driver bug.
638    #[test]
639    fn trusted_python_memory_allocation_succeeds() {
640        if !python_available() {
641            println!("skipped: python3 not found");
642            return;
643        }
644        let f = ExternalFn {
645            name: "blow_memory_trusted".into(),
646            language: Language::Python,
647            body: "x = bytearray(10**7)\nresult = len(x)".into(),
648            arg_types: vec![],
649            return_type: ValueType::Number,
650            pure: false,
651        };
652        let out = python::eval_with(&f, &[], Duration::from_secs(10), SandboxPolicy::Trusted)
653            .expect("eval");
654        assert_eq!(out, Value::Number(10_000_000.0));
655    }
656
657    /// When `bwrap` is present, a `Restricted` body reading an absolute path
658    /// outside the sandbox's bind-mounted view (here, this very repo's
659    /// `Cargo.toml`, addressed by absolute host path) must fail — proving the
660    /// filesystem restriction is real — while the SAME body under `Trusted`
661    /// (no bwrap wrapping) succeeds. Guarded on `bwrap` being on PATH.
662    #[test]
663    fn bwrap_blocks_filesystem_read_outside_sandbox() {
664        if !bwrap_available() {
665            println!("skipped: bwrap not found");
666            return;
667        }
668        if !python_available() {
669            println!("skipped: python3 not found");
670            return;
671        }
672        // Use this crate's own Cargo.toml: guaranteed to exist, guaranteed to
673        // be under $HOME (which bwrap does not bind), and not under the parts
674        // of `/` this sandbox binds read-only mirror-fashion... except `/` IS
675        // bound read-only in our bwrap policy, so instead prove the point with
676        // a path that is genuinely absent from the sandboxed view: `/etc/hostname`
677        // is bound (since `/` is ro-bound in full), so use the crate manifest
678        // under `$HOME` combined with `--unshare-net`/no special exclusion —
679        // the real isolation this policy provides is network + tmp, so assert
680        // on a target we deliberately do NOT bind: a path under a fresh tmpfs
681        // written by the *host* just before the sandboxed run, which the
682        // sandbox's own fresh `/tmp` cannot see.
683        let host_tmp_dir =
684            std::env::temp_dir().join(format!("improv_extfn_sandbox_test_{}", std::process::id()));
685        std::fs::create_dir_all(&host_tmp_dir).expect("create host tmp dir");
686        let secret_path = host_tmp_dir.join("secret.txt");
687        std::fs::write(&secret_path, "host-only-secret").expect("write secret");
688        let secret_path_str = secret_path.to_string_lossy().replace('\\', "\\\\");
689
690        let body = format!("result = open(r'{secret_path_str}').read()",);
691
692        // Restricted: bwrap's fresh tmpfs on /tmp hides the host's real /tmp,
693        // so the file is unreadable inside the sandbox.
694        let restricted_fn = ExternalFn {
695            name: "peek_restricted".into(),
696            language: Language::Python,
697            body: body.clone(),
698            arg_types: vec![],
699            return_type: ValueType::Text,
700            pure: true,
701        };
702        let restricted_err = eval(&restricted_fn, &[], Duration::from_secs(10)).unwrap_err();
703        assert!(
704            matches!(restricted_err, ExtFnError::Runtime { .. }),
705            "expected the sandboxed read to fail, got {restricted_err:?}"
706        );
707
708        // Trusted: no bwrap wrapping, the real filesystem (and real /tmp) is
709        // visible, so the same read succeeds.
710        let trusted_fn = ExternalFn {
711            name: "peek_trusted".into(),
712            language: Language::Python,
713            body,
714            arg_types: vec![],
715            return_type: ValueType::Text,
716            pure: false,
717        };
718        let out = python::eval_with(
719            &trusted_fn,
720            &[],
721            Duration::from_secs(10),
722            SandboxPolicy::Trusted,
723        )
724        .expect("trusted eval");
725        assert_eq!(out, Value::Text("host-only-secret".into()));
726
727        let _ = std::fs::remove_dir_all(&host_tmp_dir);
728    }
729
730    /// The `IMPROV_EXTFN_ALLOW_UNSANDBOXED` fallback (rlimits only, no
731    /// `bwrap`) still runs the interpreter, with a scrubbed environment.
732    #[test]
733    fn rlimits_only_fallback_still_runs() {
734        if !python_available() {
735            println!("skipped: python3 not found");
736            return;
737        }
738        std::env::set_var("IMPROV_API_TOKEN", "parent-secret");
739        let mut cmd =
740            crate::sandbox::apply_unsandboxed_for_test(std::process::Command::new("python3"));
741        cmd.args([
742            "-I",
743            "-S",
744            "-c",
745            "import os; print(os.environ.get('IMPROV_API_TOKEN', ''), os.environ['HOME'], os.environ['PATH'])",
746        ]);
747        let out = cmd.output().expect("spawn python3 under rlimits-only");
748        assert!(out.status.success(), "status: {:?}", out.status);
749        assert_eq!(
750            String::from_utf8_lossy(&out.stdout).trim(),
751            "/tmp /usr/bin:/bin"
752        );
753    }
754
755    fn py(body: &str, return_type: ValueType) -> ExternalFn {
756        ExternalFn {
757            name: "probe".into(),
758            language: Language::Python,
759            body: body.into(),
760            arg_types: vec![],
761            return_type,
762            pure: true,
763        }
764    }
765
766    /// S6: a `Restricted` body does not inherit the host environment.
767    #[test]
768    fn restricted_python_does_not_inherit_env() {
769        if !python_available() {
770            println!("skipped: python3 not found (or no sandbox)");
771            return;
772        }
773        std::env::set_var("IMPROV_API_TOKEN", "parent-secret");
774        let f = py(
775            "import os\nresult = os.environ.get('IMPROV_API_TOKEN', '')",
776            ValueType::Text,
777        );
778        let out = eval(&f, &[], Duration::from_secs(10)).expect("eval");
779        assert_eq!(out, Value::Text(String::new()));
780    }
781
782    /// S5: under bwrap neither `~` nor `/home` shows anything of the host.
783    #[test]
784    fn bwrap_hides_host_home() {
785        if !bwrap_available() || !python_available() {
786            println!("skipped: bwrap or python3 not found");
787            return;
788        }
789        let home = std::env::temp_dir().join(format!("improv_extfn_home_{}", std::process::id()));
790        std::fs::create_dir_all(&home).expect("create fake home");
791        std::fs::write(home.join("marker-secret"), "x").expect("write marker");
792        let real_home = std::env::var_os("HOME");
793        std::env::set_var("HOME", &home);
794        let f = py(
795            "import os\nh = os.path.expanduser('~')\n\
796             result = repr([h, sorted(os.listdir(h)), sorted(os.listdir('/home'))])",
797            ValueType::Text,
798        );
799        let out = eval(&f, &[], Duration::from_secs(10));
800        if let Some(h) = real_home {
801            std::env::set_var("HOME", h);
802        }
803        let _ = std::fs::remove_dir_all(&home);
804        assert_eq!(out.expect("eval"), Value::Text("['/tmp', [], []]".into()));
805    }
806
807    /// S5: under bwrap the network namespace is empty.
808    #[test]
809    fn bwrap_blocks_network() {
810        if !bwrap_available() || !python_available() {
811            println!("skipped: bwrap or python3 not found");
812            return;
813        }
814        let f = py(
815            "import socket\ns = socket.socket()\ns.settimeout(2)\n\
816             try:\n    s.connect(('1.1.1.1', 53))\n    result = 'connected'\n\
817             except OSError:\n    result = 'blocked'",
818            ValueType::Text,
819        );
820        let out = eval(&f, &[], Duration::from_secs(10)).expect("eval");
821        assert_eq!(out, Value::Text("blocked".into()));
822    }
823
824    /// S7: output past the 64 KiB pipe buffer used to deadlock until the
825    /// timeout; 200 KB now succeeds.
826    #[test]
827    fn python_large_output_succeeds() {
828        if !python_available() {
829            println!("skipped: python3 not found (or no sandbox)");
830            return;
831        }
832        let f = py("print('x' * 200000)\nresult = 1", ValueType::Number);
833        let out = eval(&f, &[], Duration::from_secs(10)).expect("eval");
834        assert_eq!(out, Value::Number(1.0));
835    }
836
837    /// S7: stdout is capped at 1 MiB; 5 MB is an "output too large" error.
838    #[test]
839    fn python_huge_output_is_rejected() {
840        if !python_available() {
841            println!("skipped: python3 not found (or no sandbox)");
842            return;
843        }
844        let f = py("print('x' * 5000000)\nresult = 1", ValueType::Number);
845        match eval(&f, &[], Duration::from_secs(10)).unwrap_err() {
846            ExtFnError::Runtime { message, .. } => {
847                assert!(message.contains("output too large"), "{message}")
848            }
849            other => panic!("expected Runtime, got {other:?}"),
850        }
851    }
852
853    /// S8: the timeout kills the whole process group, so a forked grandchild
854    /// does not outlive the call. Runs `Trusted` (no bwrap pid namespace) so
855    /// the group kill is the only thing that can stop it, and the marker file
856    /// is visible to the host.
857    #[cfg(target_os = "linux")]
858    #[test]
859    fn timeout_kills_grandchildren() {
860        let probe = python::eval_with(
861            &py("result = 1", ValueType::Number),
862            &[],
863            Duration::from_secs(5),
864            SandboxPolicy::Trusted,
865        );
866        if probe.is_err() {
867            println!("skipped: python3 not found");
868            return;
869        }
870        let marker =
871            std::env::temp_dir().join(format!("improv_extfn_grandchild_{}", std::process::id()));
872        let _ = std::fs::remove_file(&marker);
873        let body = format!(
874            "import os, time\n\
875             if os.fork() == 0:\n\
876             \x20   open({path:?}, 'w').write(str(os.getpid()))\n\
877             \x20   time.sleep(30)\n\
878             \x20   os._exit(0)\n\
879             time.sleep(30)\nresult = 1",
880            path = marker.to_string_lossy()
881        );
882        let timeout = Duration::from_secs(1);
883        let err = python::eval_with(
884            &py(&body, ValueType::Number),
885            &[],
886            timeout,
887            SandboxPolicy::Trusted,
888        )
889        .unwrap_err();
890        assert!(matches!(err, ExtFnError::Timeout(_)), "got {err:?}");
891        std::thread::sleep(Duration::from_secs(1));
892        let pid = std::fs::read_to_string(&marker).expect("grandchild wrote its pid");
893        let _ = std::fs::remove_file(&marker);
894        let proc_dir = format!("/proc/{}", pid.trim());
895        assert!(
896            !std::path::Path::new(&proc_dir).exists(),
897            "grandchild {pid} survived the timeout"
898        );
899    }
900
901    /// S20: the body reaches `exec` byte for byte. Backslashes used to be
902    /// doubled inside `r"""..."""`, which changed what the body meant.
903    #[test]
904    fn python_body_with_backslashes_and_triple_quotes() {
905        if !python_available() {
906            println!("skipped: python3 not found (or no sandbox)");
907            return;
908        }
909        let body = r#"s = """a\\b"""
910t = '\n' + "\""
911result = s + "|" + str(len(t)) + "|" + "\x41""#;
912        let out = eval(&py(body, ValueType::Text), &[], Duration::from_secs(10)).expect("eval");
913        assert_eq!(out, Value::Text(r"a\b|2|A".into()));
914    }
915}
improv_extfn/lib.rs

improv_extfn/
lib.rs
