sh_layer4/plugin_loader/
wasm.rs

1//! WebAssembly Plugin Loader
2//!
3//! Loads and executes WASM modules in a sandboxed environment.
4//!
5//! ## Features
6//! - Real wasmtime execution
7//! - Sandboxed memory and CPU limits
8//! - Capability-based security
9//! - Async plugin execution
10//! - WASI preview1 support with sandboxed filesystem
11
12use super::capabilities::CapabilitySet;
13use super::sandbox::PluginSandbox;
14use super::{Plugin, PluginContext};
15use crate::types::Layer4Result;
16use anyhow::{anyhow, Context};
17use parking_lot::RwLock;
18use std::collections::HashMap;
19use std::path::{Path, PathBuf};
20use std::sync::Arc;
21use std::time::Instant;
22use wasmtime::*;
23use wasmtime_wasi::{DirPerms, FilePerms, WasiCtx, WasiCtxBuilder};
24
25/// WASM plugin configuration
26#[derive(Debug, Clone)]
27pub struct WasmConfig {
28    /// Maximum memory in bytes (default: 16MB)
29    pub max_memory_bytes: u64,
30    /// Maximum CPU time in milliseconds (default: 5000ms)
31    pub max_cpu_time_ms: u64,
32    /// Maximum table elements (default: 10000)
33    pub max_table_elements: u32,
34    /// Enable WASI preview1 (default: true)
35    pub enable_wasi: bool,
36    /// Allow async execution (default: true)
37    pub enable_async: bool,
38}
39
40impl Default for WasmConfig {
41    fn default() -> Self {
42        Self {
43            max_memory_bytes: 16 * 1024 * 1024, // 16 MB
44            max_cpu_time_ms: 5000,              // 5 seconds
45            max_table_elements: 10000,
46            enable_wasi: true,
47            enable_async: true,
48        }
49    }
50}
51
52/// WASM plugin state (store data)
53pub struct PluginState {
54    /// Security sandbox
55    pub sandbox: PluginSandbox,
56    /// Data directory for plugin
57    pub data_dir: PathBuf,
58    /// Execution start time
59    pub start_time: Option<Instant>,
60    /// CPU time limit
61    pub cpu_limit_ms: u64,
62    /// Memory tracker
63    pub memory_used: u64,
64    /// WASI context (for sandboxed filesystem access)
65    pub wasi_ctx: Option<WasiCtx>,
66}
67
68impl PluginState {
69    fn new(sandbox: PluginSandbox, data_dir: PathBuf, cpu_limit_ms: u64) -> Self {
70        Self {
71            sandbox,
72            data_dir,
73            start_time: None,
74            cpu_limit_ms,
75            memory_used: 0,
76            wasi_ctx: None,
77        }
78    }
79
80    fn with_wasi(mut self, wasi_ctx: WasiCtx) -> Self {
81        self.wasi_ctx = Some(wasi_ctx);
82        self
83    }
84
85    fn check_cpu_limit(&self) -> Result<()> {
86        if self.cpu_limit_ms == 0 {
87            return Ok(());
88        }
89        if let Some(start) = self.start_time {
90            let elapsed = start.elapsed().as_millis() as u64;
91            if elapsed > self.cpu_limit_ms {
92                return Err(anyhow!(
93                    "CPU time limit exceeded: {}ms > {}ms",
94                    elapsed,
95                    self.cpu_limit_ms
96                ));
97            }
98        }
99        Ok(())
100    }
101}
102
103/// WASM plugin instance
104pub struct WasmPlugin {
105    /// Plugin name
106    name: String,
107    /// Plugin version
108    version: String,
109    /// Sandbox reference
110    sandbox: PluginSandbox,
111    /// Module reference
112    module: Arc<Module>,
113    /// Engine reference
114    engine: Engine,
115    /// Configuration
116    config: WasmConfig,
117}
118
119impl WasmPlugin {
120    /// Create new WASM plugin
121    fn new(
122        name: String,
123        version: String,
124        sandbox: PluginSandbox,
125        module: Module,
126        engine: Engine,
127        config: WasmConfig,
128    ) -> Self {
129        Self {
130            name,
131            version,
132            sandbox,
133            module: Arc::new(module),
134            engine,
135            config,
136        }
137    }
138
139    /// Get module reference
140    pub fn module(&self) -> &Module {
141        &self.module
142    }
143
144    /// Get engine reference
145    pub fn engine(&self) -> &Engine {
146        &self.engine
147    }
148
149    /// Create a new store for execution
150    pub fn create_store(&self, data_dir: PathBuf) -> Store<PluginState> {
151        let cpu_limit = self.config.max_cpu_time_ms;
152        let state = PluginState::new(self.sandbox.clone(), data_dir, cpu_limit);
153        Store::new(&self.engine, state)
154    }
155
156    /// Instantiate the plugin in a store
157    pub fn instantiate(&self, store: &mut Store<PluginState>) -> Result<Instance> {
158        // Start CPU timer
159        store.data_mut().start_time = Some(Instant::now());
160
161        // Create instance
162        Instance::new(store, &self.module, &[])
163            .with_context(|| format!("Failed to instantiate WASM plugin: {}", self.name))
164    }
165
166    /// Execute a function by name with JSON input/output
167    pub fn execute_func(
168        &self,
169        store: &mut Store<PluginState>,
170        instance: &Instance,
171        func_name: &str,
172        input: &serde_json::Value,
173    ) -> Result<serde_json::Value> {
174        // Check CPU limit before execution
175        store.data().check_cpu_limit()?;
176
177        // Get the function
178        let func = instance
179            .get_typed_func::<(i32, i32), (i32, i32)>(&mut *store, func_name)
180            .with_context(|| format!("Function '{}' not found in plugin", func_name))?;
181
182        // Allocate input in WASM memory
183        let input_bytes = serde_json::to_vec(input)?;
184        let input_len = input_bytes.len() as i32;
185
186        // Get memory export
187        let memory = instance
188            .get_memory(&mut *store, "memory")
189            .ok_or_else(|| anyhow!("No memory export in plugin"))?;
190
191        // Allocate space for input (simple bump allocator)
192        let input_ptr = self.allocate_in_memory(store, memory, input_len)?;
193
194        // Write input to memory
195        memory.data_mut(&mut *store)[input_ptr as usize..][..input_len as usize]
196            .copy_from_slice(&input_bytes);
197
198        // Call the function
199        let (output_ptr, output_len) = func.call(&mut *store, (input_ptr, input_len))?;
200
201        // Read output from memory
202        let data = memory.data(&store);
203        let output_slice = &data[output_ptr as usize..][..output_len as usize];
204        let output: serde_json::Value = serde_json::from_slice(output_slice)
205            .unwrap_or_else(|_| serde_json::json!({"error": "Invalid JSON output"}));
206
207        // Check CPU limit after execution
208        store.data().check_cpu_limit()?;
209
210        Ok(output)
211    }
212
213    /// Simple bump allocator for WASM memory
214    fn allocate_in_memory(
215        &self,
216        store: &mut Store<PluginState>,
217        memory: Memory,
218        size: i32,
219    ) -> Result<i32> {
220        let data = memory.data_mut(&mut *store);
221        let current_size = data.len() as i32;
222
223        // Simple allocation at end of current data
224        // In production, you'd want a proper allocator
225        let ptr = current_size;
226
227        // Grow memory if needed
228        let needed_size = ptr + size;
229        let current_pages = memory.size(&store);
230        let needed_pages = (needed_size / 65536) + 1;
231
232        if needed_pages > current_pages as i32 {
233            let grow_by = needed_pages - current_pages as i32;
234            memory
235                .grow(&mut *store, grow_by as u64)
236                .with_context(|| "Failed to grow WASM memory")?;
237        }
238
239        // Track memory usage
240        store.data_mut().memory_used += size as u64;
241
242        Ok(ptr)
243    }
244}
245
246#[async_trait::async_trait]
247impl Plugin for WasmPlugin {
248    fn name(&self) -> &str {
249        &self.name
250    }
251
252    fn version(&self) -> &str {
253        &self.version
254    }
255
256    async fn initialize(&self, _context: &PluginContext) -> Layer4Result<()> {
257        // WASM initialization happens at instantiate time
258        Ok(())
259    }
260
261    async fn execute(&self, input: &serde_json::Value) -> Layer4Result<serde_json::Value> {
262        // Check CPU limit before execution
263        self.sandbox.check_cpu_limit()?;
264
265        // Create a new store for this execution
266        let data_dir = std::env::temp_dir();
267        let mut store = self.create_store(data_dir);
268
269        // Instantiate
270        let instance = self
271            .instantiate(&mut store)
272            .map_err(|e| anyhow!("WASM instantiation failed: {}", e))?;
273
274        // Try to find and call a default entry point
275        // Look for common function names
276        for func_name in &["execute", "run", "_start", "main"] {
277            if let Ok(output) = self.execute_func(&mut store, &instance, func_name, input) {
278                return Ok(output);
279            }
280        }
281
282        // If no entry point found, return error with available info
283        Err(anyhow!(
284            "WASM plugin '{}' has no callable entry point. Expected one of: execute, run, _start, main",
285            self.name
286        ))
287    }
288
289    async fn shutdown(&self) -> Layer4Result<()> {
290        Ok(())
291    }
292}
293
294/// WASM plugin loader
295pub struct WasmLoader {
296    /// Wasmtime engine
297    engine: Engine,
298    /// Loaded modules
299    modules: RwLock<HashMap<String, Module>>,
300    /// Plugin instances
301    plugins: RwLock<HashMap<String, Arc<WasmPlugin>>>,
302    /// Default configuration
303    config: WasmConfig,
304}
305
306impl WasmLoader {
307    /// Create new WASM loader
308    pub fn new() -> Layer4Result<Self> {
309        Self::with_config(WasmConfig::default())
310    }
311
312    /// Create WASM loader with custom configuration
313    pub fn with_config(config: WasmConfig) -> Layer4Result<Self> {
314        let mut engine_config = Config::new();
315        engine_config.wasm_backtrace_details(WasmBacktraceDetails::Enable);
316        engine_config.cranelift_opt_level(OptLevel::Speed);
317
318        // Configure memory limits
319        if config.max_memory_bytes > 0 {
320            let max_pages = (config.max_memory_bytes / 65536) + 1;
321            engine_config.wasm_memory64(true);
322            engine_config.static_memory_maximum_size(max_pages * 65536);
323        }
324
325        let engine = Engine::new(&engine_config).context("Failed to create Wasmtime engine")?;
326
327        Ok(Self {
328            engine,
329            modules: RwLock::new(HashMap::new()),
330            plugins: RwLock::new(HashMap::new()),
331            config,
332        })
333    }
334
335    /// Check if file is valid WASM
336    pub fn is_valid_wasm(path: &Path) -> bool {
337        if !path.exists() || !path.is_file() {
338            return false;
339        }
340        path.extension()
341            .and_then(|ext| ext.to_str())
342            .map(|ext| ext == "wasm")
343            .unwrap_or(false)
344    }
345
346    /// Load WASM module
347    pub fn load(&self, path: &Path, capabilities: CapabilitySet) -> Layer4Result<String> {
348        let name = path
349            .file_stem()
350            .and_then(|n| n.to_str())
351            .unwrap_or("unknown")
352            .to_string();
353
354        // Compile module (validates WASM)
355        let module = Module::from_file(&self.engine, path)
356            .with_context(|| format!("Failed to compile WASM: {:?}", path))?;
357
358        // Create sandbox
359        let sandbox = PluginSandbox::new(capabilities);
360
361        // Create plugin instance
362        let plugin = WasmPlugin::new(
363            name.clone(),
364            self.extract_version(&module)
365                .unwrap_or_else(|| "0.1.0".to_string()),
366            sandbox,
367            module,
368            self.engine.clone(),
369            self.config.clone(),
370        );
371
372        // Store module and plugin
373        let module = self.modules.read().get(&name).cloned();
374        if let Some(module) = module {
375            self.modules.write().insert(name.clone(), module);
376        }
377        self.plugins.write().insert(name.clone(), Arc::new(plugin));
378
379        tracing::info!("Loaded WASM plugin: {} from {:?}", name, path);
380
381        Ok(name)
382    }
383
384    /// Extract version from module exports if available
385    fn extract_version(&self, _module: &Module) -> Option<String> {
386        // Version extraction not directly available in wasmtime API
387        // Could be stored in custom section, but that requires different approach
388        // For now, return None and use default version
389        None
390    }
391
392    /// Get loaded plugin
393    pub fn get(&self, name: &str) -> Option<Arc<WasmPlugin>> {
394        self.plugins.read().get(name).cloned()
395    }
396
397    /// Unload plugin
398    pub fn unload(&self, name: &str) -> Layer4Result<()> {
399        self.modules.write().remove(name);
400        self.plugins.write().remove(name);
401        tracing::info!("Unloaded WASM plugin: {}", name);
402        Ok(())
403    }
404
405    /// List loaded plugins
406    pub fn list(&self) -> Vec<String> {
407        self.plugins.read().keys().cloned().collect()
408    }
409
410    /// Get engine reference
411    pub fn engine(&self) -> &Engine {
412        &self.engine
413    }
414
415    /// Load and execute a WASM module in one step
416    pub fn load_and_execute(
417        &self,
418        path: &Path,
419        input: &serde_json::Value,
420        capabilities: CapabilitySet,
421    ) -> Layer4Result<serde_json::Value> {
422        let name = self.load(path, capabilities)?;
423        let plugin = self
424            .get(&name)
425            .ok_or_else(|| anyhow!("Plugin not found after loading: {}", name))?;
426
427        // Use tokio runtime for async execution
428        let rt = tokio::runtime::Runtime::new().context("Failed to create tokio runtime")?;
429
430        rt.block_on(async { plugin.execute(input).await })
431    }
432}
433
434impl Default for WasmLoader {
435    fn default() -> Self {
436        Self::new().expect("Failed to create WasmLoader")
437    }
438}
439
440/// WASI context builder for sandboxed execution
441pub struct WasiContextBuilder {
442    /// Allowed preopened directories (guest_path, host_path, dir_perms, file_perms)
443    preopens: Vec<(String, PathBuf, DirPerms, FilePerms)>,
444    /// Environment variables
445    env: HashMap<String, String>,
446    /// Arguments
447    args: Vec<String>,
448    /// Inherit stdout/stderr
449    inherit_stdio: bool,
450    /// Inherit environment from host
451    inherit_env: bool,
452}
453
454impl WasiContextBuilder {
455    /// Create new WASI context builder
456    pub fn new() -> Self {
457        Self {
458            preopens: Vec::new(),
459            env: HashMap::new(),
460            args: Vec::new(),
461            inherit_stdio: true,
462            inherit_env: false,
463        }
464    }
465
466    /// Add a preopened directory with full permissions
467    pub fn preopen(&mut self, guest_path: &str, host_path: PathBuf) -> &mut Self {
468        self.preopens.push((
469            guest_path.to_string(),
470            host_path,
471            DirPerms::all(),
472            FilePerms::all(),
473        ));
474        self
475    }
476
477    /// Add a preopened directory with read-only permissions
478    pub fn preopen_readonly(&mut self, guest_path: &str, host_path: PathBuf) -> &mut Self {
479        self.preopens.push((
480            guest_path.to_string(),
481            host_path,
482            DirPerms::READ,
483            FilePerms::READ,
484        ));
485        self
486    }
487
488    /// Add a preopened directory with custom permissions
489    pub fn preopen_with_perms(
490        &mut self,
491        guest_path: &str,
492        host_path: PathBuf,
493        dir_perms: DirPerms,
494        file_perms: FilePerms,
495    ) -> &mut Self {
496        self.preopens
497            .push((guest_path.to_string(), host_path, dir_perms, file_perms));
498        self
499    }
500
501    /// Set an environment variable
502    pub fn env(&mut self, key: &str, value: &str) -> &mut Self {
503        self.env.insert(key.to_string(), value.to_string());
504        self
505    }
506
507    /// Add multiple environment variables
508    pub fn envs(&mut self, envs: &[(impl AsRef<str>, impl AsRef<str>)]) -> &mut Self {
509        for (k, v) in envs {
510            self.env
511                .insert(k.as_ref().to_string(), v.as_ref().to_string());
512        }
513        self
514    }
515
516    /// Add an argument
517    pub fn arg(&mut self, arg: &str) -> &mut Self {
518        self.args.push(arg.to_string());
519        self
520    }
521
522    /// Add multiple arguments
523    pub fn args(&mut self, args: &[impl AsRef<str>]) -> &mut Self {
524        for arg in args {
525            self.args.push(arg.as_ref().to_string());
526        }
527        self
528    }
529
530    /// Inherit standard I/O (stdin, stdout, stderr)
531    pub fn inherit_stdio(&mut self, inherit: bool) -> &mut Self {
532        self.inherit_stdio = inherit;
533        self
534    }
535
536    /// Inherit all environment variables from host process
537    pub fn inherit_env(&mut self, inherit: bool) -> &mut Self {
538        self.inherit_env = inherit;
539        self
540    }
541
542    /// Build the WASI context
543    pub fn build(&self) -> WasiCtx {
544        let mut builder = WasiCtxBuilder::new();
545
546        // Configure stdio
547        if self.inherit_stdio {
548            builder.inherit_stdio();
549        }
550
551        // Configure environment
552        if self.inherit_env {
553            builder.inherit_env();
554        }
555
556        // Add custom environment variables
557        for (key, value) in &self.env {
558            builder.env(key, value);
559        }
560
561        // Add arguments
562        for arg in &self.args {
563            builder.arg(arg);
564        }
565
566        // Configure preopened directories
567        for (guest_path, host_path, dir_perms, file_perms) in &self.preopens {
568            builder
569                .preopened_dir(host_path, guest_path, *dir_perms, *file_perms)
570                .expect("Failed to preopen directory");
571        }
572
573        builder.build()
574    }
575}
576
577impl Default for WasiContextBuilder {
578    fn default() -> Self {
579        Self::new()
580    }
581}
582
583#[cfg(test)]
584mod tests {
585    use super::*;
586
587    #[test]
588    fn test_wasm_loader_creation() {
589        let loader = WasmLoader::new();
590        assert!(loader.is_ok());
591        let loader = loader.unwrap();
592        assert!(loader.list().is_empty());
593    }
594
595    #[test]
596    fn test_wasm_config_default() {
597        let config = WasmConfig::default();
598        assert_eq!(config.max_memory_bytes, 16 * 1024 * 1024);
599        assert_eq!(config.max_cpu_time_ms, 5000);
600        assert!(config.enable_wasi);
601    }
602
603    #[test]
604    fn test_is_valid_wasm() {
605        let tmp = tempfile::NamedTempFile::with_suffix(".wasm").unwrap();
606        assert!(WasmLoader::is_valid_wasm(tmp.path()));
607
608        let tmp_txt = tempfile::NamedTempFile::with_suffix(".txt").unwrap();
609        assert!(!WasmLoader::is_valid_wasm(tmp_txt.path()));
610    }
611
612    #[test]
613    fn test_wasm_plugin_creation() {
614        let loader = WasmLoader::new().unwrap();
615        let sandbox = PluginSandbox::sandboxed();
616        let engine = loader.engine().clone();
617        let config = WasmConfig::default();
618
619        // Create an empty module for testing
620        let module = Module::new(&engine, "(module)").unwrap();
621        let plugin = WasmPlugin::new(
622            "test".to_string(),
623            "0.1.0".to_string(),
624            sandbox,
625            module,
626            engine,
627            config,
628        );
629
630        assert_eq!(plugin.name(), "test");
631        assert_eq!(plugin.version(), "0.1.0");
632    }
633
634    #[test]
635    fn test_wasi_context_builder() {
636        let mut builder = WasiContextBuilder::new();
637        builder.env("TEST", "value");
638        builder.arg("--help");
639
640        // Build WASI context
641        let _ctx = builder.build();
642        // WasiCtx is created successfully (no need to verify contents directly)
643        // The context is now a real wasmtime-wasi WasiCtx
644    }
645
646    #[test]
647    fn test_wasi_context_builder_with_preopen() {
648        let mut builder = WasiContextBuilder::new();
649        builder.env("HOME", "/home/user");
650        builder.arg("--test");
651        builder.preopen("/tmp", std::env::temp_dir());
652
653        let _ctx = builder.build();
654        // Successfully built with preopened directory
655    }
656
657    #[test]
658    fn test_wasi_context_builder_readonly() {
659        let mut builder = WasiContextBuilder::new();
660        builder.preopen_readonly("/data", std::env::temp_dir());
661        builder.inherit_stdio(true);
662        builder.inherit_env(false);
663
664        let _ctx = builder.build();
665        // Successfully built with read-only preopen
666    }
667
668    #[tokio::test]
669    async fn test_plugin_execute_without_entry_point_returns_error() {
670        let loader = WasmLoader::new().unwrap();
671        let sandbox = PluginSandbox::sandboxed();
672        let engine = loader.engine().clone();
673        let config = WasmConfig::default();
674
675        // Create minimal module
676        let module = Module::new(&engine, "(module)").unwrap();
677        let plugin = WasmPlugin::new(
678            "test".to_string(),
679            "0.1.0".to_string(),
680            sandbox,
681            module,
682            engine,
683            config,
684        );
685
686        let input = serde_json::json!({"test": "input"});
687        let result = plugin.execute(&input).await;
688        assert!(result.is_err());
689        let error = result.unwrap_err().to_string();
690        assert!(error.contains("no callable entry point"));
691        assert!(error.contains("execute, run, _start, main"));
692    }
693}
sh_layer4/plugin_loader/wasm.rs

sh_layer4/plugin_loader/
wasm.rs
