use crate::config::paths::Paths; use rmcp::model::Tool; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::fs; use std::path::{Path, PathBuf}; use std::sync::{Arc, LazyLock, RwLock}; use tracing; use utoipa::ToSchema; const PERMISSION_FILE: &str = "permission.yaml"; static PERMISSION_MANAGER: LazyLock> = LazyLock::new(|| Arc::new(PermissionManager::new(Paths::config_dir()))); /// Enum representing the possible permission levels for a tool. #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq, ToSchema)] #[serde(rename_all = "snake_case")] pub enum PermissionLevel { AlwaysAllow, // Tool can always be used without prompt AskBefore, // Tool requires permission to be granted before use NeverAllow, // Tool is never allowed to be used } /// Struct representing the configuration of permissions, categorized by level. #[derive(Debug, Deserialize, Serialize, Default, Clone)] pub struct PermissionConfig { pub always_allow: Vec, // List of tools that are always allowed pub ask_before: Vec, // List of tools that require user consent pub never_allow: Vec, // List of tools that are never allowed } /// PermissionManager manages permission configurations for various tools. #[derive(Debug)] pub struct PermissionManager { config_path: PathBuf, permission_map: RwLock>, } // Constants representing specific permission categories const USER_PERMISSION: &str = "user"; const SMART_APPROVE_PERMISSION: &str = "smart_approve"; impl PermissionManager { pub fn new(config_dir: PathBuf) -> Self { let permission_path = config_dir.join(PERMISSION_FILE); let permission_map = if permission_path.exists() { let file_contents = fs::read_to_string(&permission_path).expect("Failed to read permission.yaml"); serde_yaml::from_str(&file_contents).unwrap_or_else(|e| { tracing::error!( "Failed to parse {}: {}. Refusing to start with corrupted permission config.", permission_path.display(), e, ); panic!( "Corrupted permission config at {}. Fix or remove the file to continue.", permission_path.display(), ); }) } else { // Consolidate directory creation for re-use in global singleton or ACP. fs::create_dir_all(&config_dir).expect("Failed to create config directory"); HashMap::new() }; PermissionManager { config_path: permission_path, permission_map: RwLock::new(permission_map), } } pub fn instance() -> Arc { Arc::clone(&PERMISSION_MANAGER) } /// Returns a list of all the names (keys) in the permission map. pub fn get_permission_names(&self) -> Vec { self.permission_map .read() .unwrap() .keys() .cloned() .collect() } /// Retrieves the user permission level for a specific tool. pub fn get_user_permission(&self, principal_name: &str) -> Option { self.get_permission(USER_PERMISSION, principal_name) } /// Retrieves the smart approve permission level for a specific tool. pub fn get_smart_approve_permission(&self, principal_name: &str) -> Option { self.get_permission(SMART_APPROVE_PERMISSION, principal_name) } /// Retrieves the config file path. pub fn get_config_path(&self) -> &Path { self.config_path.as_path() } pub fn apply_tool_annotations(&self, tools: &[Tool]) { let mut write_annotated = Vec::new(); for tool in tools { let Some(anns) = &tool.annotations else { continue; }; if anns.read_only_hint == Some(false) { write_annotated.push(tool.name.to_string()); } } if !write_annotated.is_empty() { self.bulk_update_smart_approve_permissions( &write_annotated, PermissionLevel::AskBefore, ); } } fn bulk_update_smart_approve_permissions(&self, tool_names: &[String], level: PermissionLevel) { let mut map = self.permission_map.write().unwrap(); let permission_config = map.entry(SMART_APPROVE_PERMISSION.to_string()).or_default(); for tool_name in tool_names { // Remove from all lists to avoid duplicates permission_config.always_allow.retain(|p| p != tool_name); permission_config.ask_before.retain(|p| p != tool_name); permission_config.never_allow.retain(|p| p != tool_name); // Add to the appropriate list match &level { PermissionLevel::AlwaysAllow => { permission_config.always_allow.push(tool_name.clone()) } PermissionLevel::AskBefore => permission_config.ask_before.push(tool_name.clone()), PermissionLevel::NeverAllow => { permission_config.never_allow.push(tool_name.clone()) } } } let yaml_content = serde_yaml::to_string(&*map).expect("Failed to serialize permission config"); fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml"); } /// Helper function to retrieve the permission level for a specific permission category and tool. fn get_permission(&self, name: &str, principal_name: &str) -> Option { let map = self.permission_map.read().unwrap(); // Check if the permission category exists in the map if let Some(permission_config) = map.get(name) { // Check the permission levels for the given tool if permission_config .always_allow .contains(&principal_name.to_string()) { return Some(PermissionLevel::AlwaysAllow); } else if permission_config .ask_before .contains(&principal_name.to_string()) { return Some(PermissionLevel::AskBefore); } else if permission_config .never_allow .contains(&principal_name.to_string()) { return Some(PermissionLevel::NeverAllow); } } None // Return None if no matching permission level is found } /// Updates the user permission level for a specific tool. pub fn update_user_permission(&self, principal_name: &str, level: PermissionLevel) { self.update_permission(USER_PERMISSION, principal_name, level) } /// Updates the smart approve permission level for a specific tool. pub fn update_smart_approve_permission(&self, principal_name: &str, level: PermissionLevel) { self.update_permission(SMART_APPROVE_PERMISSION, principal_name, level) } /// Helper function to update a permission level for a specific tool in a given permission category. fn update_permission(&self, name: &str, principal_name: &str, level: PermissionLevel) { let mut map = self.permission_map.write().unwrap(); // Get or create a new PermissionConfig for the specified category let permission_config = map.entry(name.to_string()).or_default(); // Remove the principal from all existing lists to avoid duplicates permission_config .always_allow .retain(|p| p != principal_name); permission_config.ask_before.retain(|p| p != principal_name); permission_config .never_allow .retain(|p| p != principal_name); // Add the principal to the appropriate list match level { PermissionLevel::AlwaysAllow => permission_config .always_allow .push(principal_name.to_string()), PermissionLevel::AskBefore => permission_config .ask_before .push(principal_name.to_string()), PermissionLevel::NeverAllow => permission_config .never_allow .push(principal_name.to_string()), } // Serialize the updated permission map and write it back to the config file let yaml_content = serde_yaml::to_string(&*map).expect("Failed to serialize permission config"); fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml"); } /// Removes all entries where the principal name starts with the given extension name. pub fn remove_extension(&self, extension_name: &str) { let mut map = self.permission_map.write().unwrap(); for permission_config in map.values_mut() { permission_config .always_allow .retain(|p| !p.starts_with(extension_name)); permission_config .ask_before .retain(|p| !p.starts_with(extension_name)); permission_config .never_allow .retain(|p| !p.starts_with(extension_name)); } let yaml_content = serde_yaml::to_string(&*map).expect("Failed to serialize permission config"); fs::write(&self.config_path, yaml_content).expect("Failed to write to permission.yaml"); } } #[cfg(test)] mod tests { use super::*; use rmcp::model::ToolAnnotations; use rmcp::object; use tempfile::TempDir; // Helper function to create a test instance of PermissionManager with a temp dir fn create_test_permission_manager() -> (PermissionManager, TempDir) { let temp_dir = TempDir::new().unwrap(); let manager = PermissionManager::new(temp_dir.path().to_path_buf()); (manager, temp_dir) } #[test] fn test_get_permission_names_empty() { let (manager, _temp_dir) = create_test_permission_manager(); assert!(manager.get_permission_names().is_empty()); } #[test] fn test_update_user_permission() { let (manager, _temp_dir) = create_test_permission_manager(); manager.update_user_permission("tool1", PermissionLevel::AlwaysAllow); let permission = manager.get_user_permission("tool1"); assert_eq!(permission, Some(PermissionLevel::AlwaysAllow)); } #[test] fn test_update_smart_approve_permission() { let (manager, _temp_dir) = create_test_permission_manager(); manager.update_smart_approve_permission("tool2", PermissionLevel::AskBefore); let permission = manager.get_smart_approve_permission("tool2"); assert_eq!(permission, Some(PermissionLevel::AskBefore)); } #[test] fn test_get_permission_not_found() { let (manager, _temp_dir) = create_test_permission_manager(); let permission = manager.get_user_permission("non_existent_tool"); assert_eq!(permission, None); } #[test] fn test_permission_levels() { let (manager, _temp_dir) = create_test_permission_manager(); manager.update_user_permission("tool4", PermissionLevel::AlwaysAllow); manager.update_user_permission("tool5", PermissionLevel::AskBefore); manager.update_user_permission("tool6", PermissionLevel::NeverAllow); // Check the permission levels assert_eq!( manager.get_user_permission("tool4"), Some(PermissionLevel::AlwaysAllow) ); assert_eq!( manager.get_user_permission("tool5"), Some(PermissionLevel::AskBefore) ); assert_eq!( manager.get_user_permission("tool6"), Some(PermissionLevel::NeverAllow) ); } #[test] fn test_permission_update_replaces_existing_level() { let (manager, _temp_dir) = create_test_permission_manager(); // Initially AlwaysAllow manager.update_user_permission("tool7", PermissionLevel::AlwaysAllow); assert_eq!( manager.get_user_permission("tool7"), Some(PermissionLevel::AlwaysAllow) ); // Now change to NeverAllow manager.update_user_permission("tool7", PermissionLevel::NeverAllow); assert_eq!( manager.get_user_permission("tool7"), Some(PermissionLevel::NeverAllow) ); // Ensure it's removed from other levels let map = manager.permission_map.read().unwrap(); let config = map.get(USER_PERMISSION).unwrap(); assert!(!config.always_allow.contains(&"tool7".to_string())); assert!(!config.ask_before.contains(&"tool7".to_string())); assert!(config.never_allow.contains(&"tool7".to_string())); } #[test] fn test_remove_extension() { let (manager, _temp_dir) = create_test_permission_manager(); manager.update_user_permission("prefix__tool1", PermissionLevel::AlwaysAllow); manager.update_user_permission("nonprefix__tool2", PermissionLevel::AlwaysAllow); manager.update_user_permission("prefix__tool3", PermissionLevel::AskBefore); // Remove entries starting with "prefix" manager.remove_extension("prefix"); let map = manager.permission_map.read().unwrap(); let config = map.get(USER_PERMISSION).unwrap(); // Verify entries with "prefix" are removed assert!(!config.always_allow.contains(&"prefix__tool1".to_string())); assert!(!config.ask_before.contains(&"prefix__tool3".to_string())); // Verify other entries remain assert!(config .always_allow .contains(&"nonprefix__tool2".to_string())); } #[test] #[should_panic(expected = "Corrupted permission config")] fn test_corrupted_permission_file_panics() { let temp_dir = TempDir::new().unwrap(); let permission_path = temp_dir.path().join(PERMISSION_FILE); fs::write(&permission_path, "{{invalid yaml: [broken").unwrap(); PermissionManager::new(temp_dir.path().to_path_buf()); } use test_case::test_case; #[test_case( vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"})) .annotate(ToolAnnotations::new().read_only(false))], Some(PermissionLevel::AskBefore); "write_annotation_caches_ask" )] #[test_case( vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"}))], None; "unannotated_left_uncached" )] #[test_case( vec![Tool::new("tool".to_string(), String::new(), object!({"type": "object"})) .annotate(ToolAnnotations::new().read_only(true))], None; "readonly_annotation_skipped" )] fn test_apply_tool_annotations(tools: Vec, expect_cache: Option) { let (manager, _temp_dir) = create_test_permission_manager(); manager.apply_tool_annotations(&tools); assert_eq!(manager.get_smart_approve_permission("tool"), expect_cache); } }