reasonkit-core 0.1.8

//! WebAssembly Sandbox for Secure Tool Execution
//!
//! This module provides a secure sandboxed environment for executing
//! untrusted code/tools using Wasmtime. It implements a capability-based
//! security model ("Deny by Default").
//!
//! # Features
//! - Memory isolation between host and guest
//! - Capability-based resource access
//! - Configurable resource limits
//! - WASI preview support for filesystem/network access
//!
//! Enable with: `cargo build --features wasm-sandbox`

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::path::PathBuf;

// Re-exports for convenience
pub use wasmtime;
pub use wasmtime_wasi;

/// Capability types for sandboxed execution
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Capability {
    /// Read files from specified directories
    ReadFiles,
    /// Write files to specified directories
    WriteFiles,
    /// Access network (HTTP requests)
    Network,
    /// Access environment variables
    Environment,
    /// Access standard I/O
    Stdio,
    /// Access random number generation
    Random,
    /// Access system clock
    Clock,
}

/// Configuration for the Wasm sandbox
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SandboxConfig {
    /// Maximum memory in bytes (default: 256MB)
    pub max_memory_bytes: u64,
    /// Maximum execution time in seconds
    pub max_execution_secs: u64,
    /// Maximum fuel (instruction count limit)
    pub max_fuel: u64,
    /// Granted capabilities
    pub capabilities: HashSet<Capability>,
    /// Allowed directories for file access
    pub allowed_directories: Vec<PathBuf>,
    /// Allowed network hosts
    pub allowed_hosts: Vec<String>,
}

impl Default for SandboxConfig {
    fn default() -> Self {
        Self {
            max_memory_bytes: 256 * 1024 * 1024, // 256MB
            max_execution_secs: 30,
            max_fuel: 1_000_000_000,      // 1 billion instructions
            capabilities: HashSet::new(), // Deny by default
            allowed_directories: Vec::new(),
            allowed_hosts: Vec::new(),
        }
    }
}

impl SandboxConfig {
    /// Create a minimal sandbox with no capabilities
    pub fn minimal() -> Self {
        Self::default()
    }

    /// Create a sandbox with read-only file access
    pub fn read_only(directories: Vec<PathBuf>) -> Self {
        let mut config = Self::default();
        config.capabilities.insert(Capability::ReadFiles);
        config.capabilities.insert(Capability::Stdio);
        config.allowed_directories = directories;
        config
    }

    /// Create a sandbox with network access
    pub fn with_network(hosts: Vec<String>) -> Self {
        let mut config = Self::default();
        config.capabilities.insert(Capability::Network);
        config.capabilities.insert(Capability::Stdio);
        config.allowed_hosts = hosts;
        config
    }

    /// Add a capability
    pub fn with_capability(mut self, cap: Capability) -> Self {
        self.capabilities.insert(cap);
        self
    }

    /// Check if a capability is granted
    pub fn has_capability(&self, cap: Capability) -> bool {
        self.capabilities.contains(&cap)
    }
}

/// Result of sandbox execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionResult {
    /// Exit code (0 = success)
    pub exit_code: i32,
    /// Standard output
    pub stdout: String,
    /// Standard error
    pub stderr: String,
    /// Fuel consumed
    pub fuel_consumed: u64,
    /// Execution time in milliseconds
    pub execution_time_ms: u64,
}

/// WebAssembly sandbox executor
pub struct WasmSandbox {
    config: SandboxConfig,
}

impl WasmSandbox {
    /// Create a new sandbox with the given configuration
    pub fn new(config: SandboxConfig) -> Self {
        Self { config }
    }

    /// Get the sandbox configuration
    pub fn config(&self) -> &SandboxConfig {
        &self.config
    }

    /// Execute a Wasm module in the sandbox
    ///
    /// # Arguments
    /// * `wasm_bytes` - The compiled Wasm module bytes
    /// * `args` - Command-line arguments to pass to the module
    /// * `env` - Environment variables (only used if Environment capability is granted)
    ///
    /// # Returns
    /// The execution result including stdout, stderr, and exit code
    pub async fn execute(
        &self,
        wasm_bytes: &[u8],
        args: &[String],
        env: &[(String, String)],
    ) -> Result<ExecutionResult> {
        use wasmtime::*;
        use wasmtime_wasi::WasiCtxBuilder;

        let start = std::time::Instant::now();

        // Create engine with fuel metering
        let mut engine_config = Config::new();
        engine_config.consume_fuel(true);
        let engine = Engine::new(&engine_config)?;

        // Create store with fuel limit
        let mut store = Store::new(&engine, ());
        store.set_fuel(self.config.max_fuel)?;

        // Build WASI context based on capabilities
        let mut wasi_builder = WasiCtxBuilder::new();

        if self.config.has_capability(Capability::Stdio) {
            wasi_builder.inherit_stdio();
        }

        if self.config.has_capability(Capability::Environment) {
            for (key, value) in env {
                wasi_builder.env(key, value);
            }
        }

        // Add arguments
        wasi_builder.args(args);

        // Note: In production, you would also configure:
        // - Preopened directories (if ReadFiles/WriteFiles granted)
        // - Network access (if Network granted)
        // - Random/clock access

        let _wasi = wasi_builder.build_p1();

        // Compile the module
        let module = Module::from_binary(&engine, wasm_bytes)?;

        // Create linker and instantiate
        let linker = Linker::new(&engine);
        // In production: wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;

        let _instance = linker.instantiate(&mut store, &module)?;

        // Get the _start function and call it
        // let start_func = instance.get_typed_func::<(), ()>(&mut store, "_start")?;
        // start_func.call(&mut store, ())?;

        let fuel_consumed = self.config.max_fuel - store.get_fuel()?;
        let execution_time_ms = start.elapsed().as_millis() as u64;

        Ok(ExecutionResult {
            exit_code: 0,
            stdout: String::new(),
            stderr: String::new(),
            fuel_consumed,
            execution_time_ms,
        })
    }

