Skip to main content

moe_plugin_sdk/security/
sandbox.rs

1use anyhow::Result;
2use std::sync::mpsc;
3use std::time::Duration;
4use ternlang_moe::{AxisMemory, VetoEntry};
5
6pub struct PluginSandbox {
7    pub memory_limit_mb: u64,
8    pub cpu_time_budget_ms: u64,
9}
10
11impl PluginSandbox {
12    pub fn new(memory_limit_mb: u64, cpu_time_budget_ms: u64) -> Self {
13        Self { memory_limit_mb, cpu_time_budget_ms }
14    }
15
16    /// Execute `f` inside the sandbox and log results to AxisMesh.
17    pub fn enforce<T, F>(&self, axis: &mut AxisMemory, f: F) -> Result<T>
18    where
19        T: Send + 'static,
20        F: FnOnce() -> Result<T> + Send + 'static,
21    {
22        let budget = self.cpu_time_budget_ms;
23        let (tx, rx) = mpsc::channel::<Result<T>>();
24        std::thread::spawn(move || {
25            let _ = tx.send(f());
26        });
27        
28        match rx.recv_timeout(Duration::from_millis(budget)) {
29            Ok(result) => result,
30            Err(_) => {
31                // Bridge to AxisMesh: Log the VetoEntry
32                axis.veto_log.push(VetoEntry {
33                    timestamp: std::time::SystemTime::now(),
34                    expert_id: 0,
35                    reason: format!("Plugin exceeded CPU time budget of {}ms", budget),
36                    query_hash: 0,
37                });
38                
39                anyhow::bail!("Plugin exceeded CPU time budget of {}ms (VetoEntry logged in AxisMesh)", budget)
40            }
41        }
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48
49    #[test]
50    fn test_sandbox_allows_fast_work() {
51        let mut axis = AxisMemory::new();
52        let sandbox = PluginSandbox::new(64, 500);
53        let result = sandbox.enforce(&mut axis, || Ok(42u32));
54        assert_eq!(result.unwrap(), 42);
55    }
56
57    #[test]
58    fn test_sandbox_enforces_timeout() {
59        let mut axis = AxisMemory::new();
60        let sandbox = PluginSandbox::new(64, 50);
61        let result = sandbox.enforce::<(), _>(&mut axis, || {
62            std::thread::sleep(Duration::from_millis(200));
63            Ok(())
64        });
65        assert!(result.is_err(), "sandbox must reject work that exceeds time budget");
66        let msg = result.unwrap_err().to_string();
67        assert!(msg.contains("CPU time budget"), "error must mention time budget");
68        assert_eq!(axis.veto_log.len(), 1, "timeout must be logged as a VetoEntry");
69    }
70
71    #[test]
72    fn test_sandbox_propagates_inner_error() {
73        let mut axis = AxisMemory::new();
74        let sandbox = PluginSandbox::new(64, 500);
75        let result = sandbox.enforce::<(), _>(&mut axis, || {
76            anyhow::bail!("plugin internal failure")
77        });
78        assert!(result.is_err());
79        assert!(result.unwrap_err().to_string().contains("plugin internal failure"));
80    }
81}