owlen/crates/owlen-core/src/mcp/permission.rs

/// 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::config::Config;
use crate::{Error, Result};
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<dyn Fn(&str, &McpToolCall) -> bool + Send + Sync>;

/// Callback for logging tool invocations
pub type LogCallback = Arc<dyn Fn(&str, &McpToolCall, &Result<McpToolResponse>) + Send + Sync>;

/// Permission-enforcing wrapper around an MCP client
pub struct PermissionLayer {
    inner: Box<dyn McpClient>,
    config: Arc<Config>,
    consent_callback: Option<ConsentCallback>,
    log_callback: Option<LogCallback>,
    allowed_tools: HashSet<String>,
}

impl PermissionLayer {
    /// Create a new permission layer wrapping the given client
    pub fn new(inner: Box<dyn McpClient>, config: Arc<Config>) -> 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<McpToolResponse>,
    ) {
        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<Vec<McpToolDescriptor>> {
        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<McpToolResponse> {
        // 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
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::LocalMcpClient;
    use crate::tools::registry::ToolRegistry;
    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(crate::ui::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(crate::ui::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());
    }
}