Skip to main content

traverse_runtime/executor/
wasm.rs

1//! Wasmtime-backed WASM executor.
2//!
3//! Executes `wasm32-wasi` capability binaries inside a sandboxed Wasmtime engine.
4//! Input is fed via WASI stdin; output is captured from WASI stdout.
5//! No ambient WASI authority is granted — all capabilities are deny-by-default.
6
7use serde::Deserialize;
8use serde_json::Value;
9use sha2::{Digest, Sha256};
10use std::collections::{HashMap, VecDeque};
11use std::fmt::Write as _;
12use std::fs;
13use std::sync::{LazyLock, Mutex};
14use wasmtime::{Config, Engine, Linker, Module, Store, StoreLimits, StoreLimitsBuilder};
15use wasmtime_wasi::WasiCtxBuilder;
16use wasmtime_wasi::p1::WasiP1Ctx;
17use wasmtime_wasi::p2::pipe::{MemoryInputPipe, MemoryOutputPipe};
18
19use super::{ArtifactType, CapabilityExecutor, ExecutorCapability, ExecutorError};
20
21/// Traverse Host ABI v1 is independently versioned from the runtime crate.
22pub const SUPPORTED_HOST_ABI_VERSION: &str = "1.0.0";
23
24const HOST_ABI_V1_WHITELIST: &str = include_str!("host_abi_v1.json");
25const DEFAULT_FUEL_BUDGET: u64 = 5_000_000;
26const DEFAULT_MEMORY_LIMIT_BYTES: usize = 8 * 1024 * 1024;
27const DEFAULT_TABLE_ELEMENT_LIMIT: usize = 1_024;
28const DEFAULT_INSTANCE_LIMIT: usize = 1;
29const DEFAULT_TABLE_LIMIT: usize = 8;
30const DEFAULT_LINEAR_MEMORY_LIMIT: usize = 1;
31const DEFAULT_MODULE_CACHE_MAX_ENTRIES: usize = 64;
32
33static HOST_ABI_V1_WHITELIST_CACHE: LazyLock<Result<HostAbiWhitelist, String>> =
34    LazyLock::new(|| {
35        serde_json::from_str::<HostAbiWhitelist>(HOST_ABI_V1_WHITELIST).map_err(|e| e.to_string())
36    });
37
38/// A host import observed in a WASM module.
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct HostAbiImport {
41    /// Imported module namespace.
42    pub module: String,
43    /// Imported function or item name.
44    pub name: String,
45}
46
47/// Successful load-time ABI validation evidence.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct HostAbiValidation {
50    /// ABI version used for whitelist validation.
51    pub abi_version: String,
52    /// All imports observed in deterministic module/name order.
53    pub imports: Vec<HostAbiImport>,
54}
55
56#[derive(Debug, Clone, Deserialize)]
57struct HostAbiWhitelist {
58    abi_version: String,
59    imports: Vec<HostAbiWhitelistImport>,
60}
61
62#[derive(Debug, Clone, Deserialize)]
63struct HostAbiWhitelistImport {
64    module: String,
65    name: String,
66}
67
68/// Return the Traverse Host ABI versions supported by this runtime.
69#[must_use]
70pub fn supported_host_abi_versions() -> &'static [&'static str] {
71    &[SUPPORTED_HOST_ABI_VERSION]
72}
73
74/// Validate a WASM binary against the declared Traverse Host ABI import whitelist.
75///
76/// # Errors
77///
78/// Returns [`ExecutorError`] when the binary is malformed, the ABI version is unsupported,
79/// or a module imports a host function outside the whitelist.
80pub fn verify_wasm_host_abi_bytes(
81    wasm_bytes: &[u8],
82    abi_version: &str,
83) -> Result<HostAbiValidation, ExecutorError> {
84    let engine = Engine::default();
85    let module = Module::from_binary(&engine, wasm_bytes).map_err(|e| {
86        ExecutorError::MalformedWasmArtifact {
87            error_code: "malformed_wasm_artifact".to_string(),
88            detail: format!("module compile: {e}"),
89        }
90    })?;
91    validate_module_imports(&module, abi_version)
92}
93
94/// Executes `.wasm32-wasi` capability binaries via Wasmtime.
95///
96/// Every invocation creates a fresh Wasmtime `Store` — no state leaks between calls.
97#[derive(Debug)]
98pub struct WasmExecutor {
99    engine: Engine,
100    limits: WasmExecutionLimits,
101    module_cache: Mutex<CompiledModuleCache>,
102}
103
104impl WasmExecutor {
105    /// Create a new [`WasmExecutor`] with a default Wasmtime engine.
106    ///
107    /// # Errors
108    ///
109    /// Returns [`ExecutorError::RuntimeSetupFailed`] if Wasmtime cannot initialise.
110    pub fn new() -> Result<Self, ExecutorError> {
111        Self::with_limits(WasmExecutionLimits::default())
112    }
113
114    /// Create a [`WasmExecutor`] with explicit per-invocation resource limits.
115    ///
116    /// # Errors
117    ///
118    /// Returns [`ExecutorError::RuntimeSetupFailed`] if Wasmtime cannot initialise.
119    pub fn with_limits(limits: WasmExecutionLimits) -> Result<Self, ExecutorError> {
120        Self::with_limits_and_cache_config(limits, WasmModuleCacheConfig::default())
121    }
122
123    /// Create a [`WasmExecutor`] with explicit resource limits and module cache bounds.
124    ///
125    /// # Errors
126    ///
127    /// Returns [`ExecutorError::RuntimeSetupFailed`] if Wasmtime cannot initialise.
128    pub fn with_limits_and_cache_config(
129        limits: WasmExecutionLimits,
130        cache_config: WasmModuleCacheConfig,
131    ) -> Result<Self, ExecutorError> {
132        let mut config = Config::new();
133        config.consume_fuel(true);
134        let engine = Engine::new(&config)
135            .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("engine config: {e}")))?;
136        Ok(Self {
137            engine,
138            limits,
139            module_cache: Mutex::new(CompiledModuleCache::new(cache_config.max_entries)),
140        })
141    }
142
143    /// Return current compiled-module cache counters.
144    #[must_use]
145    pub fn module_cache_stats(&self) -> WasmModuleCacheStats {
146        let cache = self
147            .module_cache
148            .lock()
149            .unwrap_or_else(std::sync::PoisonError::into_inner);
150        cache.stats()
151    }
152}
153
154/// Per-invocation resource limits for [`WasmExecutor`].
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156pub struct WasmExecutionLimits {
157    /// Fuel units available for guest code before it traps as a timeout.
158    pub fuel_budget: u64,
159    /// Maximum bytes for each guest linear memory.
160    pub memory_bytes: usize,
161    /// Maximum elements for each guest table.
162    pub table_elements: usize,
163    /// Maximum instances in the store.
164    pub instances: usize,
165    /// Maximum tables in the store.
166    pub tables: usize,
167    /// Maximum linear memories in the store.
168    pub memories: usize,
169}
170
171impl Default for WasmExecutionLimits {
172    fn default() -> Self {
173        Self {
174            fuel_budget: DEFAULT_FUEL_BUDGET,
175            memory_bytes: DEFAULT_MEMORY_LIMIT_BYTES,
176            table_elements: DEFAULT_TABLE_ELEMENT_LIMIT,
177            instances: DEFAULT_INSTANCE_LIMIT,
178            tables: DEFAULT_TABLE_LIMIT,
179            memories: DEFAULT_LINEAR_MEMORY_LIMIT,
180        }
181    }
182}
183
184/// Bounded compiled-module cache configuration for [`WasmExecutor`].
185#[derive(Debug, Clone, Copy, PartialEq, Eq)]
186pub struct WasmModuleCacheConfig {
187    /// Maximum number of compiled modules retained by checksum.
188    pub max_entries: usize,
189}
190
191impl Default for WasmModuleCacheConfig {
192    fn default() -> Self {
193        Self {
194            max_entries: DEFAULT_MODULE_CACHE_MAX_ENTRIES,
195        }
196    }
197}
198
199/// Snapshot of compiled-module cache counters.
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub struct WasmModuleCacheStats {
202    /// Current retained compiled modules.
203    pub entries: usize,
204    /// Number of executions served from cache.
205    pub hits: u64,
206    /// Number of executions that compiled a module before insertion.
207    pub misses: u64,
208    /// Number of deterministic oldest-entry evictions.
209    pub evictions: u64,
210}
211
212#[derive(Debug, Clone)]
213struct CachedModule {
214    module: Module,
215    validation: HostAbiValidation,
216}
217
218#[derive(Debug)]
219struct CompiledModuleCache {
220    max_entries: usize,
221    entries: HashMap<String, CachedModule>,
222    insertion_order: VecDeque<String>,
223    hits: u64,
224    misses: u64,
225    evictions: u64,
226}
227
228impl CompiledModuleCache {
229    fn new(max_entries: usize) -> Self {
230        Self {
231            max_entries: max_entries.max(1),
232            entries: HashMap::new(),
233            insertion_order: VecDeque::new(),
234            hits: 0,
235            misses: 0,
236            evictions: 0,
237        }
238    }
239
240    fn get(&mut self, checksum: &str, abi_version: &str) -> Option<CachedModule> {
241        let cached = self.entries.get(checksum)?;
242        if cached.validation.abi_version != abi_version {
243            self.misses += 1;
244            return None;
245        }
246        self.hits += 1;
247        Some(cached.clone())
248    }
249
250    fn insert(&mut self, checksum: String, cached: CachedModule) {
251        self.misses += 1;
252        while self.entries.len() >= self.max_entries {
253            if let Some(oldest) = self.insertion_order.pop_front()
254                && self.entries.remove(&oldest).is_some()
255            {
256                self.evictions += 1;
257            }
258        }
259        self.insertion_order.push_back(checksum.clone());
260        self.entries.insert(checksum, cached);
261    }
262
263    fn stats(&self) -> WasmModuleCacheStats {
264        WasmModuleCacheStats {
265            entries: self.entries.len(),
266            hits: self.hits,
267            misses: self.misses,
268            evictions: self.evictions,
269        }
270    }
271}
272
273struct WasmStoreState {
274    wasi: WasiP1Ctx,
275    limits: StoreLimits,
276}
277
278impl CapabilityExecutor for WasmExecutor {
279    fn execute(
280        &self,
281        capability: &ExecutorCapability,
282        input: &Value,
283    ) -> Result<Value, ExecutorError> {
284        if capability.artifact_type != ArtifactType::Wasm {
285            return Err(ExecutorError::UnsupportedArtifactType);
286        }
287
288        // --- Load binary ---
289        let wasm_path = capability.wasm_binary_path.as_deref().ok_or_else(|| {
290            ExecutorError::BinaryLoadFailed("no wasm_binary_path set".to_string())
291        })?;
292
293        let binary = fs::read(wasm_path).map_err(|e| {
294            ExecutorError::BinaryLoadFailed(format!("cannot read {wasm_path}: {e}"))
295        })?;
296
297        // --- Checksum validation ---
298        if let Some(expected) = capability.wasm_checksum.as_deref() {
299            let actual = sha256_hex(&binary);
300            if actual != expected {
301                return Err(ExecutorError::ChecksumMismatch {
302                    expected: expected.to_string(),
303                    actual,
304                });
305            }
306        }
307
308        let abi_version = capability
309            .host_abi_version
310            .as_deref()
311            .unwrap_or(SUPPORTED_HOST_ABI_VERSION);
312
313        self.run_wasm(&binary, input, abi_version)
314    }
315}
316
317impl WasmExecutor {
318    /// Execute pre-loaded WASM bytes with the given input.
319    ///
320    /// Exposed separately so tests can pass raw bytes without needing a file on disk.
321    ///
322    /// # Errors
323    ///
324    /// Returns [`ExecutorError`] if input serialization fails, the WASM module cannot be
325    /// compiled or linked, execution fails, or stdout is not valid JSON.
326    pub fn run_bytes(&self, wasm_bytes: &[u8], input: &Value) -> Result<Value, ExecutorError> {
327        self.run_bytes_with_host_abi(wasm_bytes, input, SUPPORTED_HOST_ABI_VERSION)
328    }
329
330    /// Execute pre-loaded WASM bytes with an explicit Traverse Host ABI version.
331    ///
332    /// # Errors
333    ///
334    /// Returns [`ExecutorError`] if ABI validation fails or execution cannot complete.
335    pub fn run_bytes_with_host_abi(
336        &self,
337        wasm_bytes: &[u8],
338        input: &Value,
339        abi_version: &str,
340    ) -> Result<Value, ExecutorError> {
341        self.run_wasm(wasm_bytes, input, abi_version)
342    }
343
344    fn run_wasm(
345        &self,
346        wasm_bytes: &[u8],
347        input: &Value,
348        abi_version: &str,
349    ) -> Result<Value, ExecutorError> {
350        let input_json = serde_json::to_string(input)
351            .map_err(|e| ExecutorError::ExecutionFailed(format!("input serialization: {e}")))?;
352
353        let cached_module = self.compiled_module(wasm_bytes, abi_version)?;
354
355        // Clone pipe reference before passing to builder — needed to read output after execution
356        let stdout_pipe = MemoryOutputPipe::new(65536);
357        let stdout_ref = stdout_pipe.clone();
358
359        // Build a WASI context: stdin = input JSON, stdout = captured buffer
360        // No filesystem, no network, no env vars — deny-by-default
361        let wasi_ctx: WasiP1Ctx = WasiCtxBuilder::new()
362            .stdin(MemoryInputPipe::new(input_json.into_bytes()))
363            .stdout(stdout_pipe)
364            .build_p1();
365
366        let mut linker: Linker<WasmStoreState> = Linker::new(&self.engine);
367        wasmtime_wasi::p1::add_to_linker_sync(&mut linker, |s| &mut s.wasi)
368            .map_err(|e| ExecutorError::RuntimeSetupFailed(e.to_string()))?;
369
370        let mut store = Store::new(
371            &self.engine,
372            WasmStoreState {
373                wasi: wasi_ctx,
374                limits: self.store_limits(),
375            },
376        );
377        store.limiter(|state| &mut state.limits);
378        store
379            .set_fuel(self.limits.fuel_budget)
380            .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("set fuel: {e}")))?;
381
382        linker
383            .module(&mut store, "", &cached_module.module)
384            .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("module link: {e}")))?;
385
386        linker
387            .get_default(&mut store, "")
388            .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("get_default: {e}")))?
389            .typed::<(), ()>(&store)
390            .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("typed: {e}")))?
391            .call(&mut store, ())
392            .map_err(|error| classify_wasm_execution_error(&error))?;
393
394        // Extract captured stdout — contents() reads the buffer without consuming it
395        let raw_output = stdout_ref.contents();
396
397        serde_json::from_slice::<Value>(&raw_output).map_err(|e| {
398            ExecutorError::OutputDeserializationFailed(format!(
399                "stdout is not valid JSON: {e} — raw: {}",
400                String::from_utf8_lossy(&raw_output)
401            ))
402        })
403    }
404
405    fn store_limits(&self) -> StoreLimits {
406        StoreLimitsBuilder::new()
407            .memory_size(self.limits.memory_bytes)
408            .table_elements(self.limits.table_elements)
409            .instances(self.limits.instances)
410            .tables(self.limits.tables)
411            .memories(self.limits.memories)
412            .trap_on_grow_failure(true)
413            .build()
414    }
415
416    fn compiled_module(
417        &self,
418        wasm_bytes: &[u8],
419        abi_version: &str,
420    ) -> Result<CachedModule, ExecutorError> {
421        let checksum = sha256_hex(wasm_bytes);
422        {
423            let mut cache = self
424                .module_cache
425                .lock()
426                .unwrap_or_else(std::sync::PoisonError::into_inner);
427            if let Some(cached) = cache.get(&checksum, abi_version) {
428                return Ok(cached);
429            }
430        }
431
432        let module = Module::from_binary(&self.engine, wasm_bytes).map_err(|e| {
433            ExecutorError::MalformedWasmArtifact {
434                error_code: "malformed_wasm_artifact".to_string(),
435                detail: format!("module compile: {e}"),
436            }
437        })?;
438        let validation = validate_module_imports(&module, abi_version)?;
439        let cached = CachedModule { module, validation };
440
441        let mut cache = self
442            .module_cache
443            .lock()
444            .unwrap_or_else(std::sync::PoisonError::into_inner);
445        cache.insert(checksum, cached.clone());
446        Ok(cached)
447    }
448}
449
450fn classify_wasm_execution_error(error: &wasmtime::Error) -> ExecutorError {
451    let display = error.to_string();
452    let debug = format!("{error:?}");
453    if display.contains("all fuel consumed by WebAssembly")
454        || debug.contains("all fuel consumed by WebAssembly")
455    {
456        return ExecutorError::Timeout(debug);
457    }
458    if display.contains("forcing trap when growing") || debug.contains("forcing trap when growing")
459    {
460        return ExecutorError::ResourceExhausted(debug);
461    }
462    ExecutorError::ExecutionFailed(display)
463}
464
465fn sha256_hex(data: &[u8]) -> String {
466    let mut hasher = Sha256::new();
467    hasher.update(data);
468    hasher
469        .finalize()
470        .iter()
471        .fold(String::new(), |mut acc, byte| {
472            let _ = write!(acc, "{byte:02x}");
473            acc
474        })
475}
476
477fn validate_module_imports(
478    module: &Module,
479    abi_version: &str,
480) -> Result<HostAbiValidation, ExecutorError> {
481    let whitelist = host_abi_whitelist(abi_version)?;
482    let mut imports = module
483        .imports()
484        .map(|import| HostAbiImport {
485            module: import.module().to_string(),
486            name: import.name().to_string(),
487        })
488        .collect::<Vec<_>>();
489    imports.sort_by(|a, b| a.module.cmp(&b.module).then_with(|| a.name.cmp(&b.name)));
490
491    for import in &imports {
492        if !whitelist
493            .imports
494            .iter()
495            .any(|allowed| allowed.module == import.module && allowed.name == import.name)
496        {
497            return Err(ExecutorError::UnauthorizedHostImport {
498                error_code: "unauthorized_host_import".to_string(),
499                abi_version: abi_version.to_string(),
500                module: import.module.clone(),
501                name: import.name.clone(),
502            });
503        }
504    }
505
506    Ok(HostAbiValidation {
507        abi_version: whitelist.abi_version,
508        imports,
509    })
510}
511
512fn host_abi_whitelist(abi_version: &str) -> Result<HostAbiWhitelist, ExecutorError> {
513    if abi_version != SUPPORTED_HOST_ABI_VERSION {
514        return Err(ExecutorError::UnsupportedAbiVersion {
515            error_code: "unsupported_abi_version".to_string(),
516            requested: abi_version.to_string(),
517            supported: supported_host_abi_versions().join(", "),
518        });
519    }
520
521    HOST_ABI_V1_WHITELIST_CACHE
522        .as_ref()
523        .cloned()
524        .map_err(|e| ExecutorError::RuntimeSetupFailed(format!("invalid ABI whitelist: {e}")))
525}