hehe_tools/sandbox/
native.rs

1use crate::error::Result;
2use crate::traits::{Tool, ToolOutput};
3use async_trait::async_trait;
4use hehe_core::Context;
5use serde_json::Value;
6use std::collections::HashSet;
7use std::path::PathBuf;
8use std::sync::Arc;
9
10#[derive(Clone, Debug)]
11pub struct SandboxConfig {
12    pub allowed_paths: HashSet<PathBuf>,
13    pub denied_paths: HashSet<PathBuf>,
14    pub allowed_hosts: HashSet<String>,
15    pub denied_hosts: HashSet<String>,
16    pub allow_shell: bool,
17    pub allow_network: bool,
18    pub max_file_size: usize,
19    pub max_output_size: usize,
20}
21
22impl Default for SandboxConfig {
23    fn default() -> Self {
24        Self {
25            allowed_paths: HashSet::new(),
26            denied_paths: HashSet::new(),
27            allowed_hosts: HashSet::new(),
28            denied_hosts: HashSet::new(),
29            allow_shell: false,
30            allow_network: true,
31            max_file_size: 10 * 1024 * 1024,
32            max_output_size: 1024 * 1024,
33        }
34    }
35}
36
37impl SandboxConfig {
38    pub fn new() -> Self {
39        Self::default()
40    }
41
42    pub fn allow_all() -> Self {
43        Self {
44            allow_shell: true,
45            allow_network: true,
46            ..Default::default()
47        }
48    }
49
50    pub fn allow_path(mut self, path: impl Into<PathBuf>) -> Self {
51        self.allowed_paths.insert(path.into());
52        self
53    }
54
55    pub fn deny_path(mut self, path: impl Into<PathBuf>) -> Self {
56        self.denied_paths.insert(path.into());
57        self
58    }
59
60    pub fn allow_host(mut self, host: impl Into<String>) -> Self {
61        self.allowed_hosts.insert(host.into());
62        self
63    }
64
65    pub fn deny_host(mut self, host: impl Into<String>) -> Self {
66        self.denied_hosts.insert(host.into());
67        self
68    }
69
70    pub fn with_shell(mut self, allow: bool) -> Self {
71        self.allow_shell = allow;
72        self
73    }
74
75    pub fn with_network(mut self, allow: bool) -> Self {
76        self.allow_network = allow;
77        self
78    }
79
80    pub fn is_path_allowed(&self, path: &PathBuf) -> bool {
81        for denied in &self.denied_paths {
82            if path.starts_with(denied) {
83                return false;
84            }
85        }
86
87        if self.allowed_paths.is_empty() {
88            return true;
89        }
90
91        for allowed in &self.allowed_paths {
92            if path.starts_with(allowed) {
93                return true;
94            }
95        }
96
97        false
98    }
99
100    pub fn is_host_allowed(&self, host: &str) -> bool {
101        if !self.allow_network {
102            return false;
103        }
104
105        if self.denied_hosts.contains(host) {
106            return false;
107        }
108
109        if self.allowed_hosts.is_empty() {
110            return true;
111        }
112
113        self.allowed_hosts.contains(host)
114    }
115}
116
117#[async_trait]
118pub trait Sandbox: Send + Sync {
119    fn config(&self) -> &SandboxConfig;
120
121    fn check_tool(&self, tool: &dyn Tool) -> Result<()>;
122
123    async fn execute(
124        &self,
125        tool: Arc<dyn Tool>,
126        ctx: &Context,
127        input: Value,
128    ) -> Result<ToolOutput>;
129}
130
131pub struct NativeSandbox {
132    config: SandboxConfig,
133}
134
135impl NativeSandbox {
136    pub fn new(config: SandboxConfig) -> Self {
137        Self { config }
138    }
139
140    pub fn permissive() -> Self {
141        Self::new(SandboxConfig::allow_all())
142    }
143}
144
145impl Default for NativeSandbox {
146    fn default() -> Self {
147        Self::new(SandboxConfig::default())
148    }
149}
150
151#[async_trait]
152impl Sandbox for NativeSandbox {
153    fn config(&self) -> &SandboxConfig {
154        &self.config
155    }
156
157    fn check_tool(&self, tool: &dyn Tool) -> Result<()> {
158        let name = tool.name();
159        
160        if tool.is_dangerous() && !self.config.allow_shell {
161            if name == "execute_shell" {
162                return Err(crate::error::ToolError::permission_denied(
163                    "Shell execution is not allowed in this sandbox",
164                ));
165            }
166        }
167
168        Ok(())
169    }
170
171    async fn execute(
172        &self,
173        tool: Arc<dyn Tool>,
174        ctx: &Context,
175        input: Value,
176    ) -> Result<ToolOutput> {
177        self.check_tool(tool.as_ref())?;
178        tool.execute(ctx, input).await
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185    use std::path::PathBuf;
186
187    #[test]
188    fn test_sandbox_config_default() {
189        let config = SandboxConfig::default();
190        assert!(!config.allow_shell);
191        assert!(config.allow_network);
192    }
193
194    #[test]
195    fn test_sandbox_config_path_check() {
196        let config = SandboxConfig::new()
197            .allow_path("/home/user/workspace")
198            .deny_path("/home/user/workspace/secrets");
199
200        assert!(config.is_path_allowed(&PathBuf::from("/home/user/workspace/project")));
201        assert!(!config.is_path_allowed(&PathBuf::from("/home/user/workspace/secrets/key")));
202        assert!(!config.is_path_allowed(&PathBuf::from("/etc/passwd")));
203    }
204
205    #[test]
206    fn test_sandbox_config_host_check() {
207        let config = SandboxConfig::new()
208            .allow_host("api.example.com")
209            .deny_host("malicious.com");
210
211        assert!(config.is_host_allowed("api.example.com"));
212        assert!(!config.is_host_allowed("malicious.com"));
213        assert!(!config.is_host_allowed("other.com"));
214    }
215
216    #[test]
217    fn test_sandbox_config_no_restrictions() {
218        let config = SandboxConfig::allow_all();
219        
220        assert!(config.is_path_allowed(&PathBuf::from("/any/path")));
221        assert!(config.is_host_allowed("any.host.com"));
222    }
223}
hehe_tools/sandbox/native.rs

hehe_tools/sandbox/
native.rs
