1use 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
21pub 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#[derive(Debug, Clone, PartialEq, Eq)]
40pub struct HostAbiImport {
41 pub module: String,
43 pub name: String,
45}
46
47#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct HostAbiValidation {
50 pub abi_version: String,
52 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#[must_use]
70pub fn supported_host_abi_versions() -> &'static [&'static str] {
71 &[SUPPORTED_HOST_ABI_VERSION]
72}
73
74pub 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#[derive(Debug)]
98pub struct WasmExecutor {
99 engine: Engine,
100 limits: WasmExecutionLimits,
101 module_cache: Mutex<CompiledModuleCache>,
102}
103
104impl WasmExecutor {
105 pub fn new() -> Result<Self, ExecutorError> {
111 Self::with_limits(WasmExecutionLimits::default())
112 }
113
114 pub fn with_limits(limits: WasmExecutionLimits) -> Result<Self, ExecutorError> {
120 Self::with_limits_and_cache_config(limits, WasmModuleCacheConfig::default())
121 }
122
123 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 #[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156pub struct WasmExecutionLimits {
157 pub fuel_budget: u64,
159 pub memory_bytes: usize,
161 pub table_elements: usize,
163 pub instances: usize,
165 pub tables: usize,
167 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
186pub struct WasmModuleCacheConfig {
187 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub struct WasmModuleCacheStats {
202 pub entries: usize,
204 pub hits: u64,
206 pub misses: u64,
208 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 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 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 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 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 let stdout_pipe = MemoryOutputPipe::new(65536);
357 let stdout_ref = stdout_pipe.clone();
358
359 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 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}