use std::collections::HashMap;
use std::panic;
use std::sync::{Arc, Mutex};
use crate::error::HawkError;
pub type Result<T> = std::result::Result<T, TalonError>;
#[derive(Debug, thiserror::Error)]
pub enum TalonError {
#[error("signature verification failed for '{0}'")]
InvalidSignature(String),
#[error("talon not found: '{0}'")]
NotFound(String),
#[error("talon already installed: '{0}'")]
AlreadyInstalled(String),
#[error("talon lifecycle error: {0}")]
Lifecycle(String),
#[error("talon panicked: {0}")]
Panicked(String),
}
impl From<TalonError> for HawkError {
fn from(e: TalonError) -> Self {
HawkError::Config(e.to_string())
}
}
#[derive(Debug, Clone)]
pub struct Capability {
pub name: String,
pub description: String,
}
#[derive(Debug, Clone, Default)]
pub struct TalonConfig {
pub settings: HashMap<String, String>,
}
#[derive(Debug, Clone, PartialEq)]
pub enum TalonStatus {
Loaded,
Unloaded,
Failed(String),
}
#[derive(Debug, Clone)]
pub struct TalonRecord {
pub name: String,
pub version: String,
pub status: TalonStatus,
pub capabilities: Vec<Capability>,
pub signature: String,
}
pub trait Talon: Send + Sync {
fn name(&self) -> &str;
fn version(&self) -> &str;
fn load(&mut self) -> Result<()>;
fn unload(&mut self) -> Result<()>;
fn configure(&mut self, config: TalonConfig) -> Result<()>;
fn capabilities(&self) -> Vec<Capability>;
}
fn expected_signature(name: &str, version: &str) -> String {
use sha2::{Digest, Sha256};
let mut h = Sha256::new();
h.update(format!("{name}:{version}").as_bytes());
hex::encode(h.finalize())
}
fn verify_signature(name: &str, version: &str, signature: &str) -> bool {
expected_signature(name, version) == signature
}
pub fn make_signature(name: &str, version: &str) -> String {
expected_signature(name, version)
}
pub struct TalonRegistry {
records: Arc<Mutex<HashMap<String, TalonRecord>>>,
instances: Arc<Mutex<HashMap<String, Box<dyn Talon>>>>,
}
impl TalonRegistry {
pub fn new() -> Self {
Self {
records: Arc::new(Mutex::new(HashMap::new())),
instances: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn install(&self, name: &str, version: &str, signature: &str, capabilities: Vec<Capability>) -> Result<()> {
if !verify_signature(name, version, signature) {
eprintln!("SECURITY WARNING: signature verification failed for talon '{name}'. Installation rejected.");
return Err(TalonError::InvalidSignature(name.to_string()));
}
let mut records = self.records.lock().unwrap();
if records.contains_key(name) {
return Err(TalonError::AlreadyInstalled(name.to_string()));
}
records.insert(name.to_string(), TalonRecord {
name: name.to_string(),
version: version.to_string(),
status: TalonStatus::Unloaded,
capabilities,
signature: signature.to_string(),
});
Ok(())
}
pub fn load(&self, name: &str) -> Result<()> {
let mut instances = self.instances.lock().unwrap();
let mut records = self.records.lock().unwrap();
let record = records.get_mut(name).ok_or_else(|| TalonError::NotFound(name.to_string()))?;
if let Some(t) = instances.get_mut(name) {
let t_ptr = t.as_mut() as *mut dyn Talon;
let result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
unsafe { &mut *t_ptr }.load()
}));
match result {
Ok(Ok(())) => { record.status = TalonStatus::Loaded; Ok(()) }
Ok(Err(e)) => {
let msg = e.to_string();
record.status = TalonStatus::Failed(msg.clone());
Err(TalonError::Lifecycle(msg))
}
Err(_) => {
let msg = "talon panicked during load".to_string();
record.status = TalonStatus::Failed(msg.clone());
Err(TalonError::Panicked(msg))
}
}
} else {
record.status = TalonStatus::Loaded;
Ok(())
}
}
pub fn unload(&self, name: &str) -> Result<()> {
let mut instances = self.instances.lock().unwrap();
let mut records = self.records.lock().unwrap();
let record = records.get_mut(name).ok_or_else(|| TalonError::NotFound(name.to_string()))?;
if let Some(t) = instances.get_mut(name) {
let t_ptr = t.as_mut() as *mut dyn Talon;
let result = panic::catch_unwind(panic::AssertUnwindSafe(|| unsafe { &mut *t_ptr }.unload()));
match result {
Ok(Ok(())) => {}
Ok(Err(e)) => {
let msg = e.to_string();
record.status = TalonStatus::Failed(msg.clone());
return Err(TalonError::Lifecycle(msg));
}
Err(_) => {
let msg = "talon panicked during unload".to_string();
record.status = TalonStatus::Failed(msg.clone());
return Err(TalonError::Panicked(msg));
}
}
}
record.status = TalonStatus::Unloaded;
Ok(())
}
pub fn list(&self) -> Vec<TalonRecord> {
self.records.lock().unwrap().values().cloned().collect()
}
pub fn get_capabilities(&self, name: &str) -> Option<Vec<Capability>> {
self.records.lock().unwrap().get(name).map(|r| r.capabilities.clone())
}
pub fn is_authorized(agent_manifest_talons: &[String], talon_name: &str) -> bool {
agent_manifest_talons.iter().any(|t| t == talon_name)
}
pub fn register_instance(&self, talon: Box<dyn Talon>) -> Result<()> {
let name = talon.name().to_string();
self.instances.lock().unwrap().insert(name, talon);
Ok(())
}
}
impl Default for TalonRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct GoodTalon { loaded: bool }
impl GoodTalon { fn new() -> Self { Self { loaded: false } } }
impl Talon for GoodTalon {
fn name(&self) -> &str { "good-talon" }
fn version(&self) -> &str { "1.0.0" }
fn load(&mut self) -> Result<()> { self.loaded = true; Ok(()) }
fn unload(&mut self) -> Result<()> { self.loaded = false; Ok(()) }
fn configure(&mut self, _: TalonConfig) -> Result<()> { Ok(()) }
fn capabilities(&self) -> Vec<Capability> {
vec![Capability { name: "browse".to_string(), description: "web browsing".to_string() }]
}
}
struct PanickingTalon;
impl Talon for PanickingTalon {
fn name(&self) -> &str { "panic-talon" }
fn version(&self) -> &str { "0.1.0" }
fn load(&mut self) -> Result<()> { panic!("intentional panic for isolation test"); }
fn unload(&mut self) -> Result<()> { Ok(()) }
fn configure(&mut self, _: TalonConfig) -> Result<()> { Ok(()) }
fn capabilities(&self) -> Vec<Capability> { vec![] }
}
struct FailingTalon;
impl Talon for FailingTalon {
fn name(&self) -> &str { "fail-talon" }
fn version(&self) -> &str { "0.1.0" }
fn load(&mut self) -> Result<()> { Err(TalonError::Lifecycle("load failed".to_string())) }
fn unload(&mut self) -> Result<()> { Ok(()) }
fn configure(&mut self, _: TalonConfig) -> Result<()> { Ok(()) }
fn capabilities(&self) -> Vec<Capability> { vec![] }
}
fn install_good(registry: &TalonRegistry) {
let sig = make_signature("good-talon", "1.0.0");
registry.install("good-talon", "1.0.0", &sig, vec![
Capability { name: "browse".to_string(), description: "web browsing".to_string() },
]).unwrap();
}
#[test]
fn load_and_unload_lifecycle() {
let registry = TalonRegistry::new();
install_good(®istry);
registry.register_instance(Box::new(GoodTalon::new())).unwrap();
registry.load("good-talon").unwrap();
assert_eq!(registry.records.lock().unwrap()["good-talon"].status, TalonStatus::Loaded);
registry.unload("good-talon").unwrap();
assert_eq!(registry.records.lock().unwrap()["good-talon"].status, TalonStatus::Unloaded);
}
#[test]
fn valid_signature_accepted() {
let registry = TalonRegistry::new();
let sig = make_signature("my-talon", "2.0.0");
assert!(registry.install("my-talon", "2.0.0", &sig, vec![]).is_ok());
}
#[test]
fn invalid_signature_rejected() {
let registry = TalonRegistry::new();
assert!(matches!(registry.install("my-talon", "2.0.0", "bad-sig", vec![]), Err(TalonError::InvalidSignature(_))));
}
#[test]
fn wrong_version_signature_rejected() {
let registry = TalonRegistry::new();
let sig = make_signature("my-talon", "1.0.0");
assert!(matches!(registry.install("my-talon", "2.0.0", &sig, vec![]), Err(TalonError::InvalidSignature(_))));
}
#[test]
fn panicking_talon_does_not_crash_process() {
let registry = TalonRegistry::new();
let sig = make_signature("panic-talon", "0.1.0");
registry.install("panic-talon", "0.1.0", &sig, vec![]).unwrap();
registry.register_instance(Box::new(PanickingTalon)).unwrap();
let result = registry.load("panic-talon");
assert!(matches!(result, Err(TalonError::Panicked(_))));
assert!(matches!(registry.records.lock().unwrap()["panic-talon"].status, TalonStatus::Failed(_)));
}
#[test]
fn failing_talon_is_marked_failed() {
let registry = TalonRegistry::new();
let sig = make_signature("fail-talon", "0.1.0");
registry.install("fail-talon", "0.1.0", &sig, vec![]).unwrap();
registry.register_instance(Box::new(FailingTalon)).unwrap();
let result = registry.load("fail-talon");
assert!(matches!(result, Err(TalonError::Lifecycle(_))));
assert!(matches!(registry.records.lock().unwrap()["fail-talon"].status, TalonStatus::Failed(_)));
}
#[test]
fn authorized_agent_can_use_talon() {
let declared = vec!["browser-talon".to_string(), "github-talon".to_string()];
assert!(TalonRegistry::is_authorized(&declared, "browser-talon"));
}
#[test]
fn unauthorized_agent_cannot_use_talon() {
let declared = vec!["github-talon".to_string()];
assert!(!TalonRegistry::is_authorized(&declared, "browser-talon"));
}
#[test]
fn empty_manifest_talons_denies_all() {
assert!(!TalonRegistry::is_authorized(&[], "any-talon"));
}
#[test]
fn get_capabilities_returns_installed_caps() {
let registry = TalonRegistry::new();
install_good(®istry);
let caps = registry.get_capabilities("good-talon").unwrap();
assert_eq!(caps.len(), 1);
assert_eq!(caps[0].name, "browse");
}
#[test]
fn get_capabilities_returns_none_for_unknown() {
let registry = TalonRegistry::new();
assert!(registry.get_capabilities("unknown").is_none());
}
#[test]
fn list_returns_all_installed() {
let registry = TalonRegistry::new();
install_good(®istry);
let sig2 = make_signature("other-talon", "0.5.0");
registry.install("other-talon", "0.5.0", &sig2, vec![]).unwrap();
assert_eq!(registry.list().len(), 2);
}
#[test]
fn duplicate_install_rejected() {
let registry = TalonRegistry::new();
install_good(®istry);
let sig = make_signature("good-talon", "1.0.0");
assert!(matches!(registry.install("good-talon", "1.0.0", &sig, vec![]), Err(TalonError::AlreadyInstalled(_))));
}
#[test]
fn load_unknown_talon_returns_not_found() {
let registry = TalonRegistry::new();
assert!(matches!(registry.load("ghost"), Err(TalonError::NotFound(_))));
}
}