aegis-delegate 0.2.0

AEGIS delegation and authority model โ€” delegation trees, session keys, scope constraints
Documentation
// AEGIS Delegate โ€” Revocation Store Trait Abstraction
//
// Reference: AEGIS Specification v1.0.0 ยง9.5
//
// Defines a pluggable storage interface for delegation revocation tracking.
// Includes an in-memory implementation for development and testing.

use std::collections::HashSet;
use std::sync::RwLock;

use async_trait::async_trait;

use openagent_aegis_core::DelegationError;

// ---------------------------------------------------------------------------
// RevocationStore Trait
// ---------------------------------------------------------------------------

/// Pluggable storage backend for tracking revoked delegation IDs.
///
/// Implementations may use in-memory storage, databases, or distributed
/// caches. All operations are async to accommodate network-backed stores.
///
/// This trait complements [`crate::revocation::RevocationRegistry`] by
/// providing an async, error-aware interface suitable for production
/// persistence backends.
#[async_trait]
pub trait RevocationStore: Send + Sync {
    /// Mark a delegation as revoked.
    ///
    /// Revoking an already-revoked delegation is idempotent and does not
    /// return an error.
    async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError>;

    /// Check whether a delegation has been revoked.
    async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError>;

    /// Revoke multiple delegations at once (batch/cascade revocation).
    ///
    /// This is typically used after cascade revocation computes the full
    /// set of transitively affected delegation IDs.
    async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError>;
}

// ---------------------------------------------------------------------------
// InMemoryRevocationStore
// ---------------------------------------------------------------------------

/// In-memory revocation store backed by a `RwLock<HashSet>`.
///
/// Suitable for development, testing, and single-instance deployments.
/// For production multi-node deployments, use a database-backed
/// implementation of [`RevocationStore`].
pub struct InMemoryRevocationStore {
    revoked: RwLock<HashSet<String>>,
}

impl InMemoryRevocationStore {
    /// Creates a new empty in-memory revocation store.
    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);
    }
}