agent_sandbox/runtime/
mod.rs

1use std::sync::{Arc, OnceLock};
2
3use wasmtime::{Config, Engine, Linker, Module, Store, StoreLimits, StoreLimitsBuilder, Trap};
4use wasmtime_wasi::WasiCtx;
5use wasmtime_wasi::p2::pipe::{MemoryInputPipe, MemoryOutputPipe};
6
7use crate::config::SandboxConfig;
8use crate::error::{Result, SandboxError};
9
10/// Result of executing a command in the sandbox.
11#[derive(Debug, Clone)]
12pub struct ExecResult {
13    pub exit_code: i32,
14    pub stdout: Vec<u8>,
15    pub stderr: Vec<u8>,
16}
17
18/// Store data combining WASI context with resource limits.
19struct SandboxState {
20    wasi: wasmtime_wasi::p1::WasiP1Ctx,
21    limits: StoreLimits,
22}
23
24/// Cached WASM engine and compiled module shared across all Sandbox instances.
25struct CachedModule {
26    engine: Engine,
27    module: Module,
28}
29
30/// Global cache for the compiled WASM module.
31/// Compiling the toolbox WASM binary is expensive, so we do it once.
32static MODULE_CACHE: OnceLock<std::result::Result<CachedModule, String>> = OnceLock::new();
33
34fn get_or_compile_module() -> Result<(&'static Engine, &'static Module)> {
35    let cached = MODULE_CACHE.get_or_init(|| {
36        let precompiled_bytes = include_bytes!(env!("TOOLBOX_CWASM_PATH"));
37
38        if precompiled_bytes.is_empty() {
39            return Err("WASM toolbox not available".to_string());
40        }
41
42        // Engine config MUST match build.rs exactly
43        let mut engine_config = Config::new();
44        engine_config.consume_fuel(true);
45
46        let engine =
47            Engine::new(&engine_config).map_err(|e| format!("engine creation failed: {e}"))?;
48
49        // SAFETY: The precompiled bytes come from our own build.rs via
50        // Engine::precompile_module() with the same engine config and wasmtime version.
51        let module = unsafe { Module::deserialize(&engine, precompiled_bytes) }
52            .map_err(|e| format!("module deserialization failed: {e}"))?;
53
54        Ok(CachedModule { engine, module })
55    });
56
57    match cached {
58        Ok(c) => Ok((&c.engine, &c.module)),
59        Err(e) => Err(SandboxError::Other(e.clone())),
60    }
61}
62
63/// The WASI runtime that manages Wasmtime engine and module compilation.
64pub struct WasiRuntime {
65    engine: &'static Engine,
66    module: &'static Module,
67    config: Arc<SandboxConfig>,
68}
69
70impl WasiRuntime {
71    /// Create a new WASI runtime with the given sandbox config.
72    /// The toolbox WASM binary is compiled once and cached globally.
73    pub fn new(config: SandboxConfig) -> Result<Self> {
74        let (engine, module) = get_or_compile_module()?;
75
76        Ok(Self {
77            engine,
78            module,
79            config: Arc::new(config),
80        })
81    }
82
83    /// Execute a command inside the WASM sandbox.
84    pub async fn exec(&self, command: &str, args: &[String]) -> Result<ExecResult> {
85        let config = self.config.clone();
86        let engine = self.engine;
87        let module = self.module;
88        let command = command.to_string();
89        let args = args.to_vec();
90        let timeout = config.timeout;
91
92        // Run in blocking thread since Wasmtime is synchronous, with a wall-clock timeout
93        let task = tokio::task::spawn_blocking(move || {
94            exec_sync(engine, module, &config, &command, &args)
95        });
96
97        match tokio::time::timeout(timeout, task).await {
98            Ok(Ok(result)) => result,
99            Ok(Err(e)) => Err(SandboxError::Other(format!("task join error: {}", e))),
100            Err(_) => Err(SandboxError::Timeout(timeout)),
101        }
102    }
103}
104
105fn exec_sync(
106    engine: &Engine,
107    module: &Module,
108    config: &SandboxConfig,
109    command: &str,
110    args: &[String],
111) -> Result<ExecResult> {
112    // Build argv: [command, ...args]
113    let mut argv: Vec<String> = vec![command.to_string()];
114    argv.extend(args.iter().cloned());
115
116    let argv_refs: Vec<&str> = argv.iter().map(|s| s.as_str()).collect();
117
118    // Set up stdout/stderr capture via MemoryOutputPipe
119    let stdout_pipe = MemoryOutputPipe::new(1024 * 1024); // 1MB capacity
120    let stderr_pipe = MemoryOutputPipe::new(1024 * 1024);
121
122    // Build WASI context using WasiCtx::builder()
123    let mut builder = WasiCtx::builder();
124    builder.args(&argv_refs);
125    builder.stdin(MemoryInputPipe::new(b"" as &[u8])); // Empty stdin — prevents blocking on host stdin
126    builder.stdout(stdout_pipe.clone());
127    builder.stderr(stderr_pipe.clone());
128
129    // Set TOOLBOX_CMD env var for BusyBox-style dispatch
130    builder.env("TOOLBOX_CMD", command);
131
132    // Set user-configured env vars
133    for (key, value) in &config.env_vars {
134        builder.env(key, value);
135    }
136
137    // Mount work directory
138    let work_dir = config.work_dir.canonicalize().map_err(|e| {
139        SandboxError::Io(std::io::Error::new(
140            std::io::ErrorKind::NotFound,
141            format!("work_dir '{}': {}", config.work_dir.display(), e),
142        ))
143    })?;
144
145    let dir = wasmtime_wasi::DirPerms::all();
146    let file = wasmtime_wasi::FilePerms::all();
147    builder.preopened_dir(&work_dir, "/work", dir, file)?;
148
149    // Mount additional directories
150    for mount in &config.mounts {
151        let host = mount.host_path.canonicalize().map_err(|e| {
152            SandboxError::Io(std::io::Error::new(
153                std::io::ErrorKind::NotFound,
154                format!("mount '{}': {}", mount.host_path.display(), e),
155            ))
156        })?;
157
158        let (d, f) = if mount.writable {
159            (
160                wasmtime_wasi::DirPerms::all(),
161                wasmtime_wasi::FilePerms::all(),
162            )
163        } else {
164            (
165                wasmtime_wasi::DirPerms::READ,
166                wasmtime_wasi::FilePerms::READ,
167            )
168        };
169
170        builder.preopened_dir(&host, &mount.guest_path, d, f)?;
171    }
172
173    // Build the WASIp1 context
174    let wasi_p1 = builder.build_p1();
175
176    // Build memory limiter
177    let limits = StoreLimitsBuilder::new()
178        .memory_size(config.memory_limit_bytes as usize)
179        .build();
180
181    let mut store = Store::new(
182        engine,
183        SandboxState {
184            wasi: wasi_p1,
185            limits,
186        },
187    );
188    store.limiter(|state| &mut state.limits);
189
190    // Set fuel limit
191    store.set_fuel(config.fuel_limit)?;
192
193    // Link WASI p1 and instantiate
194    let mut linker = Linker::new(engine);
195    wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |state: &mut SandboxState| &mut state.wasi)?;
196
197    linker.module(&mut store, "", module)?;
198
199    // Get the default function (_start) and call it
200    let func = linker
201        .get_default(&mut store, "")?
202        .typed::<(), ()>(&store)?;
203
204    let exit_code = match func.call(&mut store, ()) {
205        Ok(()) => 0,
206        Err(e) => {
207            // Check if it's a normal process exit
208            if let Some(exit) = e.downcast_ref::<wasmtime_wasi::I32Exit>() {
209                exit.0
210            } else if e.downcast_ref::<Trap>() == Some(&Trap::OutOfFuel) {
211                return Err(SandboxError::Timeout(config.timeout));
212            } else {
213                return Err(SandboxError::Runtime(e));
214            }
215        }
216    };
217
218    let stdout_bytes = stdout_pipe.contents().to_vec();
219    let stderr_bytes = stderr_pipe.contents().to_vec();
220
221    Ok(ExecResult {
222        exit_code,
223        stdout: stdout_bytes,
224        stderr: stderr_bytes,
225    })
226}
agent_sandbox/runtime/mod.rs

agent_sandbox/runtime/
mod.rs
