use crate::audit::{self, TrustLevel};
use crate::permission::{
PermissionConfig, PermissionContext, PermissionDecision, StoredPermission, StoredTrustLevel,
};
use crate::{LoadedPlugin, LoaderError, PluginLoader};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Clone)]
pub struct PluginRegistry {
inner: Arc<RwLock<RegistryInner>>,
loader: Arc<PluginLoader>,
permission: Option<Arc<PermissionConfig>>,
}
struct RegistryInner {
plugins: HashMap<String, PluginEntry>,
path_to_command: HashMap<PathBuf, String>,
}
struct PluginEntry {
plugin: LoadedPlugin,
source_path: Option<PathBuf>,
}
impl PluginRegistry {
pub fn new() -> Result<Self, LoaderError> {
Ok(Self {
inner: Arc::new(RwLock::new(RegistryInner {
plugins: HashMap::new(),
path_to_command: HashMap::new(),
})),
loader: Arc::new(PluginLoader::new()?),
permission: None,
})
}
pub fn with_loader(loader: PluginLoader) -> Self {
Self {
inner: Arc::new(RwLock::new(RegistryInner {
plugins: HashMap::new(),
path_to_command: HashMap::new(),
})),
loader: Arc::new(loader),
permission: None,
}
}
pub fn with_permissions(config: PermissionConfig) -> Result<Self, LoaderError> {
Ok(Self {
inner: Arc::new(RwLock::new(RegistryInner {
plugins: HashMap::new(),
path_to_command: HashMap::new(),
})),
loader: Arc::new(PluginLoader::new()?),
permission: Some(Arc::new(config)),
})
}
pub fn set_permissions(&mut self, config: PermissionConfig) {
self.permission = Some(Arc::new(config));
}
pub async fn load_plugin(&self, path: impl AsRef<Path>) -> Result<String, LoaderError> {
let path = path.as_ref();
let wasm_bytes = tokio::fs::read(path).await.map_err(|e| {
LoaderError::MemoryAccess(format!("Failed to read file {}: {}", path.display(), e))
})?;
let plugin = self.loader.load(&wasm_bytes)?;
let command_name = plugin.manifest.command.name.clone();
let mut inner = self.inner.write().await;
if let Some(old_cmd) = inner.path_to_command.remove(path) {
inner.plugins.remove(&old_cmd);
}
inner
.path_to_command
.insert(path.to_path_buf(), command_name.clone());
inner.plugins.insert(
command_name.clone(),
PluginEntry {
plugin,
source_path: Some(path.to_path_buf()),
},
);
tracing::info!(command = %command_name, path = %path.display(), "Plugin loaded");
Ok(command_name)
}
pub async fn register(&self, plugin: LoadedPlugin) -> String {
let command_name = plugin.manifest.command.name.clone();
let mut inner = self.inner.write().await;
inner.plugins.insert(
command_name.clone(),
PluginEntry {
plugin,
source_path: None,
},
);
tracing::info!(command = %command_name, "Plugin registered");
command_name
}
pub async fn unload_by_path(&self, path: impl AsRef<Path>) -> Option<String> {
let path = path.as_ref();
let mut inner = self.inner.write().await;
if let Some(command_name) = inner.path_to_command.remove(path) {
inner.plugins.remove(&command_name);
tracing::info!(command = %command_name, path = %path.display(), "Plugin unloaded");
Some(command_name)
} else {
None
}
}
pub async fn unload(&self, command_name: &str) -> bool {
let mut inner = self.inner.write().await;
if let Some(entry) = inner.plugins.remove(command_name) {
if let Some(path) = entry.source_path {
inner.path_to_command.remove(&path);
}
tracing::info!(command = %command_name, "Plugin unloaded");
true
} else {
false
}
}
pub async fn reload_by_path(&self, path: impl AsRef<Path>) -> Result<String, LoaderError> {
self.load_plugin(path).await
}
pub async fn list_commands(&self) -> Vec<String> {
let inner = self.inner.read().await;
inner.plugins.keys().cloned().collect()
}
pub async fn has_command(&self, command_name: &str) -> bool {
let inner = self.inner.read().await;
inner.plugins.contains_key(command_name)
}
pub async fn execute(
&self,
command_name: &str,
args: &[String],
) -> Result<sen_plugin_api::ExecuteResult, RegistryError> {
let mut inner = self.inner.write().await;
let entry = inner
.plugins
.get_mut(command_name)
.ok_or_else(|| RegistryError::CommandNotFound(command_name.to_string()))?;
if let Some(ref perm_config) = self.permission {
let capabilities = &entry.plugin.manifest.capabilities;
let _ = perm_config
.audit
.record(audit::permission_requested(command_name, capabilities));
let key =
perm_config
.store
.make_key(command_name, None, perm_config.strategy.granularity());
let stored = perm_config.store.get(&key).ok().flatten();
let ctx = PermissionContext {
plugin_name: command_name,
command_path: &[],
requested: capabilities,
granted: stored.as_ref().map(|s| &s.capabilities),
interactive: perm_config.prompt.is_interactive(),
};
let decision = if let Some(ref stored_perm) = stored {
if stored_perm.has_escalated(capabilities) {
let _ = perm_config.audit.record(audit::escalation_detected(
command_name,
&stored_perm.capabilities,
capabilities,
));
perm_config.strategy.on_escalation(&ctx)
} else {
perm_config.strategy.check(&ctx)
}
} else {
perm_config.strategy.check(&ctx)
};
match decision {
PermissionDecision::Allow => {
let _ = perm_config.audit.record(audit::permission_granted(
command_name,
capabilities,
TrustLevel::Permanent,
));
}
PermissionDecision::Deny(reason) => {
let _ = perm_config.audit.record(audit::permission_denied(
command_name,
capabilities,
&reason,
));
return Err(RegistryError::PermissionDenied {
plugin: command_name.to_string(),
reason,
});
}
PermissionDecision::Prompt => {
let prompt_result = if let Some(ref stored_perm) = stored {
perm_config.prompt.prompt_escalation(
command_name,
&stored_perm.capabilities,
capabilities,
)
} else {
perm_config.prompt.prompt(command_name, capabilities)
};
match prompt_result {
Ok(result) if result.is_allowed() => {
if result.should_persist() {
let trust_level =
result.to_trust_level().unwrap_or(StoredTrustLevel::Session);
let stored_perm =
StoredPermission::new(capabilities.clone(), trust_level);
let _ = perm_config.store.set(&key, stored_perm);
}
let audit_trust = match result.to_trust_level() {
Some(StoredTrustLevel::Permanent) => TrustLevel::Permanent,
Some(StoredTrustLevel::Session) => TrustLevel::Session,
None => TrustLevel::Once,
};
let _ = perm_config.audit.record(audit::permission_granted(
command_name,
capabilities,
audit_trust,
));
}
Ok(_) | Err(_) => {
let _ = perm_config.audit.record(audit::permission_denied(
command_name,
capabilities,
"User denied permission",
));
return Err(RegistryError::PermissionDenied {
plugin: command_name.to_string(),
reason: "User denied permission".to_string(),
});
}
}
}
PermissionDecision::AllowPartial(_reduced) => {
let _ = perm_config.audit.record(audit::permission_granted(
command_name,
capabilities,
TrustLevel::Once,
));
}
}
}
entry
.plugin
.instance
.execute(args)
.map_err(RegistryError::Execution)
}
pub async fn get_manifest(&self, command_name: &str) -> Option<sen_plugin_api::PluginManifest> {
let inner = self.inner.read().await;
inner
.plugins
.get(command_name)
.map(|e| e.plugin.manifest.clone())
}
pub async fn get_all_manifests(&self) -> Vec<sen_plugin_api::PluginManifest> {
let inner = self.inner.read().await;
inner
.plugins
.values()
.map(|e| e.plugin.manifest.clone())
.collect()
}
pub async fn len(&self) -> usize {
let inner = self.inner.read().await;
inner.plugins.len()
}
pub async fn is_empty(&self) -> bool {
let inner = self.inner.read().await;
inner.plugins.is_empty()
}
}
#[derive(Debug, thiserror::Error)]
pub enum RegistryError {
#[error("Command not found: {0}")]
CommandNotFound(String),
#[error("Plugin execution failed: {0}")]
Execution(#[source] LoaderError),
#[error("Permission denied for plugin '{plugin}': {reason}")]
PermissionDenied { plugin: String, reason: String },
}
#[cfg(test)]
mod tests {
use super::*;
use crate::audit::MemoryAuditSink;
use crate::permission::{
AutoPromptHandler, MemoryPermissionStore, PermissionPresets, PermissionStore, PromptResult,
RecordingPromptHandler,
};
const HELLO_PLUGIN_WASM: &[u8] = include_bytes!(
"../../examples/hello-plugin/target/wasm32-unknown-unknown/release/hello_plugin.wasm"
);
#[tokio::test]
async fn test_registry_register_and_execute() {
let registry = PluginRegistry::new().unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
let cmd = registry.register(plugin).await;
assert_eq!(cmd, "hello");
assert!(registry.has_command("hello").await);
let result = registry
.execute("hello", &["World".to_string()])
.await
.unwrap();
match result {
sen_plugin_api::ExecuteResult::Success(output) => {
assert_eq!(output, "Hello, World!");
}
_ => panic!("Expected success"),
}
}
#[tokio::test]
async fn test_registry_unload() {
let registry = PluginRegistry::new().unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
registry.register(plugin).await;
assert!(registry.has_command("hello").await);
registry.unload("hello").await;
assert!(!registry.has_command("hello").await);
}
#[tokio::test]
async fn test_registry_list_commands() {
let registry = PluginRegistry::new().unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
registry.register(plugin).await;
let commands = registry.list_commands().await;
assert_eq!(commands, vec!["hello"]);
}
#[tokio::test]
async fn test_registry_with_permissions_trust_all() {
let config = PermissionPresets::trust_all_dangerous();
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
registry.register(plugin).await;
let result = registry
.execute("hello", &["World".to_string()])
.await
.unwrap();
match result {
sen_plugin_api::ExecuteResult::Success(output) => {
assert_eq!(output, "Hello, World!");
}
_ => panic!("Expected success with trust_all"),
}
}
#[tokio::test]
async fn test_registry_with_permissions_testing_preset() {
let config = PermissionPresets::testing();
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
registry.register(plugin).await;
let result = registry
.execute("hello", &["World".to_string()])
.await
.unwrap();
match result {
sen_plugin_api::ExecuteResult::Success(output) => {
assert_eq!(output, "Hello, World!");
}
_ => panic!("Expected success with testing preset"),
}
}
#[tokio::test]
async fn test_registry_with_permissions_deny_on_prompt() {
let config = PermissionConfig::new(
crate::permission::DefaultPermissionStrategy,
MemoryPermissionStore::new(),
AutoPromptHandler::always_deny(),
crate::audit::NullAuditSink,
crate::permission::TrustFlagConfig::default(),
);
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let mut plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
plugin.manifest.capabilities = sen_plugin_api::Capabilities::default()
.with_stdio(sen_plugin_api::StdioCapability::stdout_only());
registry.register(plugin).await;
let result = registry.execute("hello", &["World".to_string()]).await;
match result {
Err(RegistryError::PermissionDenied { plugin, reason }) => {
assert_eq!(plugin, "hello");
assert!(reason.contains("denied"));
}
Ok(_) => panic!("Expected PermissionDenied error"),
Err(e) => panic!("Unexpected error: {:?}", e),
}
}
#[tokio::test]
async fn test_registry_with_permissions_audit_logging() {
let audit_sink = std::sync::Arc::new(MemoryAuditSink::new());
let config = PermissionConfig {
strategy: std::sync::Arc::new(crate::permission::DefaultPermissionStrategy),
store: std::sync::Arc::new(MemoryPermissionStore::new()),
prompt: std::sync::Arc::new(AutoPromptHandler::always_allow()),
audit: audit_sink.clone(),
trust_flags: crate::permission::TrustFlagConfig::default(),
};
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let mut plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
plugin.manifest.capabilities = sen_plugin_api::Capabilities::default()
.with_stdio(sen_plugin_api::StdioCapability::stdout_only());
registry.register(plugin).await;
let _ = registry.execute("hello", &["World".to_string()]).await;
let events = audit_sink.events();
assert!(!events.is_empty(), "Should have audit events");
let request_events =
audit_sink.find_by_type(crate::audit::AuditEventType::PermissionRequested);
assert!(
!request_events.is_empty(),
"Should have permission request event"
);
}
#[tokio::test]
async fn test_registry_with_permissions_stores_grant() {
let store = std::sync::Arc::new(MemoryPermissionStore::new());
let config = PermissionConfig {
strategy: std::sync::Arc::new(crate::permission::DefaultPermissionStrategy),
store: store.clone(),
prompt: std::sync::Arc::new(AutoPromptHandler::always_allow()),
audit: std::sync::Arc::new(crate::audit::NullAuditSink),
trust_flags: crate::permission::TrustFlagConfig::default(),
};
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let mut plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
plugin.manifest.capabilities = sen_plugin_api::Capabilities::default()
.with_stdio(sen_plugin_api::StdioCapability::stdout_only());
registry.register(plugin).await;
let _ = registry.execute("hello", &["World".to_string()]).await;
let stored = store.get("hello").unwrap();
assert!(stored.is_some(), "Permission should be stored after grant");
}
#[tokio::test]
async fn test_registry_with_permissions_prompt_recording() {
let prompt_handler =
std::sync::Arc::new(RecordingPromptHandler::new(PromptResult::AllowAlways));
let config = PermissionConfig {
strategy: std::sync::Arc::new(crate::permission::DefaultPermissionStrategy),
store: std::sync::Arc::new(MemoryPermissionStore::new()),
prompt: prompt_handler.clone(),
audit: std::sync::Arc::new(crate::audit::NullAuditSink),
trust_flags: crate::permission::TrustFlagConfig::default(),
};
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let mut plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
plugin.manifest.capabilities = sen_plugin_api::Capabilities::default()
.with_stdio(sen_plugin_api::StdioCapability::stdout_only());
registry.register(plugin).await;
let _ = registry.execute("hello", &["World".to_string()]).await;
assert_eq!(
prompt_handler.prompt_count(),
1,
"Should have prompted once"
);
let prompts = prompt_handler.prompts();
assert_eq!(prompts[0].plugin, "hello");
}
#[tokio::test]
async fn test_registry_without_permissions_skips_check() {
let registry = PluginRegistry::new().unwrap();
let loader = PluginLoader::new().unwrap();
let mut plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
plugin.manifest.capabilities = sen_plugin_api::Capabilities::default()
.with_fs_read(vec![sen_plugin_api::PathPattern::new("/")]);
registry.register(plugin).await;
let result = registry
.execute("hello", &["World".to_string()])
.await
.unwrap();
match result {
sen_plugin_api::ExecuteResult::Success(output) => {
assert_eq!(output, "Hello, World!");
}
_ => panic!("Expected success without permission config"),
}
}
#[tokio::test]
async fn test_registry_empty_capabilities_allowed() {
let prompt_handler = std::sync::Arc::new(RecordingPromptHandler::new(PromptResult::Deny));
let config = PermissionConfig {
strategy: std::sync::Arc::new(crate::permission::DefaultPermissionStrategy),
store: std::sync::Arc::new(MemoryPermissionStore::new()),
prompt: prompt_handler.clone(),
audit: std::sync::Arc::new(crate::audit::NullAuditSink),
trust_flags: crate::permission::TrustFlagConfig::default(),
};
let registry = PluginRegistry::with_permissions(config).unwrap();
let loader = PluginLoader::new().unwrap();
let plugin = loader.load(HELLO_PLUGIN_WASM).unwrap();
registry.register(plugin).await;
let result = registry.execute("hello", &["World".to_string()]).await;
assert!(result.is_ok(), "Empty capabilities should be allowed");
assert_eq!(
prompt_handler.prompt_count(),
0,
"Should not prompt for empty caps"
);
}
}