1use std::collections::HashMap;
40use std::sync::{Arc, PoisonError, RwLock};
41
42use async_trait::async_trait;
43
44use cheers_core::{DeviceId, MembershipTag, PrincipalId, RevocationEntry, RevocationKey, RevocationSet, Revoked, StoreError};
45use hmac::{Hmac, Mac};
46use sha2::Sha256;
47
48use crate::artifact::IssuerTrust;
49use crate::key_set::VerifyError;
50
51const MEMBERSHIP_TAG_CONTEXT: &[u8] = b"urn:cheers:revocation:membership-tag:v2";
53
54pub fn membership_tag(key: &RevocationKey, kind: &str, id: &str, principal: &PrincipalId) -> MembershipTag {
62 let mut mac = <Hmac<Sha256> as Mac>::new_from_slice(key.as_bytes()).expect("HMAC takes any key length");
63 mac.update(MEMBERSHIP_TAG_CONTEXT);
64 let principal = principal.to_string();
65 for part in [kind, id, principal.as_str()] {
66 mac.update(&(part.len() as u64).to_be_bytes());
67 mac.update(part.as_bytes());
68 }
69 MembershipTag::from_bytes(mac.finalize().into_bytes().into())
70}
71
72pub fn revocation_entry(entry: &Revoked, key_of: impl FnOnce(&str, &str) -> RevocationKey) -> RevocationEntry {
76 match entry {
77 Revoked::Jti { jti, exp } => RevocationEntry::Jti { jti: jti.clone(), exp: *exp },
78 Revoked::Device { device, at_seq } => RevocationEntry::Device { device: device.clone(), at_seq: *at_seq },
79 Revoked::Membership { kind, id, principal, at_epoch } => RevocationEntry::Membership {
80 tag: membership_tag(&key_of(kind, id), kind, id, principal),
81 at_epoch: *at_epoch,
82 },
83 }
84}
85
86#[doc(hidden)]
88pub fn test_revocation_key(kind: &str, id: &str) -> RevocationKey {
89 use sha2::Digest as _;
90 RevocationKey::from_bytes(Sha256::digest(format!("test-key/{kind}/{id}")).into())
91}
92
93#[async_trait]
102pub trait RevocationReader: Send + Sync {
103 async fn is_revoked(&self, jti: &str) -> Result<bool, StoreError>;
110
111 async fn is_device_revoked(&self, device: &DeviceId, seq: u64) -> Result<bool, StoreError>;
116
117 async fn is_membership_revoked(
127 &self,
128 key: &RevocationKey,
129 kind: &str,
130 id: &str,
131 principal: &PrincipalId,
132 snapshot_epoch: u64,
133 ) -> Result<bool, StoreError>;
134}
135
136#[async_trait]
138impl<T: RevocationReader + ?Sized> RevocationReader for Arc<T> {
139 async fn is_revoked(&self, jti: &str) -> Result<bool, StoreError> {
140 (**self).is_revoked(jti).await
141 }
142
143 async fn is_device_revoked(&self, device: &DeviceId, seq: u64) -> Result<bool, StoreError> {
144 (**self).is_device_revoked(device, seq).await
145 }
146
147 async fn is_membership_revoked(
148 &self,
149 key: &RevocationKey,
150 kind: &str,
151 id: &str,
152 principal: &PrincipalId,
153 snapshot_epoch: u64,
154 ) -> Result<bool, StoreError> {
155 (**self).is_membership_revoked(key, kind, id, principal, snapshot_epoch).await
156 }
157}
158
159#[derive(Debug, Clone, Copy, PartialEq, Eq)]
161pub enum Adopted {
162 Advanced { from: Option<u64>, to: u64 },
164 Unchanged { epoch: u64 },
167}
168
169#[derive(Debug, thiserror::Error)]
171pub enum AdoptError {
172 #[error("revocation set refused: {0}")]
175 Verify(#[from] VerifyError),
176
177 #[error("stale revocation set: epoch {offered} is below the held epoch {held}")]
180 Stale { held: u64, offered: u64 },
181}
182
183struct Held {
185 epoch: u64,
186 token: String,
187 jtis: HashMap<String, Option<i64>>,
188 devices: HashMap<DeviceId, u64>,
189 memberships: HashMap<MembershipTag, u64>,
190}
191
192impl Held {
193 fn new(set: RevocationSet, token: &str) -> Self {
194 let mut held = Self {
195 epoch: set.epoch,
196 token: token.to_owned(),
197 jtis: HashMap::new(),
198 devices: HashMap::new(),
199 memberships: HashMap::new(),
200 };
201 for entry in set.revoked {
204 match entry {
205 RevocationEntry::Jti { jti, exp } => {
206 let kept = held.jtis.entry(jti).or_insert(exp);
207 *kept = match (*kept, exp) {
208 (Some(a), Some(b)) => Some(a.max(b)),
209 _ => None,
210 };
211 }
212 RevocationEntry::Device { device, at_seq } => {
213 let kept = held.devices.entry(device).or_insert(at_seq);
214 *kept = (*kept).max(at_seq);
215 }
216 RevocationEntry::Membership { tag, at_epoch } => {
217 let kept = held.memberships.entry(tag).or_insert(at_epoch);
218 *kept = (*kept).max(at_epoch);
219 }
220 }
221 }
222 held
223 }
224}
225
226pub struct ReplicatedRevocations {
239 trust: IssuerTrust,
240 held: RwLock<Option<Arc<Held>>>,
241}
242
243impl std::fmt::Debug for ReplicatedRevocations {
244 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245 f.debug_struct("ReplicatedRevocations")
246 .field("trust", &self.trust)
247 .field("epoch", &self.epoch())
248 .finish()
249 }
250}
251
252impl ReplicatedRevocations {
253 pub fn new(trust: IssuerTrust) -> Self {
255 Self {
256 trust,
257 held: RwLock::new(None),
258 }
259 }
260
261 pub fn issuer(&self) -> &str {
262 self.trust.issuer()
263 }
264
265 pub async fn adopt(&self, token: &str) -> Result<Adopted, AdoptError> {
272 let set: RevocationSet = self.trust.verify(token).await?;
273 let offered = Held::new(set, token);
274 let mut held = self.held.write().unwrap_or_else(PoisonError::into_inner);
275 let from = held.as_ref().map(|h| h.epoch);
276 match from {
277 Some(epoch) if offered.epoch < epoch => Err(AdoptError::Stale {
278 held: epoch,
279 offered: offered.epoch,
280 }),
281 Some(epoch) if offered.epoch == epoch => Ok(Adopted::Unchanged { epoch }),
282 _ => {
283 let to = offered.epoch;
284 *held = Some(Arc::new(offered));
285 Ok(Adopted::Advanced { from, to })
286 }
287 }
288 }
289
290 pub fn epoch(&self) -> Option<u64> {
292 self.held().map(|h| h.epoch)
293 }
294
295 pub fn export(&self) -> Option<String> {
298 self.held().map(|h| h.token.clone())
299 }
300
301 fn held(&self) -> Option<Arc<Held>> {
302 self.held.read().unwrap_or_else(PoisonError::into_inner).clone()
303 }
304}
305
306#[async_trait]
307impl RevocationReader for ReplicatedRevocations {
308 async fn is_revoked(&self, jti: &str) -> Result<bool, StoreError> {
309 Ok(self.held().is_some_and(|h| h.jtis.contains_key(jti)))
313 }
314
315 async fn is_device_revoked(&self, device: &DeviceId, seq: u64) -> Result<bool, StoreError> {
316 Ok(self.held().is_some_and(|h| h.devices.get(device).is_some_and(|&at| seq < at)))
317 }
318
319 async fn is_membership_revoked(
320 &self,
321 key: &RevocationKey,
322 kind: &str,
323 id: &str,
324 principal: &PrincipalId,
325 snapshot_epoch: u64,
326 ) -> Result<bool, StoreError> {
327 let Some(held) = self.held() else {
328 return Ok(false);
329 };
330 let tag = membership_tag(key, kind, id, principal);
331 Ok(held.memberships.get(&tag).is_some_and(|&at| snapshot_epoch < at))
332 }
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338 use crate::jwks::{JwkKey, JwksDoc, KeySet};
339 use crate::key_set::KeySetVerifier;
340 use crate::public_verifier::PasetoV4PublicVerifier;
341 use cheers_core::{KeyRole, McpClaims, PrincipalId, SignedArtifact, TokenVerifier};
342 use pasetors::keys::{AsymmetricKeyPair, AsymmetricSecretKey, Generate};
343 use pasetors::version4::{PublicToken, V4};
344
345 const ISS: &str = "https://cheers.test";
346
347 fn kp() -> AsymmetricKeyPair<V4> {
348 AsymmetricKeyPair::<V4>::generate().unwrap()
349 }
350
351 fn pub_bytes(k: &AsymmetricKeyPair<V4>) -> [u8; 32] {
352 k.public.as_bytes().try_into().unwrap()
353 }
354
355 fn sign_with(secret: &AsymmetricSecretKey<V4>, kid: &str, payload: &impl serde::Serialize, ia: Option<&[u8]>) -> String {
357 let body = serde_json::to_vec(payload).unwrap();
358 let footer = format!(r#"{{"kid":"{kid}"}}"#);
359 PublicToken::sign(secret, &body, Some(footer.as_bytes()), ia).unwrap()
360 }
361
362 fn sign_set(secret: &AsymmetricSecretKey<V4>, kid: &str, set: &RevocationSet) -> String {
363 sign_with(secret, kid, set, Some(RevocationSet::IMPLICIT_ASSERTION))
364 }
365
366 fn set(epoch: u64, revoked: Vec<Revoked>) -> RevocationSet {
367 let entries = revoked.iter().map(|e| revocation_entry(e, test_revocation_key)).collect();
368 RevocationSet::new(ISS, epoch, entries)
369 }
370
371 struct Rig {
372 issuer: AsymmetricKeyPair<V4>,
373 assertion: AsymmetricKeyPair<V4>,
374 selfsig: AsymmetricKeyPair<V4>,
375 replica: ReplicatedRevocations,
376 }
377
378 fn jwks_rig() -> Rig {
379 let (issuer, assertion, selfsig) = (kp(), kp(), kp());
380 let doc = JwksDoc {
381 keys: vec![
382 JwkKey::ed25519("iss-1", &pub_bytes(&issuer), ISS, KeyRole::Issuer),
383 JwkKey::ed25519("asr-1", &pub_bytes(&assertion), "svc:issues", KeyRole::Assertion),
384 JwkKey::ed25519("self-1", &pub_bytes(&selfsig), "svc:issues", KeyRole::SelfSigner),
385 ],
386 };
387 let keys = KeySetVerifier::from_key_set(KeySet::from_doc(doc).unwrap());
388 Rig {
389 replica: ReplicatedRevocations::new(IssuerTrust::key_set(ISS, keys)),
390 issuer,
391 assertion,
392 selfsig,
393 }
394 }
395
396 fn adopt(r: &ReplicatedRevocations, token: &str) -> Result<Adopted, AdoptError> {
397 pollster::block_on(r.adopt(token))
398 }
399
400 fn jti(r: &ReplicatedRevocations, j: &str) -> bool {
401 pollster::block_on(r.is_revoked(j)).unwrap()
402 }
403
404 fn device(r: &ReplicatedRevocations, d: &str) -> bool {
405 pollster::block_on(r.is_device_revoked(&DeviceId::new(d), 0)).unwrap()
406 }
407
408 fn member(r: &ReplicatedRevocations, kind: &str, id: &str, user: &str) -> bool {
409 pollster::block_on(r.is_membership_revoked(&test_revocation_key(kind, id), kind, id, &PrincipalId::user(user), 0)).unwrap()
410 }
411
412 #[test]
413 fn empty_replica_answers_false() {
414 let r = jwks_rig();
415 assert_eq!(r.replica.epoch(), None);
416 assert_eq!(r.replica.export(), None);
417 assert!(!jti(&r.replica, "j1"));
418 assert!(!device(&r.replica, "phone"));
419 assert!(!member(&r.replica, "namespace", "ns-1", "alice"));
420 }
421
422 #[test]
423 fn adopted_set_answers_jti_device_and_membership_queries() {
424 let r = jwks_rig();
425 let s = set(
426 5,
427 vec![
428 Revoked::jti("j1", None),
429 Revoked::device("phone", 1),
430 Revoked::membership("namespace", "ns-1", PrincipalId::user("alice"), 1),
431 ],
432 );
433 assert_eq!(
434 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &s)).unwrap(),
435 Adopted::Advanced { from: None, to: 5 }
436 );
437 assert!(jti(&r.replica, "j1"));
438 assert!(!jti(&r.replica, "j2"));
439 assert!(device(&r.replica, "phone"));
440 assert!(!device(&r.replica, "laptop"));
441 assert!(member(&r.replica, "namespace", "ns-1", "alice"));
442 assert!(!member(&r.replica, "namespace", "ns-1", "bob"));
444 assert!(!member(&r.replica, "namespace", "ns-2", "alice"));
445 assert!(!member(&r.replica, "room", "ns-1", "alice"));
446 let wrong = test_revocation_key("namespace", "ns-2");
448 let alice = PrincipalId::user("alice");
449 assert!(!pollster::block_on(r.replica.is_membership_revoked(&wrong, "namespace", "ns-1", &alice, 0)).unwrap());
450 assert!(!device(&r.replica, "j1"));
452 }
453
454 #[test]
455 fn a_published_set_names_no_member_in_clear() {
456 let s = set(1, vec![Revoked::membership("namespace", "ns-secret", PrincipalId::user("alice-secret"), 3)]);
457 let wire = serde_json::to_string(&s).unwrap();
458 assert!(!wire.contains("ns-secret") && !wire.contains("alice-secret"), "{wire}");
459 assert!(wire.contains("\"tag\""));
460 }
461
462 #[test]
463 fn tags_are_length_prefixed_so_triples_do_not_collide() {
464 let key = test_revocation_key("k", "i");
465 assert_ne!(
466 membership_tag(&key, "ab", "c", &PrincipalId::user("d")),
467 membership_tag(&key, "a", "bc", &PrincipalId::user("d"))
468 );
469 }
470
471 #[test]
472 fn a_user_and_a_key_with_one_id_string_tag_differently() {
473 let key = test_revocation_key("namespace", "ns-1");
474 let k = PrincipalId::from_public_key(&[3; 32]);
475 let u = PrincipalId::user(k.id.clone());
476 assert_eq!(u.id, k.id);
477 assert_ne!(membership_tag(&key, "namespace", "ns-1", &u), membership_tag(&key, "namespace", "ns-1", &k));
478
479 let r = jwks_rig();
481 let s = set(5, vec![Revoked::membership("namespace", "ns-1", k.clone(), 7)]);
482 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &s)).unwrap();
483 let revoked = |p: &PrincipalId| pollster::block_on(r.replica.is_membership_revoked(&key, "namespace", "ns-1", p, 6)).unwrap();
484 assert!(revoked(&k));
485 assert!(!revoked(&u));
486 }
487
488 #[test]
489 fn device_and_membership_entries_mask_only_below_their_bound() {
490 let r = jwks_rig();
491 let s = set(5, vec![Revoked::device("phone", 100), Revoked::membership("namespace", "ns-1", PrincipalId::user("alice"), 7)]);
492 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &s)).unwrap();
493 let dev = |seq| pollster::block_on(r.replica.is_device_revoked(&DeviceId::new("phone"), seq)).unwrap();
494 assert!(dev(99));
495 assert!(!dev(100), "a binding minted at or after the revoke (re-enrollment) is admitted");
496 assert!(!dev(101));
497 let key = test_revocation_key("namespace", "ns-1");
498 let mem = |e| {
499 pollster::block_on(r.replica.is_membership_revoked(&key, "namespace", "ns-1", &PrincipalId::user("alice"), e)).unwrap()
500 };
501 assert!(mem(6));
502 assert!(!mem(7));
503 assert!(!mem(8));
504 }
505
506 #[test]
507 fn lower_epoch_is_refused_and_equal_epoch_is_a_no_op() {
508 let r = jwks_rig();
509 let newer = set(10, vec![Revoked::jti("j1", None), Revoked::jti("j2", None)]);
510 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &newer)).unwrap();
511
512 let older = set(9, vec![Revoked::jti("j1", None)]);
513 match adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &older)) {
514 Err(AdoptError::Stale { held: 10, offered: 9 }) => {}
515 other => panic!("expected Stale, got {other:?}"),
516 }
517 let same_epoch = set(10, vec![]);
518 assert_eq!(
519 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &same_epoch)).unwrap(),
520 Adopted::Unchanged { epoch: 10 }
521 );
522 assert_eq!(r.replica.epoch(), Some(10));
524 assert!(jti(&r.replica, "j2"));
525
526 let next = set(11, vec![Revoked::jti("j1", None)]);
528 assert_eq!(
529 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &next)).unwrap(),
530 Adopted::Advanced { from: Some(10), to: 11 }
531 );
532 assert!(!jti(&r.replica, "j2"));
533 }
534
535 #[test]
536 fn an_access_token_is_refused_as_a_set() {
537 let r = jwks_rig();
538 let token = sign_with(&r.issuer.secret, "iss-1", &set(1, vec![Revoked::jti("j1", None)]), None);
540 assert!(matches!(
541 adopt(&r.replica, &token),
542 Err(AdoptError::Verify(VerifyError::SignatureMismatch))
543 ));
544 let mcp = McpClaims::new(ISS, "https://aud.test", PrincipalId::user("alice"), 1, 2, "j", vec![]);
545 let token = sign_with(&r.issuer.secret, "iss-1", &mcp, None);
546 assert!(matches!(adopt(&r.replica, &token), Err(AdoptError::Verify(_))));
547 assert_eq!(r.replica.epoch(), None);
548 }
549
550 #[test]
551 fn a_set_is_refused_as_an_access_token() {
552 let r = jwks_rig();
553 let token = sign_set(&r.issuer.secret, "iss-1", &set(1, vec![]));
554 let keys = KeySetVerifier::from_issuer_key("iss-1", &pub_bytes(&r.issuer), ISS);
555 let res = pollster::block_on(keys.verify_mcp(&token, 0, ISS, None));
556 assert!(matches!(res, Err(VerifyError::SignatureMismatch)), "{res:?}");
557 let pinned = PasetoV4PublicVerifier::from_public_key(&pub_bytes(&r.issuer)).unwrap();
558 assert!(pinned.verify_mcp_at(&token, 0, "iss-1").is_err());
559 assert!(pinned.verify_at(&token, 0).is_err());
560 }
561
562 #[test]
563 fn only_an_issuer_role_key_may_sign_a_set() {
564 let r = jwks_rig();
565 let s = set(1, vec![]);
566 for (secret, kid) in [(&r.assertion.secret, "asr-1"), (&r.selfsig.secret, "self-1")] {
567 match adopt(&r.replica, &sign_set(secret, kid, &s)) {
568 Err(AdoptError::Verify(VerifyError::KeyRoleRejected { .. })) => {}
569 other => panic!("{kid}: expected KeyRoleRejected, got {other:?}"),
570 }
571 }
572 assert!(matches!(
574 adopt(&r.replica, &sign_set(&r.selfsig.secret, "iss-1", &s)),
575 Err(AdoptError::Verify(VerifyError::SignatureMismatch))
576 ));
577 assert_eq!(r.replica.epoch(), None);
578 }
579
580 #[test]
581 fn a_set_for_another_issuer_is_refused() {
582 let r = jwks_rig();
583 let foreign = RevocationSet::new("https://evil.test", 1, vec![]);
584 assert!(matches!(
585 adopt(&r.replica, &sign_set(&r.issuer.secret, "iss-1", &foreign)),
586 Err(AdoptError::Verify(VerifyError::BadIssuer { .. }))
587 ));
588 }
589
590 #[test]
591 fn pinned_key_trust_verifies_and_refuses_other_keys() {
592 let issuer = kp();
593 let pinned = PasetoV4PublicVerifier::from_public_key(&pub_bytes(&issuer)).unwrap();
594 let replica = ReplicatedRevocations::new(IssuerTrust::pinned(ISS, pinned));
595 let s = set(3, vec![Revoked::device("phone", 1)]);
596 assert!(matches!(
597 adopt(&replica, &sign_set(&kp().secret, "any", &s)),
598 Err(AdoptError::Verify(VerifyError::SignatureMismatch))
599 ));
600 assert!(matches!(
601 adopt(&replica, &sign_with(&issuer.secret, "any", &s, None)),
602 Err(AdoptError::Verify(VerifyError::SignatureMismatch))
603 ));
604 adopt(&replica, &sign_set(&issuer.secret, "any", &s)).unwrap();
605 assert!(device(&replica, "phone"));
606 let foreign = RevocationSet::new("https://evil.test", 4, vec![]);
607 assert!(matches!(
608 adopt(&replica, &sign_set(&issuer.secret, "any", &foreign)),
609 Err(AdoptError::Verify(VerifyError::BadIssuer { .. }))
610 ));
611 }
612
613 #[test]
614 fn export_import_round_trips_through_a_fresh_replica() {
615 let r = jwks_rig();
616 let s = set(
617 42,
618 vec![Revoked::jti("j1", None), Revoked::device("phone", 1), Revoked::membership("band", "b1", PrincipalId::user("carol"), 1)],
619 );
620 let token = sign_set(&r.issuer.secret, "iss-1", &s);
621 adopt(&r.replica, &token).unwrap();
622 let bytes = r.replica.export().expect("held set exports");
623 assert_eq!(bytes, token);
624
625 let pinned = PasetoV4PublicVerifier::from_public_key(&pub_bytes(&r.issuer)).unwrap();
627 let restarted = ReplicatedRevocations::new(IssuerTrust::pinned(ISS, pinned));
628 assert_eq!(
629 adopt(&restarted, &bytes).unwrap(),
630 Adopted::Advanced { from: None, to: 42 }
631 );
632 assert!(jti(&restarted, "j1"));
633 assert!(device(&restarted, "phone"));
634 assert!(member(&restarted, "band", "b1", "carol"));
635 let mut tampered = bytes.clone();
637 tampered.pop();
638 let fresh = ReplicatedRevocations::new(IssuerTrust::key_set(
639 ISS,
640 KeySetVerifier::from_issuer_key("iss-1", &pub_bytes(&r.issuer), ISS),
641 ));
642 assert!(adopt(&fresh, &tampered).is_err());
643 assert_eq!(fresh.epoch(), None);
644 }
645
646 #[test]
647 fn concurrent_adopts_and_reads_settle_on_the_highest_epoch() {
648 let r = Arc::new(jwks_rig());
649 let tokens: Vec<String> = (1..=32u64)
651 .map(|e| {
652 let revoked = (1..=e).map(|i| Revoked::jti(format!("j{i}"), None)).collect();
653 sign_set(&r.issuer.secret, "iss-1", &set(e, revoked))
654 })
655 .collect();
656 let handles: Vec<_> = tokens
657 .into_iter()
658 .rev()
659 .map(|t| {
660 let r = r.clone();
661 std::thread::spawn(move || {
662 let _ = adopt(&r.replica, &t);
663 let e = r.replica.epoch().expect("something is held after an adopt");
666 assert!(jti(&r.replica, &format!("j{e}")));
667 })
668 })
669 .collect();
670 for h in handles {
671 h.join().unwrap();
672 }
673 assert_eq!(r.replica.epoch(), Some(32));
674 assert!(jti(&r.replica, "j32"));
675 assert!(!jti(&r.replica, "j33"));
676 }
677
678 #[test]
679 fn trait_is_dyn_compatible_and_shareable() {
680 fn _reader(_: &dyn RevocationReader) {}
681 fn _shared(r: Arc<ReplicatedRevocations>) -> impl RevocationReader {
682 r
683 }
684 }
685}