1use std::sync::atomic::{AtomicU8, Ordering};
7
8use crate::TrustLevel;
9
10use crate::executor::{ToolCall, ToolError, ToolExecutor, ToolOutput};
11use crate::permissions::{PermissionAction, PermissionPolicy};
12use crate::registry::ToolDef;
13
14const QUARANTINE_DENIED: &[&str] = &["bash", "file_write", "web_scrape"];
16
17fn trust_to_u8(level: TrustLevel) -> u8 {
18 match level {
19 TrustLevel::Trusted => 0,
20 TrustLevel::Verified => 1,
21 TrustLevel::Quarantined => 2,
22 TrustLevel::Blocked => 3,
23 }
24}
25
26fn u8_to_trust(v: u8) -> TrustLevel {
27 match v {
28 0 => TrustLevel::Trusted,
29 1 => TrustLevel::Verified,
30 2 => TrustLevel::Quarantined,
31 _ => TrustLevel::Blocked,
32 }
33}
34
35pub struct TrustGateExecutor<T: ToolExecutor> {
37 inner: T,
38 policy: PermissionPolicy,
39 effective_trust: AtomicU8,
40}
41
42impl<T: ToolExecutor + std::fmt::Debug> std::fmt::Debug for TrustGateExecutor<T> {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 f.debug_struct("TrustGateExecutor")
45 .field("inner", &self.inner)
46 .field("policy", &self.policy)
47 .field("effective_trust", &self.effective_trust())
48 .finish()
49 }
50}
51
52impl<T: ToolExecutor> TrustGateExecutor<T> {
53 #[must_use]
54 pub fn new(inner: T, policy: PermissionPolicy) -> Self {
55 Self {
56 inner,
57 policy,
58 effective_trust: AtomicU8::new(trust_to_u8(TrustLevel::Trusted)),
59 }
60 }
61
62 pub fn set_effective_trust(&self, level: TrustLevel) {
63 self.effective_trust
64 .store(trust_to_u8(level), Ordering::Relaxed);
65 }
66
67 #[must_use]
68 pub fn effective_trust(&self) -> TrustLevel {
69 u8_to_trust(self.effective_trust.load(Ordering::Relaxed))
70 }
71
72 fn check_trust(&self, tool_id: &str, input: &str) -> Result<(), ToolError> {
73 match self.effective_trust() {
74 TrustLevel::Blocked => {
75 return Err(ToolError::Blocked {
76 command: "all tools blocked (trust=blocked)".to_owned(),
77 });
78 }
79 TrustLevel::Quarantined => {
80 if QUARANTINE_DENIED.contains(&tool_id) {
81 return Err(ToolError::Blocked {
82 command: format!("{tool_id} denied (trust=quarantined)"),
83 });
84 }
85 }
86 TrustLevel::Trusted | TrustLevel::Verified => {}
87 }
88
89 match self.policy.check(tool_id, input) {
90 PermissionAction::Allow => Ok(()),
91 PermissionAction::Ask => Err(ToolError::ConfirmationRequired {
92 command: input.to_owned(),
93 }),
94 PermissionAction::Deny => Err(ToolError::Blocked {
95 command: input.to_owned(),
96 }),
97 }
98 }
99}
100
101impl<T: ToolExecutor> ToolExecutor for TrustGateExecutor<T> {
102 async fn execute(&self, response: &str) -> Result<Option<ToolOutput>, ToolError> {
103 match self.effective_trust() {
104 TrustLevel::Blocked | TrustLevel::Quarantined => {
105 return Err(ToolError::Blocked {
106 command: format!(
107 "tool execution denied (trust={})",
108 format!("{:?}", self.effective_trust()).to_lowercase()
109 ),
110 });
111 }
112 TrustLevel::Trusted | TrustLevel::Verified => {}
113 }
114 self.inner.execute(response).await
115 }
116
117 async fn execute_confirmed(&self, response: &str) -> Result<Option<ToolOutput>, ToolError> {
118 match self.effective_trust() {
119 TrustLevel::Blocked | TrustLevel::Quarantined => {
120 return Err(ToolError::Blocked {
121 command: format!(
122 "tool execution denied (trust={})",
123 format!("{:?}", self.effective_trust()).to_lowercase()
124 ),
125 });
126 }
127 TrustLevel::Trusted | TrustLevel::Verified => {}
128 }
129 self.inner.execute_confirmed(response).await
130 }
131
132 fn tool_definitions(&self) -> Vec<ToolDef> {
133 self.inner.tool_definitions()
134 }
135
136 async fn execute_tool_call(&self, call: &ToolCall) -> Result<Option<ToolOutput>, ToolError> {
137 let input = call
138 .params
139 .get("command")
140 .and_then(|v| v.as_str())
141 .unwrap_or("");
142 self.check_trust(&call.tool_id, input)?;
143 self.inner.execute_tool_call(call).await
144 }
145
146 fn set_skill_env(&self, env: Option<std::collections::HashMap<String, String>>) {
147 self.inner.set_skill_env(env);
148 }
149
150 fn set_effective_trust(&self, level: crate::TrustLevel) {
151 self.effective_trust
152 .store(trust_to_u8(level), Ordering::Relaxed);
153 }
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159
160 #[derive(Debug)]
161 struct MockExecutor;
162 impl ToolExecutor for MockExecutor {
163 async fn execute(&self, _: &str) -> Result<Option<ToolOutput>, ToolError> {
164 Ok(None)
165 }
166 async fn execute_tool_call(
167 &self,
168 call: &ToolCall,
169 ) -> Result<Option<ToolOutput>, ToolError> {
170 Ok(Some(ToolOutput {
171 tool_name: call.tool_id.clone(),
172 summary: "ok".into(),
173 blocks_executed: 1,
174 filter_stats: None,
175 diff: None,
176 streamed: false,
177 terminal_id: None,
178 locations: None,
179 raw_response: None,
180 }))
181 }
182 }
183
184 fn make_call(tool_id: &str) -> ToolCall {
185 ToolCall {
186 tool_id: tool_id.into(),
187 params: serde_json::Map::new(),
188 }
189 }
190
191 fn make_call_with_cmd(tool_id: &str, cmd: &str) -> ToolCall {
192 let mut params = serde_json::Map::new();
193 params.insert("command".into(), serde_json::Value::String(cmd.into()));
194 ToolCall {
195 tool_id: tool_id.into(),
196 params,
197 }
198 }
199
200 #[tokio::test]
201 async fn trusted_allows_all() {
202 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
203 gate.set_effective_trust(TrustLevel::Trusted);
204
205 let result = gate.execute_tool_call(&make_call("bash")).await;
206 assert!(matches!(
208 result,
209 Err(ToolError::ConfirmationRequired { .. })
210 ));
211 }
212
213 #[tokio::test]
214 async fn quarantined_denies_bash() {
215 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
216 gate.set_effective_trust(TrustLevel::Quarantined);
217
218 let result = gate.execute_tool_call(&make_call("bash")).await;
219 assert!(matches!(result, Err(ToolError::Blocked { .. })));
220 }
221
222 #[tokio::test]
223 async fn quarantined_denies_file_write() {
224 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
225 gate.set_effective_trust(TrustLevel::Quarantined);
226
227 let result = gate.execute_tool_call(&make_call("file_write")).await;
228 assert!(matches!(result, Err(ToolError::Blocked { .. })));
229 }
230
231 #[tokio::test]
232 async fn quarantined_allows_file_read() {
233 let policy = crate::permissions::PermissionPolicy::from_legacy(&[], &[]);
234 let gate = TrustGateExecutor::new(MockExecutor, policy);
235 gate.set_effective_trust(TrustLevel::Quarantined);
236
237 let result = gate.execute_tool_call(&make_call("file_read")).await;
238 assert!(matches!(
240 result,
241 Err(ToolError::ConfirmationRequired { .. })
242 ));
243 }
244
245 #[tokio::test]
246 async fn blocked_denies_everything() {
247 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
248 gate.set_effective_trust(TrustLevel::Blocked);
249
250 let result = gate.execute_tool_call(&make_call("file_read")).await;
251 assert!(matches!(result, Err(ToolError::Blocked { .. })));
252 }
253
254 #[tokio::test]
255 async fn policy_deny_overrides_trust() {
256 let policy = crate::permissions::PermissionPolicy::from_legacy(&["sudo".into()], &[]);
257 let gate = TrustGateExecutor::new(MockExecutor, policy);
258 gate.set_effective_trust(TrustLevel::Trusted);
259
260 let result = gate
261 .execute_tool_call(&make_call_with_cmd("bash", "sudo rm"))
262 .await;
263 assert!(matches!(result, Err(ToolError::Blocked { .. })));
264 }
265
266 #[tokio::test]
267 async fn blocked_denies_execute() {
268 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
269 gate.set_effective_trust(TrustLevel::Blocked);
270
271 let result = gate.execute("some response").await;
272 assert!(matches!(result, Err(ToolError::Blocked { .. })));
273 }
274
275 #[tokio::test]
276 async fn blocked_denies_execute_confirmed() {
277 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
278 gate.set_effective_trust(TrustLevel::Blocked);
279
280 let result = gate.execute_confirmed("some response").await;
281 assert!(matches!(result, Err(ToolError::Blocked { .. })));
282 }
283
284 #[tokio::test]
285 async fn trusted_allows_execute() {
286 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
287 gate.set_effective_trust(TrustLevel::Trusted);
288
289 let result = gate.execute("some response").await;
290 assert!(result.is_ok());
291 }
292
293 #[tokio::test]
294 async fn verified_with_allow_policy_succeeds() {
295 let policy = crate::permissions::PermissionPolicy::from_legacy(&[], &[]);
296 let gate = TrustGateExecutor::new(MockExecutor, policy);
297 gate.set_effective_trust(TrustLevel::Verified);
298
299 let result = gate
300 .execute_tool_call(&make_call_with_cmd("bash", "echo hi"))
301 .await
302 .unwrap();
303 assert!(result.is_some());
304 }
305
306 #[tokio::test]
307 async fn quarantined_denies_web_scrape() {
308 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
309 gate.set_effective_trust(TrustLevel::Quarantined);
310
311 let result = gate.execute_tool_call(&make_call("web_scrape")).await;
312 assert!(matches!(result, Err(ToolError::Blocked { .. })));
313 }
314
315 #[derive(Debug)]
316 struct EnvCapture {
317 captured: std::sync::Mutex<Option<std::collections::HashMap<String, String>>>,
318 }
319 impl EnvCapture {
320 fn new() -> Self {
321 Self {
322 captured: std::sync::Mutex::new(None),
323 }
324 }
325 }
326 impl ToolExecutor for EnvCapture {
327 async fn execute(&self, _: &str) -> Result<Option<ToolOutput>, ToolError> {
328 Ok(None)
329 }
330 async fn execute_tool_call(&self, _: &ToolCall) -> Result<Option<ToolOutput>, ToolError> {
331 Ok(None)
332 }
333 fn set_skill_env(&self, env: Option<std::collections::HashMap<String, String>>) {
334 *self.captured.lock().unwrap() = env;
335 }
336 }
337
338 #[test]
339 fn set_skill_env_forwarded_to_inner() {
340 let inner = EnvCapture::new();
341 let gate = TrustGateExecutor::new(inner, PermissionPolicy::default());
342
343 let mut env = std::collections::HashMap::new();
344 env.insert("MY_VAR".to_owned(), "42".to_owned());
345 gate.set_skill_env(Some(env.clone()));
346
347 let captured = gate.inner.captured.lock().unwrap();
348 assert_eq!(*captured, Some(env));
349 }
350
351 #[test]
352 fn set_effective_trust_interior_mutability() {
353 let gate = TrustGateExecutor::new(MockExecutor, PermissionPolicy::default());
354 assert_eq!(gate.effective_trust(), TrustLevel::Trusted);
355
356 gate.set_effective_trust(TrustLevel::Quarantined);
357 assert_eq!(gate.effective_trust(), TrustLevel::Quarantined);
358
359 gate.set_effective_trust(TrustLevel::Blocked);
360 assert_eq!(gate.effective_trust(), TrustLevel::Blocked);
361
362 gate.set_effective_trust(TrustLevel::Trusted);
363 assert_eq!(gate.effective_trust(), TrustLevel::Trusted);
364 }
365}