920 lines
29 KiB
Rust
920 lines
29 KiB
Rust
//! 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<Value>,
|
|
pub enum_values: Option<Vec<String>>,
|
|
pub validation_regex: Option<Regex>,
|
|
pub min_value: Option<f64>,
|
|
pub max_value: Option<f64>,
|
|
pub min_length: Option<usize>,
|
|
pub max_length: Option<usize>,
|
|
}
|
|
|
|
impl ToolParameter {
|
|
pub fn new(name: impl Into<String>, 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<String>) -> 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<String>,
|
|
pub allowed_directories: HashSet<String>,
|
|
pub max_execution_time: Option<Duration>,
|
|
pub require_explicit_permissions: bool,
|
|
pub audit_all_operations: bool,
|
|
}
|
|
|
|
impl SecurityPolicy {
|
|
pub fn new(name: impl Into<String>) -> 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<String>) -> &mut Self {
|
|
self.blocked_commands.insert(cmd.into());
|
|
self
|
|
}
|
|
|
|
pub fn add_allowed_directory(&mut self, dir: impl Into<String>) -> &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<u64>, // bytes
|
|
pub max_execution_time: Option<Duration>,
|
|
pub max_output_size: Option<usize>, // bytes
|
|
pub max_cpu_percent: Option<f64>,
|
|
pub max_open_files: Option<usize>,
|
|
pub max_processes: Option<usize>,
|
|
}
|
|
|
|
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<String>,
|
|
pub environment_variables: HashMap<String, String>,
|
|
pub user: Option<String>,
|
|
pub group: Option<String>,
|
|
pub chroot_path: Option<String>,
|
|
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<String>) -> &mut Self {
|
|
self.working_directory = Some(path.into());
|
|
self
|
|
}
|
|
|
|
pub fn add_environment_variable(
|
|
&mut self,
|
|
key: impl Into<String>,
|
|
value: impl Into<String>,
|
|
) -> &mut Self {
|
|
self.environment_variables.insert(key.into(), value.into());
|
|
self
|
|
}
|
|
|
|
pub fn set_user(&mut self, user: impl Into<String>) -> &mut Self {
|
|
self.user = Some(user.into());
|
|
self
|
|
}
|
|
|
|
pub fn set_group(&mut self, group: impl Into<String>) -> &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<u64>,
|
|
pub exit_code: Option<i32>,
|
|
pub errors: Vec<String>,
|
|
pub warnings: Vec<String>,
|
|
pub metadata: HashMap<String, Value>,
|
|
}
|
|
|
|
impl ToolResult {
|
|
pub fn success(output: impl Into<String>) -> 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<String>) -> 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<String>, 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<String, Value>,
|
|
pub result: Option<ToolResult>,
|
|
pub timestamp: Instant,
|
|
pub execution_time: Duration,
|
|
pub user_context: Option<String>,
|
|
pub violations: Vec<SandboxViolation>,
|
|
}
|
|
|
|
/// Tool definition with comprehensive metadata
|
|
#[derive(Debug, Clone)]
|
|
pub struct ToolDefinition {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub version: String,
|
|
pub parameters: Vec<ToolParameter>,
|
|
pub dependencies: Vec<String>,
|
|
pub category: String,
|
|
pub tags: Vec<String>,
|
|
pub author: Option<String>,
|
|
pub license: Option<String>,
|
|
pub documentation_url: Option<String>,
|
|
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<String>,
|
|
description: impl Into<String>,
|
|
parameters: Vec<ToolParameter>,
|
|
) -> 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<String>) -> Self {
|
|
self.version = version.into();
|
|
self
|
|
}
|
|
|
|
pub fn with_dependencies(mut self, deps: Vec<String>) -> Self {
|
|
self.dependencies = deps;
|
|
self
|
|
}
|
|
|
|
pub fn with_category(mut self, category: impl Into<String>) -> 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<String, Value>) -> 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<String, Value>,
|
|
context: &ExecutionContext,
|
|
) -> Result<ToolResult>;
|
|
}
|
|
|
|
/// 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<String>,
|
|
}
|
|
|
|
/// Main tool executor with comprehensive security and management features
|
|
pub struct ToolExecutor {
|
|
name: String,
|
|
tools: Arc<DashMap<String, IndexMap<String, ToolDefinition>>>, // name -> version -> definition
|
|
security_policy: Arc<RwLock<SecurityPolicy>>,
|
|
resource_limits: Arc<RwLock<ResourceLimits>>,
|
|
execution_environment: Arc<RwLock<ExecutionEnvironment>>,
|
|
global_timeout: Arc<RwLock<Option<Duration>>>,
|
|
sandbox_enabled: bool,
|
|
audit_enabled: bool,
|
|
audit_log: Arc<RwLock<Vec<AuditEntry>>>,
|
|
executor_impl: Arc<dyn ToolExecution + Send + Sync>,
|
|
}
|
|
|
|
/// Default tool execution implementation
|
|
pub struct DefaultToolExecutor;
|
|
|
|
#[async_trait]
|
|
impl ToolExecution for DefaultToolExecutor {
|
|
async fn execute_tool(
|
|
&self,
|
|
tool: &ToolDefinition,
|
|
parameters: HashMap<String, Value>,
|
|
context: &ExecutionContext,
|
|
) -> Result<ToolResult> {
|
|
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<i32> = (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<i32> = 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::<i32>(),
|
|
"avg" => {
|
|
if data.is_empty() {
|
|
0
|
|
} else {
|
|
data.iter().sum::<i32>() / 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<String>) -> 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<String, Value>,
|
|
) -> Result<ToolResult> {
|
|
// 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<String, Value>,
|
|
) -> Result<ToolResult> {
|
|
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<Vec<AuditEntry>> {
|
|
Ok(self.audit_log.read().clone())
|
|
}
|
|
|
|
pub fn list_tools(&self) -> Vec<String> {
|
|
self.tools.iter().map(|entry| entry.key().clone()).collect()
|
|
}
|
|
|
|
pub fn get_tool_info(&self, name: &str) -> Option<Vec<ToolDefinition>> {
|
|
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()
|
|
}
|
|
}
|