Skip to main content

av_identity/
revocation.rs

1//! JTI-based token revocation.
2//!
3//! Revoked token identifiers are stored in a backend that supports TTL-based
4//! expiry. Retention must include the verifier's clock leeway, including
5//! its final accepted second.
6
7use parking_lot::RwLock;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11/// Signature-verified identity needed to check an ancestor's revocation.
12#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct RevocationIdentity {
14    /// Original issuer. Required so independently issued JTIs cannot collide.
15    pub iss: String,
16    /// Token identifier within this issuer's namespace.
17    pub jti: String,
18    /// Issuing agent instance.
19    pub instance_uid: String,
20    /// Original issued-at time, retained across token exchange.
21    pub iat: u64,
22}
23
24impl From<&crate::NhiClaims> for RevocationIdentity {
25    fn from(claims: &crate::NhiClaims) -> Self {
26        Self {
27            iss: claims.iss.clone(),
28            jti: claims.jti.clone(),
29            instance_uid: claims.instance_uid.clone(),
30            iat: claims.iat,
31        }
32    }
33}
34
35/// The rule which revoked a validated token.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum RevokedBy {
38    /// Its unique token identifier was revoked.
39    Jti,
40    /// Its instance was revoked through this issued-at cutoff (inclusive).
41    Instance {
42        /// Last refused issued-at second.
43        cutoff: u64,
44    },
45}
46
47/// Whether a revoke request created a new entry.
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum RevokeOutcome {
50    /// This request first recorded the revocation.
51    Revoked,
52    /// The token was already revoked; retention was extended if necessary.
53    AlreadyRevoked,
54}
55
56/// Local, nonblocking health of a revocation dependency.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct RevocationHealth {
59    /// Whether dependency reads are currently admitted by the circuit breaker.
60    pub available: bool,
61    /// Number of cached entries known to be revoked; never live-token entries.
62    pub local_entries: usize,
63    /// Elapsed monotonic seconds since the last successful dependency call.
64    pub seconds_since_last_success: Option<u64>,
65}
66
67/// Trait for pluggable revocation backends. The harness wires this to
68/// the `StateStore` (Redis for distributed deployments, in-memory otherwise).
69pub trait RevocationStore: Send + Sync {
70    /// Read local health without contacting storage. `None` means the backend
71    /// does not expose a fallible dependency or availability guard.
72    fn health(&self) -> Option<RevocationHealth> {
73        None
74    }
75
76    /// Mark a JTI as revoked across issuers for an authorized operator. Holder
77    /// requests must use `try_revoke_token`. `expires_at` is the epoch second at which the
78    /// entry is last required (the token's `exp` plus verifier leeway). An error
79    /// means the caller cannot confirm that the write succeeded.
80    fn revoke(&self, jti: &str, expires_at: u64) -> Result<(), String>;
81
82    /// `Ok(true)` if an operator-wide JTI rule is present. Identity validation
83    /// must use `check_token` to include issuer-scoped rules. An error means the store
84    /// could not answer; the validator then refuses the token.
85    fn is_revoked(&self, jti: &str) -> Result<bool, String>;
86
87    /// Reclaim entries past their retention period when the backend needs an
88    /// explicit sweep. Distributed stores normally use their native TTL.
89    fn sweep_expired(&self, _now_s: u64) {}
90
91    /// Atomically record a revocation and report whether it is new.
92    /// Backends overriding this method can deduplicate audit records.
93    fn try_revoke(&self, jti: &str, expires_at: u64) -> Result<RevokeOutcome, String> {
94        self.revoke(jti, expires_at)?;
95        Ok(RevokeOutcome::Revoked)
96    }
97
98    /// Revoke tokens of an instance issued at or before `issued_at_or_before`.
99    /// Repeated writes must never lower either the cutoff or retention time.
100    fn revoke_instance(
101        &self,
102        _instance_uid: &str,
103        _issued_at_or_before: u64,
104        _expires_at: u64,
105    ) -> Result<(), String> {
106        Err("instance revocation is not supported by this backend".into())
107    }
108
109    /// Revoke a holder-authenticated token only within its verified issuer.
110    /// This must use a separate key domain from operator-wide JTI rules.
111    /// Backends without scoped storage fail closed instead of widening the
112    /// holder's authority to every issuer.
113    fn try_revoke_token(&self, _issuer: &str, _jti: &str, _expires_at: u64) -> Result<RevokeOutcome, String> {
114        Err("issuer-scoped token revocation is not supported by this backend".into())
115    }
116
117    /// Check an issuer-scoped holder rule and the explicit operator-wide rules.
118    /// The default preserves global checks for older implementations, which
119    /// cannot accept scoped writes until they override both methods.
120    fn check_token(
121        &self,
122        _issuer: &str,
123        jti: &str,
124        instance_uid: &str,
125        iat: u64,
126    ) -> Result<Option<RevokedBy>, String> {
127        self.check(jti, instance_uid, iat)
128    }
129
130    /// Check both the operator-wide token identifier and instance cutoff.
131    fn check(&self, jti: &str, _instance_uid: &str, _iat: u64) -> Result<Option<RevokedBy>, String> {
132        Ok(self.is_revoked(jti)?.then_some(RevokedBy::Jti))
133    }
134}
135
136/// In-memory revocation store with lazy expiry sweeps. Suitable for
137/// single-instance deployments and testing.
138pub struct InMemoryRevocationStore {
139    entries: RwLock<HashMap<String, u64>>,
140    issued_entries: RwLock<HashMap<(String, String), u64>>,
141    instances: RwLock<HashMap<String, (u64, u64)>>,
142}
143
144impl InMemoryRevocationStore {
145    /// Create an empty store.
146    pub fn new() -> Self {
147        Self {
148            entries: RwLock::new(HashMap::new()),
149            issued_entries: RwLock::new(HashMap::new()),
150            instances: RwLock::new(HashMap::new()),
151        }
152    }
153
154    /// Remove entries whose `expires_at` is in the past.
155    pub fn sweep(&self, now_s: u64) {
156        self.entries.write().retain(|_, exp| *exp >= now_s);
157        self.issued_entries.write().retain(|_, exp| *exp >= now_s);
158        self.instances.write().retain(|_, (_, exp)| *exp >= now_s);
159    }
160
161    /// Number of entries currently stored (for metrics/testing).
162    pub fn len(&self) -> usize {
163        self.entries.read().len() + self.issued_entries.read().len() + self.instances.read().len()
164    }
165
166    /// True when the store is empty.
167    pub fn is_empty(&self) -> bool {
168        self.entries.read().is_empty()
169            && self.issued_entries.read().is_empty()
170            && self.instances.read().is_empty()
171    }
172}
173
174impl Default for InMemoryRevocationStore {
175    fn default() -> Self {
176        Self::new()
177    }
178}
179
180impl RevocationStore for InMemoryRevocationStore {
181    fn sweep_expired(&self, now_s: u64) {
182        self.sweep(now_s);
183    }
184
185    fn revoke(&self, jti: &str, expires_at: u64) -> Result<(), String> {
186        self.try_revoke(jti, expires_at).map(|_| ())
187    }
188
189    fn try_revoke(&self, jti: &str, expires_at: u64) -> Result<RevokeOutcome, String> {
190        let mut entries = self.entries.write();
191        match entries.entry(jti.to_owned()) {
192            std::collections::hash_map::Entry::Vacant(entry) => {
193                entry.insert(expires_at);
194                Ok(RevokeOutcome::Revoked)
195            }
196            std::collections::hash_map::Entry::Occupied(mut entry) => {
197                *entry.get_mut() = (*entry.get()).max(expires_at);
198                Ok(RevokeOutcome::AlreadyRevoked)
199            }
200        }
201    }
202
203    fn is_revoked(&self, jti: &str) -> Result<bool, String> {
204        Ok(self.entries.read().contains_key(jti))
205    }
206
207    fn try_revoke_token(&self, issuer: &str, jti: &str, expires_at: u64) -> Result<RevokeOutcome, String> {
208        let mut entries = self.issued_entries.write();
209        match entries.entry((issuer.to_owned(), jti.to_owned())) {
210            std::collections::hash_map::Entry::Vacant(entry) => {
211                entry.insert(expires_at);
212                Ok(RevokeOutcome::Revoked)
213            }
214            std::collections::hash_map::Entry::Occupied(mut entry) => {
215                *entry.get_mut() = (*entry.get()).max(expires_at);
216                Ok(RevokeOutcome::AlreadyRevoked)
217            }
218        }
219    }
220
221    fn check_token(&self, issuer: &str, jti: &str, uid: &str, iat: u64) -> Result<Option<RevokedBy>, String> {
222        if self
223            .issued_entries
224            .read()
225            .contains_key(&(issuer.to_owned(), jti.to_owned()))
226        {
227            return Ok(Some(RevokedBy::Jti));
228        }
229        self.check(jti, uid, iat)
230    }
231
232    fn revoke_instance(&self, instance_uid: &str, cutoff: u64, expires_at: u64) -> Result<(), String> {
233        let mut entries = self.instances.write();
234        let entry = entries
235            .entry(instance_uid.to_owned())
236            .or_insert((cutoff, expires_at));
237        entry.0 = entry.0.max(cutoff);
238        entry.1 = entry.1.max(expires_at);
239        Ok(())
240    }
241
242    fn check(&self, jti: &str, instance_uid: &str, iat: u64) -> Result<Option<RevokedBy>, String> {
243        if self.is_revoked(jti)? {
244            return Ok(Some(RevokedBy::Jti));
245        }
246        Ok(self
247            .instances
248            .read()
249            .get(instance_uid)
250            .filter(|(cutoff, _)| iat <= *cutoff)
251            .map(|(cutoff, _)| RevokedBy::Instance { cutoff: *cutoff }))
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    #![allow(clippy::unwrap_used)]
258
259    use super::*;
260
261    #[test]
262    fn revoke_and_query() {
263        let store = InMemoryRevocationStore::new();
264        assert!(!store.is_revoked("jti-1").unwrap());
265        store.revoke("jti-1", 1000).unwrap();
266        assert!(store.is_revoked("jti-1").unwrap());
267        assert!(!store.is_revoked("jti-2").unwrap());
268    }
269
270    #[test]
271    fn sweep_removes_expired() {
272        let store = InMemoryRevocationStore::new();
273        store.revoke("old", 100).unwrap();
274        store.revoke("fresh", 2000).unwrap();
275        assert_eq!(store.len(), 2);
276        store.sweep(500);
277        assert_eq!(store.len(), 1);
278        assert!(!store.is_revoked("old").unwrap());
279        assert!(store.is_revoked("fresh").unwrap());
280    }
281
282    #[test]
283    fn re_revoking_never_shortens_the_retention_window() {
284        let store = InMemoryRevocationStore::new();
285        store.revoke("jti", 2000).unwrap();
286        store.revoke("jti", 100).unwrap();
287        store.sweep(500);
288        assert!(
289            store.is_revoked("jti").unwrap(),
290            "a second revoke with an earlier expiry must not let the sweep drop the entry early"
291        );
292    }
293
294    #[test]
295    fn empty_store() {
296        let store = InMemoryRevocationStore::new();
297        assert!(store.is_empty());
298        store.revoke("x", 100).unwrap();
299        assert!(!store.is_empty());
300    }
301}