1use std::borrow::Cow;
67use std::collections::{BTreeMap, HashMap};
68use std::sync::{Arc, Mutex, PoisonError};
69use std::time::{SystemTime, UNIX_EPOCH};
70
71use async_trait::async_trait;
72
73use cheers_core::{DeviceId, Error, PrincipalId, RevocationKey, RevocationSet, Revoked, StoreError};
74use cheers_verify::{revocation_entry, RevocationReader};
75
76use crate::codec::PasetoV4SecretMinter;
77use crate::ownership::OwnershipStore;
78
79#[derive(Debug, Clone, PartialEq, Eq)]
81pub struct RevocationSnapshot {
82 pub epoch: u64,
83 pub revoked: Vec<Revoked>,
84}
85
86pub fn next_epoch(current: u64, now: i64) -> u64 {
88 current.saturating_add(1).max(u64::try_from(now).unwrap_or(0))
89}
90
91#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct RevokedColumns<'a> {
99 pub kind: &'static str,
100 pub subject: Cow<'a, str>,
101 pub resource_kind: &'a str,
102 pub resource_id: &'a str,
103 pub bound: Option<i64>,
104 pub expires_at: Option<i64>,
105}
106
107pub fn revoked_columns(entry: &Revoked) -> RevokedColumns<'_> {
110 let clamp = |b: u64| i64::try_from(b).unwrap_or(i64::MAX);
111 match entry {
112 Revoked::Jti { jti, exp } => RevokedColumns {
113 kind: "jti",
114 subject: Cow::Borrowed(jti),
115 resource_kind: "",
116 resource_id: "",
117 bound: None,
118 expires_at: *exp,
119 },
120 Revoked::Device { device, at_seq } => RevokedColumns {
121 kind: "device",
122 subject: Cow::Borrowed(device.as_str()),
123 resource_kind: "",
124 resource_id: "",
125 bound: Some(clamp(*at_seq)),
126 expires_at: None,
127 },
128 Revoked::Membership { kind, id, principal, at_epoch } => RevokedColumns {
129 kind: "membership",
130 subject: Cow::Owned(principal.to_string()),
131 resource_kind: kind,
132 resource_id: id,
133 bound: Some(clamp(*at_epoch)),
134 expires_at: None,
135 },
136 }
137}
138
139pub fn revoked_from_columns(
141 kind: &str,
142 subject: String,
143 resource_kind: String,
144 resource_id: String,
145 bound: Option<i64>,
146 expires_at: Option<i64>,
147) -> Result<Revoked, StoreError> {
148 let bound = || -> Result<u64, StoreError> {
149 let b = bound.ok_or_else(|| StoreError::Backend(format!("{kind} revocation {subject:?} has no bound")))?;
150 u64::try_from(b).map_err(|_| StoreError::Backend(format!("negative revocation bound {b}")))
151 };
152 match kind {
153 "jti" => Ok(Revoked::jti(subject, expires_at)),
154 "device" => Ok(Revoked::device(DeviceId::new(subject.clone()), bound()?)),
155 "membership" => {
156 let principal = subject
157 .parse::<PrincipalId>()
158 .map_err(|e| StoreError::Backend(format!("membership revocation subject {subject:?}: {e}")))?;
159 Ok(Revoked::membership(resource_kind, resource_id, principal, bound()?))
160 }
161 other => Err(StoreError::Backend(format!("unknown revocation kind {other:?}"))),
162 }
163}
164
165pub fn epoch_from_sql(epoch: i64) -> Result<u64, StoreError> {
167 u64::try_from(epoch).map_err(|_| StoreError::Backend(format!("negative revocation epoch {epoch}")))
168}
169
170#[async_trait]
174pub trait RevocationWriter: Send + Sync {
175 async fn revoke(&self, entry: &Revoked) -> Result<(), StoreError>;
183
184 async fn snapshot(&self) -> Result<RevocationSnapshot, StoreError>;
189}
190
191#[derive(Debug, Clone, PartialEq, Eq)]
193pub struct SignedRevocationSet {
194 pub epoch: u64,
195 pub token: String,
197}
198
199pub struct RevocationPublisher<W> {
206 store: W,
207 keys: Arc<dyn OwnershipStore>,
208 minter: PasetoV4SecretMinter,
209 issuer: String,
210 kid: String,
211}
212
213impl<W> std::fmt::Debug for RevocationPublisher<W> {
214 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
215 f.debug_struct("RevocationPublisher")
216 .field("issuer", &self.issuer)
217 .field("kid", &self.kid)
218 .finish_non_exhaustive()
219 }
220}
221
222impl<W: RevocationWriter> RevocationPublisher<W> {
223 pub fn new(
224 store: W,
225 keys: Arc<dyn OwnershipStore>,
226 minter: PasetoV4SecretMinter,
227 issuer: impl Into<String>,
228 kid: impl Into<String>,
229 ) -> Self {
230 Self {
231 store,
232 keys,
233 minter,
234 issuer: issuer.into(),
235 kid: kid.into(),
236 }
237 }
238
239 pub fn issuer(&self) -> &str {
240 &self.issuer
241 }
242
243 pub fn store(&self) -> &W {
244 &self.store
245 }
246
247 pub async fn current(&self, now: i64) -> Result<SignedRevocationSet, Error> {
255 let mut snapshot = self.store.snapshot().await?;
256 snapshot.revoked.retain(|entry| !entry.is_lapsed_at(now));
257 let mut keys: BTreeMap<(String, String), RevocationKey> = BTreeMap::new();
258 for entry in &snapshot.revoked {
259 if let Revoked::Membership { kind, id, .. } = entry {
260 if !keys.contains_key(&(kind.clone(), id.clone())) {
261 let key = self.keys.revocation_key(kind, id).await?;
262 keys.insert((kind.clone(), id.clone()), key);
263 }
264 }
265 }
266 let entries = snapshot
267 .revoked
268 .iter()
269 .map(|entry| revocation_entry(entry, |kind, id| keys[&(kind.to_owned(), id.to_owned())].clone()))
270 .collect();
271 let set = RevocationSet::new(self.issuer.clone(), snapshot.epoch, entries);
272 let token = self.minter.mint_artifact(&set, &self.kid)?;
273 Ok(SignedRevocationSet {
274 epoch: set.epoch,
275 token,
276 })
277 }
278}
279
280#[derive(Debug, Clone, Default)]
284pub struct MemoryRevocationStore(Arc<Mutex<MemoryLog>>);
285
286type IdentityKey = (u8, String, String, String);
288
289fn identity_key(entry: &Revoked) -> IdentityKey {
290 let (v, a, b, c) = entry.identity();
291 (v, a.to_owned(), b.to_owned(), c.into_owned())
292}
293
294#[derive(Debug, Default)]
295struct MemoryLog {
296 epoch: u64,
297 entries: HashMap<IdentityKey, Revoked>,
299}
300
301impl MemoryRevocationStore {
302 pub fn new() -> Self {
303 Self::default()
304 }
305
306 fn get(&self, probe: &Revoked) -> Option<Revoked> {
308 let log = self.0.lock().unwrap_or_else(PoisonError::into_inner);
309 log.entries.get(&identity_key(probe)).cloned()
310 }
311
312 pub fn gc(&self, now: i64) -> u64 {
315 let mut log = self.0.lock().unwrap_or_else(PoisonError::into_inner);
316 let before = log.entries.len();
317 log.entries.retain(|_, entry| !entry.is_lapsed_at(now));
318 let removed = (before - log.entries.len()) as u64;
319 if removed > 0 {
320 log.epoch = next_epoch(log.epoch, now_unix());
321 }
322 removed
323 }
324}
325
326#[async_trait]
327impl RevocationWriter for MemoryRevocationStore {
328 async fn revoke(&self, entry: &Revoked) -> Result<(), StoreError> {
329 let mut log = self.0.lock().unwrap_or_else(PoisonError::into_inner);
330 let key = identity_key(entry);
331 let raises = log
332 .entries
333 .get(&key)
334 .is_none_or(|held| entry.cmp_bound(held) == std::cmp::Ordering::Greater);
335 if raises {
336 log.entries.insert(key, entry.clone());
337 log.epoch = next_epoch(log.epoch, now_unix());
338 }
339 Ok(())
340 }
341
342 async fn snapshot(&self) -> Result<RevocationSnapshot, StoreError> {
343 let log = self.0.lock().unwrap_or_else(PoisonError::into_inner);
344 Ok(RevocationSnapshot {
345 epoch: log.epoch,
346 revoked: log.entries.values().cloned().collect(),
347 })
348 }
349}
350
351#[async_trait]
352impl RevocationReader for MemoryRevocationStore {
353 async fn is_revoked(&self, jti: &str) -> Result<bool, StoreError> {
354 Ok(self.get(&Revoked::jti(jti, None)).is_some())
355 }
356
357 async fn is_device_revoked(&self, device: &DeviceId, seq: u64) -> Result<bool, StoreError> {
358 Ok(matches!(
359 self.get(&Revoked::device(device.clone(), 0)),
360 Some(Revoked::Device { at_seq, .. }) if seq < at_seq
361 ))
362 }
363
364 async fn is_membership_revoked(
366 &self,
367 _key: &RevocationKey,
368 kind: &str,
369 id: &str,
370 principal: &PrincipalId,
371 snapshot_epoch: u64,
372 ) -> Result<bool, StoreError> {
373 Ok(matches!(
374 self.get(&Revoked::membership(kind, id, principal.clone(), 0)),
375 Some(Revoked::Membership { at_epoch, .. }) if snapshot_epoch < at_epoch
376 ))
377 }
378}
379
380pub(crate) fn now_unix() -> i64 {
381 SystemTime::now()
382 .duration_since(UNIX_EPOCH)
383 .map(|d| d.as_secs() as i64)
384 .unwrap_or(0)
385}
386
387#[cfg(test)]
388mod tests {
389 use super::*;
390 use cheers_core::{KeyRole, RevocationEntry, SignedArtifact};
391 use cheers_verify::{
392 AdoptError, Adopted, IssuerTrust, JwkKey, JwksDoc, KeySet, KeySetVerifier, ReplicatedRevocations,
393 };
394
395 const ISS: &str = "https://cheers.test";
396
397 #[test]
398 fn next_epoch_is_strictly_monotonic_with_a_clock_floor() {
399 assert_eq!(next_epoch(0, 1_000), 1_000);
400 assert_eq!(next_epoch(1_000, 1_000), 1_001);
402 assert_eq!(next_epoch(5_000, 1_000), 5_001);
403 assert_eq!(next_epoch(5_000, 10), 5_001);
405 assert_eq!(next_epoch(5_000, -1), 5_001);
406 assert_eq!(next_epoch(u64::MAX, 0), u64::MAX);
407 }
408
409 #[test]
410 fn revoke_keeps_the_greatest_bound_and_only_raises_advance_the_epoch() {
411 let store = MemoryRevocationStore::default();
412 pollster::block_on(async {
413 assert_eq!(store.snapshot().await.unwrap().epoch, 0);
414 store.revoke(&Revoked::jti("tok-1", Some(100))).await.unwrap();
415 let e1 = store.snapshot().await.unwrap().epoch;
416 assert!(e1 > 0);
417 store.revoke(&Revoked::jti("tok-1", Some(100))).await.unwrap();
419 store.revoke(&Revoked::jti("tok-1", Some(50))).await.unwrap();
420 assert_eq!(store.snapshot().await.unwrap().epoch, e1);
421 store.revoke(&Revoked::jti("tok-1", None)).await.unwrap();
423 let e2 = store.snapshot().await.unwrap().epoch;
424 assert!(e2 > e1);
425 store.revoke(&Revoked::jti("tok-1", Some(i64::MAX))).await.unwrap();
426 assert_eq!(store.snapshot().await.unwrap().epoch, e2);
427
428 store.revoke(&Revoked::device("phone", 10)).await.unwrap();
429 let e3 = store.snapshot().await.unwrap().epoch;
430 store.revoke(&Revoked::device("phone", 9)).await.unwrap();
431 assert_eq!(store.snapshot().await.unwrap().epoch, e3);
432 store.revoke(&Revoked::device("phone", 20)).await.unwrap();
433 assert!(store.snapshot().await.unwrap().epoch > e3);
434 let phone = DeviceId::new("phone");
435 assert!(store.is_device_revoked(&phone, 19).await.unwrap());
436 assert!(!store.is_device_revoked(&phone, 20).await.unwrap());
437
438 store.revoke(&Revoked::membership("namespace", "ns-1", PrincipalId::user("alice"), 7)).await.unwrap();
439 let alice = PrincipalId::user("alice");
440 let k = cheers_verify::test_revocation_key("namespace", "ns-1");
441 assert!(store.is_membership_revoked(&k, "namespace", "ns-1", &alice, 6).await.unwrap());
442 assert!(!store.is_membership_revoked(&k, "namespace", "ns-1", &alice, 7).await.unwrap());
443 assert!(!store.is_membership_revoked(&k, "namespace", "ns-1", &PrincipalId::user("bob"), 0).await.unwrap());
444
445 let snap = store.snapshot().await.unwrap();
446 assert_eq!(snap.revoked.len(), 3);
447 assert!(store.is_revoked("tok-1").await.unwrap());
448 });
449 }
450
451 #[test]
452 fn gc_drops_lapsed_jtis_only_and_advances_the_epoch_iff_it_removed_something() {
453 let store = MemoryRevocationStore::default();
454 pollster::block_on(async {
455 store.revoke(&Revoked::jti("short", Some(50))).await.unwrap();
456 store.revoke(&Revoked::jti("forever", None)).await.unwrap();
457 store.revoke(&Revoked::device("phone", 1)).await.unwrap();
458 let before = store.snapshot().await.unwrap().epoch;
459 assert_eq!(store.gc(49), 0);
460 assert_eq!(store.snapshot().await.unwrap().epoch, before);
461 assert_eq!(store.gc(50), 1);
462 let after = store.snapshot().await.unwrap();
463 assert!(after.epoch > before);
464 assert_eq!(after.revoked.len(), 2);
465 assert!(!store.is_revoked("short").await.unwrap());
466 assert!(store.is_revoked("forever").await.unwrap());
467 });
468 }
469
470 #[test]
471 fn the_published_set_omits_jtis_lapsed_at_mint_time() {
472 let (publisher, keys) = publisher();
473 pollster::block_on(async {
474 publisher.store().revoke(&Revoked::jti("short", Some(50))).await.unwrap();
475 publisher.store().revoke(&Revoked::jti("forever", None)).await.unwrap();
476 let names = |token: String| {
477 let keys = &keys;
478 async move {
479 let set: RevocationSet = keys.verify_artifact(&token).await.unwrap();
480 set.revoked
481 }
482 };
483 assert_eq!(
484 names(publisher.current(49).await.unwrap().token).await,
485 vec![RevocationEntry::Jti { jti: "forever".into(), exp: None }, RevocationEntry::Jti { jti: "short".into(), exp: Some(50) }]
486 );
487 assert_eq!(
488 names(publisher.current(50).await.unwrap().token).await,
489 vec![RevocationEntry::Jti { jti: "forever".into(), exp: None }]
490 );
491 });
492 }
493
494 fn publisher() -> (RevocationPublisher<MemoryRevocationStore>, KeySetVerifier) {
495 let (minter, verifier) = PasetoV4SecretMinter::generate().unwrap();
496 let public: [u8; 32] = verifier.public_key().as_bytes().try_into().unwrap();
497 let keys = KeySetVerifier::from_issuer_key("iss-1", &public, ISS);
498 let ownership: Arc<dyn OwnershipStore> = Arc::new(crate::MemoryOwnershipStore::new());
499 (RevocationPublisher::new(MemoryRevocationStore::default(), ownership, minter, ISS, "iss-1"), keys)
500 }
501
502 #[test]
503 fn a_published_membership_is_a_tag_a_key_holder_can_test() {
504 let (minter, verifier) = PasetoV4SecretMinter::generate().unwrap();
505 let public: [u8; 32] = verifier.public_key().as_bytes().try_into().unwrap();
506 let keys = KeySetVerifier::from_issuer_key("iss-1", &public, ISS);
507 let ownership = crate::MemoryOwnershipStore::new();
508 let publisher = RevocationPublisher::new(
509 MemoryRevocationStore::default(),
510 Arc::new(ownership.clone()),
511 minter,
512 ISS,
513 "iss-1",
514 );
515 let replica = ReplicatedRevocations::new(IssuerTrust::key_set(ISS, keys.clone()));
516 pollster::block_on(async {
517 publisher.store().revoke(&Revoked::membership("room", "r-secret", PrincipalId::user("alice-secret"), 9)).await.unwrap();
518 let signed = publisher.current(0).await.unwrap();
519 let set: RevocationSet = keys.verify_artifact(&signed.token).await.unwrap();
520 let wire = serde_json::to_string(&set).unwrap();
521 assert!(!wire.contains("r-secret") && !wire.contains("alice-secret"), "{wire}");
522 replica.adopt(&signed.token).await.unwrap();
523 let key = ownership.revocation_key("room", "r-secret").await.unwrap();
524 assert_eq!(ownership.revocation_key("room", "r-secret").await.unwrap(), key);
526 let alice = PrincipalId::user("alice-secret");
527 assert!(replica.is_membership_revoked(&key, "room", "r-secret", &alice, 8).await.unwrap());
528 assert!(!replica.is_membership_revoked(&key, "room", "r-secret", &alice, 9).await.unwrap());
529 let other = ownership.revocation_key("room", "r-other").await.unwrap();
530 assert_ne!(other, key);
531 assert!(!replica.is_membership_revoked(&other, "room", "r-secret", &alice, 8).await.unwrap());
532 });
533 }
534
535 #[test]
536 fn published_sets_verify_at_a_replica_in_epoch_order() {
537 let (publisher, keys) = publisher();
538 let replica = ReplicatedRevocations::new(IssuerTrust::key_set(ISS, keys));
539 pollster::block_on(async {
540 let empty = publisher.current(0).await.unwrap();
541 assert_eq!(empty.epoch, 0);
542 assert_eq!(replica.adopt(&empty.token).await.unwrap(), Adopted::Advanced { from: None, to: 0 });
543
544 publisher.store().revoke(&Revoked::device("phone", 1)).await.unwrap();
545 let first = publisher.current(0).await.unwrap();
546 assert!(first.epoch > empty.epoch);
547 assert_eq!(publisher.current(0).await.unwrap(), first);
549 replica.adopt(&first.token).await.unwrap();
550 assert!(replica.is_device_revoked(&DeviceId::new("phone"), 0).await.unwrap());
551
552 publisher.store().revoke(&Revoked::jti("j1", None)).await.unwrap();
553 let second = publisher.current(0).await.unwrap();
554 replica.adopt(&second.token).await.unwrap();
555 assert!(replica.is_revoked("j1").await.unwrap());
556 assert!(matches!(replica.adopt(&first.token).await, Err(AdoptError::Stale { .. })));
558 });
559 }
560
561 #[test]
562 fn mint_artifact_is_domain_separated_from_access_tokens() {
563 let (minter, verifier) = PasetoV4SecretMinter::generate().unwrap();
564 let public: [u8; 32] = verifier.public_key().as_bytes().try_into().unwrap();
565 let keys = KeySetVerifier::from_issuer_key("iss-1", &public, ISS);
566 let set = RevocationSet::new(ISS, 1, vec![RevocationEntry::Jti { jti: "j1".into(), exp: None }]);
567 let token = minter.mint_artifact(&set, "iss-1").unwrap();
568 assert_eq!(pollster::block_on(keys.verify_artifact::<RevocationSet>(&token)).unwrap(), set);
569 assert_eq!(verifier.verify_artifact::<RevocationSet>(&token).unwrap(), set);
570
571 assert!(cheers_core::TokenVerifier::verify_at(&verifier, &token, 0).is_err());
573 assert!(verifier.verify_mcp_at(&token, 0, "iss-1").is_err());
574 assert!(pollster::block_on(keys.verify_mcp(&token, 0, ISS, None)).is_err());
575 let mcp = cheers_core::McpClaims::new(ISS, "aud", cheers_core::PrincipalId::user("alice"), 1, 2, "j", vec![]);
577 let access = minter.mint_mcp(&mcp, "iss-1").unwrap();
578 assert!(verifier.verify_artifact::<RevocationSet>(&access).is_err());
579 assert!(pollster::block_on(keys.verify_artifact::<RevocationSet>(&access)).is_err());
580 }
581
582 #[test]
583 fn a_set_signed_by_a_non_issuer_role_key_is_refused() {
584 let (minter, verifier) = PasetoV4SecretMinter::generate().unwrap();
585 let public: [u8; 32] = verifier.public_key().as_bytes().try_into().unwrap();
586 let doc = JwksDoc {
587 keys: vec![JwkKey::ed25519("svc-1", &public, ISS, KeyRole::Assertion)],
588 };
589 let keys = KeySetVerifier::from_key_set(KeySet::from_doc(doc).unwrap());
590 let token = minter.mint_artifact(&RevocationSet::new(ISS, 1, vec![]), "svc-1").unwrap();
591 assert!(pollster::block_on(keys.verify_artifact::<RevocationSet>(&token)).is_err());
592 assert!(!RevocationSet::IMPLICIT_ASSERTION.is_empty());
593 }
594
595 #[test]
596 fn trait_is_dyn_compatible() {
597 fn _writer(_: &dyn RevocationWriter) {}
598 }
599}