1use std::collections::HashSet;
9use std::sync::RwLock;
10
11use async_trait::async_trait;
12
13use openagent_aegis_core::DelegationError;
14
15#[async_trait]
28pub trait RevocationStore: Send + Sync {
29 async fn revoke(&self, delegation_id: &str) -> Result<(), DelegationError>;
34
35 async fn is_revoked(&self, delegation_id: &str) -> Result<bool, DelegationError>;
37
38 async fn revoke_batch(&self, ids: &[String]) -> Result<(), DelegationError>;
43}
44
45pub struct InMemoryRevocationStore {
55 revoked: RwLock<HashSet<String>>,
56}
57
58impl InMemoryRevocationStore {
59 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}