1use std::collections::HashMap;
2use std::net::IpAddr;
3use std::path::Path;
4use std::sync::Arc;
5use std::time::Instant;
6
7use parking_lot::RwLock;
8use tracing::debug;
9
10use super::config::*;
11use super::path::PathPolicyResolver;
12use ai_agents_core::{
13 CommandBindingKind, CommandPolicyBinding, DomainPolicyBinding, PathAccessMode,
14 PathPolicyBinding, Result, ResultLimitBinding, ResultLimitKind, ToolCallClassification,
15 ToolExecutionLimits, ToolPolicyBindings, ToolSafetyMetadata,
16};
17use serde_json::Value;
18
19#[derive(Debug, Default)]
20struct ToolCallTracker {
21 calls: HashMap<String, Vec<Instant>>,
22}
23
24impl ToolCallTracker {
25 fn admit(&mut self, tool_id: &str, rate_limit: Option<u32>) -> bool {
26 let Some(rate_limit) = rate_limit else {
27 return true;
28 };
29 let now = Instant::now();
30 let window = std::time::Duration::from_secs(60);
31 let calls = self.calls.entry(tool_id.to_string()).or_default();
32 calls.retain(|timestamp| now.duration_since(*timestamp) < window);
33 if calls.len() >= rate_limit as usize {
34 return false;
35 }
36 calls.push(now);
37 true
38 }
39
40 fn reset(&mut self) {
41 self.calls.clear();
42 }
43}
44
45#[derive(Debug, Clone)]
46pub struct ToolSecurityEngine {
47 config: ToolSecurityConfig,
48 tool_call_tracker: Arc<RwLock<ToolCallTracker>>,
49 policy_version: u64,
50}
51
52impl ToolSecurityEngine {
53 pub fn new(config: ToolSecurityConfig) -> Self {
55 Self::try_new(config).expect("invalid tool security configuration")
56 }
57
58 pub fn try_new(config: ToolSecurityConfig) -> Result<Self> {
60 Self::try_new_with_policy_version(config, 1)
61 }
62
63 pub fn new_with_policy_version(config: ToolSecurityConfig, policy_version: u64) -> Self {
65 Self::try_new_with_policy_version(config, policy_version)
66 .expect("invalid tool security configuration")
67 }
68
69 pub fn try_new_with_policy_version(
71 config: ToolSecurityConfig,
72 policy_version: u64,
73 ) -> Result<Self> {
74 config.validate()?;
75 Ok(Self {
76 config,
77 tool_call_tracker: Arc::new(RwLock::new(ToolCallTracker::default())),
78 policy_version,
79 })
80 }
81
82 pub fn validate(&self) -> Result<()> {
84 self.config.validate()
85 }
86
87 pub fn config(&self) -> &ToolSecurityConfig {
88 &self.config
89 }
90
91 pub fn policy_version(&self) -> u64 {
92 self.policy_version
93 }
94
95 pub fn prepare_tool_arguments(&self, tool_id: &str, args: &Value) -> Value {
96 let bindings = legacy_policy_bindings(tool_id);
97 self.prepare_tool_arguments_with_bindings(tool_id, args, &bindings)
98 }
99
100 pub fn prepare_tool_arguments_with_bindings(
101 &self,
102 tool_id: &str,
103 args: &Value,
104 bindings: &ToolPolicyBindings,
105 ) -> Value {
106 if !self.config.enabled {
107 return args.clone();
108 }
109 let mut prepared = args.clone();
110 let Some(tool_config) = self.config.tools.get(tool_id) else {
111 return prepared;
112 };
113 normalize_default_path_arguments(bindings, &mut prepared);
114 apply_policy_caps(tool_config, bindings, &mut prepared);
115 prepared
116 }
117
118 pub fn attach_internal_tool_policy(&self, _tool_id: &str, args: &Value) -> Value {
119 args.clone()
120 }
121
122 pub fn get_tool_output_cap(
123 &self,
124 tool_id: &str,
125 classification_cap: Option<usize>,
126 ) -> Option<usize> {
127 let policy_cap = self
128 .config
129 .enabled
130 .then(|| self.config.tools.get(tool_id))
131 .flatten()
132 .and_then(|config| config.max_output_chars);
133 min_optional_usize(classification_cap, policy_cap)
134 }
135
136 pub fn effective_limits(
137 &self,
138 tool_id: &str,
139 safety: &ToolSafetyMetadata,
140 classification: &ToolCallClassification,
141 ) -> ToolExecutionLimits {
142 let policy = self
143 .config
144 .enabled
145 .then(|| self.config.tools.get(tool_id))
146 .flatten();
147 let policy_timeout_ms = self.get_tool_timeout(tool_id);
148 ToolExecutionLimits {
149 timeout_ms: Some(
150 classification
151 .timeout_ms
152 .map_or(policy_timeout_ms, |timeout_ms| {
153 timeout_ms.min(policy_timeout_ms)
154 }),
155 ),
156 max_output_chars: min_optional_usize(
157 classification.max_output_chars,
158 policy.and_then(|config| config.max_output_chars),
159 ),
160 max_result_chars: safety.max_result_size_chars,
161 max_results: policy.and_then(|config| config.max_results),
162 max_file_size_bytes: policy.and_then(|config| config.max_file_size_bytes),
163 max_response_bytes: policy.and_then(|config| config.max_response_bytes),
164 max_redirects: policy.and_then(|config| config.max_redirects),
165 max_replacements: policy.and_then(|config| config.max_replacements),
166 max_changed_files: policy.and_then(|config| config.max_changed_files),
167 max_changed_lines: policy.and_then(|config| config.max_changed_lines),
168 }
169 }
170
171 pub fn policy_snapshot(&self, tool_id: &str) -> Value {
172 if !self.config.enabled {
173 return Value::Null;
174 }
175 self.config
176 .tools
177 .get(tool_id)
178 .and_then(|config| serde_json::to_value(config).ok())
179 .unwrap_or(Value::Null)
180 }
181
182 pub fn classification_approval_message(
184 &self,
185 tool_id: &str,
186 classification: &ToolCallClassification,
187 ) -> Option<String> {
188 if !classification.requires_approval || classification.read_only {
189 return None;
190 }
191 if !matches!(
192 classification.operation,
193 ai_agents_core::ToolOperationKind::Write
194 | ai_agents_core::ToolOperationKind::Edit
195 | ai_agents_core::ToolOperationKind::Delete
196 | ai_agents_core::ToolOperationKind::Patch
197 | ai_agents_core::ToolOperationKind::Command
198 ) {
199 return None;
200 }
201 let tool_config = self
202 .config
203 .enabled
204 .then(|| self.config.tools.get(tool_id))
205 .flatten();
206 if tool_config.is_some_and(|config| config.allow_without_confirmation) {
207 return None;
208 }
209 Some(format!(
210 "Confirm {} operation for tool '{}' ?",
211 format!("{:?}", classification.operation).to_ascii_lowercase(),
212 tool_id
213 ))
214 }
215
216 pub fn custom_config(&self, tool_id: &str) -> Value {
217 if !self.config.enabled {
218 return Value::Null;
219 }
220 self.config
221 .tools
222 .get(tool_id)
223 .map(|config| Value::Object(config.config.clone().into_iter().collect()))
224 .unwrap_or(Value::Null)
225 }
226
227 pub async fn check_tool_execution(
228 &self,
229 tool_id: &str,
230 args: &serde_json::Value,
231 ) -> Result<SecurityCheckResult> {
232 let bindings = legacy_policy_bindings(tool_id);
233 self.check_tool_execution_with_bindings(tool_id, args, &bindings)
234 .await
235 }
236
237 pub async fn check_tool_execution_with_bindings(
238 &self,
239 tool_id: &str,
240 args: &serde_json::Value,
241 bindings: &ToolPolicyBindings,
242 ) -> Result<SecurityCheckResult> {
243 let validation = self
244 .validate_tool_execution_with_bindings(tool_id, args, bindings)
245 .await?;
246 if validation.is_allowed() {
247 let admission = self.admit_tool_execution(tool_id);
248 if !admission.is_allowed() {
249 return Ok(admission);
250 }
251 }
252 Ok(validation)
253 }
254
255 pub async fn validate_tool_execution_with_bindings(
257 &self,
258 tool_id: &str,
259 args: &serde_json::Value,
260 bindings: &ToolPolicyBindings,
261 ) -> Result<SecurityCheckResult> {
262 if !self.config.enabled {
263 return Ok(SecurityCheckResult::Allow);
264 }
265
266 let tool_config = match self.config.tools.get(tool_id) {
267 Some(config) => config,
268 None if self.config.fail_closed => {
269 return Ok(SecurityCheckResult::Block {
270 reason: format!("Tool '{}' has no explicit security policy", tool_id),
271 });
272 }
273 None => {
274 debug!(tool_id = %tool_id, "Tool execution allowed by legacy open policy");
275 return Ok(SecurityCheckResult::Allow);
276 }
277 };
278
279 if !tool_config.enabled {
280 return Ok(SecurityCheckResult::Unavailable {
281 reason: format!("Tool '{}' is disabled", tool_id),
282 });
283 }
284
285 if let Some(result) =
286 validate_policy_bindings(tool_id, tool_config, bindings, self.config.fail_closed)
287 {
288 return Ok(result);
289 }
290
291 if let Some(result) = self.check_domain_policy(tool_id, tool_config, args, bindings) {
292 return Ok(result);
293 }
294
295 if let Some(result) = self.check_path_policy(tool_id, tool_config, args, bindings) {
296 return Ok(result);
297 }
298
299 if let Some(result) = self.check_operation_policy(tool_id, tool_config, args, bindings) {
300 return Ok(result);
301 }
302
303 if let Some(result) = self.check_command_policy(tool_id, tool_config, args, bindings) {
304 return Ok(result);
305 }
306
307 if tool_config.require_confirmation {
308 let message = tool_config
309 .confirmation_message
310 .clone()
311 .unwrap_or_else(|| format!("Confirm execution of tool '{}' ?", tool_id));
312 return Ok(SecurityCheckResult::RequireConfirmation { message });
313 }
314
315 debug!(tool_id = %tool_id, "Tool execution allowed by policy validation");
316 Ok(SecurityCheckResult::Allow)
317 }
318
319 pub fn admit_tool_execution(&self, tool_id: &str) -> SecurityCheckResult {
320 if !self.config.enabled {
321 return SecurityCheckResult::Allow;
322 }
323 let tool_config = match self.config.tools.get(tool_id) {
324 Some(config) => config,
325 None if self.config.fail_closed => {
326 return SecurityCheckResult::Block {
327 reason: format!("Tool '{}' has no explicit security policy", tool_id),
328 };
329 }
330 None => {
331 self.tool_call_tracker.write().admit(tool_id, None);
332 return SecurityCheckResult::Allow;
333 }
334 };
335 if !tool_config.enabled {
336 return SecurityCheckResult::Unavailable {
337 reason: format!("Tool '{}' is disabled", tool_id),
338 };
339 }
340 let mut tracker = self.tool_call_tracker.write();
341 if !tracker.admit(tool_id, tool_config.rate_limit) {
342 let rate_limit = tool_config.rate_limit.unwrap_or_default();
343 return SecurityCheckResult::Block {
344 reason: format!(
345 "Rate limit exceeded for tool '{}': {} calls per minute",
346 tool_id, rate_limit
347 ),
348 };
349 }
350 debug!(tool_id = %tool_id, "Tool execution admitted");
351 SecurityCheckResult::Allow
352 }
353
354 pub fn check_command_execution(
355 &self,
356 tool_id: &str,
357 command: &str,
358 args: &[String],
359 ) -> SecurityCheckResult {
360 if !self.config.enabled {
361 return SecurityCheckResult::Allow;
362 }
363 let Some(tool_config) = self.config.tools.get(tool_id) else {
364 return if self.config.fail_closed {
365 SecurityCheckResult::Block {
366 reason: format!("Tool '{}' has no explicit command policy", tool_id),
367 }
368 } else {
369 SecurityCheckResult::Allow
370 };
371 };
372 let value = serde_json::json!({
373 "command": command,
374 "argv": std::iter::once(command.to_string()).chain(args.iter().cloned()).collect::<Vec<_>>()
375 });
376 let bindings = legacy_policy_bindings(tool_id);
377 self.check_command_policy(tool_id, tool_config, &value, &bindings)
378 .unwrap_or(SecurityCheckResult::Allow)
379 }
380
381 fn check_domain_policy(
382 &self,
383 tool_id: &str,
384 tool_config: &ToolPolicyConfig,
385 args: &serde_json::Value,
386 bindings: &ToolPolicyBindings,
387 ) -> Option<SecurityCheckResult> {
388 let values = bound_domain_values(args, bindings);
389 if values.is_empty() {
390 return missing_bound_value_result(
391 tool_id,
392 "domain",
393 domain_policy_configured(tool_config),
394 self.config.fail_closed,
395 );
396 }
397 for value in values {
398 let parsed = if value.is_url {
399 match reqwest::Url::parse(&value.value) {
400 Ok(parsed) => parsed,
401 Err(_) => {
402 return Some(SecurityCheckResult::Block {
403 reason: format!("URL is invalid for tool '{}'", tool_id),
404 });
405 }
406 }
407 } else {
408 match reqwest::Url::parse(&format!("https://{}", value.value)) {
409 Ok(parsed) => parsed,
410 Err(_) => {
411 return Some(SecurityCheckResult::Block {
412 reason: format!("Domain is invalid for tool '{}'", tool_id),
413 });
414 }
415 }
416 };
417 let host = parsed.host_str().map(normalize_host)?;
418
419 if !tool_config.allowed_schemes.is_empty()
420 && !tool_config
421 .allowed_schemes
422 .iter()
423 .any(|scheme| scheme.eq_ignore_ascii_case(parsed.scheme()))
424 {
425 return Some(SecurityCheckResult::Block {
426 reason: format!(
427 "URL scheme '{}' is not allowed for tool '{}'",
428 parsed.scheme(),
429 tool_id
430 ),
431 });
432 }
433
434 if !tool_config.allowed_ports.is_empty() {
435 let port = parsed.port_or_known_default().unwrap_or(0);
436 if !tool_config.allowed_ports.contains(&port) {
437 return Some(SecurityCheckResult::Block {
438 reason: format!(
439 "URL port '{}' is not allowed for tool '{}'",
440 port, tool_id
441 ),
442 });
443 }
444 }
445
446 if tool_config.blocked_private_networks && host_is_private_or_local(&host) {
447 return Some(SecurityCheckResult::Block {
448 reason: format!(
449 "Private, localhost, link-local, or metadata host is blocked for tool '{}'",
450 tool_id
451 ),
452 });
453 }
454
455 let denied = tool_config
456 .blocked_domains
457 .iter()
458 .chain(tool_config.domains.deny.iter());
459 for pattern in denied {
460 if host_matches(pattern, &host) {
461 return Some(SecurityCheckResult::Block {
462 reason: format!("Domain '{}' is blocked for tool '{}'", pattern, tool_id),
463 });
464 }
465 }
466
467 for pattern in &tool_config.domains.unavailable {
468 if host_matches(pattern, &host) {
469 return Some(SecurityCheckResult::Unavailable {
470 reason: format!(
471 "Domain '{}' is unavailable for tool '{}'",
472 pattern, tool_id
473 ),
474 });
475 }
476 }
477
478 for pattern in &tool_config.domains.requires_approval {
479 if host_matches(pattern, &host) {
480 return Some(SecurityCheckResult::RequireConfirmation {
481 message: format!(
482 "Confirm access to domain '{}' for tool '{}' ?",
483 host, tool_id
484 ),
485 });
486 }
487 }
488
489 let allowed: Vec<&String> = tool_config
490 .allowed_domains
491 .iter()
492 .chain(tool_config.domains.allow.iter())
493 .collect();
494 if !allowed.is_empty() && !allowed.iter().any(|pattern| host_matches(pattern, &host)) {
495 return Some(SecurityCheckResult::Block {
496 reason: format!("URL domain not in allowed list for tool '{}'", tool_id),
497 });
498 }
499 }
500
501 None
502 }
503
504 fn check_path_policy(
505 &self,
506 tool_id: &str,
507 tool_config: &ToolPolicyConfig,
508 args: &serde_json::Value,
509 bindings: &ToolPolicyBindings,
510 ) -> Option<SecurityCheckResult> {
511 let values = bound_path_values(args, bindings);
512 if values.is_empty() {
513 return missing_bound_value_result(
514 tool_id,
515 "path",
516 path_policy_configured(tool_config),
517 self.config.fail_closed,
518 );
519 }
520 let resolver = match PathPolicyResolver::new() {
521 Ok(resolver) => resolver,
522 Err(error) => return Some(path_resolution_block(tool_id, error)),
523 };
524 for value in values {
525 if let Err(error) = resolver.resolve_path(Path::new(&value.path)) {
526 return Some(path_resolution_block(tool_id, error));
527 }
528
529 for pattern in tool_config
530 .blocked_paths
531 .iter()
532 .chain(tool_config.paths.deny.iter())
533 {
534 match resolver.matches_restriction(Path::new(&value.path), Path::new(pattern)) {
535 Ok(true) => {
536 return Some(SecurityCheckResult::Block {
537 reason: format!("Path is blocked for tool '{}'", tool_id),
538 });
539 }
540 Ok(false) => {}
541 Err(error) => return Some(path_resolution_block(tool_id, error)),
542 }
543 }
544
545 for pattern in tool_config.paths.unavailable.iter() {
546 match resolver.matches_restriction(Path::new(&value.path), Path::new(pattern)) {
547 Ok(true) => {
548 return Some(SecurityCheckResult::Unavailable {
549 reason: format!("Path is unavailable for tool '{}'", tool_id),
550 });
551 }
552 Ok(false) => {}
553 Err(error) => return Some(path_resolution_block(tool_id, error)),
554 }
555 }
556
557 for pattern in tool_config.paths.requires_approval.iter() {
558 match resolver.matches_restriction(Path::new(&value.path), Path::new(pattern)) {
559 Ok(true) => {
560 return Some(SecurityCheckResult::RequireConfirmation {
561 message: format!(
562 "Confirm access to path '{}' for tool '{}' ?",
563 value.path, tool_id
564 ),
565 });
566 }
567 Ok(false) => {}
568 Err(error) => return Some(path_resolution_block(tool_id, error)),
569 }
570 }
571
572 if !matches!(value.kind, ai_agents_core::PathBindingKind::Cwd)
573 && matches!(
574 value.mode,
575 PathAccessMode::Write | PathAccessMode::ReadWrite
576 )
577 && !has_write_allowlist(tool_config)
578 {
579 let dry_run = args
580 .get("dry_run")
581 .and_then(Value::as_bool)
582 .unwrap_or(false);
583 if matches!(tool_config.no_write_policy, NoWritePolicyBehavior::Deny) || !dry_run {
584 return Some(SecurityCheckResult::Block {
585 reason: format!(
586 "Tool '{}' cannot mutate paths without an explicit write_paths policy",
587 tool_id
588 ),
589 });
590 }
591 }
592
593 let allowed = allowed_paths_for_value(tool_config, &value);
594 if matches!(value.kind, ai_agents_core::PathBindingKind::Cwd) && allowed.is_empty() {
595 return Some(SecurityCheckResult::Block {
596 reason: format!(
597 "Tool '{}' requires an explicit working_dirs policy for command cwd",
598 tool_id
599 ),
600 });
601 }
602 if !allowed.is_empty() {
603 let mut matches_allowed = false;
604 for pattern in allowed {
605 match resolver.is_allowed(Path::new(&value.path), Path::new(pattern)) {
606 Ok(true) => {
607 matches_allowed = true;
608 break;
609 }
610 Ok(false) => {}
611 Err(error) => return Some(path_resolution_block(tool_id, error)),
612 }
613 }
614 if !matches_allowed {
615 return Some(SecurityCheckResult::Block {
616 reason: format!("Path not in allowed list for tool '{}'", tool_id),
617 });
618 }
619 }
620 }
621
622 None
623 }
624
625 fn check_operation_policy(
626 &self,
627 tool_id: &str,
628 tool_config: &ToolPolicyConfig,
629 args: &serde_json::Value,
630 bindings: &ToolPolicyBindings,
631 ) -> Option<SecurityCheckResult> {
632 let operations = bound_operation_values(args, bindings);
633 if operations.is_empty() {
634 return missing_bound_value_result(
635 tool_id,
636 "operation",
637 operation_policy_configured(tool_config),
638 self.config.fail_closed,
639 );
640 }
641 for operation in operations {
642 if contains_casefold(&tool_config.operations.deny, &operation) {
643 return Some(SecurityCheckResult::Block {
644 reason: format!(
645 "Operation '{}' is blocked for tool '{}'",
646 operation, tool_id
647 ),
648 });
649 }
650 if contains_casefold(&tool_config.operations.unavailable, &operation) {
651 return Some(SecurityCheckResult::Unavailable {
652 reason: format!(
653 "Operation '{}' is unavailable for tool '{}'",
654 operation, tool_id
655 ),
656 });
657 }
658 if contains_casefold(&tool_config.operations.requires_approval, &operation) {
659 return Some(SecurityCheckResult::RequireConfirmation {
660 message: format!("Confirm operation '{}' for tool '{}' ?", operation, tool_id),
661 });
662 }
663 if !tool_config.operations.allow.is_empty()
664 && !contains_casefold(&tool_config.operations.allow, &operation)
665 {
666 return Some(SecurityCheckResult::Block {
667 reason: format!(
668 "Operation '{}' is not allowed for tool '{}'",
669 operation, tool_id
670 ),
671 });
672 }
673 }
674
675 None
676 }
677
678 fn check_command_policy(
679 &self,
680 tool_id: &str,
681 tool_config: &ToolPolicyConfig,
682 args: &serde_json::Value,
683 bindings: &ToolPolicyBindings,
684 ) -> Option<SecurityCheckResult> {
685 let commands = bound_command_values(args, bindings);
686 if commands.is_empty() {
687 return missing_bound_value_result(
688 tool_id,
689 "command",
690 command_policy_configured(tool_config),
691 self.config.fail_closed,
692 );
693 }
694 for command in commands {
695 let display = command.display();
696 let command_name = command.command_name();
697
698 if command.is_string
699 && command_denies_shell(tool_config)
700 && contains_shell_syntax(&display)
701 {
702 return Some(SecurityCheckResult::Block {
703 reason: format!(
704 "Command '{}' uses shell syntax denied for tool '{}'",
705 display, tool_id
706 ),
707 });
708 }
709 if contains_casefold(&tool_config.commands.deny, &display)
710 || contains_casefold(&tool_config.commands.deny, &command_name)
711 {
712 return Some(SecurityCheckResult::Block {
713 reason: format!("Command '{}' is blocked for tool '{}'", display, tool_id),
714 });
715 }
716 if contains_casefold(&tool_config.commands.unavailable, &display)
717 || contains_casefold(&tool_config.commands.unavailable, &command_name)
718 {
719 return Some(SecurityCheckResult::Unavailable {
720 reason: format!(
721 "Command '{}' is unavailable for tool '{}'",
722 display, tool_id
723 ),
724 });
725 }
726 if contains_casefold(&tool_config.commands.requires_approval, &display)
727 || contains_casefold(&tool_config.commands.requires_approval, &command_name)
728 {
729 return Some(SecurityCheckResult::RequireConfirmation {
730 message: format!("Confirm command '{}' for tool '{}' ?", display, tool_id),
731 });
732 }
733 let has_exact_allowlist = command_has_exact_allowlist(tool_config);
734 if command_requires_exact_allowlist(tool_id) && !has_exact_allowlist {
735 return Some(SecurityCheckResult::Block {
736 reason: format!(
737 "Tool '{}' requires allowed_commands or command_templates before execution",
738 tool_id
739 ),
740 });
741 }
742 if has_exact_allowlist {
743 if !command_matches_allowed(tool_config, &command.argv) {
744 if command_allows_escalation(tool_config) {
745 return Some(SecurityCheckResult::RequireConfirmation {
746 message: format!(
747 "Confirm command '{}' outside the exact allowlist for tool '{}' ?",
748 display, tool_id
749 ),
750 });
751 }
752 return Some(SecurityCheckResult::Block {
753 reason: format!(
754 "Command '{}' is not in the exact argv allowlist for tool '{}'",
755 display, tool_id
756 ),
757 });
758 }
759 continue;
760 }
761 if !tool_config.commands.allow.is_empty()
762 && !contains_casefold(&tool_config.commands.allow, &display)
763 && !contains_casefold(&tool_config.commands.allow, &command_name)
764 {
765 return Some(SecurityCheckResult::Block {
766 reason: format!(
767 "Command '{}' is not allowed for tool '{}'",
768 display, tool_id
769 ),
770 });
771 }
772 }
773
774 None
775 }
776
777 pub fn get_tool_timeout(&self, tool_id: &str) -> u64 {
778 self.config
779 .tools
780 .get(tool_id)
781 .and_then(|c| c.timeout_ms)
782 .unwrap_or(self.config.default_timeout_ms)
783 }
784
785 pub fn reset_session(&self) {
786 self.tool_call_tracker.write().reset();
787 }
788}
789
790impl Default for ToolSecurityEngine {
791 fn default() -> Self {
792 Self::new(ToolSecurityConfig::default())
793 }
794}
795
796fn normalize_default_path_arguments(bindings: &ToolPolicyBindings, args: &mut Value) {
797 for binding in &bindings.path_fields {
798 let Some(default_path) = binding.default_path.as_deref() else {
799 continue;
800 };
801 if value_at_path(args, &binding.field).is_none() {
802 set_root_value(
803 args,
804 &binding.field,
805 Value::String(default_path.to_string()),
806 );
807 }
808 }
809}
810
811fn apply_policy_caps(config: &ToolPolicyConfig, bindings: &ToolPolicyBindings, args: &mut Value) {
812 let Some(obj) = args.as_object_mut() else {
813 return;
814 };
815 for binding in &bindings.result_limit_fields {
816 match binding.kind {
817 ResultLimitKind::MaxResults | ResultLimitKind::Pagination => {
818 apply_usize_cap(obj, &binding.field, config.max_results);
819 }
820 ResultLimitKind::MaxLines => {
821 apply_usize_cap(obj, &binding.field, config.max_results);
822 }
823 ResultLimitKind::MaxOutputChars => {
824 apply_usize_cap(obj, &binding.field, config.max_output_chars);
825 }
826 ResultLimitKind::MaxFileSizeBytes => {
827 apply_u64_cap(obj, &binding.field, config.max_file_size_bytes);
828 }
829 ResultLimitKind::MaxResponseBytes => {
830 apply_usize_cap(obj, &binding.field, config.max_response_bytes);
831 }
832 ResultLimitKind::MaxRedirects => {
833 apply_usize_cap(obj, &binding.field, config.max_redirects);
834 }
835 ResultLimitKind::MaxReplacements => {
836 apply_usize_cap(obj, &binding.field, config.max_replacements);
837 }
838 ResultLimitKind::MaxChangedFiles => {
839 apply_usize_cap(obj, &binding.field, config.max_changed_files);
840 }
841 ResultLimitKind::MaxChangedLines => {
842 apply_usize_cap(obj, &binding.field, config.max_changed_lines);
843 }
844 }
845 }
846}
847
848fn legacy_policy_bindings(tool_id: &str) -> ToolPolicyBindings {
849 match tool_id {
850 "glob" => ToolPolicyBindings {
851 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
852 result_limit_fields: vec![ResultLimitBinding::new(
853 "max_results",
854 ResultLimitKind::MaxResults,
855 )],
856 ..Default::default()
857 },
858 "grep" => ToolPolicyBindings {
859 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
860 result_limit_fields: vec![
861 ResultLimitBinding::new("max_results", ResultLimitKind::MaxResults),
862 ResultLimitBinding::new("max_file_size_bytes", ResultLimitKind::MaxFileSizeBytes),
863 ResultLimitBinding::new("max_output_chars", ResultLimitKind::MaxOutputChars),
864 ],
865 ..Default::default()
866 },
867 "file_read" => ToolPolicyBindings {
868 path_fields: vec![PathPolicyBinding::read("path")],
869 result_limit_fields: vec![
870 ResultLimitBinding::new("max_bytes", ResultLimitKind::MaxFileSizeBytes),
871 ResultLimitBinding::new("max_lines", ResultLimitKind::MaxLines),
872 ],
873 ..Default::default()
874 },
875 "file_list" => ToolPolicyBindings {
876 path_fields: vec![PathPolicyBinding::read("path")],
877 result_limit_fields: vec![ResultLimitBinding::new(
878 "max_results",
879 ResultLimitKind::MaxResults,
880 )],
881 ..Default::default()
882 },
883 "file_info" => ToolPolicyBindings {
884 path_fields: vec![PathPolicyBinding::read("path")],
885 ..Default::default()
886 },
887 "git_status" => ToolPolicyBindings {
888 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
889 result_limit_fields: vec![ResultLimitBinding::new(
890 "max_results",
891 ResultLimitKind::MaxResults,
892 )],
893 ..Default::default()
894 },
895 "git_diff" => ToolPolicyBindings {
896 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
897 result_limit_fields: vec![ResultLimitBinding::new(
898 "max_output_chars",
899 ResultLimitKind::MaxOutputChars,
900 )],
901 ..Default::default()
902 },
903 "diagnostics" => ToolPolicyBindings {
904 path_fields: vec![PathPolicyBinding::read("path").with_default_path(".")],
905 result_limit_fields: vec![ResultLimitBinding::new(
906 "max_results",
907 ResultLimitKind::MaxResults,
908 )],
909 ..Default::default()
910 },
911 "web_fetch" => ToolPolicyBindings {
912 domain_fields: vec![DomainPolicyBinding::url("url")],
913 result_limit_fields: vec![
914 ResultLimitBinding::new("max_chars", ResultLimitKind::MaxOutputChars),
915 ResultLimitBinding::new("max_response_bytes", ResultLimitKind::MaxResponseBytes),
916 ResultLimitBinding::new("max_redirects", ResultLimitKind::MaxRedirects),
917 ],
918 ..Default::default()
919 },
920 "http" => ToolPolicyBindings {
921 domain_fields: vec![DomainPolicyBinding::url("url")],
922 operation_fields: vec!["method".to_string()],
923 ..Default::default()
924 },
925 "file" => ToolPolicyBindings {
926 path_fields: vec![PathPolicyBinding::read_write("path")],
927 operation_fields: vec!["operation".to_string()],
928 ..Default::default()
929 },
930 "file_write" => ToolPolicyBindings {
931 path_fields: vec![PathPolicyBinding::write("path")],
932 result_limit_fields: vec![
933 ResultLimitBinding::new("max_changed_files", ResultLimitKind::MaxChangedFiles),
934 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
935 ],
936 ..Default::default()
937 },
938 "file_edit" => ToolPolicyBindings {
939 path_fields: vec![PathPolicyBinding::write("path")],
940 result_limit_fields: vec![
941 ResultLimitBinding::new("max_replacements", ResultLimitKind::MaxReplacements),
942 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
943 ],
944 ..Default::default()
945 },
946 "patch" => ToolPolicyBindings {
947 path_fields: vec![
948 ai_agents_core::PathPolicyBinding::new(
949 "base_path",
950 PathAccessMode::Write,
951 ai_agents_core::PathBindingKind::PatchBase,
952 )
953 .with_default_path("."),
954 ],
955 result_limit_fields: vec![
956 ResultLimitBinding::new("max_changed_files", ResultLimitKind::MaxChangedFiles),
957 ResultLimitBinding::new("max_changed_lines", ResultLimitKind::MaxChangedLines),
958 ],
959 ..Default::default()
960 },
961 "copy_path" => ToolPolicyBindings {
962 path_fields: vec![
963 PathPolicyBinding::read("source_path"),
964 PathPolicyBinding::write("destination_path"),
965 ],
966 ..Default::default()
967 },
968 "move_path" => ToolPolicyBindings {
969 path_fields: vec![
970 PathPolicyBinding::read_write("source_path"),
971 PathPolicyBinding::write("destination_path"),
972 ],
973 ..Default::default()
974 },
975 "delete_path" => ToolPolicyBindings {
976 path_fields: vec![PathPolicyBinding::write("path")],
977 ..Default::default()
978 },
979 "command" => ToolPolicyBindings {
980 command_fields: vec![
981 CommandPolicyBinding::command("command"),
982 CommandPolicyBinding::argv("argv"),
983 CommandPolicyBinding::env("env"),
984 ],
985 path_fields: vec![
986 ai_agents_core::PathPolicyBinding::new(
987 "cwd",
988 PathAccessMode::ReadWrite,
989 ai_agents_core::PathBindingKind::Cwd,
990 )
991 .with_default_path("."),
992 ],
993 result_limit_fields: vec![ResultLimitBinding::new(
994 "max_output_chars",
995 ResultLimitKind::MaxOutputChars,
996 )],
997 ..Default::default()
998 },
999 _ => ToolPolicyBindings::default(),
1000 }
1001}
1002
1003#[derive(Debug, Clone)]
1004struct BoundPathValue {
1005 path: String,
1006 mode: PathAccessMode,
1007 kind: ai_agents_core::PathBindingKind,
1008}
1009
1010#[derive(Debug, Clone)]
1011struct BoundDomainValue {
1012 value: String,
1013 is_url: bool,
1014}
1015
1016#[derive(Debug, Clone)]
1017struct BoundCommandValue {
1018 argv: Vec<String>,
1019 is_string: bool,
1020}
1021
1022impl BoundCommandValue {
1023 fn display(&self) -> String {
1024 self.argv.join(" ")
1025 }
1026
1027 fn command_name(&self) -> String {
1028 self.argv.first().cloned().unwrap_or_default()
1029 }
1030}
1031
1032fn validate_policy_bindings(
1033 tool_id: &str,
1034 config: &ToolPolicyConfig,
1035 bindings: &ToolPolicyBindings,
1036 fail_closed: bool,
1037) -> Option<SecurityCheckResult> {
1038 if !fail_closed {
1039 return None;
1040 }
1041 if path_policy_configured(config) && !bindings.has_path_bindings() {
1042 return Some(SecurityCheckResult::Block {
1043 reason: format!(
1044 "path policy configured for {} but tool exposes no path policy bindings",
1045 tool_id
1046 ),
1047 });
1048 }
1049 if domain_policy_configured(config) && !bindings.has_domain_bindings() {
1050 return Some(SecurityCheckResult::Block {
1051 reason: format!(
1052 "domain policy configured for {} but tool exposes no domain policy bindings",
1053 tool_id
1054 ),
1055 });
1056 }
1057 if command_policy_configured(config) && !bindings.has_command_bindings() {
1058 return Some(SecurityCheckResult::Block {
1059 reason: format!(
1060 "command policy configured for {} but tool exposes no command policy bindings",
1061 tool_id
1062 ),
1063 });
1064 }
1065 if operation_policy_configured(config) && !bindings.has_operation_bindings() {
1066 return Some(SecurityCheckResult::Block {
1067 reason: format!(
1068 "operation policy configured for {} but tool exposes no operation policy bindings",
1069 tool_id
1070 ),
1071 });
1072 }
1073 if result_limit_policy_configured(config) && !bindings.has_result_limit_bindings() {
1074 return Some(SecurityCheckResult::Block {
1075 reason: format!(
1076 "result-limit policy configured for {} but tool exposes no result-limit policy bindings",
1077 tool_id
1078 ),
1079 });
1080 }
1081 None
1082}
1083
1084fn missing_bound_value_result(
1085 tool_id: &str,
1086 policy_kind: &str,
1087 configured: bool,
1088 fail_closed: bool,
1089) -> Option<SecurityCheckResult> {
1090 if configured && fail_closed {
1091 Some(SecurityCheckResult::Block {
1092 reason: format!(
1093 "{} policy configured for {} but no bound {} argument was present",
1094 policy_kind, tool_id, policy_kind
1095 ),
1096 })
1097 } else {
1098 None
1099 }
1100}
1101
1102fn path_policy_configured(config: &ToolPolicyConfig) -> bool {
1103 !config.allowed_paths.is_empty()
1104 || !config.read_paths.is_empty()
1105 || !config.write_paths.is_empty()
1106 || !config.working_dirs.is_empty()
1107 || !config.commands.working_dirs.is_empty()
1108 || !config.blocked_paths.is_empty()
1109 || !config.paths.allow.is_empty()
1110 || !config.paths.deny.is_empty()
1111 || !config.paths.requires_approval.is_empty()
1112 || !config.paths.unavailable.is_empty()
1113}
1114
1115fn domain_policy_configured(config: &ToolPolicyConfig) -> bool {
1116 !config.allowed_domains.is_empty()
1117 || !config.blocked_domains.is_empty()
1118 || !config.allowed_schemes.is_empty()
1119 || !config.allowed_ports.is_empty()
1120 || !config.domains.allow.is_empty()
1121 || !config.domains.deny.is_empty()
1122 || !config.domains.requires_approval.is_empty()
1123 || !config.domains.unavailable.is_empty()
1124}
1125
1126fn command_policy_configured(config: &ToolPolicyConfig) -> bool {
1127 !config.commands.allow.is_empty()
1128 || !config.commands.deny.is_empty()
1129 || !config.commands.requires_approval.is_empty()
1130 || !config.commands.unavailable.is_empty()
1131 || !config.commands.allowed_commands.is_empty()
1132 || !config.commands.templates.is_empty()
1133 || !config.allowed_commands.is_empty()
1134 || !config.command_templates.is_empty()
1135 || !config.env_passthrough.is_empty()
1136 || !config.commands.env_passthrough.is_empty()
1137}
1138
1139fn operation_policy_configured(config: &ToolPolicyConfig) -> bool {
1140 !config.operations.allow.is_empty()
1141 || !config.operations.deny.is_empty()
1142 || !config.operations.requires_approval.is_empty()
1143 || !config.operations.unavailable.is_empty()
1144}
1145
1146fn result_limit_policy_configured(config: &ToolPolicyConfig) -> bool {
1147 config.max_file_size_bytes.is_some()
1149 || config.max_results.is_some()
1150 || config.max_response_bytes.is_some()
1151 || config.max_redirects.is_some()
1152 || config.max_replacements.is_some()
1153 || config.max_changed_files.is_some()
1154 || config.max_changed_lines.is_some()
1155}
1156
1157fn bound_path_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundPathValue> {
1158 let mut values = Vec::new();
1159 for binding in &bindings.path_fields {
1160 collect_path_binding_values(args, binding, &mut values);
1161 }
1162 values
1163}
1164
1165fn collect_path_binding_values(
1166 args: &Value,
1167 binding: &PathPolicyBinding,
1168 values: &mut Vec<BoundPathValue>,
1169) {
1170 let value = value_at_path(args, &binding.field).cloned().or_else(|| {
1171 binding
1172 .default_path
1173 .as_ref()
1174 .map(|path| Value::String(path.clone()))
1175 });
1176 let Some(value) = value else {
1177 return;
1178 };
1179 match value {
1180 Value::String(path) => values.push(BoundPathValue {
1181 path,
1182 mode: effective_path_mode(args, binding),
1183 kind: binding.kind,
1184 }),
1185 Value::Array(items) => {
1186 for item in items {
1187 if let Some(path) = item.as_str() {
1188 values.push(BoundPathValue {
1189 path: path.to_string(),
1190 mode: effective_path_mode(args, binding),
1191 kind: binding.kind,
1192 });
1193 }
1194 }
1195 }
1196 _ => {}
1197 }
1198}
1199
1200fn bound_domain_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundDomainValue> {
1201 let mut values = Vec::new();
1202 for binding in &bindings.domain_fields {
1203 collect_domain_binding_values(args, binding, &mut values);
1204 }
1205 values
1206}
1207
1208fn collect_domain_binding_values(
1209 args: &Value,
1210 binding: &DomainPolicyBinding,
1211 values: &mut Vec<BoundDomainValue>,
1212) {
1213 let Some(value) = value_at_path(args, &binding.field) else {
1214 return;
1215 };
1216 match value {
1217 Value::String(value) => values.push(BoundDomainValue {
1218 value: value.clone(),
1219 is_url: binding.is_url,
1220 }),
1221 Value::Array(items) => {
1222 for item in items {
1223 if let Some(value) = item.as_str() {
1224 values.push(BoundDomainValue {
1225 value: value.to_string(),
1226 is_url: binding.is_url,
1227 });
1228 }
1229 }
1230 }
1231 _ => {}
1232 }
1233}
1234
1235fn bound_operation_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<String> {
1236 bindings
1237 .operation_fields
1238 .iter()
1239 .filter_map(|field| value_at_path(args, field).and_then(Value::as_str))
1240 .map(|value| value.trim().to_ascii_lowercase())
1241 .collect()
1242}
1243
1244fn bound_command_values(args: &Value, bindings: &ToolPolicyBindings) -> Vec<BoundCommandValue> {
1245 let mut values = Vec::new();
1246 for binding in &bindings.command_fields {
1247 collect_command_binding_values(args, binding, &mut values);
1248 }
1249 values
1250}
1251
1252fn collect_command_binding_values(
1253 args: &Value,
1254 binding: &CommandPolicyBinding,
1255 values: &mut Vec<BoundCommandValue>,
1256) {
1257 let Some(value) = value_at_path(args, &binding.field) else {
1258 return;
1259 };
1260 match binding.kind {
1261 CommandBindingKind::CommandString => {
1262 if let Some(command) = value.as_str() {
1263 if let Some(argv) = parse_command_words(command) {
1264 values.push(BoundCommandValue {
1265 argv,
1266 is_string: true,
1267 });
1268 } else {
1269 values.push(BoundCommandValue {
1270 argv: vec![command.to_string()],
1271 is_string: true,
1272 });
1273 }
1274 }
1275 }
1276 CommandBindingKind::Argv => {
1277 if let Some(argv) = value.as_array().map(|items| {
1278 items
1279 .iter()
1280 .filter_map(Value::as_str)
1281 .map(str::to_string)
1282 .collect::<Vec<_>>()
1283 }) && !argv.is_empty()
1284 {
1285 values.push(BoundCommandValue {
1286 argv,
1287 is_string: false,
1288 });
1289 }
1290 }
1291 CommandBindingKind::Cwd
1292 | CommandBindingKind::TemplateVariable
1293 | CommandBindingKind::Env => {}
1294 }
1295}
1296
1297fn allowed_paths_for_value<'a>(
1298 config: &'a ToolPolicyConfig,
1299 value: &BoundPathValue,
1300) -> Vec<&'a String> {
1301 if matches!(value.kind, ai_agents_core::PathBindingKind::Cwd) {
1302 return config
1303 .working_dirs
1304 .iter()
1305 .chain(config.commands.working_dirs.iter())
1306 .collect();
1307 }
1308 let mut allowed: Vec<&String> = config
1309 .allowed_paths
1310 .iter()
1311 .chain(config.paths.allow.iter())
1312 .collect();
1313 match value.mode {
1314 PathAccessMode::Read => allowed.extend(config.read_paths.iter()),
1315 PathAccessMode::Write | PathAccessMode::ReadWrite => {
1316 allowed.extend(config.write_paths.iter());
1317 }
1318 }
1319 allowed
1320}
1321
1322fn has_write_allowlist(config: &ToolPolicyConfig) -> bool {
1323 !config.write_paths.is_empty()
1324 || !config.allowed_paths.is_empty()
1325 || !config.paths.allow.is_empty()
1326}
1327
1328fn effective_path_mode(args: &Value, binding: &PathPolicyBinding) -> PathAccessMode {
1329 if !matches!(binding.mode, PathAccessMode::ReadWrite) {
1330 return binding.mode;
1331 }
1332 let operation = args
1333 .get("operation")
1334 .and_then(Value::as_str)
1335 .unwrap_or_default()
1336 .to_ascii_lowercase();
1337 match operation.as_str() {
1338 "read" | "exists" | "list" | "info" => PathAccessMode::Read,
1339 "write" | "append" | "mkdir" | "delete" | "edit" | "patch" => PathAccessMode::Write,
1340 _ => binding.mode,
1341 }
1342}
1343
1344fn value_at_path<'a>(value: &'a Value, field: &str) -> Option<&'a Value> {
1345 let mut current = value;
1346 for segment in field.split('.') {
1347 if segment.is_empty() {
1348 return None;
1349 }
1350 current = current.get(segment)?;
1351 }
1352 Some(current)
1353}
1354
1355fn set_root_value(args: &mut Value, field: &str, value: Value) {
1356 if field.contains('.') {
1357 return;
1358 }
1359 let Some(obj) = args.as_object_mut() else {
1360 return;
1361 };
1362 obj.insert(field.to_string(), value);
1363}
1364
1365fn apply_usize_cap(obj: &mut serde_json::Map<String, Value>, key: &str, cap: Option<usize>) {
1366 let Some(cap) = cap else {
1367 return;
1368 };
1369 let effective = obj
1370 .get(key)
1371 .and_then(|value| value.as_u64())
1372 .map(|value| value.min(cap as u64) as usize)
1373 .unwrap_or(cap);
1374 obj.insert(key.to_string(), Value::from(effective));
1375}
1376
1377fn apply_u64_cap(obj: &mut serde_json::Map<String, Value>, key: &str, cap: Option<u64>) {
1378 let Some(cap) = cap else {
1379 return;
1380 };
1381 let effective = obj
1382 .get(key)
1383 .and_then(|value| value.as_u64())
1384 .map(|value| value.min(cap))
1385 .unwrap_or(cap);
1386 obj.insert(key.to_string(), Value::from(effective));
1387}
1388
1389fn min_optional_usize(left: Option<usize>, right: Option<usize>) -> Option<usize> {
1390 match (left, right) {
1391 (Some(left), Some(right)) => Some(left.min(right)),
1392 (Some(left), None) => Some(left),
1393 (None, Some(right)) => Some(right),
1394 (None, None) => None,
1395 }
1396}
1397
1398fn normalize_host(host: &str) -> String {
1399 host.trim_end_matches('.').to_ascii_lowercase()
1400}
1401
1402fn host_matches(pattern: &str, host: &str) -> bool {
1403 let pattern = normalize_host(pattern.trim_start_matches("*."));
1404 host == pattern || host.ends_with(&format!(".{}", pattern))
1405}
1406
1407fn path_resolution_block(
1408 tool_id: &str,
1409 error: super::path::PathResolutionError,
1410) -> SecurityCheckResult {
1411 SecurityCheckResult::Block {
1412 reason: format!(
1413 "Path policy resolution failed for tool '{}': {}",
1414 tool_id, error
1415 ),
1416 }
1417}
1418
1419fn contains_casefold(values: &[String], needle: &str) -> bool {
1420 values
1421 .iter()
1422 .any(|value| value.eq_ignore_ascii_case(needle))
1423}
1424
1425fn command_denies_shell(config: &ToolPolicyConfig) -> bool {
1426 config.deny_shell || config.commands.deny_shell
1427}
1428
1429fn command_allows_escalation(config: &ToolPolicyConfig) -> bool {
1430 config.allow_command_escalation || config.commands.allow_escalation
1431}
1432
1433fn command_requires_exact_allowlist(tool_id: &str) -> bool {
1434 tool_id == "command"
1435}
1436
1437fn command_has_exact_allowlist(config: &ToolPolicyConfig) -> bool {
1438 !config.allowed_commands.is_empty()
1439 || !config.commands.allowed_commands.is_empty()
1440 || !config.command_templates.is_empty()
1441 || !config.commands.templates.is_empty()
1442}
1443
1444fn command_matches_allowed(config: &ToolPolicyConfig, argv: &[String]) -> bool {
1445 config
1446 .allowed_commands
1447 .iter()
1448 .chain(config.commands.allowed_commands.iter())
1449 .any(|rule| rule.argv == argv)
1450 || config
1451 .command_templates
1452 .iter()
1453 .chain(config.commands.templates.iter())
1454 .any(|template| command_matches_template(&template.argv, argv))
1455}
1456
1457fn command_matches_template(template: &[String], argv: &[String]) -> bool {
1458 template.len() == argv.len()
1459 && template.iter().zip(argv.iter()).all(|(expected, actual)| {
1460 (expected.starts_with('{') && expected.ends_with('}')) || expected == actual
1461 })
1462}
1463
1464fn contains_shell_syntax(value: &str) -> bool {
1465 const DENIED: &[char] = &[';', '&', '|', '<', '>', '`', '$', '\n', '\r'];
1466 value.chars().any(|ch| DENIED.contains(&ch))
1467 || value.contains("$(")
1468 || value.contains("${")
1469 || value.contains("<(")
1470 || value.contains(">(")
1471}
1472
1473fn parse_command_words(value: &str) -> Option<Vec<String>> {
1474 let mut words = Vec::new();
1475 let mut current = String::new();
1476 let mut quote: Option<char> = None;
1477 for ch in value.chars() {
1478 match (quote, ch) {
1479 (Some(q), c) if c == q => quote = None,
1480 (Some(_), c) => current.push(c),
1481 (None, '\'' | '"') => quote = Some(ch),
1482 (None, c) if c.is_whitespace() => {
1483 if !current.is_empty() {
1484 words.push(std::mem::take(&mut current));
1485 }
1486 }
1487 (None, c) => current.push(c),
1488 }
1489 }
1490 if quote.is_some() {
1491 return None;
1492 }
1493 if !current.is_empty() {
1494 words.push(current);
1495 }
1496 (!words.is_empty()).then_some(words)
1497}
1498
1499fn host_is_private_or_local(host: &str) -> bool {
1500 if matches!(
1501 host,
1502 "localhost"
1503 | "metadata"
1504 | "metadata.google.internal"
1505 | "169.254.169.254"
1506 | "100.100.100.200"
1507 ) || host.ends_with(".localhost")
1508 {
1509 return true;
1510 }
1511 match host.parse::<IpAddr>() {
1512 Ok(IpAddr::V4(ip)) => {
1513 ip.is_private()
1514 || ip.is_loopback()
1515 || ip.is_link_local()
1516 || ip.is_multicast()
1517 || ip.is_documentation()
1518 || ip.octets() == [169, 254, 169, 254]
1519 }
1520 Ok(IpAddr::V6(ip)) => {
1521 ip.is_loopback()
1522 || ip.is_unspecified()
1523 || ip.is_multicast()
1524 || ip.segments()[0] & 0xfe00 == 0xfc00
1525 || ip.segments()[0] & 0xffc0 == 0xfe80
1526 }
1527 Err(_) => false,
1528 }
1529}
1530
1531#[cfg(test)]
1532mod tests {
1533 use super::*;
1534
1535 fn enabled_security_config() -> ToolSecurityConfig {
1536 ToolSecurityConfig {
1537 enabled: true,
1538 ..Default::default()
1539 }
1540 }
1541
1542 #[test]
1543 fn test_default_engine() {
1544 let engine = ToolSecurityEngine::default();
1545 assert!(!engine.config().enabled);
1546 }
1547
1548 #[tokio::test]
1549 async fn test_tool_domain_blocking() {
1550 let mut config = enabled_security_config();
1551
1552 let http_config = ToolPolicyConfig {
1553 blocked_domains: vec!["evil.com".to_string()],
1554 ..Default::default()
1555 };
1556 config.tools.insert("http".to_string(), http_config);
1557
1558 let engine = ToolSecurityEngine::new(config);
1559
1560 let args = serde_json::json!({"url": "https://evil.com/api"});
1561 let result = engine.check_tool_execution("http", &args).await.unwrap();
1562 assert!(result.is_blocked());
1563
1564 let args = serde_json::json!({"url": "https://not-evil.com/api"});
1565 let result = engine.check_tool_execution("http", &args).await.unwrap();
1566 assert!(result.is_allowed());
1567 }
1568
1569 #[tokio::test]
1570 async fn test_tool_allowed_domains() {
1571 let mut config = enabled_security_config();
1572
1573 let http_config = ToolPolicyConfig {
1574 allowed_domains: vec!["api.example.com".to_string()],
1575 ..Default::default()
1576 };
1577 config.tools.insert("http".to_string(), http_config);
1578
1579 let engine = ToolSecurityEngine::new(config);
1580
1581 let args = serde_json::json!({"url": "https://api.example.com/v1"});
1582 let result = engine.check_tool_execution("http", &args).await.unwrap();
1583 assert!(result.is_allowed());
1584
1585 let args = serde_json::json!({"url": "https://other.com/api"});
1586 let result = engine.check_tool_execution("http", &args).await.unwrap();
1587 assert!(result.is_blocked());
1588 }
1589
1590 #[tokio::test]
1591 async fn test_tool_disabled() {
1592 let mut config = enabled_security_config();
1593
1594 let tool_config = ToolPolicyConfig {
1595 enabled: false,
1596 ..Default::default()
1597 };
1598 config.tools.insert("dangerous".to_string(), tool_config);
1599
1600 let engine = ToolSecurityEngine::new(config);
1601
1602 let result = engine
1603 .check_tool_execution("dangerous", &serde_json::json!({}))
1604 .await
1605 .unwrap();
1606 assert!(result.is_blocked());
1607 assert!(result.is_unavailable());
1608 }
1609
1610 #[tokio::test]
1611 async fn test_tool_confirmation_required() {
1612 let mut config = enabled_security_config();
1613
1614 let tool_config = ToolPolicyConfig {
1615 require_confirmation: true,
1616 confirmation_message: Some("Are you sure?".to_string()),
1617 ..Default::default()
1618 };
1619 config.tools.insert("delete".to_string(), tool_config);
1620
1621 let engine = ToolSecurityEngine::new(config);
1622
1623 let result = engine
1624 .check_tool_execution("delete", &serde_json::json!({}))
1625 .await
1626 .unwrap();
1627
1628 match result {
1629 SecurityCheckResult::RequireConfirmation { message } => {
1630 assert_eq!(message, "Are you sure?");
1631 }
1632 _ => panic!("Expected RequireConfirmation"),
1633 }
1634 }
1635
1636 #[test]
1637 fn test_get_tool_timeout() {
1638 let mut config = ToolSecurityConfig {
1639 default_timeout_ms: 5000,
1640 ..Default::default()
1641 };
1642
1643 let tool_config = ToolPolicyConfig {
1644 timeout_ms: Some(10000),
1645 ..Default::default()
1646 };
1647 config.tools.insert("slow".to_string(), tool_config);
1648
1649 let engine = ToolSecurityEngine::new(config);
1650
1651 assert_eq!(engine.get_tool_timeout("slow"), 10000);
1652 assert_eq!(engine.get_tool_timeout("other"), 5000);
1653 }
1654
1655 #[test]
1656 fn call_classification_timeout_only_lowers_policy_timeout() {
1657 let engine = ToolSecurityEngine::new(ToolSecurityConfig {
1658 default_timeout_ms: 5_000,
1659 ..Default::default()
1660 });
1661 let safety = ToolSafetyMetadata::compute();
1662 let mut classification = ToolCallClassification::from_metadata(&safety);
1663
1664 classification.timeout_ms = Some(1_000);
1665 assert_eq!(
1666 engine
1667 .effective_limits("custom", &safety, &classification)
1668 .timeout_ms,
1669 Some(1_000)
1670 );
1671
1672 classification.timeout_ms = Some(10_000);
1673 assert_eq!(
1674 engine
1675 .effective_limits("custom", &safety, &classification)
1676 .timeout_ms,
1677 Some(5_000)
1678 );
1679 }
1680
1681 #[tokio::test]
1682 async fn test_path_restrictions() {
1683 let directory = tempfile::tempdir().unwrap();
1684 let allowed_root = directory.path().join("allowed");
1685 std::fs::create_dir(&allowed_root).unwrap();
1686 let allowed_path = allowed_root.join("test.txt");
1687 let denied_path = directory.path().join("denied/test.txt");
1688 let mut config = enabled_security_config();
1689
1690 let tool_config = ToolPolicyConfig {
1691 allowed_paths: vec![allowed_root.to_string_lossy().into_owned()],
1692 ..Default::default()
1693 };
1694 config.tools.insert("file_write".to_string(), tool_config);
1695
1696 let engine = ToolSecurityEngine::new(config);
1697
1698 let args = serde_json::json!({"path": allowed_path});
1699 let result = engine
1700 .check_tool_execution("file_write", &args)
1701 .await
1702 .unwrap();
1703 assert!(result.is_allowed(), "{result:?}");
1704
1705 let args = serde_json::json!({"path": denied_path});
1706 let result = engine
1707 .check_tool_execution("file_write", &args)
1708 .await
1709 .unwrap();
1710 assert!(result.is_blocked(), "{result:?}");
1711 }
1712
1713 #[tokio::test]
1714 async fn test_operation_policy() {
1715 let mut config = enabled_security_config();
1716 let tool_config = ToolPolicyConfig {
1717 operations: OperationPolicyConfig {
1718 deny: vec!["delete".to_string()],
1719 requires_approval: vec!["write".to_string()],
1720 ..Default::default()
1721 },
1722 ..Default::default()
1723 };
1724 config.tools.insert("file".to_string(), tool_config);
1725 let engine = ToolSecurityEngine::new(config);
1726
1727 let result = engine
1728 .check_tool_execution("file", &serde_json::json!({"operation": "delete"}))
1729 .await
1730 .unwrap();
1731 assert!(result.is_blocked());
1732
1733 let result = engine
1734 .check_tool_execution("file", &serde_json::json!({"operation": "write"}))
1735 .await
1736 .unwrap();
1737 assert!(result.requires_approval());
1738 }
1739
1740 #[tokio::test]
1741 async fn omitted_optional_path_uses_default_for_policy() {
1742 let mut config = enabled_security_config();
1743 let tool_config = ToolPolicyConfig {
1744 read_paths: vec!["./crates".to_string()],
1745 ..Default::default()
1746 };
1747 config.tools.insert("grep".to_string(), tool_config);
1748 let engine = ToolSecurityEngine::new(config);
1749
1750 let result = engine
1751 .check_tool_execution("grep", &serde_json::json!({"pattern": "Tool"}))
1752 .await
1753 .unwrap();
1754 assert!(result.is_blocked());
1755
1756 let prepared =
1757 engine.prepare_tool_arguments("grep", &serde_json::json!({"pattern": "Tool"}));
1758 assert_eq!(prepared.get("path").and_then(Value::as_str), Some("."));
1759 }
1760
1761 #[tokio::test]
1762 async fn fail_closed_requires_path_bindings_for_custom_tools() {
1763 let mut config = ToolSecurityConfig {
1764 fail_closed: true,
1765 ..enabled_security_config()
1766 };
1767 let tool_config = ToolPolicyConfig {
1768 read_paths: vec!["./allowed".to_string()],
1769 ..Default::default()
1770 };
1771 config
1772 .tools
1773 .insert("custom_search".to_string(), tool_config);
1774 let engine = ToolSecurityEngine::new(config);
1775
1776 let result = engine
1777 .check_tool_execution_with_bindings(
1778 "custom_search",
1779 &serde_json::json!({"path": "./allowed/file.txt"}),
1780 &ToolPolicyBindings::default(),
1781 )
1782 .await
1783 .unwrap();
1784
1785 assert!(result.is_blocked());
1786 assert!(
1787 result
1788 .reason()
1789 .unwrap_or_default()
1790 .contains("tool exposes no path policy bindings")
1791 );
1792 }
1793
1794 #[tokio::test]
1795 async fn custom_path_bindings_enforce_blocked_paths() {
1796 let mut config = ToolSecurityConfig {
1797 fail_closed: true,
1798 ..enabled_security_config()
1799 };
1800 let tool_config = ToolPolicyConfig {
1801 read_paths: vec!["./allowed".to_string()],
1802 blocked_paths: vec!["./allowed/private".to_string()],
1803 ..Default::default()
1804 };
1805 config
1806 .tools
1807 .insert("custom_search".to_string(), tool_config);
1808 let engine = ToolSecurityEngine::new(config);
1809 let bindings = ToolPolicyBindings {
1810 path_fields: vec![PathPolicyBinding::read("root")],
1811 ..Default::default()
1812 };
1813
1814 let allowed = engine
1815 .check_tool_execution_with_bindings(
1816 "custom_search",
1817 &serde_json::json!({"root": "./allowed/src"}),
1818 &bindings,
1819 )
1820 .await
1821 .unwrap();
1822 assert!(allowed.is_allowed());
1823
1824 let blocked = engine
1825 .check_tool_execution_with_bindings(
1826 "custom_search",
1827 &serde_json::json!({"root": "./allowed/private/secrets.txt"}),
1828 &bindings,
1829 )
1830 .await
1831 .unwrap();
1832 assert!(blocked.is_blocked());
1833 }
1834
1835 #[cfg(unix)]
1836 #[tokio::test]
1837 async fn blocked_path_cannot_be_reached_through_symlink_alias() {
1838 use std::os::unix::fs::symlink;
1839
1840 let root = tempfile::tempdir().unwrap();
1841 let private = root.path().join("private");
1842 let public = root.path().join("public");
1843 std::fs::create_dir_all(&private).unwrap();
1844 std::fs::create_dir_all(&public).unwrap();
1845 symlink(&private, public.join("alias")).unwrap();
1846
1847 let mut config = ToolSecurityConfig {
1848 fail_closed: true,
1849 ..enabled_security_config()
1850 };
1851 let tool_config = ToolPolicyConfig {
1852 read_paths: vec![root.path().to_string_lossy().into_owned()],
1853 blocked_paths: vec![private.to_string_lossy().into_owned()],
1854 ..Default::default()
1855 };
1856 config
1857 .tools
1858 .insert("custom_search".to_string(), tool_config);
1859 let engine = ToolSecurityEngine::new(config);
1860 let bindings = ToolPolicyBindings {
1861 path_fields: vec![PathPolicyBinding::read("root")],
1862 ..Default::default()
1863 };
1864
1865 let result = engine
1866 .check_tool_execution_with_bindings(
1867 "custom_search",
1868 &serde_json::json!({"root": public.join("alias/secret.txt")}),
1869 &bindings,
1870 )
1871 .await
1872 .unwrap();
1873
1874 assert!(result.is_blocked());
1875 }
1876
1877 #[cfg(unix)]
1878 #[tokio::test]
1879 async fn path_restrictions_keep_results_through_symlink_aliases() {
1880 use std::os::unix::fs::symlink;
1881
1882 let root = tempfile::tempdir().unwrap();
1883 let workspace = root.path().join("workspace");
1884 let restricted = root.path().join("restricted");
1885 std::fs::create_dir(&workspace).unwrap();
1886 std::fs::create_dir(&restricted).unwrap();
1887 symlink(&restricted, workspace.join("alias")).unwrap();
1888 let candidate = workspace.join("alias/secret.txt");
1889 let bindings = ToolPolicyBindings {
1890 path_fields: vec![PathPolicyBinding::read("root")],
1891 ..Default::default()
1892 };
1893
1894 let mut denied_config = enabled_security_config();
1895 denied_config.tools.insert(
1896 "custom_search".to_string(),
1897 ToolPolicyConfig {
1898 read_paths: vec![restricted.to_string_lossy().into_owned()],
1899 blocked_paths: vec![restricted.to_string_lossy().into_owned()],
1900 ..Default::default()
1901 },
1902 );
1903 let denied = ToolSecurityEngine::new(denied_config)
1904 .check_tool_execution_with_bindings(
1905 "custom_search",
1906 &serde_json::json!({"root": candidate}),
1907 &bindings,
1908 )
1909 .await
1910 .unwrap();
1911 assert!(denied.is_blocked());
1912 assert!(!denied.is_unavailable());
1913
1914 let mut unavailable_config = enabled_security_config();
1915 unavailable_config.tools.insert(
1916 "custom_search".to_string(),
1917 ToolPolicyConfig {
1918 read_paths: vec![restricted.to_string_lossy().into_owned()],
1919 paths: PathPolicyConfig {
1920 unavailable: vec![restricted.to_string_lossy().into_owned()],
1921 ..Default::default()
1922 },
1923 ..Default::default()
1924 },
1925 );
1926 let unavailable = ToolSecurityEngine::new(unavailable_config)
1927 .check_tool_execution_with_bindings(
1928 "custom_search",
1929 &serde_json::json!({"root": candidate}),
1930 &bindings,
1931 )
1932 .await
1933 .unwrap();
1934 assert!(unavailable.is_unavailable());
1935
1936 let mut approval_config = enabled_security_config();
1937 approval_config.tools.insert(
1938 "custom_search".to_string(),
1939 ToolPolicyConfig {
1940 read_paths: vec![restricted.to_string_lossy().into_owned()],
1941 paths: PathPolicyConfig {
1942 requires_approval: vec![restricted.to_string_lossy().into_owned()],
1943 ..Default::default()
1944 },
1945 ..Default::default()
1946 },
1947 );
1948 let approval = ToolSecurityEngine::new(approval_config)
1949 .check_tool_execution_with_bindings(
1950 "custom_search",
1951 &serde_json::json!({"root": candidate}),
1952 &bindings,
1953 )
1954 .await
1955 .unwrap();
1956 assert!(approval.requires_approval());
1957 }
1958
1959 #[tokio::test]
1960 async fn copy_and_move_bindings_enforce_source_and_destination_roots() {
1961 let root = tempfile::tempdir().unwrap();
1962 let readable = root.path().join("readable");
1963 let writable = root.path().join("writable");
1964 let outside = root.path().join("outside");
1965 std::fs::create_dir(&readable).unwrap();
1966 std::fs::create_dir(&writable).unwrap();
1967 std::fs::create_dir(&outside).unwrap();
1968
1969 let mut copy_config = enabled_security_config();
1970 copy_config.tools.insert(
1971 "copy_path".to_string(),
1972 ToolPolicyConfig {
1973 read_paths: vec![readable.to_string_lossy().into_owned()],
1974 write_paths: vec![writable.to_string_lossy().into_owned()],
1975 ..Default::default()
1976 },
1977 );
1978 let copy_engine = ToolSecurityEngine::new(copy_config);
1979 let copy_bindings = legacy_policy_bindings("copy_path");
1980 let copy_allowed = copy_engine
1981 .check_tool_execution_with_bindings(
1982 "copy_path",
1983 &serde_json::json!({
1984 "source_path": readable.join("source.txt"),
1985 "destination_path": writable.join("destination.txt")
1986 }),
1987 ©_bindings,
1988 )
1989 .await
1990 .unwrap();
1991 assert!(copy_allowed.is_allowed());
1992
1993 let copy_source_blocked = copy_engine
1994 .check_tool_execution_with_bindings(
1995 "copy_path",
1996 &serde_json::json!({
1997 "source_path": outside.join("source.txt"),
1998 "destination_path": writable.join("destination.txt")
1999 }),
2000 ©_bindings,
2001 )
2002 .await
2003 .unwrap();
2004 assert!(copy_source_blocked.is_blocked());
2005
2006 let copy_destination_blocked = copy_engine
2007 .check_tool_execution_with_bindings(
2008 "copy_path",
2009 &serde_json::json!({
2010 "source_path": readable.join("source.txt"),
2011 "destination_path": outside.join("destination.txt")
2012 }),
2013 ©_bindings,
2014 )
2015 .await
2016 .unwrap();
2017 assert!(copy_destination_blocked.is_blocked());
2018
2019 let mut move_config = enabled_security_config();
2020 move_config.tools.insert(
2021 "move_path".to_string(),
2022 ToolPolicyConfig {
2023 read_paths: vec![readable.to_string_lossy().into_owned()],
2024 write_paths: vec![writable.to_string_lossy().into_owned()],
2025 ..Default::default()
2026 },
2027 );
2028 let move_engine = ToolSecurityEngine::new(move_config);
2029 let move_bindings = legacy_policy_bindings("move_path");
2030 let move_allowed = move_engine
2031 .check_tool_execution_with_bindings(
2032 "move_path",
2033 &serde_json::json!({
2034 "source_path": writable.join("source.txt"),
2035 "destination_path": writable.join("destination.txt")
2036 }),
2037 &move_bindings,
2038 )
2039 .await
2040 .unwrap();
2041 assert!(move_allowed.is_allowed());
2042
2043 let move_source_blocked = move_engine
2044 .check_tool_execution_with_bindings(
2045 "move_path",
2046 &serde_json::json!({
2047 "source_path": readable.join("source.txt"),
2048 "destination_path": writable.join("destination.txt")
2049 }),
2050 &move_bindings,
2051 )
2052 .await
2053 .unwrap();
2054 assert!(move_source_blocked.is_blocked());
2055
2056 let move_destination_blocked = move_engine
2057 .check_tool_execution_with_bindings(
2058 "move_path",
2059 &serde_json::json!({
2060 "source_path": writable.join("source.txt"),
2061 "destination_path": outside.join("destination.txt")
2062 }),
2063 &move_bindings,
2064 )
2065 .await
2066 .unwrap();
2067 assert!(move_destination_blocked.is_blocked());
2068 }
2069
2070 #[test]
2071 fn custom_config_is_exposed_separately() {
2072 let mut config = enabled_security_config();
2073 let mut tool_config = ToolPolicyConfig::default();
2074 tool_config
2075 .config
2076 .insert("backend".to_string(), serde_json::json!("tantivy"));
2077 config.tools.insert("my_search".to_string(), tool_config);
2078 let engine = ToolSecurityEngine::new(config);
2079
2080 assert_eq!(engine.custom_config("my_search")["backend"], "tantivy");
2081 }
2082
2083 #[test]
2084 fn policy_caps_are_applied_as_upper_bounds() {
2085 let mut config = enabled_security_config();
2086 let tool_config = ToolPolicyConfig {
2087 max_results: Some(5),
2088 max_file_size_bytes: Some(1024),
2089 max_output_chars: Some(1000),
2090 ..Default::default()
2091 };
2092 config.tools.insert("grep".to_string(), tool_config);
2093 let engine = ToolSecurityEngine::new(config);
2094
2095 let prepared = engine.prepare_tool_arguments(
2096 "grep",
2097 &serde_json::json!({
2098 "pattern": "Tool",
2099 "path": ".",
2100 "max_results": 50,
2101 "max_file_size_bytes": 8192,
2102 "max_output_chars": 20000
2103 }),
2104 );
2105 assert_eq!(prepared.get("max_results").and_then(Value::as_u64), Some(5));
2106 let lower = engine.prepare_tool_arguments(
2107 "grep",
2108 &serde_json::json!({"pattern": "Tool", "max_results": 3}),
2109 );
2110 assert_eq!(lower.get("max_results").and_then(Value::as_u64), Some(3));
2111 assert_eq!(
2112 prepared.get("max_file_size_bytes").and_then(Value::as_u64),
2113 Some(1024)
2114 );
2115 assert_eq!(
2116 prepared.get("max_output_chars").and_then(Value::as_u64),
2117 Some(1000)
2118 );
2119 }
2120
2121 #[test]
2122 fn invalid_result_limit_is_rejected_before_engine_construction() {
2123 let mut config = enabled_security_config();
2124 config.tools.insert(
2125 "web_search".to_string(),
2126 ToolPolicyConfig {
2127 max_results: Some(0),
2128 ..Default::default()
2129 },
2130 );
2131 let error = ToolSecurityEngine::try_new(config).unwrap_err();
2132 assert!(
2133 error
2134 .to_string()
2135 .contains("max_results must be greater than 0")
2136 );
2137 }
2138
2139 #[tokio::test]
2140 async fn fail_closed_blocks_missing_result_limit_bindings() {
2141 let mut config = ToolSecurityConfig {
2142 fail_closed: true,
2143 ..enabled_security_config()
2144 };
2145 let tool_config = ToolPolicyConfig {
2146 max_results: Some(5),
2147 ..Default::default()
2148 };
2149 config
2150 .tools
2151 .insert("custom_search".to_string(), tool_config);
2152 let engine = ToolSecurityEngine::new(config);
2153
2154 let result = engine
2155 .check_tool_execution_with_bindings(
2156 "custom_search",
2157 &serde_json::json!({"query": "rust"}),
2158 &ToolPolicyBindings::default(),
2159 )
2160 .await
2161 .unwrap();
2162
2163 assert!(result.is_blocked());
2164 assert!(result.reason().unwrap().contains("result-limit policy"));
2165 }
2166
2167 #[tokio::test]
2168 async fn fail_closed_allows_configured_result_limit_bindings() {
2169 let mut config = ToolSecurityConfig {
2170 fail_closed: true,
2171 ..enabled_security_config()
2172 };
2173 let tool_config = ToolPolicyConfig {
2174 max_results: Some(5),
2175 ..Default::default()
2176 };
2177 config
2178 .tools
2179 .insert("custom_search".to_string(), tool_config);
2180 let engine = ToolSecurityEngine::new(config);
2181 let bindings = ToolPolicyBindings {
2182 result_limit_fields: vec![ResultLimitBinding::new(
2183 "limit",
2184 ResultLimitKind::MaxResults,
2185 )],
2186 ..Default::default()
2187 };
2188
2189 let result = engine
2190 .check_tool_execution_with_bindings(
2191 "custom_search",
2192 &serde_json::json!({"query": "rust", "limit": 10}),
2193 &bindings,
2194 )
2195 .await
2196 .unwrap();
2197
2198 assert!(result.is_allowed());
2199 }
2200
2201 #[tokio::test]
2202 async fn read_paths_do_not_authorize_file_write() {
2203 let mut config = enabled_security_config();
2204 let tool_config = ToolPolicyConfig {
2205 read_paths: vec!["./workspace".to_string()],
2206 no_write_policy: NoWritePolicyBehavior::Deny,
2207 ..Default::default()
2208 };
2209 config.tools.insert("file_write".to_string(), tool_config);
2210 let engine = ToolSecurityEngine::new(config);
2211
2212 let result = engine
2213 .check_tool_execution_with_bindings(
2214 "file_write",
2215 &serde_json::json!({"path": "./workspace/out.txt", "dry_run": false}),
2216 &legacy_policy_bindings("file_write"),
2217 )
2218 .await
2219 .unwrap();
2220
2221 assert!(result.is_blocked());
2222 }
2223
2224 #[tokio::test]
2225 async fn command_cwd_requires_working_dir_allowlist() {
2226 let mut config = enabled_security_config();
2227 let tool_config = ToolPolicyConfig {
2228 read_paths: vec![".".to_string()],
2229 allowed_commands: vec![CommandRuleConfig {
2230 argv: vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
2231 }],
2232 ..Default::default()
2233 };
2234 config.tools.insert("command".to_string(), tool_config);
2235 let engine = ToolSecurityEngine::new(config);
2236
2237 let result = engine
2238 .check_tool_execution_with_bindings(
2239 "command",
2240 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
2241 &legacy_policy_bindings("command"),
2242 )
2243 .await
2244 .unwrap();
2245
2246 assert!(result.is_blocked());
2247 }
2248
2249 #[tokio::test]
2250 async fn command_requires_exact_argv_allowlist() {
2251 let mut config = enabled_security_config();
2252 let tool_config = ToolPolicyConfig {
2253 allow_without_confirmation: true,
2254 working_dirs: vec![".".to_string()],
2255 ..Default::default()
2256 };
2257 config.tools.insert("command".to_string(), tool_config);
2258 let engine = ToolSecurityEngine::new(config);
2259
2260 let result = engine
2261 .check_tool_execution_with_bindings(
2262 "command",
2263 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
2264 &legacy_policy_bindings("command"),
2265 )
2266 .await
2267 .unwrap();
2268
2269 assert!(result.is_blocked());
2270 assert!(
2271 result
2272 .reason()
2273 .unwrap()
2274 .contains("requires allowed_commands or command_templates")
2275 );
2276 }
2277
2278 #[tokio::test]
2279 async fn command_exact_argv_allowlist_is_enforced() {
2280 let mut config = enabled_security_config();
2281 let tool_config = ToolPolicyConfig {
2282 allowed_commands: vec![CommandRuleConfig {
2283 argv: vec!["cargo".to_string(), "fmt".to_string(), "--all".to_string()],
2284 }],
2285 working_dirs: vec![".".to_string()],
2286 ..Default::default()
2287 };
2288 config.tools.insert("command".to_string(), tool_config);
2289 let engine = ToolSecurityEngine::new(config);
2290
2291 let allowed = engine
2292 .check_tool_execution_with_bindings(
2293 "command",
2294 &serde_json::json!({"argv": ["cargo", "fmt", "--all"], "cwd": "."}),
2295 &legacy_policy_bindings("command"),
2296 )
2297 .await
2298 .unwrap();
2299 assert!(allowed.is_allowed());
2300
2301 let blocked = engine
2302 .check_tool_execution_with_bindings(
2303 "command",
2304 &serde_json::json!({"argv": ["cargo", "test"], "cwd": "."}),
2305 &legacy_policy_bindings("command"),
2306 )
2307 .await
2308 .unwrap();
2309 assert!(blocked.is_blocked());
2310 }
2311
2312 #[tokio::test]
2313 async fn validation_does_not_consume_rate_limit_admission() {
2314 let mut config = enabled_security_config();
2315 let tool_config = ToolPolicyConfig {
2316 rate_limit: Some(1),
2317 ..Default::default()
2318 };
2319 config.tools.insert("limited".to_string(), tool_config);
2320 let engine = ToolSecurityEngine::new(config);
2321 let bindings = legacy_policy_bindings("limited");
2322
2323 for _ in 0..3 {
2324 let result = engine
2325 .validate_tool_execution_with_bindings("limited", &serde_json::json!({}), &bindings)
2326 .await
2327 .unwrap();
2328 assert!(result.is_allowed());
2329 }
2330 assert!(engine.admit_tool_execution("limited").is_allowed());
2331 assert!(engine.admit_tool_execution("limited").is_blocked());
2332 }
2333
2334 #[tokio::test]
2335 async fn public_check_preserves_rate_limit_admission() {
2336 let mut config = enabled_security_config();
2337 let tool_config = ToolPolicyConfig {
2338 rate_limit: Some(1),
2339 ..Default::default()
2340 };
2341 config.tools.insert("limited".to_string(), tool_config);
2342 let engine = ToolSecurityEngine::new(config);
2343
2344 let first = engine
2345 .check_tool_execution("limited", &serde_json::json!({}))
2346 .await
2347 .unwrap();
2348 let second = engine
2349 .check_tool_execution("limited", &serde_json::json!({}))
2350 .await
2351 .unwrap();
2352
2353 assert!(first.is_allowed());
2354 assert!(second.is_blocked());
2355 }
2356
2357 #[test]
2358 fn concurrent_rate_limit_admission_is_atomic() {
2359 let mut config = enabled_security_config();
2360 let tool_config = ToolPolicyConfig {
2361 rate_limit: Some(1),
2362 ..Default::default()
2363 };
2364 config.tools.insert("limited".to_string(), tool_config);
2365 let engine = Arc::new(ToolSecurityEngine::new_with_policy_version(config, 17));
2366 let barrier = Arc::new(std::sync::Barrier::new(8));
2367 let handles = (0..8)
2368 .map(|_| {
2369 let engine = Arc::clone(&engine);
2370 let barrier = Arc::clone(&barrier);
2371 std::thread::spawn(move || {
2372 barrier.wait();
2373 engine.admit_tool_execution("limited").is_allowed()
2374 })
2375 })
2376 .collect::<Vec<_>>();
2377 let admitted = handles
2378 .into_iter()
2379 .map(|handle| handle.join().unwrap())
2380 .filter(|admitted| *admitted)
2381 .count();
2382
2383 assert_eq!(admitted, 1);
2384 assert_eq!(engine.policy_version(), 17);
2385 }
2386
2387 #[tokio::test]
2388 async fn omitted_dry_run_is_treated_as_actual_mutation() {
2389 let mut config = enabled_security_config();
2390 for tool_id in ["file_edit", "copy_path"] {
2391 let tool_config = ToolPolicyConfig {
2392 no_write_policy: NoWritePolicyBehavior::DryRunOnly,
2393 ..Default::default()
2394 };
2395 config.tools.insert(tool_id.to_string(), tool_config);
2396 }
2397 let engine = ToolSecurityEngine::new(config);
2398
2399 let edit = engine
2400 .check_tool_execution_with_bindings(
2401 "file_edit",
2402 &serde_json::json!({"path": "./note.txt"}),
2403 &legacy_policy_bindings("file_edit"),
2404 )
2405 .await
2406 .unwrap();
2407 assert!(edit.is_blocked());
2408
2409 let copy = engine
2410 .check_tool_execution_with_bindings(
2411 "copy_path",
2412 &serde_json::json!({
2413 "source_path": "./source.txt",
2414 "destination_path": "./destination.txt"
2415 }),
2416 &legacy_policy_bindings("copy_path"),
2417 )
2418 .await
2419 .unwrap();
2420 assert!(copy.is_blocked());
2421 }
2422
2423 #[tokio::test]
2424 async fn no_write_policy_dry_run_only_allows_dry_run() {
2425 let mut config = enabled_security_config();
2426 let tool_config = ToolPolicyConfig {
2427 no_write_policy: NoWritePolicyBehavior::DryRunOnly,
2428 ..Default::default()
2429 };
2430 config.tools.insert("file_edit".to_string(), tool_config);
2431 let engine = ToolSecurityEngine::new(config);
2432
2433 let dry_run = engine
2434 .check_tool_execution_with_bindings(
2435 "file_edit",
2436 &serde_json::json!({"path": "./note.txt", "dry_run": true}),
2437 &legacy_policy_bindings("file_edit"),
2438 )
2439 .await
2440 .unwrap();
2441 assert!(dry_run.is_allowed());
2442
2443 let actual = engine
2444 .check_tool_execution_with_bindings(
2445 "file_edit",
2446 &serde_json::json!({"path": "./note.txt", "dry_run": false}),
2447 &legacy_policy_bindings("file_edit"),
2448 )
2449 .await
2450 .unwrap();
2451 assert!(actual.is_blocked());
2452 }
2453}