1use parking_lot::RwLock;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
13pub struct RevocationIdentity {
14 pub iss: String,
16 pub jti: String,
18 pub instance_uid: String,
20 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum RevokedBy {
38 Jti,
40 Instance {
42 cutoff: u64,
44 },
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum RevokeOutcome {
50 Revoked,
52 AlreadyRevoked,
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct RevocationHealth {
59 pub available: bool,
61 pub local_entries: usize,
63 pub seconds_since_last_success: Option<u64>,
65}
66
67pub trait RevocationStore: Send + Sync {
70 fn health(&self) -> Option<RevocationHealth> {
73 None
74 }
75
76 fn revoke(&self, jti: &str, expires_at: u64) -> Result<(), String>;
81
82 fn is_revoked(&self, jti: &str) -> Result<bool, String>;
86
87 fn sweep_expired(&self, _now_s: u64) {}
90
91 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 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 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 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 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
136pub 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 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 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 pub fn len(&self) -> usize {
163 self.entries.read().len() + self.issued_entries.read().len() + self.instances.read().len()
164 }
165
166 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}