arete_auth/
multi_key.rs

1use crate::claims::AuthContext;
2use crate::error::VerifyError;
3use crate::keys::VerifyingKey;
4use crate::token::TokenVerifier;
5use std::collections::HashMap;
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8use tokio::sync::RwLock;
9
10/// A key with its metadata for rotation
11#[derive(Clone)]
12pub struct RotationKey {
13    /// The verifying key
14    pub key: VerifyingKey,
15    /// Key ID for JWKS compatibility
16    pub key_id: String,
17    /// When this key was added
18    pub added_at: Instant,
19    /// Optional: when this key should be removed (for grace period rotation)
20    pub expires_at: Option<Instant>,
21    /// Whether this is the primary (current) key
22    pub is_primary: bool,
23}
24
25impl RotationKey {
26    /// Create a new primary key
27    pub fn primary(key: VerifyingKey, key_id: impl Into<String>) -> Self {
28        Self {
29            key,
30            key_id: key_id.into(),
31            added_at: Instant::now(),
32            expires_at: None,
33            is_primary: true,
34        }
35    }
36
37    /// Create a secondary (rotating out) key with expiration
38    pub fn secondary(key: VerifyingKey, key_id: impl Into<String>, grace_period: Duration) -> Self {
39        Self {
40            key,
41            key_id: key_id.into(),
42            added_at: Instant::now(),
43            expires_at: Some(Instant::now() + grace_period),
44            is_primary: false,
45        }
46    }
47
48    /// Check if this key has expired
49    pub fn is_expired(&self) -> bool {
50        self.expires_at
51            .map(|exp| Instant::now() > exp)
52            .unwrap_or(false)
53    }
54}
55
56/// Multi-key verifier supporting graceful key rotation
57///
58/// This verifier maintains multiple keys and attempts verification with each
59/// until one succeeds. This allows zero-downtime key rotation:
60///
61/// 1. Generate new key pair
62/// 2. Add new key as primary, mark old key as secondary with grace period
63/// 3. Update JWKS to include both keys
64/// 4. After grace period, remove old key
65///
66/// # Example
67/// ```rust
68/// use arete_auth::{MultiKeyVerifier, RotationKey, SigningKey};
69/// use std::time::Duration;
70///
71/// // Generate key pairs
72/// let old_signing_key = SigningKey::generate();
73/// let old_verifying_key = old_signing_key.verifying_key();
74/// let new_signing_key = SigningKey::generate();
75/// let new_verifying_key = new_signing_key.verifying_key();
76///
77/// // Create rotation keys
78/// let old_key = RotationKey::secondary(old_verifying_key, "key-1", Duration::from_secs(86400));
79/// let new_key = RotationKey::primary(new_verifying_key, "key-2");
80///
81/// let verifier = MultiKeyVerifier::new(vec![old_key, new_key], "issuer", "audience")
82///     .with_cleanup_interval(Duration::from_secs(3600));
83/// ```
84/// Build the per-key verifier used for one rotation candidate, carrying the
85/// accepted-audience set and origin policy through unchanged.
86fn build_inner_verifier(
87    key: crate::keys::VerifyingKey,
88    issuer: &str,
89    audiences: &crate::AudienceSet,
90    require_origin: bool,
91    claim_seal_keys: Option<&crate::sealed::ClaimSealKeys>,
92) -> TokenVerifier {
93    let verifier = match audiences.as_single() {
94        Some(single) => TokenVerifier::new(key, issuer, single),
95        None => TokenVerifier::with_audiences(key, issuer, audiences.iter())
96            .expect("an AudienceSet is never empty"),
97    };
98    let verifier = if require_origin {
99        verifier.with_origin_validation()
100    } else {
101        verifier
102    };
103    match claim_seal_keys {
104        Some(keys) => verifier.with_claim_seal_keys(keys.clone()),
105        None => verifier,
106    }
107}
108
109pub struct MultiKeyVerifier {
110    keys: Arc<RwLock<HashMap<String, RotationKey>>>,
111    issuer: String,
112    audiences: crate::AudienceSet,
113    require_origin: bool,
114    cleanup_interval: Duration,
115    last_cleanup: Arc<RwLock<Instant>>,
116    claim_seal_keys: Option<crate::sealed::ClaimSealKeys>,
117}
118
119impl MultiKeyVerifier {
120    /// Create a new multi-key verifier
121    pub fn new(
122        keys: Vec<RotationKey>,
123        issuer: impl Into<String>,
124        audience: impl Into<String>,
125    ) -> Self {
126        let key_map: HashMap<String, RotationKey> =
127            keys.into_iter().map(|k| (k.key_id.clone(), k)).collect();
128
129        Self {
130            keys: Arc::new(RwLock::new(key_map)),
131            issuer: issuer.into(),
132            audiences: crate::AudienceSet::single(audience),
133            require_origin: false,
134            claim_seal_keys: None,
135            cleanup_interval: Duration::from_secs(3600), // 1 hour default
136            last_cleanup: Arc::new(RwLock::new(Instant::now())),
137        }
138    }
139
140    /// Create from a single key (for backward compatibility)
141    pub fn from_single_key(
142        key: VerifyingKey,
143        key_id: impl Into<String>,
144        issuer: impl Into<String>,
145        audience: impl Into<String>,
146    ) -> Self {
147        Self::new(vec![RotationKey::primary(key, key_id)], issuer, audience)
148    }
149
150    /// Require origin validation
151    pub fn with_origin_validation(mut self) -> Self {
152        self.require_origin = true;
153        self
154    }
155
156    /// Open sealed claims with these keys. Without them, a token that carries
157    /// sealed claims is refused.
158    pub fn with_claim_seal_keys(mut self, keys: crate::sealed::ClaimSealKeys) -> Self {
159        self.claim_seal_keys = Some(keys);
160        self
161    }
162
163    /// Set cleanup interval for expired keys
164    pub fn with_cleanup_interval(mut self, interval: Duration) -> Self {
165        self.cleanup_interval = interval;
166        self
167    }
168
169    /// Add a new key to the verifier
170    pub async fn add_key(&self, key: RotationKey) {
171        let mut keys = self.keys.write().await;
172
173        // If adding a primary key, demote existing primary to secondary
174        if key.is_primary {
175            for existing in keys.values_mut() {
176                if existing.is_primary {
177                    existing.is_primary = false;
178                    // Set grace period for old primary
179                    existing.expires_at = Some(Instant::now() + Duration::from_secs(86400));
180                    // 24 hours
181                }
182            }
183        }
184
185        keys.insert(key.key_id.clone(), key);
186    }
187
188    /// Remove a key by ID
189    pub async fn remove_key(&self, key_id: &str) {
190        let mut keys = self.keys.write().await;
191        keys.remove(key_id);
192    }
193
194    /// Get all key IDs
195    pub async fn key_ids(&self) -> Vec<String> {
196        let keys = self.keys.read().await;
197        keys.keys().cloned().collect()
198    }
199
200    /// Get primary key ID
201    pub async fn primary_key_id(&self) -> Option<String> {
202        let keys = self.keys.read().await;
203        keys.values()
204            .find(|k| k.is_primary)
205            .map(|k| k.key_id.clone())
206    }
207
208    /// Clean up expired keys
209    async fn cleanup_expired_keys(&self) {
210        let should_cleanup = {
211            let last = self.last_cleanup.read().await;
212            last.elapsed() >= self.cleanup_interval
213        };
214
215        if !should_cleanup {
216            return;
217        }
218
219        let mut keys = self.keys.write().await;
220        let expired: Vec<String> = keys
221            .iter()
222            .filter(|(_, k)| k.is_expired())
223            .map(|(id, _)| id.clone())
224            .collect();
225
226        for key_id in expired {
227            keys.remove(&key_id);
228        }
229
230        // Update last cleanup time
231        let mut last = self.last_cleanup.write().await;
232        *last = Instant::now();
233    }
234
235    /// Verify a token against all keys
236    pub async fn verify(
237        &self,
238        token: &str,
239        expected_origin: Option<&str>,
240        expected_client_ip: Option<&str>,
241    ) -> Result<AuthContext, VerifyError> {
242        // Clean up expired keys periodically
243        self.cleanup_expired_keys().await;
244
245        let keys = self.keys.read().await;
246
247        if keys.is_empty() {
248            return Err(VerifyError::KeyNotFound("no keys configured".to_string()));
249        }
250
251        let mut last_error = None;
252
253        // Try primary key first, then secondary keys
254        let mut key_order: Vec<&RotationKey> = keys.values().collect();
255        key_order.sort_by_key(|k| !k.is_primary); // Primary first
256
257        for key_entry in key_order {
258            if key_entry.is_expired() {
259                continue;
260            }
261
262            let verifier = build_inner_verifier(
263                key_entry.key.clone(),
264                &self.issuer,
265                &self.audiences,
266                self.require_origin,
267                self.claim_seal_keys.as_ref(),
268            );
269
270            match verifier.verify(token, expected_origin, expected_client_ip) {
271                Ok(ctx) => {
272                    return Ok(ctx);
273                }
274                Err(VerifyError::InvalidSignature) => {
275                    // Wrong key, try next
276                    last_error = Some(VerifyError::InvalidSignature);
277                    continue;
278                }
279                Err(e) => {
280                    // Other errors (expired, invalid format, etc.) - don't try other keys
281                    return Err(e);
282                }
283            }
284        }
285
286        // All keys failed
287        Err(last_error.unwrap_or(VerifyError::InvalidSignature))
288    }
289
290    /// Verify without cleaning up (for high-throughput scenarios)
291    pub async fn verify_fast(
292        &self,
293        token: &str,
294        expected_origin: Option<&str>,
295        expected_client_ip: Option<&str>,
296    ) -> Result<AuthContext, VerifyError> {
297        let keys = self.keys.read().await;
298
299        if keys.is_empty() {
300            return Err(VerifyError::KeyNotFound("no keys configured".to_string()));
301        }
302
303        let mut last_error = None;
304
305        // Try primary key first, then secondary keys
306        let mut key_order: Vec<&RotationKey> = keys.values().collect();
307        key_order.sort_by_key(|k| !k.is_primary);
308
309        for key_entry in key_order {
310            if key_entry.is_expired() {
311                continue;
312            }
313
314            let verifier = build_inner_verifier(
315                key_entry.key.clone(),
316                &self.issuer,
317                &self.audiences,
318                self.require_origin,
319                self.claim_seal_keys.as_ref(),
320            );
321
322            match verifier.verify(token, expected_origin, expected_client_ip) {
323                Ok(ctx) => return Ok(ctx),
324                Err(VerifyError::InvalidSignature) => {
325                    last_error = Some(VerifyError::InvalidSignature);
326                    continue;
327                }
328                Err(e) => return Err(e),
329            }
330        }
331
332        Err(last_error.unwrap_or(VerifyError::InvalidSignature))
333    }
334}
335
336/// Builder for constructing a MultiKeyVerifier with rotation support
337pub struct MultiKeyVerifierBuilder {
338    keys: Vec<RotationKey>,
339    issuer: String,
340    audience: String,
341    require_origin: bool,
342    cleanup_interval: Duration,
343}
344
345impl MultiKeyVerifierBuilder {
346    /// Create a new builder
347    pub fn new(issuer: impl Into<String>, audience: impl Into<String>) -> Self {
348        Self {
349            keys: Vec::new(),
350            issuer: issuer.into(),
351            audience: audience.into(),
352            require_origin: false,
353            cleanup_interval: Duration::from_secs(3600),
354        }
355    }
356
357    /// Add a primary key
358    pub fn with_primary_key(mut self, key: VerifyingKey, key_id: impl Into<String>) -> Self {
359        self.keys.push(RotationKey::primary(key, key_id));
360        self
361    }
362
363    /// Add a secondary key with grace period
364    pub fn with_secondary_key(
365        mut self,
366        key: VerifyingKey,
367        key_id: impl Into<String>,
368        grace_period: Duration,
369    ) -> Self {
370        self.keys
371            .push(RotationKey::secondary(key, key_id, grace_period));
372        self
373    }
374
375    /// Require origin validation
376    pub fn with_origin_validation(mut self) -> Self {
377        self.require_origin = true;
378        self
379    }
380
381    /// Set cleanup interval
382    pub fn with_cleanup_interval(mut self, interval: Duration) -> Self {
383        self.cleanup_interval = interval;
384        self
385    }
386
387    /// Build the verifier
388    pub fn build(self) -> MultiKeyVerifier {
389        let mut verifier = MultiKeyVerifier::new(self.keys, self.issuer, self.audience);
390        if self.require_origin {
391            verifier = verifier.with_origin_validation();
392        }
393        verifier.with_cleanup_interval(self.cleanup_interval)
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use crate::claims::{KeyClass, SessionClaims};
401    use crate::keys::SigningKey;
402    use crate::token::TokenSigner;
403
404    #[tokio::test]
405    async fn test_multi_key_verifier_single_key() {
406        let signing_key = SigningKey::generate();
407        let verifying_key = signing_key.verifying_key();
408
409        let signer = TokenSigner::new(signing_key, "test-issuer");
410        let verifier = MultiKeyVerifier::from_single_key(
411            verifying_key,
412            "key-1",
413            "test-issuer",
414            "test-audience",
415        );
416
417        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
418            .with_scope("read")
419            .with_metering_key("meter-123")
420            .with_key_class(KeyClass::Publishable)
421            .build();
422
423        let token = signer.sign(claims).unwrap();
424        let context = verifier.verify(&token, None, None).await.unwrap();
425
426        assert_eq!(context.subject, "test-subject");
427        assert_eq!(verifier.primary_key_id().await, Some("key-1".to_string()));
428    }
429
430    #[tokio::test]
431    async fn test_key_rotation() {
432        // Create old key pair
433        let old_signing_key = SigningKey::generate();
434        let old_verifying_key = old_signing_key.verifying_key();
435        let old_signer = TokenSigner::new(old_signing_key, "test-issuer");
436
437        // Create new key pair
438        let new_signing_key = SigningKey::generate();
439        let new_verifying_key = new_signing_key.verifying_key();
440        let new_signer = TokenSigner::new(new_signing_key, "test-issuer");
441
442        // Start with old key as primary
443        let old_key = RotationKey::primary(old_verifying_key.clone(), "key-old");
444        let verifier = MultiKeyVerifier::new(vec![old_key], "test-issuer", "test-audience");
445
446        // Sign token with old key
447        let old_claims = SessionClaims::builder("test-issuer", "subject-1", "test-audience")
448            .with_scope("read")
449            .with_metering_key("meter-1")
450            .with_key_class(KeyClass::Publishable)
451            .build();
452        let old_token = old_signer.sign(old_claims).unwrap();
453
454        // Verify old token works
455        let ctx = verifier.verify(&old_token, None, None).await.unwrap();
456        assert_eq!(ctx.subject, "subject-1");
457
458        // Rotate: add new key as primary (old key becomes secondary)
459        let new_key = RotationKey::primary(new_verifying_key, "key-new");
460        verifier.add_key(new_key).await;
461
462        // Verify old token still works (grace period)
463        let ctx = verifier.verify(&old_token, None, None).await.unwrap();
464        assert_eq!(ctx.subject, "subject-1");
465
466        // Sign and verify new token
467        let new_claims = SessionClaims::builder("test-issuer", "subject-2", "test-audience")
468            .with_scope("read")
469            .with_metering_key("meter-2")
470            .with_key_class(KeyClass::Publishable)
471            .build();
472        let new_token = new_signer.sign(new_claims).unwrap();
473
474        let ctx = verifier.verify(&new_token, None, None).await.unwrap();
475        assert_eq!(ctx.subject, "subject-2");
476
477        // Check that new key is now primary
478        assert_eq!(verifier.primary_key_id().await, Some("key-new".to_string()));
479
480        // Both keys should be present
481        let key_ids = verifier.key_ids().await;
482        assert!(key_ids.contains(&"key-old".to_string()));
483        assert!(key_ids.contains(&"key-new".to_string()));
484    }
485
486    #[tokio::test]
487    async fn test_verifier_builder() {
488        let signing_key = SigningKey::generate();
489        let verifying_key = signing_key.verifying_key();
490
491        let verifier = MultiKeyVerifierBuilder::new("test-issuer", "test-audience")
492            .with_primary_key(verifying_key, "key-1")
493            .with_origin_validation()
494            .build();
495
496        let signer = TokenSigner::new(signing_key, "test-issuer");
497        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
498            .with_scope("read")
499            .with_origin("https://trusted.example.com")
500            .with_key_class(KeyClass::Secret)
501            .build();
502
503        let token = signer.sign(claims).unwrap();
504        let ctx = verifier
505            .verify(&token, Some("https://trusted.example.com"), None)
506            .await
507            .unwrap();
508        assert_eq!(ctx.subject, "test-subject");
509    }
510
511    #[tokio::test]
512    async fn test_invalid_signature_with_multiple_keys() {
513        // Create two different key pairs
514        let key1_signing = SigningKey::generate();
515
516        let key2_signing = SigningKey::generate();
517
518        let signer = TokenSigner::new(key1_signing, "test-issuer");
519
520        // Create verifier with only key2
521        let verifier = MultiKeyVerifier::from_single_key(
522            key2_signing.verifying_key(),
523            "key-2",
524            "test-issuer",
525            "test-audience",
526        );
527
528        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
529            .with_scope("read")
530            .with_key_class(KeyClass::Publishable)
531            .build();
532
533        let token = signer.sign(claims).unwrap();
534
535        // Should fail because token was signed with key1, verifier only has key2
536        let result = verifier.verify(&token, None, None).await;
537        assert!(matches!(result, Err(VerifyError::InvalidSignature)));
538    }
539
540    #[tokio::test]
541    async fn test_jwks_key_rotation_grace_period() {
542        use crate::token::{Jwk, Jwks};
543        use base64::Engine;
544
545        // Create old key pair with specific key ID
546        let old_signing_key = SigningKey::generate();
547        let old_verifying_key = old_signing_key.verifying_key();
548        let old_kid = old_verifying_key.key_id();
549        let old_signer = TokenSigner::new(old_signing_key, "test-issuer");
550
551        // Create new key pair with specific key ID
552        let new_signing_key = SigningKey::generate();
553        let new_verifying_key = new_signing_key.verifying_key();
554        let new_kid = new_verifying_key.key_id();
555        let new_signer = TokenSigner::new(new_signing_key, "test-issuer");
556
557        // Create JWKS with both keys using their actual key IDs
558        let old_key_b64 =
559            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(old_verifying_key.to_bytes());
560        let new_key_b64 =
561            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(new_verifying_key.to_bytes());
562
563        let jwks = Jwks {
564            keys: vec![
565                Jwk {
566                    kty: "OKP".to_string(),
567                    use_: Some("sig".to_string()),
568                    kid: old_kid,
569                    x: old_key_b64,
570                },
571                Jwk {
572                    kty: "OKP".to_string(),
573                    use_: Some("sig".to_string()),
574                    kid: new_kid,
575                    x: new_key_b64,
576                },
577            ],
578        };
579
580        // Create verifier from JWKS
581        let verifier =
582            crate::verifier::AsyncVerifier::with_jwks(jwks, "test-issuer", "test-audience");
583
584        // Sign and verify token with old key
585        let old_claims = SessionClaims::builder("test-issuer", "subject-old", "test-audience")
586            .with_scope("read")
587            .with_key_class(KeyClass::Secret)
588            .build();
589        let old_token = old_signer.sign(old_claims).unwrap();
590
591        // Old token should still verify during rotation
592        let ctx = verifier.verify(&old_token, None, None).await.unwrap();
593        assert_eq!(ctx.subject, "subject-old");
594
595        // Sign and verify token with new key
596        let new_claims = SessionClaims::builder("test-issuer", "subject-new", "test-audience")
597            .with_scope("read")
598            .with_key_class(KeyClass::Secret)
599            .build();
600        let new_token = new_signer.sign(new_claims).unwrap();
601
602        // New token should also verify
603        let ctx = verifier.verify(&new_token, None, None).await.unwrap();
604        assert_eq!(ctx.subject, "subject-new");
605    }
606
607    #[tokio::test]
608    async fn test_jwks_key_not_found() {
609        use crate::token::{Jwk, Jwks};
610        use base64::Engine;
611
612        // Create a key pair
613        let signing_key = SigningKey::generate();
614        let _verifying_key = signing_key.verifying_key();
615        let signer = TokenSigner::new(signing_key, "test-issuer");
616
617        // Create JWKS with a different key (not the one used for signing)
618        let different_key = SigningKey::generate();
619        let different_verifying_key = different_key.verifying_key();
620        let different_key_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD
621            .encode(different_verifying_key.to_bytes());
622
623        let jwks = Jwks {
624            keys: vec![Jwk {
625                kty: "OKP".to_string(),
626                use_: Some("sig".to_string()),
627                kid: "different-key".to_string(),
628                x: different_key_b64,
629            }],
630        };
631
632        let verifier =
633            crate::verifier::AsyncVerifier::with_jwks(jwks, "test-issuer", "test-audience");
634
635        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
636            .with_scope("read")
637            .with_key_class(KeyClass::Secret)
638            .build();
639        let token = signer.sign(claims).unwrap();
640
641        // Should fail with key not found
642        let result = verifier.verify(&token, None, None).await;
643        assert!(matches!(result, Err(VerifyError::KeyNotFound(_))));
644    }
645
646    #[tokio::test]
647    async fn test_jwks_with_origin_validation() {
648        use crate::token::{Jwk, Jwks};
649        use base64::Engine;
650
651        let signing_key = SigningKey::generate();
652        let verifying_key = signing_key.verifying_key();
653        let kid = verifying_key.key_id();
654        let signer = TokenSigner::new(signing_key, "test-issuer");
655
656        let key_b64 =
657            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(verifying_key.to_bytes());
658
659        let jwks = Jwks {
660            keys: vec![Jwk {
661                kty: "OKP".to_string(),
662                use_: Some("sig".to_string()),
663                kid,
664                x: key_b64,
665            }],
666        };
667
668        // Create verifier with origin validation
669        let verifier =
670            crate::verifier::AsyncVerifier::with_jwks(jwks, "test-issuer", "test-audience")
671                .with_origin_validation();
672
673        // Token with matching origin
674        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
675            .with_scope("read")
676            .with_key_class(KeyClass::Secret)
677            .with_origin("https://trusted.example.com")
678            .build();
679        let token = signer.sign(claims).unwrap();
680
681        // Should succeed with matching origin
682        let ctx = verifier
683            .verify(&token, Some("https://trusted.example.com"), None)
684            .await
685            .unwrap();
686        assert_eq!(ctx.subject, "test-subject");
687
688        // Should fail with wrong origin
689        let result = verifier
690            .verify(&token, Some("https://evil.example.com"), None)
691            .await;
692        assert!(matches!(result, Err(VerifyError::OriginMismatch { .. })));
693    }
694}
arete_auth/multi_key.rs

arete_auth/
multi_key.rs
