1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::path::{Component, Path, PathBuf};
4use std::sync::Arc;
5use std::time::Instant;
6
7use parking_lot::RwLock;
8use tracing::debug;
9
10use super::config::*;
11use ai_agents_core::{
12 CommandBindingKind, CommandPolicyBinding, DomainPolicyBinding, PathAccessMode,
13 PathPolicyBinding, Result, ResultLimitBinding, ResultLimitKind, ToolCallClassification,
14 ToolExecutionLimits, ToolPolicyBindings, ToolSafetyMetadata,
15};
16use serde_json::Value;
17
18#[derive(Debug, Default)]
19struct ToolCallTracker {
20 calls: HashMap<String, Vec<Instant>>,
21}
22
23impl ToolCallTracker {
24 fn record_call(&mut self, tool_id: &str) {
25 self.calls
26 .entry(tool_id.to_string())
27 .or_default()
28 .push(Instant::now());
29 }
30
31 fn get_calls_in_window(&self, tool_id: &str, window_seconds: u64) -> usize {
32 let now = Instant::now();
33 let window = std::time::Duration::from_secs(window_seconds);
34
35 self.calls
36 .get(tool_id)
37 .map(|calls| {
38 calls
39 .iter()
40 .filter(|t| now.duration_since(**t) < window)
41 .count()
42 })
43 .unwrap_or(0)
44 }
45
46 fn reset(&mut self) {
47 self.calls.clear();
48 }
49}
50
51#[derive(Debug, Clone)]
52pub struct ToolSecurityEngine {
53 config: ToolSecurityConfig,
54 tool_call_tracker: Arc<RwLock<ToolCallTracker>>,
55 policy_version: u64,
56}
57
58impl ToolSecurityEngine {
59 pub fn new(config: ToolSecurityConfig) -> Self {
60 Self {
61 config,
62 tool_call_tracker: Arc::new(RwLock::new(ToolCallTracker::default())),
63 policy_version: 1,
64 }
65 }
66
67 pub fn config(&self) -> &ToolSecurityConfig {
68 &self.config
69 }
70
71 pub fn policy_version(&self) -> u64 {
72 self.policy_version
73 }
74
75 pub fn prepare_tool_arguments(&self, tool_id: &str, args: &Value) -> Value {
76 let bindings = legacy_policy_bindings(tool_id);
77 self.prepare_tool_arguments_with_bindings(tool_id, args, &bindings)
78 }
79
80 pub fn prepare_tool_arguments_with_bindings(
81 &self,
82 tool_id: &str,
83 args: &Value,
84 bindings: &ToolPolicyBindings,
85 ) -> Value {
86 if !self.config.enabled {
87 return args.clone();
88 }
89 let mut prepared = args.clone();
90 let Some(tool_config) = self.config.tools.get(tool_id) else {
91 return prepared;
92 };
93 normalize_default_path_arguments(bindings, &mut prepared);
94 apply_policy_caps(tool_config, bindings, &mut prepared);
95 prepared
96 }
97
98 pub fn attach_internal_tool_policy(&self, _tool_id: &str, args: &Value) -> Value {
99 args.clone()
100 }
101
102 pub fn get_tool_output_cap(
103 &self,
104 tool_id: &str,
105 classification_cap: Option<usize>,
106 ) -> Option<usize> {
107 let policy_cap = self
108 .config
109 .enabled
110 .then(|| self.config.tools.get(tool_id))
111 .flatten()
112 .and_then(|config| config.max_output_chars);
113 min_optional_usize(classification_cap, policy_cap)
114 }
115
116 pub fn effective_limits(
117 &self,
118 tool_id: &str,
119 safety: &ToolSafetyMetadata,
120 classification: &ToolCallClassification,
121 ) -> ToolExecutionLimits {
122 let policy = self
123 .config
124 .enabled
125 .then(|| self.config.tools.get(tool_id))
126 .flatten();
127 ToolExecutionLimits {
128 timeout_ms: Some(self.get_tool_timeout(tool_id)),
129 max_output_chars: min_optional_usize(
130 classification.max_output_chars,
131 policy.and_then(|config| config.max_output_chars),
132 ),
133 max_result_chars: safety.max_result_size_chars,
134 max_results: policy.and_then(|config| config.max_results),
135 max_file_size_bytes: policy.and_then(|config| config.max_file_size_bytes),
136 max_response_bytes: policy.and_then(|config| config.max_response_bytes),
137 max_redirects: policy.and_then(|config| config.max_redirects),
138 max_replacements: policy.and_then(|config| config.max_replacements),
139 max_changed_files: policy.and_then(|config| config.max_changed_files),
140 max_changed_lines: policy.and_then(|config| config.max_changed_lines),
141 }
142 }
143
144 pub fn policy_snapshot(&self, tool_id: &str) -> Value {
145 if !self.config.enabled {
146 return Value::Null;
147 }
148 self.config
149 .tools
150 .get(tool_id)
151 .and_then(|config| serde_json::to_value(config).ok())
152 .unwrap_or(Value::Null)
153 }
154
155 pub fn classification_approval_message(
157 &self,
158 tool_id: &str,
159 classification: &ToolCallClassification,
160 ) -> Option<String> {
161 if !classification.requires_approval || classification.read_only {
162 return None;
163 }
164 if !matches!(
165 classification.operation,
166 ai_agents_core::ToolOperationKind::Write
167 | ai_agents_core::ToolOperationKind::Edit
168 | ai_agents_core::ToolOperationKind::Delete
169 | ai_agents_core::ToolOperationKind::Patch
170 | ai_agents_core::ToolOperationKind::Command
171 ) {
172 return None;
173 }
174 let tool_config = self
175 .config
176 .enabled
177 .then(|| self.config.tools.get(tool_id))
178 .flatten();
179 if tool_config.is_some_and(|config| config.allow_without_confirmation) {
180 return None;
181 }
182 Some(format!(
183 "Confirm {} operation for tool '{}' ?",
184 format!("{:?}", classification.operation).to_ascii_lowercase(),
185 tool_id
186 ))
187 }
188
189 pub fn custom_config(&self, tool_id: &str) -> Value {
190 if !self.config.enabled {
191 return Value::Null;
192 }
193 self.config
194 .tools
195 .get(tool_id)
196 .map(|config| Value::Object(config.config.clone().into_iter().collect()))
197 .unwrap_or(Value::Null)
198 }
199
200 pub async fn check_tool_execution(
201 &self,
202 tool_id: &str,
203 args: &serde_json::Value,
204 ) -> Result<SecurityCheckResult> {
205 let bindings = legacy_policy_bindings(tool_id);
206 self.check_tool_execution_with_bindings(tool_id, args, &bindings)
207 .await
208 }
209
210 pub async fn check_tool_execution_with_bindings(
211 &self,
212 tool_id: &str,
213 args: &serde_json::Value,
214 bindings: &ToolPolicyBindings,
215 ) -> Result<SecurityCheckResult> {
216 if !self.config.enabled {
217 return Ok(SecurityCheckResult::Allow);
218 }
219
220 let tool_config = match self.config.tools.get(tool_id) {
221 Some(config) => config,
222 None if self.config.fail_closed => {
223 return Ok(SecurityCheckResult::Block {
224 reason: format!("Tool '{}' has no explicit security policy", tool_id),
225 });
226 }
227 None => {
228 self.tool_call_tracker.write().record_call(tool_id);
229 debug!(tool_id = %tool_id, "Tool execution allowed by legacy open policy");
230 return Ok(SecurityCheckResult::Allow);
231 }
232 };
233
234 if !tool_config.enabled {
235 return Ok(SecurityCheckResult::Unavailable {
236 reason: format!("Tool '{}' is disabled", tool_id),
237 });
238 }
239
240 if let Some(rate_limit) = tool_config.rate_limit {
241 let calls = self
242 .tool_call_tracker
243 .read()
244 .get_calls_in_window(tool_id, 60);
245 if calls >= rate_limit as usize {
246 return Ok(SecurityCheckResult::Block {
247 reason: format!(
248 "Rate limit exceeded for tool '{}': {} calls per minute",
249 tool_id, rate_limit
250 ),
251 });
252 }
253 }
254
255 if let Some(result) =
256 validate_policy_bindings(tool_id, tool_config, bindings, self.config.fail_closed)
257 {
258 return Ok(result);
259 }
260
261 if let Some(result) = self.check_domain_policy(tool_id, tool_config, args, bindings) {
262 return Ok(result);
263 }
264
265 if let Some(result) = self.check_path_policy(tool_id, tool_config, args, bindings) {
266 return Ok(result);
267 }
268
269 if let Some(result) = self.check_operation_policy(tool_id, tool_config, args, bindings) {
270 return Ok(result);
271 }
272
273 if let Some(result) = self.check_command_policy(tool_id, tool_config, args, bindings) {
274 return Ok(result);
275 }
276
277 if tool_config.require_confirmation {
278 let message = tool_config
279 .confirmation_message
280 .clone()
281 .unwrap_or_else(|| format!("Confirm execution of tool '{}' ?", tool_id));
282 return Ok(SecurityCheckResult::RequireConfirmation { message });
283 }
284
285 self.tool_call_tracker.write().record_call(tool_id);
286 debug!(tool_id = %tool_id, "Tool execution allowed");
287
288 Ok(SecurityCheckResult::Allow)
289 }
290
291 pub fn check_command_execution(
292 &self,
293 tool_id: &str,
294 command: &str,
295 args: &[String],
296 ) -> SecurityCheckResult {
297 if !self.config.enabled {
298 return SecurityCheckResult::Allow;
299 }
300 let Some(tool_config) = self.config.tools.get(tool_id) else {
301 return if self.config.fail_closed {
302 SecurityCheckResult::Block {
303 reason: format!("Tool '{}' has no explicit command policy", tool_id),
304 }
305 } else {
306 SecurityCheckResult::Allow
307 };
308 };
309 let value = serde_json::json!({
310 "command": command,
311 "argv": std::iter::once(command.to_string()).chain(args.iter().cloned()).collect::<Vec<_>>()
312 });
313 let bindings = legacy_policy_bindings(tool_id);
314 self.check_command_policy(tool_id, tool_config, &value, &bindings)
315 .unwrap_or(SecurityCheckResult::Allow)
316 }
317
318 fn check_domain_policy(
319 &self,
320 tool_id: &str,
321 tool_config: &ToolPolicyConfig,
322 args: &serde_json::Value,
323 bindings: &ToolPolicyBindings,
324 ) -> Option<SecurityCheckResult> {
325 let values = bound_domain_values(args, bindings);
326 if values.is_empty() {
327 return missing_bound_value_result(
328 tool_id,
329 "domain",
330 domain_policy_configured(tool_config),
331 self.config.fail_closed,
332 );
333 }
334 for value in values {
335 let parsed = if value.is_url {
336 match reqwest::Url::parse(&value.value) {
337 Ok(parsed) => parsed,
338 Err(_) => {
339 return Some(SecurityCheckResult::Block {
340 reason: format!("URL is invalid for tool '{}'", tool_id),
341 });
342 }
343 }
344 } else {
345 match reqwest::Url::parse(&format!("https://{}", value.value)) {
346 Ok(parsed) => parsed,
347 Err(_) => {
348 return Some(SecurityCheckResult::Block {
349 reason: format!("Domain is invalid for tool '{}'", tool_id),
350 });
351 }
352 }
353 };
354 let host = parsed.host_str().map(normalize_host)?;
355
356 if !tool_config.allowed_schemes.is_empty()
357 && !tool_config
358 .allowed_schemes
359 .iter()
360 .any(|scheme| scheme.eq_ignore_ascii_case(parsed.scheme()))
361 {
362 return Some(SecurityCheckResult::Block {
363 reason: format!(
364 "URL scheme '{}' is not allowed for tool '{}'",
365 parsed.scheme(),
366 tool_id
367 ),
368 });
369 }
370
371 if !tool_config.allowed_ports.is_empty() {
372 let port = parsed.port_or_known_default().unwrap_or(0);
373 if !tool_config.allowed_ports.contains(&port) {
374 return Some(SecurityCheckResult::Block {
375 reason: format!(
376 "URL port '{}' is not allowed for tool '{}'",
377 port, tool_id
378 ),
379 });
380 }
381 }
382
383 if tool_config.blocked_private_networks && host_is_private_or_local(&host) {
384 return Some(SecurityCheckResult::Block {
385 reason: format!(
386 "Private, localhost, link-local, or metadata host is blocked for tool '{}'",
387 tool_id
388 ),
389 });
390 }
391
392 let denied = tool_config
393 .blocked_domains
394 .iter()
395 .chain(tool_config.domains.deny.iter());
396 for pattern in denied {
397 if host_matches(pattern, &host) {
398 return Some(SecurityCheckResult::Block {
399 reason: format!("Domain '{}' is blocked for tool '{}'", pattern, tool_id),
400 });
401 }
402 }
403
404 for pattern in &tool_config.domains.unavailable {
405 if host_matches(pattern, &host) {
406 return Some(SecurityCheckResult::Unavailable {
407 reason: format!(
408 "Domain '{}' is unavailable for tool '{}'",
409 pattern, tool_id
410 ),
411 });
412 }
413 }
414
415 for pattern in &tool_config.domains.requires_approval {
416 if host_matches(pattern, &host) {
417 return Some(SecurityCheckResult::RequireConfirmation {
418 message: format!(
419 "Confirm access to domain '{}' for tool '{}' ?",
420 host, tool_id
421 ),
422 });
423 }
424 }
425
426 let allowed: Vec<&String> = tool_config
427 .allowed_domains
428 .iter()
429 .chain(tool_config.domains.allow.iter())
430 .collect();
431 if !allowed.is_empty() && !allowed.iter().any(|pattern| host_matches(pattern, &host)) {
432 return Some(SecurityCheckResult::Block {
433 reason: format!("URL domain not in allowed list for tool '{}'", tool_id),
434 });
435 }
436 }
437
438 None
439 }
440
441 fn check_path_policy(
442 &self,
443 tool_id: &str,
444 tool_config: &ToolPolicyConfig,
445 args: &serde_json::Value,
446 bindings: &ToolPolicyBindings,
447 ) -> Option<SecurityCheckResult> {
448 let values = bound_path_values(args, bindings);
449 if values.is_empty() {
450 return missing_bound_value_result(
451 tool_id,
452 "path",
453 path_policy_configured(tool_config),
454 self.config.fail_closed,
455 );
456 }
457 for value in values {
458 let normalized = normalize_path(&value.path);
459
460 for pattern in tool_config
461 .blocked_paths
462 .iter()
463 .chain(tool_config.paths.deny.iter())
464 {
465 if path_matches_bound(pattern, &value.path, &normalized) {
466 return Some(SecurityCheckResult::Block {
467 reason: format!("Path is blocked for tool '{}'", tool_id),
468 });
469 }
470 }
471
472 for pattern in tool_config.paths.unavailable.iter() {
473 if path_matches_bound(pattern, &value.path, &normalized) {
474 return Some(SecurityCheckResult::Unavailable {
475 reason: format!("Path is unavailable for tool '{}'", tool_id),
476 });
477 }
478 }
479
480 for pattern in tool_config.paths.requires_approval.iter() {
481 if path_matches_bound(pattern, &value.path, &normalized) {
482 return Some(SecurityCheckResult::RequireConfirmation {
483 message: format!(
484 "Confirm access to path '{}' for tool '{}' ?",
485 value.path, tool_id
486 ),
487 });
488 }
489 }
490
491 if !matches!(value.kind, ai_agents_core::PathBindingKind::Cwd)
492 && matches!(
493 value.mode,
494 PathAccessMode::Write | PathAccessMode::ReadWrite
495 )
496 && !has_write_allowlist(tool_config)
497 {
498 let dry_run = args
499 .get("dry_run")
500 .and_then(Value::as_bool)
501 .unwrap_or_else(|| default_dry_run_for_tool(tool_id));
502 if matches!(tool_config.no_write_policy, NoWritePolicyBehavior::Deny) || !dry_run {
503 return Some(SecurityCheckResult::Block {
504 reason: format!(
505 "Tool '{}' cannot mutate paths without an explicit write_paths policy",
506 tool_id
507 ),
508 });
509 }
510 }
511
512 let allowed = allowed_paths_for_value(tool_config, &value);
513 if matches!(value.kind, ai_agents_core::PathBindingKind::Cwd) && allowed.is_empty() {
514 return Some(SecurityCheckResult::Block {
515 reason: format!(
516 "Tool '{}' requires an explicit working_dirs policy for command cwd",
517 tool_id
518 ),
519 });
520 }
521 if !allowed.is_empty()
522 && !allowed
523 .iter()
524 .any(|pattern| path_matches_bound(pattern, &value.path, &normalized))
525 {
526 return Some(SecurityCheckResult::Block {
527 reason: format!("Path not in allowed list for tool '{}'", tool_id),
528 });
529 }
530 }
531
532 None
533 }
534
535 fn check_operation_policy(
536 &self,
537 tool_id: &str,
538 tool_config: &ToolPolicyConfig,
539 args: &serde_json::Value,
540 bindings: &ToolPolicyBindings,
541 ) -> Option<SecurityCheckResult> {
542 let operations = bound_operation_values(args, bindings);
543 if operations.is_empty() {
544 return missing_bound_value_result(
545 tool_id,
546 "operation",
547 operation_policy_configured(tool_config),
548 self.config.fail_closed,
549 );
550 }
551 for operation in operations {
552 if contains_casefold(&tool_config.operations.deny, &operation) {
553 return Some(SecurityCheckResult::Block {
554 reason: format!(
555 "Operation '{}' is blocked for tool '{}'",
556 operation, tool_id
557 ),
558 });
559 }
560 if contains_casefold(&tool_config.operations.unavailable, &operation) {
561 return Some(SecurityCheckResult::Unavailable {
562 reason: format!(
563 "Operation '{}' is unavailable for tool '{}'",
564 operation, tool_id
565 ),
566 });
567 }
568 if contains_casefold(&tool_config.operations.requires_approval, &operation) {
569 return Some(SecurityCheckResult::RequireConfirmation {
570 message: format!("Confirm operation '{}' for tool '{}' ?", operation, tool_id),
571 });
572 }
573 if !tool_config.operations.allow.is_empty()
574 && !contains_casefold(&tool_config.operations.allow, &operation)
575 {
576 return Some(SecurityCheckResult::Block {
577 reason: format!(
578 "Operation '{}' is not allowed for tool '{}'",
579 operation, tool_id
580 ),
581 });
582 }
583 }
584
585 None
586 }
587
588 fn check_command_policy(
589 &self,
590 tool_id: &str,
591 tool_config: &ToolPolicyConfig,
592 args: &serde_json::Value,
593 bindings: &ToolPolicyBindings,
594 ) -> Option<SecurityCheckResult> {
595 let commands = bound_command_values(args, bindings);
596 if commands.is_empty() {
597 return missing_bound_value_result(
598 tool_id,
599 "command",
600 command_policy_configured(tool_config),
601 self.config.fail_closed,
602 );
603 }
604 for command in commands {
605 let display = command.display();
606 let command_name = command.command_name();
607
608 if command.is_string
609 && command_denies_shell(tool_config)
610 && contains_shell_syntax(&display)
611 {
612 return Some(SecurityCheckResult::Block {
613 reason: format!(
614 "Command '{}' uses shell syntax denied for tool '{}'",
615 display, tool_id
616 ),
617 });
618 }
619 if contains_casefold(&tool_config.commands.deny, &display)
620 || contains_casefold(&tool_config.commands.deny, &command_name)
621 {
622 return Some(SecurityCheckResult::Block {
623 reason: format!("Command '{}' is blocked for tool '{}'", display, tool_id),
624 });
625 }
626 if contains_casefold(&tool_config.commands.unavailable, &display)
627 || contains_casefold(&tool_config.commands.unavailable, &command_name)
628 {
629 return Some(SecurityCheckResult::Unavailable {
630 reason: format!(
631 "Command '{}' is unavailable for tool '{}'",
632 display, tool_id
633 ),
634 });
635 }
636 if contains_casefold(&tool_config.commands.requires_approval, &display)
637 || contains_casefold(&tool_config.commands.requires_approval, &command_name)
638 {
639 return Some(SecurityCheckResult::RequireConfirmation {
640 message: format!("Confirm command '{}' for tool '{}' ?", display, tool_id),
641 });
642 }
643 let has_exact_allowlist = command_has_exact_allowlist(tool_config);
644 if command_requires_exact_allowlist(tool_id) && !has_exact_allowlist {
645 return Some(SecurityCheckResult::Block {
646 reason: format!(
647 "Tool '{}' requires allowed_commands or command_templates before execution",
648 tool_id
649 ),
650 });
651 }
652 if has_exact_allowlist {
653 if !command_matches_allowed(tool_config, &command.argv) {
654 if command_allows_escalation(tool_config) {
655 return Some(SecurityCheckResult::RequireConfirmation {
656 message: format!(
657 "Confirm command '{}' outside the exact allowlist for tool '{}' ?",
658 display, tool_id
659 ),
660 });
661 }
662 return Some(SecurityCheckResult::Block {
663 reason: format!(
664 "Command '{}' is not in the exact argv allowlist for tool '{}'",
665 display, tool_id
666 ),
667 });
668 }
669 continue;
670 }
671 if !tool_config.commands.allow.is_empty()
672 && !contains_casefold(&tool_config.commands.allow, &display)
673 && !contains_casefold(&tool_config.commands.allow, &command_name)
674 {
675 return Some(SecurityCheckResult::Block {
676 reason: format!(
677 "Command '{}' is not allowed for tool '{}'",
678 display, tool_id
679 ),
680 });
681 }
682 }
683
684 None
685 }
686
687 pub fn get_tool_timeout(&self, tool_id: &str) -> u64 {
688 self.config
689 .tools
690 .get(tool_id)
691 .and_then(|c| c.timeout_ms)
692 .unwrap_or(self.config.default_timeout_ms)
693 }
694
695 pub fn reset_session(&self) {
696 self.tool_call_tracker.write().reset();
697 }
698}
699
700impl Default for ToolSecurityEngine {
701 fn default() -> Self {
702 Self::new(ToolSecurityConfig::default())
703 }
704}
705
706fn normalize_default_path_arguments(bindings: &ToolPolicyBindings, args: &mut Value) {
707 for binding in &bindings.path_fields {
708 let Some(default_path) = binding.default_path.as_deref() else {
709 continue;
710 };
711 if value_at_path(args, &binding.field).is_none() {
712 set_root_value(
713 args,
714 &binding.field,
715 Value::String(default_path.to_string()),
716 );
717 }
718 }
719}
720
721fn apply_policy_caps(config: &ToolPolicyConfig, bindings: &ToolPolicyBindings, args: &mut Value) {
722 let Some(obj) = args.as_object_mut() else {
723 return;
724 };
725 for binding in &bindings.result_limit_fields {
726 match binding.kind {
727 ResultLimitKind::MaxResults | ResultLimitKind::Pagination => {
728 apply_usize_cap(obj, &binding.field, config.max_results);
729 }
730 ResultLimitKind::MaxLines => {
731 apply_usize_cap(obj, &binding.field, config.max_results);
732 }
733 ResultLimitKind::MaxOutputChars => {
734 apply_usize_cap(obj, &binding.field, config.max_output_chars);
735 }
736 ResultLimitKind::MaxFileSizeBytes => {
737 apply_u64_cap(obj, &binding.field, config.max_file_size_bytes);
738 }
739 ResultLimitKind::MaxResponseBytes => {
740 apply_usize_cap(obj, &binding.field, config.max_response_bytes);
741 }
742 ResultLimitKind::MaxRedirects => {
743 apply_usize_cap(obj, &binding.field, config.max_redirects);
744 }
745 ResultLimitKind::MaxReplacements => {
746 apply_usize_cap(obj, &binding.field, config.max_replacements);
747 }
748 ResultLimitKind::MaxChangedFiles => {
749 apply_usize_cap(obj, &binding.field, config.max_changed_files);
750 }
751 ResultLimitKind::MaxChangedLines => {
752 apply_usize_cap(obj, &binding.field, config.max_changed_lines);
753 }
754 }
755 }
756}
757
758fn legacy_policy_bindings(tool_id: &str) -> ToolPolicyBindings {
759 match tool_id {
760 "glob" => ToolPolicyBindings {
761 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
762 result_limit_fields: vec![ResultLimitBinding::new(
763 "max_results",
764 ResultLimitKind::MaxResults,
765 )],
766 ..Default::default()
767 },
768 "grep" => ToolPolicyBindings {
769 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
770 result_limit_fields: vec![
771 ResultLimitBinding::new("max_results", ResultLimitKind::MaxResults),
772 ResultLimitBinding::new("max_file_size_bytes", ResultLimitKind::MaxFileSizeBytes),
773 ResultLimitBinding::new("max_output_chars", ResultLimitKind::MaxOutputChars),
774 ],
775 ..Default::default()
776 },
777 "file_read" => ToolPolicyBindings {
778 path_fields: vec![PathPolicyBinding::read("path")],
779 result_limit_fields: vec![
780 ResultLimitBinding::new("max_bytes", ResultLimitKind::MaxFileSizeBytes),
781 ResultLimitBinding::new("max_lines", ResultLimitKind::MaxLines),
782 ],
783 ..Default::default()
784 },
785 "file_list" => ToolPolicyBindings {
786 path_fields: vec![PathPolicyBinding::read("path")],
787 result_limit_fields: vec![ResultLimitBinding::new(
788 "max_results",
789 ResultLimitKind::MaxResults,
790 )],
791 ..Default::default()
792 },
793 "file_info" => ToolPolicyBindings {
794 path_fields: vec![PathPolicyBinding::read("path")],
795 ..Default::default()
796 },
797 "git_status" => ToolPolicyBindings {
798 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
799 result_limit_fields: vec![ResultLimitBinding::new(
800 "max_results",
801 ResultLimitKind::MaxResults,
802 )],
803 ..Default::default()
804 },
805 "git_diff" => ToolPolicyBindings {
806 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
807 result_limit_fields: vec![ResultLimitBinding::new(
808 "max_output_chars",
809 ResultLimitKind::MaxOutputChars,
810 )],
811 ..Default::default()
812 },
813 "diagnostics" => ToolPolicyBindings {
814 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
815 result_limit_fields: vec![ResultLimitBinding::new(
816 "max_results",
817 ResultLimitKind::MaxResults,
818 )],
819 ..Default::default()
820 },
821 "web_fetch" => ToolPolicyBindings {
822 domain_fields: vec![DomainPolicyBinding::url("url")],
823 result_limit_fields: vec![
824 ResultLimitBinding::new("max_chars", ResultLimitKind::MaxOutputChars),
825 ResultLimitBinding::new("max_response_bytes", ResultLimitKind::MaxResponseBytes),
826 ResultLimitBinding::new("max_redirects", ResultLimitKind::MaxRedirects),
827 ],
828 ..Default::default()
829 },
830 "http" => ToolPolicyBindings {
831 domain_fields: vec![DomainPolicyBinding::url("url")],
832 operation_fields: vec!["method".to_string()],
833 ..Default::default()
834 },
835 "file" => ToolPolicyBindings {
836 path_fields: vec![PathPolicyBinding::read_write("path")],
837 operation_fields: vec!["operation".to_string()],
838 ..Default::default()
839 },
840 "file_write" => ToolPolicyBindings {
841 path_fields: vec![PathPolicyBinding::write("path")],
842 result_limit_fields: vec![
843 ResultLimitBinding::new("max_changed_files", ResultLimitKind::MaxChangedFiles),
844 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
845 ],
846 ..Default::default()
847 },
848 "file_edit" => ToolPolicyBindings {
849 path_fields: vec![PathPolicyBinding::write("path")],
850 result_limit_fields: vec![
851 ResultLimitBinding::new("max_replacements", ResultLimitKind::MaxReplacements),
852 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
853 ],
854 ..Default::default()
855 },
856 "patch" => ToolPolicyBindings {
857 path_fields: vec![
858 ai_agents_core::PathPolicyBinding::new(
859 "base_path",
860 PathAccessMode::Write,
861 ai_agents_core::PathBindingKind::PatchBase,
862 )
863 .with_default_path("."),
864 ],
865 result_limit_fields: vec![
866 ResultLimitBinding::new("max_changed_files", ResultLimitKind::MaxChangedFiles),
867 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
868 ],
869 ..Default::default()
870 },
871 "command" => ToolPolicyBindings {
872 command_fields: vec![
873 CommandPolicyBinding::command("command"),
874 CommandPolicyBinding::argv("argv"),
875 CommandPolicyBinding::env("env"),
876 ],
877 path_fields: vec![
878 ai_agents_core::PathPolicyBinding::new(
879 "cwd",
880 PathAccessMode::ReadWrite,
881 ai_agents_core::PathBindingKind::Cwd,
882 )
883 .with_default_path("."),
884 ],
885 result_limit_fields: vec![ResultLimitBinding::new(
886 "max_output_chars",
887 ResultLimitKind::MaxOutputChars,
888 )],
889 ..Default::default()
890 },
891 _ => ToolPolicyBindings::default(),
892 }
893}
894
895#[derive(Debug, Clone)]
896struct BoundPathValue {
897 path: String,
898 mode: PathAccessMode,
899 kind: ai_agents_core::PathBindingKind,
900}
901
902#[derive(Debug, Clone)]
903struct BoundDomainValue {
904 value: String,
905 is_url: bool,
906}
907
908#[derive(Debug, Clone)]
909struct BoundCommandValue {
910 argv: Vec<String>,
911 is_string: bool,
912}
913
914impl BoundCommandValue {
915 fn display(&self) -> String {
916 self.argv.join(" ")
917 }
918
919 fn command_name(&self) -> String {
920 self.argv.first().cloned().unwrap_or_default()
921 }
922}
923
924fn validate_policy_bindings(
925 tool_id: &str,
926 config: &ToolPolicyConfig,
927 bindings: &ToolPolicyBindings,
928 fail_closed: bool,
929) -> Option<SecurityCheckResult> {
930 if !fail_closed {
931 return None;
932 }
933 if path_policy_configured(config) && !bindings.has_path_bindings() {
934 return Some(SecurityCheckResult::Block {
935 reason: format!(
936 "path policy configured for {} but tool exposes no path policy bindings",
937 tool_id
938 ),
939 });
940 }
941 if domain_policy_configured(config) && !bindings.has_domain_bindings() {
942 return Some(SecurityCheckResult::Block {
943 reason: format!(
944 "domain policy configured for {} but tool exposes no domain policy bindings",
945 tool_id
946 ),
947 });
948 }
949 if command_policy_configured(config) && !bindings.has_command_bindings() {
950 return Some(SecurityCheckResult::Block {
951 reason: format!(
952 "command policy configured for {} but tool exposes no command policy bindings",
953 tool_id
954 ),
955 });
956 }
957 if operation_policy_configured(config) && !bindings.has_operation_bindings() {
958 return Some(SecurityCheckResult::Block {
959 reason: format!(
960 "operation policy configured for {} but tool exposes no operation policy bindings",
961 tool_id
962 ),
963 });
964 }
965 if result_limit_policy_configured(config) && !bindings.has_result_limit_bindings() {
966 return Some(SecurityCheckResult::Block {
967 reason: format!(
968 "result-limit policy configured for {} but tool exposes no result-limit policy bindings",
969 tool_id
970 ),
971 });
972 }
973 None
974}
975
976fn missing_bound_value_result(
977 tool_id: &str,
978 policy_kind: &str,
979 configured: bool,
980 fail_closed: bool,
981) -> Option<SecurityCheckResult> {
982 if configured && fail_closed {
983 Some(SecurityCheckResult::Block {
984 reason: format!(
985 "{} policy configured for {} but no bound {} argument was present",
986 policy_kind, tool_id, policy_kind
987 ),
988 })
989 } else {
990 None
991 }
992}
993
994fn path_policy_configured(config: &ToolPolicyConfig) -> bool {
995 !config.allowed_paths.is_empty()
996 || !config.read_paths.is_empty()
997 || !config.write_paths.is_empty()
998 || !config.working_dirs.is_empty()
999 || !config.commands.working_dirs.is_empty()
1000 || !config.blocked_paths.is_empty()
1001 || !config.paths.allow.is_empty()
1002 || !config.paths.deny.is_empty()
1003 || !config.paths.requires_approval.is_empty()
1004 || !config.paths.unavailable.is_empty()
1005}
1006
1007fn domain_policy_configured(config: &ToolPolicyConfig) -> bool {
1008 !config.allowed_domains.is_empty()
1009 || !config.blocked_domains.is_empty()
1010 || !config.allowed_schemes.is_empty()
1011 || !config.allowed_ports.is_empty()
1012 || !config.domains.allow.is_empty()
1013 || !config.domains.deny.is_empty()
1014 || !config.domains.requires_approval.is_empty()
1015 || !config.domains.unavailable.is_empty()
1016}
1017
1018fn command_policy_configured(config: &ToolPolicyConfig) -> bool {
1019 !config.commands.allow.is_empty()
1020 || !config.commands.deny.is_empty()
1021 || !config.commands.requires_approval.is_empty()
1022 || !config.commands.unavailable.is_empty()
1023 || !config.commands.allowed_commands.is_empty()
1024 || !config.commands.templates.is_empty()
1025 || !config.allowed_commands.is_empty()
1026 || !config.command_templates.is_empty()
1027 || !config.env_passthrough.is_empty()
1028 || !config.commands.env_passthrough.is_empty()
1029}
1030
1031fn operation_policy_configured(config: &ToolPolicyConfig) -> bool {
1032 !config.operations.allow.is_empty()
1033 || !config.operations.deny.is_empty()
1034 || !config.operations.requires_approval.is_empty()
1035 || !config.operations.unavailable.is_empty()
1036}
1037
1038fn result_limit_policy_configured(config: &ToolPolicyConfig) -> bool {
1039 config.max_file_size_bytes.is_some()
1041 || config.max_results.is_some()
1042 || config.max_response_bytes.is_some()
1043 || config.max_redirects.is_some()
1044 || config.max_replacements.is_some()
1045 || config.max_changed_files.is_some()
1046 || config.max_changed_lines.is_some()
1047}
1048
1049fn bound_path_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundPathValue> {
1050 let mut values = Vec::new();
1051 for binding in &bindings.path_fields {
1052 collect_path_binding_values(args, binding, &mut values);
1053 }
1054 values
1055}
1056
1057fn collect_path_binding_values(
1058 args: &Value,
1059 binding: &PathPolicyBinding,
1060 values: &mut Vec<BoundPathValue>,
1061) {
1062 let value = value_at_path(args, &binding.field).cloned().or_else(|| {
1063 binding
1064 .default_path
1065 .as_ref()
1066 .map(|path| Value::String(path.clone()))
1067 });
1068 let Some(value) = value else {
1069 return;
1070 };
1071 match value {
1072 Value::String(path) => values.push(BoundPathValue {
1073 path,
1074 mode: effective_path_mode(args, binding),
1075 kind: binding.kind,
1076 }),
1077 Value::Array(items) => {
1078 for item in items {
1079 if let Some(path) = item.as_str() {
1080 values.push(BoundPathValue {
1081 path: path.to_string(),
1082 mode: effective_path_mode(args, binding),
1083 kind: binding.kind,
1084 });
1085 }
1086 }
1087 }
1088 _ => {}
1089 }
1090}
1091
1092fn bound_domain_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundDomainValue> {
1093 let mut values = Vec::new();
1094 for binding in &bindings.domain_fields {
1095 collect_domain_binding_values(args, binding, &mut values);
1096 }
1097 values
1098}
1099
1100fn collect_domain_binding_values(
1101 args: &Value,
1102 binding: &DomainPolicyBinding,
1103 values: &mut Vec<BoundDomainValue>,
1104) {
1105 let Some(value) = value_at_path(args, &binding.field) else {
1106 return;
1107 };
1108 match value {
1109 Value::String(value) => values.push(BoundDomainValue {
1110 value: value.clone(),
1111 is_url: binding.is_url,
1112 }),
1113 Value::Array(items) => {
1114 for item in items {
1115 if let Some(value) = item.as_str() {
1116 values.push(BoundDomainValue {
1117 value: value.to_string(),
1118 is_url: binding.is_url,
1119 });
1120 }
1121 }
1122 }
1123 _ => {}
1124 }
1125}
1126
1127fn bound_operation_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<String> {
1128 bindings
1129 .operation_fields
1130 .iter()
1131 .filter_map(|field| value_at_path(args, field).and_then(Value::as_str))
1132 .map(|value| value.trim().to_ascii_lowercase())
1133 .collect()
1134}
1135
1136fn bound_command_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundCommandValue> {
1137 let mut values = Vec::new();
1138 for binding in &bindings.command_fields {
1139 collect_command_binding_values(args, binding, &mut values);
1140 }
1141 values
1142}
1143
1144fn collect_command_binding_values(
1145 args: &Value,
1146 binding: &CommandPolicyBinding,
1147 values: &mut Vec<BoundCommandValue>,
1148) {
1149 let Some(value) = value_at_path(args, &binding.field) else {
1150 return;
1151 };
1152 match binding.kind {
1153 CommandBindingKind::CommandString => {
1154 if let Some(command) = value.as_str() {
1155 if let Some(argv) = parse_command_words(command) {
1156 values.push(BoundCommandValue {
1157 argv,
1158 is_string: true,
1159 });
1160 } else {
1161 values.push(BoundCommandValue {
1162 argv: vec![command.to_string()],
1163 is_string: true,
1164 });
1165 }
1166 }
1167 }
1168 CommandBindingKind::Argv => {
1169 if let Some(argv) = value.as_array().map(|items| {
1170 items
1171 .iter()
1172 .filter_map(Value::as_str)
1173 .map(str::to_string)
1174 .collect::<Vec<_>>()
1175 }) {
1176 if !argv.is_empty() {
1177 values.push(BoundCommandValue {
1178 argv,
1179 is_string: false,
1180 });
1181 }
1182 }
1183 }
1184 CommandBindingKind::Cwd
1185 | CommandBindingKind::TemplateVariable
1186 | CommandBindingKind::Env => {}
1187 }
1188}
1189
1190fn allowed_paths_for_value<'a>(
1191 config: &'a ToolPolicyConfig,
1192 value: &BoundPathValue,
1193) -> Vec<&'a String> {
1194 if matches!(value.kind, ai_agents_core::PathBindingKind::Cwd) {
1195 return config
1196 .working_dirs
1197 .iter()
1198 .chain(config.commands.working_dirs.iter())
1199 .collect();
1200 }
1201 let mut allowed: Vec<&String> = config
1202 .allowed_paths
1203 .iter()
1204 .chain(config.paths.allow.iter())
1205 .collect();
1206 match value.mode {
1207 PathAccessMode::Read => allowed.extend(config.read_paths.iter()),
1208 PathAccessMode::Write | PathAccessMode::ReadWrite => {
1209 allowed.extend(config.write_paths.iter());
1210 }
1211 }
1212 allowed
1213}
1214
1215fn has_write_allowlist(config: &ToolPolicyConfig) -> bool {
1216 !config.write_paths.is_empty()
1217 || !config.allowed_paths.is_empty()
1218 || !config.paths.allow.is_empty()
1219}
1220
1221fn effective_path_mode(args: &Value, binding: &PathPolicyBinding) -> PathAccessMode {
1222 if !matches!(binding.mode, PathAccessMode::ReadWrite) {
1223 return binding.mode;
1224 }
1225 let operation = args
1226 .get("operation")
1227 .and_then(Value::as_str)
1228 .unwrap_or_default()
1229 .to_ascii_lowercase();
1230 match operation.as_str() {
1231 "read" | "exists" | "list" | "info" => PathAccessMode::Read,
1232 "write" | "append" | "mkdir" | "delete" | "edit" | "patch" => PathAccessMode::Write,
1233 _ => binding.mode,
1234 }
1235}
1236
1237fn value_at_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
1238 let mut current = value;
1239 for segment in field.split('.') {
1240 if segment.is_empty() {
1241 return None;
1242 }
1243 current = current.get(segment)?;
1244 }
1245 Some(current)
1246}
1247
1248fn set_root_value(args: &mut Value, field: &str, value: Value) {
1249 if field.contains('.') {
1250 return;
1251 }
1252 let Some(obj) = args.as_object_mut() else {
1253 return;
1254 };
1255 obj.insert(field.to_string(), value);
1256}
1257
1258fn apply_usize_cap(obj: &mut serde_json::Map<String, Value>, key: &str, cap: Option<usize>) {
1259 let Some(cap) = cap else {
1260 return;
1261 };
1262 let effective = obj
1263 .get(key)
1264 .and_then(|value| value.as_u64())
1265 .map(|value| value.min(cap as u64) as usize)
1266 .unwrap_or(cap);
1267 obj.insert(key.to_string(), Value::from(effective));
1268}
1269
1270fn apply_u64_cap(obj: &mut serde_json::Map<String, Value>, key: &str, cap: Option<u64>) {
1271 let Some(cap) = cap else {
1272 return;
1273 };
1274 let effective = obj
1275 .get(key)
1276 .and_then(|value| value.as_u64())
1277 .map(|value| value.min(cap))
1278 .unwrap_or(cap);
1279 obj.insert(key.to_string(), Value::from(effective));
1280}
1281
1282fn min_optional_usize(left: Option<usize>, right: Option<usize>) -> Option<usize> {
1283 match (left, right) {
1284 (Some(left), Some(right)) => Some(left.min(right)),
1285 (Some(left), None) => Some(left),
1286 (None, Some(right)) => Some(right),
1287 (None, None) => None,
1288 }
1289}
1290
1291fn normalize_host(host: &str) -> String {
1292 host.trim_end_matches('.').to_ascii_lowercase()
1293}
1294
1295fn host_matches(pattern: &str, host: &str) -> bool {
1296 let pattern = normalize_host(pattern.trim_start_matches("*."));
1297 host == pattern || host.ends_with(&format!(".{}", pattern))
1298}
1299
1300fn normalize_path(path: &str) -> PathBuf {
1301 let mut normalized = PathBuf::new();
1302 for component in Path::new(path).components() {
1303 match component {
1304 Component::CurDir => {}
1305 Component::ParentDir => {
1306 normalized.pop();
1307 }
1308 other => normalized.push(other.as_os_str()),
1309 }
1310 }
1311 normalized
1312}
1313
1314fn path_matches(pattern: &str, path: &Path) -> bool {
1315 let pattern = normalize_path(pattern);
1316 path.starts_with(pattern)
1317}
1318
1319fn path_matches_bound(pattern: &str, raw_path: &str, normalized: &Path) -> bool {
1320 if !path_matches(pattern, normalized) {
1321 return false;
1322 }
1323 let raw = Path::new(raw_path);
1324 let pattern_path = Path::new(pattern);
1325 let Ok(resolved) = resolve_existing_or_parent(raw) else {
1326 return false;
1327 };
1328 if pattern_path.exists() {
1329 let Ok(resolved_pattern) = pattern_path.canonicalize() else {
1330 return false;
1331 };
1332 return resolved.starts_with(resolved_pattern.components().collect::<PathBuf>());
1333 }
1334 true
1335}
1336
1337fn resolve_existing_or_parent(path: &Path) -> std::io::Result<PathBuf> {
1338 if path.exists() {
1339 return path.canonicalize();
1340 }
1341 let normalized = if path.is_absolute() {
1342 path.components().collect::<PathBuf>()
1343 } else {
1344 std::env::current_dir()?.join(path).components().collect()
1345 };
1346 let mut ancestor = normalized.as_path();
1347 let mut missing = Vec::new();
1348 while !ancestor.exists() {
1349 let Some(name) = ancestor.file_name() else {
1350 break;
1351 };
1352 missing.push(name.to_os_string());
1353 ancestor = ancestor.parent().unwrap_or_else(|| Path::new("."));
1354 }
1355 let mut resolved = ancestor.canonicalize()?;
1356 for component in missing.iter().rev() {
1357 resolved.push(component);
1358 }
1359 Ok(resolved.components().collect())
1360}
1361
1362fn default_dry_run_for_tool(tool_id: &str) -> bool {
1363 tool_id == "patch"
1364}
1365
1366fn contains_casefold(values: &[String], needle: &str) -> bool {
1367 values
1368 .iter()
1369 .any(|value| value.eq_ignore_ascii_case(needle))
1370}
1371
1372fn command_denies_shell(config: &ToolPolicyConfig) -> bool {
1373 config.deny_shell || config.commands.deny_shell
1374}
1375
1376fn command_allows_escalation(config: &ToolPolicyConfig) -> bool {
1377 config.allow_command_escalation || config.commands.allow_escalation
1378}
1379
1380fn command_requires_exact_allowlist(tool_id: &str) -> bool {
1381 tool_id == "command"
1382}
1383
1384fn command_has_exact_allowlist(config: &ToolPolicyConfig) -> bool {
1385 !config.allowed_commands.is_empty()
1386 || !config.commands.allowed_commands.is_empty()
1387 || !config.command_templates.is_empty()
1388 || !config.commands.templates.is_empty()
1389}
1390
1391fn command_matches_allowed(config: &ToolPolicyConfig, argv: &[String]) -> bool {
1392 config
1393 .allowed_commands
1394 .iter()
1395 .chain(config.commands.allowed_commands.iter())
1396 .any(|rule| rule.argv == argv)
1397 || config
1398 .command_templates
1399 .iter()
1400 .chain(config.commands.templates.iter())
1401 .any(|template| command_matches_template(&template.argv, argv))
1402}
1403
1404fn command_matches_template(template: &[String], argv: &[String]) -> bool {
1405 template.len() == argv.len()
1406 && template.iter().zip(argv.iter()).all(|(expected, actual)| {
1407 (expected.starts_with('{') && expected.ends_with('}')) || expected == actual
1408 })
1409}
1410
1411fn contains_shell_syntax(value: &str) -> bool {
1412 const DENIED: &[char] = &[';', '&', '|', '<', '>', '`', '$', '\n', '\r'];
1413 value.chars().any(|ch| DENIED.contains(&ch))
1414 || value.contains("$(")
1415 || value.contains("${")
1416 || value.contains("<(")
1417 || value.contains(">(")
1418}
1419
1420fn parse_command_words(value: &str) -> Option<Vec<String>> {
1421 let mut words = Vec::new();
1422 let mut current = String::new();
1423 let mut quote: Option<char> = None;
1424 for ch in value.chars() {
1425 match (quote, ch) {
1426 (Some(q), c) if c == q => quote = None,
1427 (Some(_), c) => current.push(c),
1428 (None, '\'' | '"') => quote = Some(ch),
1429 (None, c) if c.is_whitespace() => {
1430 if !current.is_empty() {
1431 words.push(std::mem::take(&mut current));
1432 }
1433 }
1434 (None, c) => current.push(c),
1435 }
1436 }
1437 if quote.is_some() {
1438 return None;
1439 }
1440 if !current.is_empty() {
1441 words.push(current);
1442 }
1443 (!words.is_empty()).then_some(words)
1444}
1445
1446fn host_is_private_or_local(host: &str) -> bool {
1447 if matches!(
1448 host,
1449 "localhost"
1450 | "metadata"
1451 | "metadata.google.internal"
1452 | "169.254.169.254"
1453 | "100.100.100.200"
1454 ) || host.ends_with(".localhost")
1455 {
1456 return true;
1457 }
1458 match host.parse::<IpAddr>() {
1459 Ok(IpAddr::V4(ip)) => {
1460 ip.is_private()
1461 || ip.is_loopback()
1462 || ip.is_link_local()
1463 || ip.is_multicast()
1464 || ip.is_documentation()
1465 || ip.octets() == [169, 254, 169, 254]
1466 }
1467 Ok(IpAddr::V6(ip)) => {
1468 ip.is_loopback()
1469 || ip.is_unspecified()
1470 || ip.is_multicast()
1471 || ip.segments()[0] & 0xfe00 == 0xfc00
1472 || ip.segments()[0] & 0xffc0 == 0xfe80
1473 }
1474 Err(_) => false,
1475 }
1476}
1477
1478#[cfg(test)]
1479mod tests {
1480 use super::*;
1481
1482 #[test]
1483 fn test_default_engine() {
1484 let engine = ToolSecurityEngine::default();
1485 assert!(!engine.config().enabled);
1486 }
1487
1488 #[tokio::test]
1489 async fn test_tool_domain_blocking() {
1490 let mut config = ToolSecurityConfig::default();
1491 config.enabled = true;
1492
1493 let mut http_config = ToolPolicyConfig::default();
1494 http_config.blocked_domains = vec!["evil.com".to_string()];
1495 config.tools.insert("http".to_string(), http_config);
1496
1497 let engine = ToolSecurityEngine::new(config);
1498
1499 let args = serde_json::json!({"url": "https://evil.com/api"});
1500 let result = engine.check_tool_execution("http", &args).await.unwrap();
1501 assert!(result.is_blocked());
1502
1503 let args = serde_json::json!({"url": "https://not-evil.com/api"});
1504 let result = engine.check_tool_execution("http", &args).await.unwrap();
1505 assert!(result.is_allowed());
1506 }
1507
1508 #[tokio::test]
1509 async fn test_tool_allowed_domains() {
1510 let mut config = ToolSecurityConfig::default();
1511 config.enabled = true;
1512
1513 let mut http_config = ToolPolicyConfig::default();
1514 http_config.allowed_domains = vec!["api.example.com".to_string()];
1515 config.tools.insert("http".to_string(), http_config);
1516
1517 let engine = ToolSecurityEngine::new(config);
1518
1519 let args = serde_json::json!({"url": "https://api.example.com/v1"});
1520 let result = engine.check_tool_execution("http", &args).await.unwrap();
1521 assert!(result.is_allowed());
1522
1523 let args = serde_json::json!({"url": "https://other.com/api"});
1524 let result = engine.check_tool_execution("http", &args).await.unwrap();
1525 assert!(result.is_blocked());
1526 }
1527
1528 #[tokio::test]
1529 async fn test_tool_disabled() {
1530 let mut config = ToolSecurityConfig::default();
1531 config.enabled = true;
1532
1533 let mut tool_config = ToolPolicyConfig::default();
1534 tool_config.enabled = false;
1535 config.tools.insert("dangerous".to_string(), tool_config);
1536
1537 let engine = ToolSecurityEngine::new(config);
1538
1539 let result = engine
1540 .check_tool_execution("dangerous", &serde_json::json!({}))
1541 .await
1542 .unwrap();
1543 assert!(result.is_blocked());
1544 assert!(result.is_unavailable());
1545 }
1546
1547 #[tokio::test]
1548 async fn test_tool_confirmation_required() {
1549 let mut config = ToolSecurityConfig::default();
1550 config.enabled = true;
1551
1552 let mut tool_config = ToolPolicyConfig::default();
1553 tool_config.require_confirmation = true;
1554 tool_config.confirmation_message = Some("Are you sure?".to_string());
1555 config.tools.insert("delete".to_string(), tool_config);
1556
1557 let engine = ToolSecurityEngine::new(config);
1558
1559 let result = engine
1560 .check_tool_execution("delete", &serde_json::json!({}))
1561 .await
1562 .unwrap();
1563
1564 match result {
1565 SecurityCheckResult::RequireConfirmation { message } => {
1566 assert_eq!(message, "Are you sure?");
1567 }
1568 _ => panic!("Expected RequireConfirmation"),
1569 }
1570 }
1571
1572 #[test]
1573 fn test_get_tool_timeout() {
1574 let mut config = ToolSecurityConfig::default();
1575 config.default_timeout_ms = 5000;
1576
1577 let mut tool_config = ToolPolicyConfig::default();
1578 tool_config.timeout_ms = Some(10000);
1579 config.tools.insert("slow".to_string(), tool_config);
1580
1581 let engine = ToolSecurityEngine::new(config);
1582
1583 assert_eq!(engine.get_tool_timeout("slow"), 10000);
1584 assert_eq!(engine.get_tool_timeout("other"), 5000);
1585 }
1586
1587 #[tokio::test]
1588 async fn test_path_restrictions() {
1589 let mut config = ToolSecurityConfig::default();
1590 config.enabled = true;
1591
1592 let mut tool_config = ToolPolicyConfig::default();
1593 tool_config.allowed_paths = vec!["/tmp/".to_string(), "/home/user/".to_string()];
1594 config.tools.insert("file_write".to_string(), tool_config);
1595
1596 let engine = ToolSecurityEngine::new(config);
1597
1598 let args = serde_json::json!({"path": "/tmp/test.txt"});
1599 let result = engine
1600 .check_tool_execution("file_write", &args)
1601 .await
1602 .unwrap();
1603 assert!(result.is_allowed());
1604
1605 let args = serde_json::json!({"path": "/etc/passwd"});
1606 let result = engine
1607 .check_tool_execution("file_write", &args)
1608 .await
1609 .unwrap();
1610 assert!(result.is_blocked());
1611 }
1612
1613 #[tokio::test]
1614 async fn test_operation_policy() {
1615 let mut config = ToolSecurityConfig::default();
1616 config.enabled = true;
1617 let mut tool_config = ToolPolicyConfig::default();
1618 tool_config.operations.deny = vec!["delete".to_string()];
1619 tool_config.operations.requires_approval = vec!["write".to_string()];
1620 config.tools.insert("file".to_string(), tool_config);
1621 let engine = ToolSecurityEngine::new(config);
1622
1623 let result = engine
1624 .check_tool_execution("file", &serde_json::json!({"operation": "delete"}))
1625 .await
1626 .unwrap();
1627 assert!(result.is_blocked());
1628
1629 let result = engine
1630 .check_tool_execution("file", &serde_json::json!({"operation": "write"}))
1631 .await
1632 .unwrap();
1633 assert!(result.requires_approval());
1634 }
1635
1636 #[tokio::test]
1637 async fn omitted_optional_path_uses_default_for_policy() {
1638 let mut config = ToolSecurityConfig::default();
1639 config.enabled = true;
1640 let mut tool_config = ToolPolicyConfig::default();
1641 tool_config.read_paths = vec!["./crates".to_string()];
1642 config.tools.insert("grep".to_string(), tool_config);
1643 let engine = ToolSecurityEngine::new(config);
1644
1645 let result = engine
1646 .check_tool_execution("grep", &serde_json::json!({"pattern": "Tool"}))
1647 .await
1648 .unwrap();
1649 assert!(result.is_blocked());
1650
1651 let prepared =
1652 engine.prepare_tool_arguments("grep", &serde_json::json!({"pattern": "Tool"}));
1653 assert_eq!(prepared.get("path").and_then(Value::as_str), Some("."));
1654 }
1655
1656 #[tokio::test]
1657 async fn fail_closed_requires_path_bindings_for_custom_tools() {
1658 let mut config = ToolSecurityConfig::default();
1659 config.enabled = true;
1660 config.fail_closed = true;
1661 let mut tool_config = ToolPolicyConfig::default();
1662 tool_config.read_paths = vec!["./allowed".to_string()];
1663 config
1664 .tools
1665 .insert("custom_search".to_string(), tool_config);
1666 let engine = ToolSecurityEngine::new(config);
1667
1668 let result = engine
1669 .check_tool_execution_with_bindings(
1670 "custom_search",
1671 &serde_json::json!({"path": "./allowed/file.txt"}),
1672 &ToolPolicyBindings::default(),
1673 )
1674 .await
1675 .unwrap();
1676
1677 assert!(result.is_blocked());
1678 assert!(
1679 result
1680 .reason()
1681 .unwrap_or_default()
1682 .contains("tool exposes no path policy bindings")
1683 );
1684 }
1685
1686 #[tokio::test]
1687 async fn custom_path_bindings_enforce_blocked_paths() {
1688 let mut config = ToolSecurityConfig::default();
1689 config.enabled = true;
1690 config.fail_closed = true;
1691 let mut tool_config = ToolPolicyConfig::default();
1692 tool_config.read_paths = vec!["./allowed".to_string()];
1693 tool_config.blocked_paths = vec!["./allowed/private".to_string()];
1694 config
1695 .tools
1696 .insert("custom_search".to_string(), tool_config);
1697 let engine = ToolSecurityEngine::new(config);
1698 let bindings = ToolPolicyBindings {
1699 path_fields: vec![PathPolicyBinding::read("root")],
1700 ..Default::default()
1701 };
1702
1703 let allowed = engine
1704 .check_tool_execution_with_bindings(
1705 "custom_search",
1706 &serde_json::json!({"root": "./allowed/src"}),
1707 &bindings,
1708 )
1709 .await
1710 .unwrap();
1711 assert!(allowed.is_allowed());
1712
1713 let blocked = engine
1714 .check_tool_execution_with_bindings(
1715 "custom_search",
1716 &serde_json::json!({"root": "./allowed/private/secrets.txt"}),
1717 &bindings,
1718 )
1719 .await
1720 .unwrap();
1721 assert!(blocked.is_blocked());
1722 }
1723
1724 #[test]
1725 fn custom_config_is_exposed_separately() {
1726 let mut config = ToolSecurityConfig::default();
1727 config.enabled = true;
1728 let mut tool_config = ToolPolicyConfig::default();
1729 tool_config
1730 .config
1731 .insert("backend".to_string(), serde_json::json!("tantivy"));
1732 config.tools.insert("my_search".to_string(), tool_config);
1733 let engine = ToolSecurityEngine::new(config);
1734
1735 assert_eq!(engine.custom_config("my_search")["backend"], "tantivy");
1736 }
1737
1738 #[test]
1739 fn policy_caps_are_applied_as_upper_bounds() {
1740 let mut config = ToolSecurityConfig::default();
1741 config.enabled = true;
1742 let mut tool_config = ToolPolicyConfig::default();
1743 tool_config.max_results = Some(5);
1744 tool_config.max_file_size_bytes = Some(1024);
1745 tool_config.max_output_chars = Some(1000);
1746 config.tools.insert("grep".to_string(), tool_config);
1747 let engine = ToolSecurityEngine::new(config);
1748
1749 let prepared = engine.prepare_tool_arguments(
1750 "grep",
1751 &serde_json::json!({
1752 "pattern": "Tool",
1753 "path": ".",
1754 "max_results": 50,
1755 "max_file_size_bytes": 8192,
1756 "max_output_chars": 20000
1757 }),
1758 );
1759 assert_eq!(prepared.get("max_results").and_then(Value::as_u64), Some(5));
1760 assert_eq!(
1761 prepared.get("max_file_size_bytes").and_then(Value::as_u64),
1762 Some(1024)
1763 );
1764 assert_eq!(
1765 prepared.get("max_output_chars").and_then(Value::as_u64),
1766 Some(1000)
1767 );
1768 }
1769
1770 #[tokio::test]
1771 async fn fail_closed_blocks_missing_result_limit_bindings() {
1772 let mut config = ToolSecurityConfig::default();
1773 config.enabled = true;
1774 config.fail_closed = true;
1775 let mut tool_config = ToolPolicyConfig::default();
1776 tool_config.max_results = Some(5);
1777 config
1778 .tools
1779 .insert("custom_search".to_string(), tool_config);
1780 let engine = ToolSecurityEngine::new(config);
1781
1782 let result = engine
1783 .check_tool_execution_with_bindings(
1784 "custom_search",
1785 &serde_json::json!({"query": "rust"}),
1786 &ToolPolicyBindings::default(),
1787 )
1788 .await
1789 .unwrap();
1790
1791 assert!(result.is_blocked());
1792 assert!(result.reason().unwrap().contains("result-limit policy"));
1793 }
1794
1795 #[tokio::test]
1796 async fn fail_closed_allows_configured_result_limit_bindings() {
1797 let mut config = ToolSecurityConfig::default();
1798 config.enabled = true;
1799 config.fail_closed = true;
1800 let mut tool_config = ToolPolicyConfig::default();
1801 tool_config.max_results = Some(5);
1802 config
1803 .tools
1804 .insert("custom_search".to_string(), tool_config);
1805 let engine = ToolSecurityEngine::new(config);
1806 let bindings = ToolPolicyBindings {
1807 result_limit_fields: vec![ResultLimitBinding::new(
1808 "limit",
1809 ResultLimitKind::MaxResults,
1810 )],
1811 ..Default::default()
1812 };
1813
1814 let result = engine
1815 .check_tool_execution_with_bindings(
1816 "custom_search",
1817 &serde_json::json!({"query": "rust", "limit": 10}),
1818 &bindings,
1819 )
1820 .await
1821 .unwrap();
1822
1823 assert!(result.is_allowed());
1824 }
1825
1826 #[tokio::test]
1827 async fn read_paths_do_not_authorize_file_write() {
1828 let mut config = ToolSecurityConfig::default();
1829 config.enabled = true;
1830 let mut tool_config = ToolPolicyConfig::default();
1831 tool_config.read_paths = vec!["./workspace".to_string()];
1832 tool_config.no_write_policy = NoWritePolicyBehavior::Deny;
1833 config.tools.insert("file_write".to_string(), tool_config);
1834 let engine = ToolSecurityEngine::new(config);
1835
1836 let result = engine
1837 .check_tool_execution_with_bindings(
1838 "file_write",
1839 &serde_json::json!({"path": "./workspace/out.txt", "dry_run": false}),
1840 &legacy_policy_bindings("file_write"),
1841 )
1842 .await
1843 .unwrap();
1844
1845 assert!(result.is_blocked());
1846 }
1847
1848 #[tokio::test]
1849 async fn command_cwd_requires_working_dir_allowlist() {
1850 let mut config = ToolSecurityConfig::default();
1851 config.enabled = true;
1852 let mut tool_config = ToolPolicyConfig::default();
1853 tool_config.read_paths = vec![".".to_string()];
1854 tool_config.allowed_commands = vec![CommandRuleConfig {
1855 argv: vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
1856 }];
1857 config.tools.insert("command".to_string(), tool_config);
1858 let engine = ToolSecurityEngine::new(config);
1859
1860 let result = engine
1861 .check_tool_execution_with_bindings(
1862 "command",
1863 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
1864 &legacy_policy_bindings("command"),
1865 )
1866 .await
1867 .unwrap();
1868
1869 assert!(result.is_blocked());
1870 }
1871
1872 #[tokio::test]
1873 async fn command_requires_exact_argv_allowlist() {
1874 let mut config = ToolSecurityConfig::default();
1875 config.enabled = true;
1876 let mut tool_config = ToolPolicyConfig::default();
1877 tool_config.allow_without_confirmation = true;
1878 tool_config.working_dirs = vec![".".to_string()];
1879 config.tools.insert("command".to_string(), tool_config);
1880 let engine = ToolSecurityEngine::new(config);
1881
1882 let result = engine
1883 .check_tool_execution_with_bindings(
1884 "command",
1885 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
1886 &legacy_policy_bindings("command"),
1887 )
1888 .await
1889 .unwrap();
1890
1891 assert!(result.is_blocked());
1892 assert!(
1893 result
1894 .reason()
1895 .unwrap()
1896 .contains("requires allowed_commands or command_templates")
1897 );
1898 }
1899
1900 #[tokio::test]
1901 async fn command_exact_argv_allowlist_is_enforced() {
1902 let mut config = ToolSecurityConfig::default();
1903 config.enabled = true;
1904 let mut tool_config = ToolPolicyConfig::default();
1905 tool_config.allowed_commands = vec![CommandRuleConfig {
1906 argv: vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
1907 }];
1908 tool_config.working_dirs = vec![".".to_string()];
1909 config.tools.insert("command".to_string(), tool_config);
1910 let engine = ToolSecurityEngine::new(config);
1911
1912 let allowed = engine
1913 .check_tool_execution_with_bindings(
1914 "command",
1915 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
1916 &legacy_policy_bindings("command"),
1917 )
1918 .await
1919 .unwrap();
1920 assert!(allowed.is_allowed());
1921
1922 let blocked = engine
1923 .check_tool_execution_with_bindings(
1924 "command",
1925 &serde_json::json!({"argv": ["cargo", "test"], "cwd": "."}),
1926 &legacy_policy_bindings("command"),
1927 )
1928 .await
1929 .unwrap();
1930 assert!(blocked.is_blocked());
1931 }
1932
1933 #[tokio::test]
1934 async fn no_write_policy_dry_run_only_allows_dry_run() {
1935 let mut config = ToolSecurityConfig::default();
1936 config.enabled = true;
1937 let mut tool_config = ToolPolicyConfig::default();
1938 tool_config.no_write_policy = NoWritePolicyBehavior::DryRunOnly;
1939 config.tools.insert("file_edit".to_string(), tool_config);
1940 let engine = ToolSecurityEngine::new(config);
1941
1942 let dry_run = engine
1943 .check_tool_execution_with_bindings(
1944 "file_edit",
1945 &serde_json::json!({"path": "./note.txt", "dry_run": true}),
1946 &legacy_policy_bindings("file_edit"),
1947 )
1948 .await
1949 .unwrap();
1950 assert!(dry_run.is_allowed());
1951
1952 let actual = engine
1953 .check_tool_execution_with_bindings(
1954 "file_edit",
1955 &serde_json::json!({"path": "./note.txt", "dry_run": false}),
1956 &legacy_policy_bindings("file_edit"),
1957 )
1958 .await
1959 .unwrap();
1960 assert!(actual.is_blocked());
1961 }
1962}