Skip to main content

aegis_delegate/
store.rs

1// AEGIS Delegate — Revocation Store Trait Abstraction
2//
3// Reference: AEGIS Specification v1.0.0 §9.5
4//
5// Defines a pluggable storage interface for delegation revocation tracking.
6// Includes an in-memory implementation for development and testing.
7
8use std::collections::HashSet;
9use std::sync::RwLock;
10
11use async_trait::async_trait;
12
13use openagent_aegis_core::DelegationError;
14
15// ---------------------------------------------------------------------------
16// RevocationStore Trait
17// ---------------------------------------------------------------------------
18
19/// Pluggable storage backend for tracking revoked delegation IDs.
20///
21/// Implementations may use in-memory storage, databases, or distributed
22/// caches. All operations are async to accommodate network-backed stores.
23///
24/// This trait complements [`crate::revocation::RevocationRegistry`] by
25/// providing an async, error-aware interface suitable for production
26/// persistence backends.
27#[async_trait]
28pub trait RevocationStore: Send + Sync {
29    /// Mark a delegation as revoked.
30    ///
31    /// Revoking an already-revoked delegation is idempotent and does not
32    /// return an error.
33    async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError>;
34
35    /// Check whether a delegation has been revoked.
36    async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError>;
37
38    /// Revoke multiple delegations at once (batch/cascade revocation).
39    ///
40    /// This is typically used after cascade revocation computes the full
41    /// set of transitively affected delegation IDs.
42    async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError>;
43}
44
45// ---------------------------------------------------------------------------
46// InMemoryRevocationStore
47// ---------------------------------------------------------------------------
48
49/// In-memory revocation store backed by a `RwLock<HashSet>`.
50///
51/// Suitable for development, testing, and single-instance deployments.
52/// For production multi-node deployments, use a database-backed
53/// implementation of [`RevocationStore`].
54pub struct InMemoryRevocationStore {
55    revoked: RwLock<HashSet<String>>,
56}
57
58impl InMemoryRevocationStore {
59    /// Creates a new empty in-memory revocation store.
60    pub fn new() -> Self {
61        Self {
62            revoked: RwLock::new(HashSet::new()),
63        }
64    }
65}
66
67impl Default for InMemoryRevocationStore {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73#[async_trait]
74impl RevocationStore for InMemoryRevocationStore {
75    async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError> {
76        let mut revoked = self
77            .revoked
78            .write()
79            .map_err(|e| DelegationError::InvalidProof {
80                reason: format!("revocation store lock poisoned: {e}"),
81            })?;
82        revoked.insert(delegation_id.to_string());
83        Ok(())
84    }
85
86    async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError> {
87        let revoked = self
88            .revoked
89            .read()
90            .map_err(|e| DelegationError::InvalidProof {
91                reason: format!("revocation store lock poisoned: {e}"),
92            })?;
93        Ok(revoked.contains(delegation_id))
94    }
95
96    async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError> {
97        let mut revoked = self
98            .revoked
99            .write()
100            .map_err(|e| DelegationError::InvalidProof {
101                reason: format!("revocation store lock poisoned: {e}"),
102            })?;
103        for id in ids {
104            revoked.insert(id.clone());
105        }
106        Ok(())
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[tokio::test]
115    async fn revoke_and_check() {
116        let store = InMemoryRevocationStore::new();
117
118        let is_rev = store.is_revoked("del-1").await;
119        assert!(is_rev.is_ok());
120        assert_eq!(is_rev.unwrap_or(true), false);
121
122        let result = store.revoke("del-1").await;
123        assert!(result.is_ok());
124
125        let is_rev = store.is_revoked("del-1").await;
126        assert!(is_rev.is_ok());
127        assert_eq!(is_rev.unwrap_or(false), true);
128    }
129
130    #[tokio::test]
131    async fn revoke_is_idempotent() {
132        let store = InMemoryRevocationStore::new();
133        let r1 = store.revoke("del-1").await;
134        assert!(r1.is_ok());
135        let r2 = store.revoke("del-1").await;
136        assert!(r2.is_ok());
137
138        let is_rev = store.is_revoked("del-1").await;
139        assert!(is_rev.is_ok());
140        assert_eq!(is_rev.unwrap_or(false), true);
141    }
142
143    #[tokio::test]
144    async fn revoke_batch_marks_all() {
145        let store = InMemoryRevocationStore::new();
146        let ids = vec!["d1".to_string(), "d2".to_string(), "d3".to_string()];
147
148        let result = store.revoke_batch(&ids).await;
149        assert!(result.is_ok());
150
151        for id in &ids {
152            let is_rev = store.is_revoked(id).await;
153            assert!(is_rev.is_ok());
154            assert_eq!(is_rev.unwrap_or(false), true);
155        }
156
157        let is_rev = store.is_revoked("d4").await;
158        assert!(is_rev.is_ok());
159        assert_eq!(is_rev.unwrap_or(true), false);
160    }
161
162    #[tokio::test]
163    async fn revoke_batch_empty_is_ok() {
164        let store = InMemoryRevocationStore::new();
165        let result = store.revoke_batch(&[]).await;
166        assert!(result.is_ok());
167    }
168
169    #[tokio::test]
170    async fn default_constructor() {
171        let store = InMemoryRevocationStore::default();
172        let is_rev = store.is_revoked("any").await;
173        assert!(is_rev.is_ok());
174        assert_eq!(is_rev.unwrap_or(true), false);
175    }
176}