cortexai-wasm 0.1.0

//! Tests for the WASM tool sandbox
//!
//! These tests verify the sandbox functionality for running
//! untrusted WASM modules with capability-based restrictions.

#![cfg(feature = "sandbox")]

use std::time::Duration;

use cortexai_core::tool::{ExecutionContext, Tool};
use cortexai_core::types::AgentId;
use cortexai_wasm::sandbox::{SandboxConfig, SandboxedTool, ToolSandbox};
use cortexai_wasm::sandbox_error::SandboxError;

#[test]
fn sandbox_config_has_correct_defaults() {
    let config = SandboxConfig::default();

    assert_eq!(config.max_memory_bytes(), 64 * 1024 * 1024); // 64MB
    assert_eq!(config.max_execution_time(), Duration::from_secs(30));
    assert!(config.allowed_paths().is_empty());
    assert!(!config.allow_network());
    assert!(config.max_fuel().is_none());
}

#[test]
fn sandbox_config_builder_sets_custom_values() {
    let config = SandboxConfig::builder()
        .max_memory_bytes(128 * 1024 * 1024)
        .max_execution_time(Duration::from_secs(10))
        .allowed_path("/guest/data".to_string(), "/host/data".to_string())
        .allow_network(true)
        .max_fuel(1_000_000)
        .build();

    assert_eq!(config.max_memory_bytes(), 128 * 1024 * 1024);
    assert_eq!(config.max_execution_time(), Duration::from_secs(10));
    assert_eq!(config.allowed_paths().len(), 1);
    assert_eq!(config.allowed_paths()[0].0, "/guest/data");
    assert_eq!(config.allowed_paths()[0].1, "/host/data");
    assert!(config.allow_network());
    assert_eq!(config.max_fuel(), Some(1_000_000));
}

#[test]
fn sandbox_error_variants_have_meaningful_messages() {
    let err = SandboxError::CompilationFailed("invalid wasm".to_string());
    assert!(err.to_string().contains("invalid wasm"));

    let err = SandboxError::ExecutionFailed("trap occurred".to_string());
    assert!(err.to_string().contains("trap occurred"));

    let err = SandboxError::MemoryLimitExceeded { limit_bytes: 64 * 1024 * 1024 };
    assert!(err.to_string().contains("67108864"));

    let err = SandboxError::TimeoutExceeded {
        limit: Duration::from_secs(30),
    };
    assert!(err.to_string().contains("30s"));

    let err = SandboxError::FuelExhausted { fuel_limit: 1000 };
    assert!(err.to_string().contains("1000"));
}

#[test]
fn tool_sandbox_can_be_created_with_config() {
    let config = SandboxConfig::builder()
        .max_memory_bytes(32 * 1024 * 1024)
        .build();
    let sandbox = ToolSandbox::new(config.clone());
    assert_eq!(sandbox.config().max_memory_bytes(), 32 * 1024 * 1024);
}

#[test]
fn tool_sandbox_rejects_invalid_wasm() {
    let sandbox = ToolSandbox::new(SandboxConfig::default());
    let result = sandbox.load_module(b"not valid wasm");
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(matches!(err, SandboxError::CompilationFailed(_)));
}

/// WAT module that exports an `add` function: (i32, i32) -> i32
const ADD_WAT: &str = r#"
(module
  (func $add (export "add") (param i32 i32) (result i32)
    local.get 0
    local.get 1
    i32.add
  )
)
"#;

#[test]
fn tool_sandbox_loads_and_executes_wat_module() {
    let sandbox = ToolSandbox::new(SandboxConfig::default());
    let module = sandbox.load_module(ADD_WAT.as_bytes()).unwrap();
    let result = sandbox.execute(&module, "add", &[3i32.to_le_bytes().as_slice(), 4i32.to_le_bytes().as_slice()].concat()).unwrap();

    // Result should be the i32 value 7, little-endian
    assert_eq!(result.len(), 4);
    let value = i32::from_le_bytes(result[0..4].try_into().unwrap());
    assert_eq!(value, 7);
}

