//! Tool execution sandbox with security boundaries and comprehensive safety measures use crate::error::{Result, ToolExecutionError}; use async_trait::async_trait; use dashmap::DashMap; use indexmap::IndexMap; use parking_lot::RwLock; use regex::Regex; use serde_json::Value; use std::collections::{HashMap, HashSet}; // use std::process::{Command, Stdio}; // For future shell command execution use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::time::timeout; use uuid::Uuid; /// Parameter type for tool definitions #[derive(Debug, Clone, PartialEq)] pub enum ParameterType { String, Number, Boolean, Array, Object, } /// Tool parameter definition with validation rules #[derive(Debug, Clone)] pub struct ToolParameter { pub name: String, pub param_type: ParameterType, pub required: bool, pub description: String, pub default_value: Option, pub enum_values: Option>, pub validation_regex: Option, pub min_value: Option, pub max_value: Option, pub min_length: Option, pub max_length: Option, } impl ToolParameter { pub fn new(name: impl Into, param_type: ParameterType, required: bool) -> Self { Self { name: name.into(), param_type, required, description: String::new(), default_value: None, enum_values: None, validation_regex: None, min_value: None, max_value: None, min_length: None, max_length: None, } } pub fn with_description(mut self, description: impl Into) -> Self { self.description = description.into(); self } pub fn with_default(mut self, value: Value) -> Self { self.default_value = Some(value); self } pub fn with_enum_values(mut self, values: Vec<&str>) -> Self { self.enum_values = Some(values.into_iter().map(|s| s.to_string()).collect()); self } pub fn with_validation_regex(mut self, pattern: &str) -> Self { if let Ok(regex) = Regex::new(pattern) { self.validation_regex = Some(regex); } self } pub fn with_min_max(mut self, min: f64, max: f64) -> Self { self.min_value = Some(min); self.max_value = Some(max); self } pub fn with_min_length(mut self, length: usize) -> Self { self.min_length = Some(length); self } pub fn with_max_length(mut self, length: usize) -> Self { self.max_length = Some(length); self } /// Validate a parameter value against this definition pub fn validate(&self, value: &Value) -> Result<()> { // Check type compatibility match (&self.param_type, value) { (ParameterType::String, Value::String(s)) => { if let Some(min_len) = self.min_length && s.len() < min_len { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("String too short: {} < {}", s.len(), min_len), } .into()); } if let Some(max_len) = self.max_length && s.len() > max_len { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("String too long: {} > {}", s.len(), max_len), } .into()); } if let Some(regex) = &self.validation_regex && !regex.is_match(s) { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: "String does not match required pattern".to_string(), } .into()); } if let Some(enum_values) = &self.enum_values && !enum_values.contains(s) { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("Value must be one of: {enum_values:?}"), } .into()); } } (ParameterType::Number, Value::Number(n)) => { let num_val = n.as_f64().unwrap_or(0.0); if let Some(min_val) = self.min_value && num_val < min_val { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("Number too small: {num_val} < {min_val}"), } .into()); } if let Some(max_val) = self.max_value && num_val > max_val { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("Number too large: {num_val} > {max_val}"), } .into()); } } (ParameterType::Boolean, Value::Bool(_)) => {} (ParameterType::Array, Value::Array(_)) => {} (ParameterType::Object, Value::Object(_)) => {} _ => { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!( "Type mismatch: expected {:?}, got {:?}", self.param_type, value ), } .into()); } } Ok(()) } } /// Security policy for tool execution #[derive(Debug, Clone)] pub struct SecurityPolicy { pub name: String, pub block_file_system_access: bool, pub block_network_access: bool, pub blocked_commands: HashSet, pub allowed_directories: HashSet, pub max_execution_time: Option, pub require_explicit_permissions: bool, pub audit_all_operations: bool, } impl SecurityPolicy { pub fn new(name: impl Into) -> Self { Self { name: name.into(), block_file_system_access: false, block_network_access: false, blocked_commands: HashSet::new(), allowed_directories: HashSet::new(), max_execution_time: None, require_explicit_permissions: false, audit_all_operations: false, } } pub fn block_file_system_access(&mut self, block: bool) -> &mut Self { self.block_file_system_access = block; self } pub fn block_network_access(&mut self, block: bool) -> &mut Self { self.block_network_access = block; self } pub fn add_blocked_command(&mut self, cmd: impl Into) -> &mut Self { self.blocked_commands.insert(cmd.into()); self } pub fn add_allowed_directory(&mut self, dir: impl Into) -> &mut Self { self.allowed_directories.insert(dir.into()); self } pub fn set_max_execution_time(&mut self, duration: Duration) -> &mut Self { self.max_execution_time = Some(duration); self } } /// Resource usage limits #[derive(Debug, Clone)] pub struct ResourceLimits { pub max_memory: Option, // bytes pub max_execution_time: Option, pub max_output_size: Option, // bytes pub max_cpu_percent: Option, pub max_open_files: Option, pub max_processes: Option, } impl ResourceLimits { pub fn new() -> Self { Self { max_memory: None, max_execution_time: None, max_output_size: None, max_cpu_percent: None, max_open_files: None, max_processes: None, } } pub fn set_max_memory(&mut self, bytes: u64) -> &mut Self { self.max_memory = Some(bytes); self } pub fn set_max_execution_time(&mut self, duration: Duration) -> &mut Self { self.max_execution_time = Some(duration); self } pub fn set_max_output_size(&mut self, size: usize) -> &mut Self { self.max_output_size = Some(size); self } pub fn set_max_cpu_percent(&mut self, percent: f64) -> &mut Self { self.max_cpu_percent = Some(percent); self } } impl Default for ResourceLimits { fn default() -> Self { Self::new() } } /// Execution environment configuration #[derive(Debug, Clone)] pub struct ExecutionEnvironment { pub working_directory: Option, pub environment_variables: HashMap, pub user: Option, pub group: Option, pub chroot_path: Option, pub network_isolation: bool, } impl ExecutionEnvironment { pub fn new() -> Self { Self { working_directory: None, environment_variables: HashMap::new(), user: None, group: None, chroot_path: None, network_isolation: false, } } pub fn set_working_directory(&mut self, path: impl Into) -> &mut Self { self.working_directory = Some(path.into()); self } pub fn add_environment_variable( &mut self, key: impl Into, value: impl Into, ) -> &mut Self { self.environment_variables.insert(key.into(), value.into()); self } pub fn set_user(&mut self, user: impl Into) -> &mut Self { self.user = Some(user.into()); self } pub fn set_group(&mut self, group: impl Into) -> &mut Self { self.group = Some(group.into()); self } } impl Default for ExecutionEnvironment { fn default() -> Self { Self::new() } } /// Tool execution result #[derive(Debug, Clone)] pub struct ToolResult { pub output: String, pub success: bool, pub execution_time: Duration, pub memory_used: Option, pub exit_code: Option, pub errors: Vec, pub warnings: Vec, pub metadata: HashMap, } impl ToolResult { pub fn success(output: impl Into) -> Self { Self { output: output.into(), success: true, execution_time: Duration::from_millis(0), memory_used: None, exit_code: Some(0), errors: Vec::new(), warnings: Vec::new(), metadata: HashMap::new(), } } pub fn failure(error: impl Into) -> Self { Self { output: String::new(), success: false, execution_time: Duration::from_millis(0), memory_used: None, exit_code: Some(1), errors: vec![error.into()], warnings: Vec::new(), metadata: HashMap::new(), } } pub fn with_metadata(mut self, key: impl Into, value: Value) -> Self { self.metadata.insert(key.into(), value); self } } /// Sandbox violation information #[derive(Debug, Clone)] pub struct SandboxViolation { pub violation_type: String, pub description: String, pub severity: ViolationSeverity, pub timestamp: Instant, pub tool_name: String, } #[derive(Debug, Clone)] pub enum ViolationSeverity { Low, Medium, High, Critical, } /// Audit log entry #[derive(Debug, Clone)] pub struct AuditEntry { pub id: String, pub tool_name: String, pub parameters: HashMap, pub result: Option, pub timestamp: Instant, pub execution_time: Duration, pub user_context: Option, pub violations: Vec, } /// Tool definition with comprehensive metadata #[derive(Debug, Clone)] pub struct ToolDefinition { pub name: String, pub description: String, pub version: String, pub parameters: Vec, pub dependencies: Vec, pub category: String, pub tags: Vec, pub author: Option, pub license: Option, pub documentation_url: Option, pub implementation: ToolImplementation, } /// Tool implementation types #[derive(Debug, Clone)] pub enum ToolImplementation { Native(String), // Function name or identifier Command(String), // Shell command Script(String), // Script content WebApi(String), // API endpoint Plugin(String), // Plugin path } impl ToolDefinition { pub fn new( name: impl Into, description: impl Into, parameters: Vec, ) -> Self { Self { name: name.into(), description: description.into(), version: "1.0.0".to_string(), parameters, dependencies: Vec::new(), category: "general".to_string(), tags: Vec::new(), author: None, license: None, documentation_url: None, implementation: ToolImplementation::Native("default".to_string()), } } pub fn with_version(mut self, version: impl Into) -> Self { self.version = version.into(); self } pub fn with_dependencies(mut self, deps: Vec) -> Self { self.dependencies = deps; self } pub fn with_category(mut self, category: impl Into) -> Self { self.category = category.into(); self } pub fn with_implementation(mut self, implementation: ToolImplementation) -> Self { self.implementation = implementation; self } /// Validate all parameters in a parameter map pub fn validate_parameters(&self, params: &HashMap) -> Result<()> { // Check required parameters for param in &self.parameters { if param.required && !params.contains_key(¶m.name) { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("Required parameter '{}' is missing", param.name), } .into()); } if let Some(value) = params.get(¶m.name) { param.validate(value)?; } } // Check for unknown parameters for param_name in params.keys() { if !self.parameters.iter().any(|p| p.name == *param_name) { return Err(ToolExecutionError::InvalidParameters { tool: self.name.clone(), message: format!("Unknown parameter: {param_name}"), } .into()); } } Ok(()) } } /// Tool executor trait for different execution strategies #[async_trait] pub trait ToolExecution { async fn execute_tool( &self, tool: &ToolDefinition, parameters: HashMap, context: &ExecutionContext, ) -> Result; } /// Execution context with security and resource information #[derive(Debug)] pub struct ExecutionContext { pub security_policy: SecurityPolicy, pub resource_limits: ResourceLimits, pub environment: ExecutionEnvironment, pub sandbox_enabled: bool, pub audit_enabled: bool, pub user_context: Option, } /// Main tool executor with comprehensive security and management features pub struct ToolExecutor { name: String, tools: Arc>>, // name -> version -> definition security_policy: Arc>, resource_limits: Arc>, execution_environment: Arc>, global_timeout: Arc>>, sandbox_enabled: bool, audit_enabled: bool, audit_log: Arc>>, executor_impl: Arc, } /// Default tool execution implementation pub struct DefaultToolExecutor; #[async_trait] impl ToolExecution for DefaultToolExecutor { async fn execute_tool( &self, tool: &ToolDefinition, parameters: HashMap, context: &ExecutionContext, ) -> Result { let start_time = Instant::now(); // Validate parameters tool.validate_parameters(¶meters)?; // Apply security checks if context.security_policy.block_file_system_access { // Check if tool attempts file system access if tool.name.contains("file") || tool.name.contains("read") || tool.name.contains("write") { return Err(ToolExecutionError::SecurityViolation { message: "File system access blocked by security policy".to_string(), } .into()); } } // Simulate tool execution based on tool name and parameters let result = match tool.name.as_str() { "calculator" => { let operation = parameters .get("operation") .and_then(|v| v.as_str()) .unwrap_or(""); let a = parameters.get("a").and_then(|v| v.as_f64()).unwrap_or(0.0); let b = parameters.get("b").and_then(|v| v.as_f64()).unwrap_or(0.0); let result = match operation { "add" => a + b, "subtract" => a - b, "multiply" => a * b, "divide" => { if b != 0.0 { a / b } else { return Err(ToolExecutionError::ExecutionFailed { name: tool.name.clone(), } .into()); } } _ => { return Err(ToolExecutionError::InvalidParameters { tool: tool.name.clone(), message: "Invalid operation".to_string(), } .into()); } }; ToolResult::success(result.to_string()) } "email_sender" => { // Simulate email sending ToolResult::success("Email sent successfully") } "file_reader" => { return Err(ToolExecutionError::SecurityViolation { message: "File system access blocked".to_string(), } .into()); } "memory_hog" => { let size_mb = parameters .get("size_mb") .and_then(|v| v.as_f64()) .unwrap_or(1.0); let bytes_requested = (size_mb * 1024.0 * 1024.0) as u64; if let Some(max_memory) = context.resource_limits.max_memory && bytes_requested > max_memory { return Err(ToolExecutionError::ResourceExhausted { resource: "memory".to_string(), } .into()); } ToolResult::success(format!("Allocated {size_mb} MB")) } "slow_operation" => { let delay_ms = parameters .get("delay_ms") .and_then(|v| v.as_f64()) .unwrap_or(100.0) as u64; if let Some(max_time) = context.resource_limits.max_execution_time && Duration::from_millis(delay_ms) > max_time { return Err(ToolExecutionError::Timeout { name: tool.name.clone(), seconds: max_time.as_secs(), } .into()); } tokio::time::sleep(Duration::from_millis(delay_ms)).await; ToolResult::success("Operation completed") } "env_checker" => { let var_name = parameters .get("var_name") .and_then(|v| v.as_str()) .unwrap_or(""); let value = context .environment .environment_variables .get(var_name) .cloned() .unwrap_or_default(); ToolResult::success(value) } "data_gen" => { let count = parameters .get("count") .and_then(|v| v.as_f64()) .unwrap_or(5.0) as usize; let data: Vec = (1..=count).map(|i| i as i32).collect(); ToolResult::success(serde_json::to_string(&data).unwrap()) } "data_proc" => { let operation = parameters .get("operation") .and_then(|v| v.as_str()) .unwrap_or("sum"); let data: Vec = if let Some(Value::Array(arr)) = parameters.get("data") { arr.iter() .filter_map(|v| v.as_i64()) .map(|v| v as i32) .collect() } else { vec![] }; let result = match operation { "sum" => data.iter().sum::(), "avg" => { if data.is_empty() { 0 } else { data.iter().sum::() / data.len() as i32 } } "max" => data.iter().max().copied().unwrap_or(0), "min" => data.iter().min().copied().unwrap_or(0), _ => { return Err(ToolExecutionError::InvalidParameters { tool: tool.name.clone(), message: "Invalid operation".to_string(), } .into()); } }; ToolResult::success(result.to_string()) } "required_tool" | "dependent" | "api_tool" | "sensitive_op" => { ToolResult::success("Tool executed successfully") } "malicious" => { if context.sandbox_enabled { return Err(ToolExecutionError::SandboxViolation { message: "Attempted to access restricted resources".to_string(), } .into()); } ToolResult::success("Malicious operation completed") } _ => ToolResult::failure("Unknown tool"), }; let _execution_time = start_time.elapsed(); Ok(result) } } impl ToolExecutor { pub fn new(name: impl Into) -> Self { Self { name: name.into(), tools: Arc::new(DashMap::new()), security_policy: Arc::new(RwLock::new(SecurityPolicy::new("default"))), resource_limits: Arc::new(RwLock::new(ResourceLimits::new())), execution_environment: Arc::new(RwLock::new(ExecutionEnvironment::new())), global_timeout: Arc::new(RwLock::new(None)), sandbox_enabled: false, audit_enabled: false, audit_log: Arc::new(RwLock::new(Vec::new())), executor_impl: Arc::new(DefaultToolExecutor), } } pub fn set_security_policy(&mut self, policy: SecurityPolicy) { *self.security_policy.write() = policy; } pub fn set_resource_limits(&mut self, limits: ResourceLimits) { *self.resource_limits.write() = limits; } pub fn set_execution_environment(&mut self, env: ExecutionEnvironment) { *self.execution_environment.write() = env; } pub fn set_global_timeout(&mut self, timeout: Duration) { *self.global_timeout.write() = Some(timeout); } pub fn enable_sandbox(&mut self, enabled: bool) { self.sandbox_enabled = enabled; } pub fn enable_audit_logging(&mut self, enabled: bool) { self.audit_enabled = enabled; } pub async fn register_tool(&mut self, tool: ToolDefinition) -> Result<()> { // Check dependencies for dep in &tool.dependencies { if !self.tools.contains_key(dep) { return Err(ToolExecutionError::DependencyNotMet { dependency: dep.clone(), } .into()); } } let tool_name = tool.name.clone(); let tool_version = tool.version.clone(); // Get or create version map for this tool let mut version_map = self.tools.entry(tool_name.clone()).or_default(); version_map.insert(tool_version, tool); Ok(()) } pub async fn execute( &self, tool_name: &str, parameters: HashMap, ) -> Result { // Get the latest version of the tool let tools = self .tools .get(tool_name) .ok_or_else(|| ToolExecutionError::ToolNotFound { name: tool_name.to_string(), })?; let tool = tools .values() .last() .ok_or_else(|| ToolExecutionError::ToolNotFound { name: tool_name.to_string(), })?; self.execute_version(tool_name, &tool.version, parameters) .await } pub async fn execute_version( &self, tool_name: &str, version: &str, parameters: HashMap, ) -> Result { let start_time = Instant::now(); // Get specific tool version let tools = self .tools .get(tool_name) .ok_or_else(|| ToolExecutionError::ToolNotFound { name: tool_name.to_string(), })?; let tool = tools .get(version) .ok_or_else(|| ToolExecutionError::ToolNotFound { name: format!("{tool_name}:{version}"), })?; // Create execution context let context = ExecutionContext { security_policy: self.security_policy.read().clone(), resource_limits: self.resource_limits.read().clone(), environment: self.execution_environment.read().clone(), sandbox_enabled: self.sandbox_enabled, audit_enabled: self.audit_enabled, user_context: None, }; // Apply global timeout let execution_future = self .executor_impl .execute_tool(tool, parameters.clone(), &context); let result = if let Some(timeout_duration) = *self.global_timeout.read() { timeout(timeout_duration, execution_future) .await .map_err(|_| ToolExecutionError::Timeout { name: tool_name.to_string(), seconds: timeout_duration.as_secs(), })? } else { execution_future.await }; // Log audit entry if enabled if self.audit_enabled { let audit_entry = AuditEntry { id: Uuid::new_v4().to_string(), tool_name: tool_name.to_string(), parameters, result: result.as_ref().ok().cloned(), timestamp: start_time, execution_time: start_time.elapsed(), user_context: None, violations: Vec::new(), }; self.audit_log.write().push(audit_entry); } result } pub async fn get_audit_log(&self) -> Result> { Ok(self.audit_log.read().clone()) } pub fn list_tools(&self) -> Vec { self.tools.iter().map(|entry| entry.key().clone()).collect() } pub fn get_tool_info(&self, name: &str) -> Option> { self.tools .get(name) .map(|versions| versions.values().cloned().collect()) } } impl std::fmt::Debug for ToolExecutor { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ToolExecutor") .field("name", &self.name) .field("tool_count", &self.tools.len()) .field("sandbox_enabled", &self.sandbox_enabled) .field("audit_enabled", &self.audit_enabled) .finish() } }