Files
rustytorch/crates/models/rtx-llm-tools/src/tools.rs
T
2026-03-04 00:08:42 +00:00

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(&param.name) {
return Err(ToolExecutionError::InvalidParameters {
tool: self.name.clone(),
message: format!("Required parameter '{}' is missing", param.name),
}
.into());
}
if let Some(value) = params.get(&param.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(&parameters)?;
// 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()
}
}