use std::collections::HashSet;
use std::sync::RwLock;
use async_trait::async_trait;
use openagent_aegis_core::DelegationError;
#[async_trait]
pub trait RevocationStore: Send + Sync {
async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError>;
async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError>;
async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError>;
}
pub struct InMemoryRevocationStore {
revoked: RwLock<HashSet<String>>,
}
impl InMemoryRevocationStore {
pub fn new() -> Self {
Self {
revoked: RwLock::new(HashSet::new()),
}
}
}
impl Default for InMemoryRevocationStore {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl RevocationStore for InMemoryRevocationStore {
async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError> {
let mut revoked = self
.revoked
.write()
.map_err(|e| DelegationError::InvalidProof {
reason: format!("revocation store lock poisoned: {e}"),
})?;
revoked.insert(delegation_id.to_string());
Ok(())
}
async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError> {
let revoked = self
.revoked
.read()
.map_err(|e| DelegationError::InvalidProof {
reason: format!("revocation store lock poisoned: {e}"),
})?;
Ok(revoked.contains(delegation_id))
}
async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError> {
let mut revoked = self
.revoked
.write()
.map_err(|e| DelegationError::InvalidProof {
reason: format!("revocation store lock poisoned: {e}"),
})?;
for id in ids {
revoked.insert(id.clone());
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn revoke_and_check() {
let store = InMemoryRevocationStore::new();
let is_rev = store.is_revoked("del-1").await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(true), false);
let result = store.revoke("del-1").await;
assert!(result.is_ok());
let is_rev = store.is_revoked("del-1").await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(false), true);
}
#[tokio::test]
async fn revoke_is_idempotent() {
let store = InMemoryRevocationStore::new();
let r1 = store.revoke("del-1").await;
assert!(r1.is_ok());
let r2 = store.revoke("del-1").await;
assert!(r2.is_ok());
let is_rev = store.is_revoked("del-1").await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(false), true);
}
#[tokio::test]
async fn revoke_batch_marks_all() {
let store = InMemoryRevocationStore::new();
let ids = vec!["d1".to_string(), "d2".to_string(), "d3".to_string()];
let result = store.revoke_batch(&ids).await;
assert!(result.is_ok());
for id in &ids {
let is_rev = store.is_revoked(id).await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(false), true);
}
let is_rev = store.is_revoked("d4").await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(true), false);
}
#[tokio::test]
async fn revoke_batch_empty_is_ok() {
let store = InMemoryRevocationStore::new();
let result = store.revoke_batch(&[]).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn default_constructor() {
let store = InMemoryRevocationStore::default();
let is_rev = store.is_revoked("any").await;
assert!(is_rev.is_ok());
assert_eq!(is_rev.unwrap_or(true), false);
}
}