/// WAT module with an infinite loop to test fuel exhaustion.
const LOOP_WAT: &str = r#"
(module
  (func $spin (export "spin") (result i32)
    (local i32)
    (loop $loop
      (local.set 0 (i32.add (local.get 0) (i32.const 1)))
      (br $loop)
    )
    local.get 0
  )
)
"#;

#[test]
fn tool_sandbox_exhausts_fuel_on_infinite_loop() {
    let config = SandboxConfig::builder()
        .max_fuel(1000)
        .build();
    let sandbox = ToolSandbox::new(config);
    let module = sandbox.load_module(LOOP_WAT.as_bytes()).unwrap();
    let result = sandbox.execute(&module, "spin", &[]);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        matches!(err, SandboxError::FuelExhausted { fuel_limit: 1000 }),
        "Expected FuelExhausted, got: {:?}",
        err
    );
}

/// WAT module that acts as a tool: has __tool_schema and __tool_execute exports.
/// __tool_schema returns a pointer+length to a JSON schema string in memory.
/// __tool_execute accepts a pointer+length to JSON args, returns pointer+length to JSON result.
///
/// For simplicity, this module stores the schema and result in linear memory and
/// returns the offset and length.
/// Helper to build WAT module bytes with embedded JSON data.
fn build_tool_wat() -> Vec<u8> {
    let schema = serde_json::json!({
        "name": "add_tool",
        "description": "Adds two numbers",
        "parameters": {},
        "dangerous": false,
        "metadata": {},
        "required_scopes": []
    });
    let schema_bytes = serde_json::to_vec(&schema).unwrap();
    let schema_len = schema_bytes.len();

    let result_bytes = br#"{"result":42}"#;
    let result_len = result_bytes.len();
    let result_offset = 512;

    // Encode data bytes as WAT hex escape sequences
    let schema_hex = bytes_to_wat_hex(&schema_bytes);
    let result_hex = bytes_to_wat_hex(result_bytes);

    let wat = format!(
        r#"(module
  (memory (export "memory") 1)
  (data (i32.const 0) "{schema_hex}")
  (func (export "__tool_schema") (result i32 i32)
    i32.const 0
    i32.const {schema_len}
  )
  (data (i32.const {result_offset}) "{result_hex}")
  (func (export "__tool_execute") (param i32 i32) (result i32 i32)
    i32.const {result_offset}
    i32.const {result_len}
  )
)"#
    );

    wat.into_bytes()
}

/// Convert bytes to WAT hex escape string (e.g., \00\ff).
fn bytes_to_wat_hex(data: &[u8]) -> String {
    data.iter().map(|b| format!("\\{:02x}", b)).collect()
}

#[tokio::test]
async fn sandboxed_tool_returns_schema_from_wasm_module() {
    let wat = build_tool_wat();
    let sandbox = ToolSandbox::new(SandboxConfig::default());
    let module = sandbox.load_module(&wat).unwrap();
    let tool = SandboxedTool::new(sandbox, module).unwrap();

    let schema = tool.schema();
    assert_eq!(schema.name, "add_tool");
    assert_eq!(schema.description, "Adds two numbers");
}

#[tokio::test]
async fn sandboxed_tool_executes_via_tool_trait() {
    let wat = build_tool_wat();
    let sandbox = ToolSandbox::new(SandboxConfig::default());
    let module = sandbox.load_module(&wat).unwrap();
    let tool = SandboxedTool::new(sandbox, module).unwrap();

    let ctx = ExecutionContext::new(AgentId::new("test-agent"));
    let args = serde_json::json!({"a": 1, "b": 2});
    let result: serde_json::Value = tool.execute(&ctx, args).await.unwrap();
    assert_eq!(result["result"], 42);
}

#[test]
fn sandboxed_tool_rejects_module_without_schema_export() {
    let sandbox = ToolSandbox::new(SandboxConfig::default());
    // The ADD_WAT module has no __tool_schema export
    let module = sandbox.load_module(ADD_WAT.as_bytes()).unwrap();
    let result = SandboxedTool::new(sandbox, module);
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert!(
        matches!(err, SandboxError::ExecutionFailed(_)),
        "Expected ExecutionFailed, got: {:?}",
        err
    );
}