/// Permission and Safety Layer for MCP /// /// This module provides runtime enforcement of security policies for tool execution. /// It wraps MCP clients to filter/whitelist tool calls, log invocations, and prompt for consent. use super::client::McpClient; use super::{McpToolCall, McpToolDescriptor, McpToolResponse}; use crate::{Error, Result}; use crate::{config::Config, mode::Mode}; use async_trait::async_trait; use std::collections::HashSet; use std::sync::Arc; /// Callback for requesting user consent for dangerous operations pub type ConsentCallback = Arc bool + Send + Sync>; /// Callback for logging tool invocations pub type LogCallback = Arc) + Send + Sync>; /// Permission-enforcing wrapper around an MCP client pub struct PermissionLayer { inner: Box, config: Arc, consent_callback: Option, log_callback: Option, allowed_tools: HashSet, } impl PermissionLayer { /// Create a new permission layer wrapping the given client pub fn new(inner: Box, config: Arc) -> Self { let allowed_tools = config.security.allowed_tools.iter().cloned().collect(); Self { inner, config, consent_callback: None, log_callback: None, allowed_tools, } } /// Set a callback for requesting user consent pub fn with_consent_callback(mut self, callback: ConsentCallback) -> Self { self.consent_callback = Some(callback); self } /// Set a callback for logging tool invocations pub fn with_log_callback(mut self, callback: LogCallback) -> Self { self.log_callback = Some(callback); self } /// Check if a tool requires dangerous filesystem operations fn requires_dangerous_filesystem(&self, tool_name: &str) -> bool { matches!( tool_name, "resources/write" | "resources/delete" | "file_write" | "file_delete" ) } /// Check if a tool is allowed by security policy fn is_tool_allowed(&self, tool_descriptor: &McpToolDescriptor) -> bool { // Check if tool requires filesystem access for fs_perm in &tool_descriptor.requires_filesystem { if !self.allowed_tools.contains(fs_perm) { return false; } } // Check if tool requires network access if tool_descriptor.requires_network && !self.allowed_tools.contains("web_search") { return false; } true } /// Request user consent for a tool call fn request_consent(&self, tool_name: &str, call: &McpToolCall) -> bool { if let Some(ref callback) = self.consent_callback { callback(tool_name, call) } else { // If no callback is set, deny dangerous operations by default !self.requires_dangerous_filesystem(tool_name) } } /// Log a tool invocation fn log_invocation( &self, tool_name: &str, call: &McpToolCall, result: &Result, ) { if let Some(ref callback) = self.log_callback { callback(tool_name, call, result); } else { // Default logging to stderr match result { Ok(resp) => { eprintln!( "[MCP] Tool '{}' executed successfully ({}ms)", tool_name, resp.duration_ms ); } Err(e) => { eprintln!("[MCP] Tool '{}' failed: {}", tool_name, e); } } } } } #[async_trait] impl McpClient for PermissionLayer { async fn list_tools(&self) -> Result> { let tools = self.inner.list_tools().await?; // Filter tools based on security policy Ok(tools .into_iter() .filter(|tool| self.is_tool_allowed(tool)) .collect()) } async fn call_tool(&self, call: McpToolCall) -> Result { // Check if tool requires consent if self.requires_dangerous_filesystem(&call.name) && self.config.privacy.require_consent_per_session && !self.request_consent(&call.name, &call) { let result = Err(Error::PermissionDenied(format!( "User denied consent for tool '{}'", call.name ))); self.log_invocation(&call.name, &call, &result); return result; } // Execute the tool call let result = self.inner.call_tool(call.clone()).await; // Log the invocation self.log_invocation(&call.name, &call, &result); result } async fn set_mode(&self, mode: Mode) -> Result<()> { self.inner.set_mode(mode).await } } #[cfg(test)] mod tests { use super::*; use crate::mcp::LocalMcpClient; use crate::tools::registry::ToolRegistry; use crate::ui::NoOpUiController; use crate::validation::SchemaValidator; use std::sync::atomic::{AtomicBool, Ordering}; #[tokio::test] async fn test_permission_layer_filters_dangerous_tools() { let config = Arc::new(Config::default()); let ui = Arc::new(NoOpUiController); let registry = Arc::new(ToolRegistry::new( Arc::new(tokio::sync::Mutex::new((*config).clone())), ui, )); let validator = Arc::new(SchemaValidator::new()); let client = Box::new(LocalMcpClient::new(registry, validator)); let mut config_mut = (*config).clone(); // Disallow file operations config_mut.security.allowed_tools = vec!["web_search".to_string()]; let permission_layer = PermissionLayer::new(client, Arc::new(config_mut)); let tools = permission_layer.list_tools().await.unwrap(); // Should not include file_write or file_delete tools assert!(!tools.iter().any(|t| t.name.contains("write"))); assert!(!tools.iter().any(|t| t.name.contains("delete"))); } #[tokio::test] async fn test_consent_callback_is_invoked() { let config = Arc::new(Config::default()); let ui = Arc::new(NoOpUiController); let registry = Arc::new(ToolRegistry::new( Arc::new(tokio::sync::Mutex::new((*config).clone())), ui, )); let validator = Arc::new(SchemaValidator::new()); let client = Box::new(LocalMcpClient::new(registry, validator)); let consent_called = Arc::new(AtomicBool::new(false)); let consent_called_clone = consent_called.clone(); let consent_callback: ConsentCallback = Arc::new(move |_tool, _call| { consent_called_clone.store(true, Ordering::SeqCst); false // Deny }); let mut config_mut = (*config).clone(); config_mut.privacy.require_consent_per_session = true; let permission_layer = PermissionLayer::new(client, Arc::new(config_mut)) .with_consent_callback(consent_callback); let call = McpToolCall { name: "resources/write".to_string(), arguments: serde_json::json!({"path": "test.txt", "content": "hello"}), }; let result = permission_layer.call_tool(call).await; assert!(consent_called.load(Ordering::SeqCst)); assert!(result.is_err()); } }