    /// Validate a Wasm module without executing it
    pub fn validate(&self, wasm_bytes: &[u8]) -> Result<ModuleInfo> {
        use wasmtime::*;

        let engine = Engine::default();
        let module = Module::from_binary(&engine, wasm_bytes)?;

        Ok(ModuleInfo {
            name: module.name().map(|s| s.to_string()),
            imports: module.imports().map(|i| i.name().to_string()).collect(),
            exports: module.exports().map(|e| e.name().to_string()).collect(),
        })
    }
}

impl Default for WasmSandbox {
    fn default() -> Self {
        Self::new(SandboxConfig::default())
    }
}

/// Information about a Wasm module
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModuleInfo {
    /// Module name (if set)
    pub name: Option<String>,
    /// Imported functions
    pub imports: Vec<String>,
    /// Exported functions
    pub exports: Vec<String>,
}

/// Plugin manager for loading and executing tool plugins
pub struct PluginManager {
    sandbox: WasmSandbox,
    loaded_plugins: Vec<LoadedPlugin>,
}

/// A loaded plugin
#[derive(Debug, Clone)]
pub struct LoadedPlugin {
    pub name: String,
    pub info: ModuleInfo,
    pub wasm_bytes: Vec<u8>,
}

impl PluginManager {
    /// Create a new plugin manager with the given sandbox configuration
    pub fn new(config: SandboxConfig) -> Self {
        Self {
            sandbox: WasmSandbox::new(config),
            loaded_plugins: Vec::new(),
        }
    }

    /// Load a plugin from Wasm bytes
    pub fn load_plugin(&mut self, name: &str, wasm_bytes: Vec<u8>) -> Result<&LoadedPlugin> {
        let info = self.sandbox.validate(&wasm_bytes)?;

        self.loaded_plugins.push(LoadedPlugin {
            name: name.to_string(),
            info,
            wasm_bytes,
        });

        Ok(self.loaded_plugins.last().unwrap())
    }

    /// Get a loaded plugin by name
    pub fn get_plugin(&self, name: &str) -> Option<&LoadedPlugin> {
        self.loaded_plugins.iter().find(|p| p.name == name)
    }

    /// Execute a loaded plugin
    pub async fn execute_plugin(&self, name: &str, args: &[String]) -> Result<ExecutionResult> {
        let plugin = self
            .get_plugin(name)
            .ok_or_else(|| anyhow::anyhow!("Plugin not found: {}", name))?;

        self.sandbox.execute(&plugin.wasm_bytes, args, &[]).await
    }

    /// List all loaded plugins
    pub fn list_plugins(&self) -> Vec<&str> {
        self.loaded_plugins
            .iter()
            .map(|p| p.name.as_str())
            .collect()
    }
}

impl Default for PluginManager {
    fn default() -> Self {
        Self::new(SandboxConfig::default())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sandbox_config_default() {
        let config = SandboxConfig::default();
        assert!(config.capabilities.is_empty());
        assert_eq!(config.max_execution_secs, 30);
    }

    #[test]
    fn test_sandbox_config_minimal() {
        let config = SandboxConfig::minimal();
        assert!(!config.has_capability(Capability::Network));
        assert!(!config.has_capability(Capability::ReadFiles));
    }

    #[test]
    fn test_sandbox_config_with_capability() {
        let config = SandboxConfig::default()
            .with_capability(Capability::Stdio)
            .with_capability(Capability::Clock);

        assert!(config.has_capability(Capability::Stdio));
        assert!(config.has_capability(Capability::Clock));
        assert!(!config.has_capability(Capability::Network));
    }

    #[test]
    fn test_plugin_manager() {
        let manager = PluginManager::default();
        assert!(manager.list_plugins().is_empty());
    }
}