moe_plugin_sdk/security/
sandbox.rs1use 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 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 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}