Skip to main content

arete_auth/
verifier.rs

1use crate::claims::AuthContext;
2use crate::error::VerifyError;
3use crate::keys::VerifyingKey;
4use crate::sealed::ClaimSealKeys;
5use crate::token::{JwksVerifier, TokenVerifier};
6use std::sync::Arc;
7use std::time::{Duration, Instant};
8use tokio::sync::RwLock;
9
10/// Cached JWKS with expiration
11#[derive(Clone)]
12struct CachedJwks {
13    verifier: JwksVerifier,
14    fetched_at: Instant,
15}
16
17/// Async verifier with JWKS caching support
18pub struct AsyncVerifier {
19    inner: VerifierInner,
20    jwks_url: Option<String>,
21    cache_duration: Duration,
22    cached_jwks: Arc<RwLock<Option<CachedJwks>>>,
23    /// Issuer for JWKS-based verification
24    issuer: String,
25    /// Audience for JWKS-based verification
26    audiences: crate::AudienceSet,
27    require_origin: bool,
28    claim_seal_keys: Option<ClaimSealKeys>,
29}
30
31enum VerifierInner {
32    Static(TokenVerifier),
33    Jwks(JwksVerifier),
34    /// Keys are fetched from `jwks_url`; there is no local key to fall back to.
35    #[cfg(feature = "jwks")]
36    Remote,
37}
38
39impl AsyncVerifier {
40    /// Create a verifier with a static key
41    pub fn with_static_key(
42        key: VerifyingKey,
43        issuer: impl Into<String>,
44        audience: impl Into<String>,
45    ) -> Self {
46        let issuer_str = issuer.into();
47        let audience_str = audience.into();
48        Self {
49            inner: VerifierInner::Static(TokenVerifier::new(
50                key,
51                issuer_str.clone(),
52                audience_str.clone(),
53            )),
54            jwks_url: None,
55            cache_duration: Duration::from_secs(3600), // 1 hour default
56            cached_jwks: Arc::new(RwLock::new(None)),
57            issuer: issuer_str,
58            audiences: crate::AudienceSet::single(audience_str),
59            require_origin: false,
60            claim_seal_keys: None,
61        }
62    }
63
64    /// Create a verifier with JWKS
65    pub fn with_jwks(
66        jwks: crate::token::Jwks,
67        issuer: impl Into<String>,
68        audience: impl Into<String>,
69    ) -> Self {
70        let issuer_str = issuer.into();
71        let audience_str = audience.into();
72        Self {
73            inner: VerifierInner::Jwks(JwksVerifier::new(
74                jwks,
75                issuer_str.clone(),
76                audience_str.clone(),
77            )),
78            jwks_url: None,
79            cache_duration: Duration::from_secs(3600),
80            cached_jwks: Arc::new(RwLock::new(None)),
81            issuer: issuer_str,
82            audiences: crate::AudienceSet::single(audience_str),
83            require_origin: false,
84            claim_seal_keys: None,
85        }
86    }
87
88    /// Create a verifier that fetches JWKS from a URL
89    #[cfg(feature = "jwks")]
90    pub fn with_jwks_url(
91        url: impl Into<String>,
92        issuer: impl Into<String>,
93        audience: impl Into<String>,
94    ) -> Self {
95        Self {
96            inner: VerifierInner::Remote,
97            jwks_url: Some(url.into()),
98            issuer: issuer.into(),
99            audiences: crate::AudienceSet::single(audience.into()),
100            cache_duration: Duration::from_secs(3600),
101            cached_jwks: Arc::new(RwLock::new(None)),
102            require_origin: false,
103            claim_seal_keys: None,
104        }
105    }
106
107    /// Require origin validation on verified tokens.
108    pub fn with_origin_validation(mut self) -> Self {
109        self.require_origin = true;
110        self.inner = match self.inner {
111            VerifierInner::Static(verifier) => {
112                VerifierInner::Static(verifier.with_origin_validation())
113            }
114            VerifierInner::Jwks(verifier) => VerifierInner::Jwks(verifier.with_origin_validation()),
115            #[cfg(feature = "jwks")]
116            VerifierInner::Remote => VerifierInner::Remote,
117        };
118        self.cached_jwks = Arc::new(RwLock::new(None));
119        self
120    }
121
122    /// Open sealed claims with these keys. Without them, a token that carries
123    /// sealed claims is refused.
124    pub fn with_claim_seal_keys(mut self, keys: ClaimSealKeys) -> Self {
125        self.inner = match self.inner {
126            VerifierInner::Static(verifier) => {
127                VerifierInner::Static(verifier.with_claim_seal_keys(keys.clone()))
128            }
129            VerifierInner::Jwks(verifier) => {
130                VerifierInner::Jwks(verifier.with_claim_seal_keys(keys.clone()))
131            }
132            #[cfg(feature = "jwks")]
133            VerifierInner::Remote => VerifierInner::Remote,
134        };
135        self.claim_seal_keys = Some(keys);
136        // A verifier cached before this call would still open with the old keys.
137        self.cached_jwks = Arc::new(RwLock::new(None));
138        self
139    }
140
141    /// Set cache duration for JWKS
142    pub fn with_cache_duration(mut self, duration: Duration) -> Self {
143        self.cache_duration = duration;
144        self
145    }
146
147    /// Verify a token with automatic JWKS fetching and caching
148    #[cfg(feature = "jwks")]
149    pub async fn verify(
150        &self,
151        token: &str,
152        expected_origin: Option<&str>,
153        expected_client_ip: Option<&str>,
154    ) -> Result<AuthContext, VerifyError> {
155        match &self.inner {
156            VerifierInner::Static(verifier) => {
157                verifier.verify(token, expected_origin, expected_client_ip)
158            }
159            VerifierInner::Jwks(verifier) => {
160                verifier.verify(token, expected_origin, expected_client_ip)
161            }
162            // A URL-backed verifier only trusts keys fetched from its JWKS URL.
163            VerifierInner::Remote => {
164                self.verify_with_cache(token, expected_origin, expected_client_ip)
165                    .await
166            }
167        }
168    }
169
170    /// Verify a token (non-JWKS version)
171    #[cfg(not(feature = "jwks"))]
172    pub fn verify(
173        &self,
174        token: &str,
175        expected_origin: Option<&str>,
176        expected_client_ip: Option<&str>,
177    ) -> Result<AuthContext, VerifyError> {
178        match &self.inner {
179            VerifierInner::Static(verifier) => {
180                verifier.verify(token, expected_origin, expected_client_ip)
181            }
182            VerifierInner::Jwks(verifier) => {
183                verifier.verify(token, expected_origin, expected_client_ip)
184            }
185        }
186    }
187
188    /// Refresh JWKS cache from the configured URL
189    #[cfg(feature = "jwks")]
190    pub async fn refresh_cache(&self) -> Result<(), VerifyError> {
191        if let Some(ref jwks_url) = self.jwks_url {
192            // Fetch JWKS from URL
193            let jwks = crate::token::JwksVerifier::fetch_jwks(jwks_url)
194                .await
195                .map_err(|e| VerifyError::InvalidFormat(format!("Failed to fetch JWKS: {}", e)))?;
196
197            // Create new verifier with fetched JWKS
198            let verifier =
199                JwksVerifier::with_audience_set(jwks, &self.issuer, self.audiences.clone());
200            let verifier = if self.require_origin {
201                verifier.with_origin_validation()
202            } else {
203                verifier
204            };
205            let verifier = match &self.claim_seal_keys {
206                Some(keys) => verifier.with_claim_seal_keys(keys.clone()),
207                None => verifier,
208            };
209
210            // Update cache
211            let mut cached = self.cached_jwks.write().await;
212            *cached = Some(CachedJwks {
213                verifier,
214                fetched_at: Instant::now(),
215            });
216        }
217        Ok(())
218    }
219
220    /// Get cached verifier if available and not expired
221    async fn get_cached_verifier(&self) -> Option<JwksVerifier> {
222        let cached = self.cached_jwks.read().await;
223        if let Some(ref cached_jwks) = *cached {
224            if cached_jwks.fetched_at.elapsed() < self.cache_duration {
225                return Some(cached_jwks.verifier.clone());
226            }
227        }
228        None
229    }
230
231    /// Verify a token with automatic JWKS caching
232    #[cfg(feature = "jwks")]
233    pub async fn verify_with_cache(
234        &self,
235        token: &str,
236        expected_origin: Option<&str>,
237        expected_client_ip: Option<&str>,
238    ) -> Result<AuthContext, VerifyError> {
239        // Try cached verifier first
240        if let Some(verifier) = self.get_cached_verifier().await {
241            match verifier.verify(token, expected_origin, expected_client_ip) {
242                Ok(ctx) => return Ok(ctx),
243                Err(VerifyError::KeyNotFound(_)) => {
244                    // Key not found in cache, refresh and retry
245                }
246                Err(e) => return Err(e),
247            }
248        }
249
250        // Refresh cache and try again
251        self.refresh_cache().await?;
252
253        if let Some(verifier) = self.get_cached_verifier().await {
254            return verifier.verify(token, expected_origin, expected_client_ip);
255        }
256
257        // No cache: verifiers with a local key fall back to it; a URL-backed
258        // verifier has none and must fail closed.
259        match &self.inner {
260            VerifierInner::Static(verifier) => {
261                verifier.verify(token, expected_origin, expected_client_ip)
262            }
263            VerifierInner::Jwks(verifier) => {
264                verifier.verify(token, expected_origin, expected_client_ip)
265            }
266            VerifierInner::Remote => Err(VerifyError::InvalidFormat(
267                "JWKS cache unavailable after refresh".to_string(),
268            )),
269        }
270    }
271}
272
273/// Simple synchronous verifier for use in non-async contexts
274pub struct SimpleVerifier {
275    inner: TokenVerifier,
276}
277
278impl SimpleVerifier {
279    /// Create a new simple verifier
280    pub fn new(key: VerifyingKey, issuer: impl Into<String>, audience: impl Into<String>) -> Self {
281        Self {
282            inner: TokenVerifier::new(key, issuer, audience),
283        }
284    }
285
286    /// Verify a token synchronously
287    pub fn verify(
288        &self,
289        token: &str,
290        expected_origin: Option<&str>,
291        expected_client_ip: Option<&str>,
292    ) -> Result<AuthContext, VerifyError> {
293        self.inner
294            .verify(token, expected_origin, expected_client_ip)
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use crate::claims::{KeyClass, SessionClaims};
302    use crate::keys::SigningKey;
303    use crate::token::TokenSigner;
304    use base64::Engine;
305
306    #[cfg(feature = "jwks")]
307    use tokio::io::{AsyncReadExt, AsyncWriteExt};
308
309    #[tokio::test]
310    async fn test_async_verifier_with_static_key() {
311        let signing_key = SigningKey::generate();
312        let verifying_key = signing_key.verifying_key();
313
314        let signer = TokenSigner::new(signing_key, "test-issuer");
315        let verifier =
316            AsyncVerifier::with_static_key(verifying_key, "test-issuer", "test-audience");
317
318        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
319            .with_scope("read")
320            .with_metering_key("meter-123")
321            .with_key_class(KeyClass::Publishable)
322            .build();
323
324        let token = signer.sign(claims).unwrap();
325        let context = verifier.verify(&token, None, None).await.unwrap();
326
327        assert_eq!(context.subject, "test-subject");
328    }
329
330    #[test]
331    fn test_simple_verifier() {
332        let signing_key = SigningKey::generate();
333        let verifying_key = signing_key.verifying_key();
334
335        let signer = TokenSigner::new(signing_key, "test-issuer");
336        let verifier = SimpleVerifier::new(verifying_key, "test-issuer", "test-audience");
337
338        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
339            .with_scope("read")
340            .with_metering_key("meter-123")
341            .with_key_class(KeyClass::Publishable)
342            .build();
343
344        let token = signer.sign(claims).unwrap();
345        let context = verifier.verify(&token, None, None).unwrap();
346
347        assert_eq!(context.subject, "test-subject");
348        assert_eq!(context.metering_key, "meter-123");
349    }
350
351    #[cfg(feature = "jwks")]
352    #[test]
353    fn test_verify_with_cache_returns_explicit_error_when_cache_stays_empty() {
354        tokio::runtime::Runtime::new().unwrap().block_on(async {
355            let signing_key = SigningKey::generate();
356            let verifying_key = signing_key.verifying_key();
357            let signer = TokenSigner::new(signing_key, "test-issuer");
358
359            let jwks = serde_json::json!({
360                "keys": [{
361                    "kty": "OKP",
362                    "use": "sig",
363                    "kid": verifying_key.key_id(),
364                    "x": base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(verifying_key.to_bytes()),
365                }]
366            })
367            .to_string();
368
369            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
370            let addr = listener.local_addr().unwrap();
371            let response_body = jwks.clone();
372            tokio::spawn(async move {
373                let (mut socket, _) = listener.accept().await.unwrap();
374                let mut buffer = [0u8; 1024];
375                let _ = socket.read(&mut buffer).await;
376
377                let response = format!(
378                    "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
379                    response_body.len(),
380                    response_body
381                );
382                socket.write_all(response.as_bytes()).await.unwrap();
383            });
384
385            let verifier = AsyncVerifier::with_jwks_url(
386                format!("http://{addr}/jwks"),
387                "test-issuer",
388                "test-audience",
389            )
390            .with_cache_duration(Duration::ZERO);
391
392            let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
393                .with_scope("read")
394                .with_metering_key("meter-123")
395                .with_key_class(KeyClass::Publishable)
396                .build();
397            let token = signer.sign(claims).unwrap();
398
399            let result = verifier.verify_with_cache(&token, None, None).await;
400            assert!(matches!(
401                result,
402                Err(VerifyError::InvalidFormat(ref msg)) if msg == "JWKS cache unavailable after refresh"
403            ));
404        });
405    }
406
407    fn b64(bytes: impl AsRef<[u8]>) -> String {
408        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
409    }
410
411    fn jwks_json(keys: &[&VerifyingKey]) -> String {
412        let keys: Vec<_> = keys
413            .iter()
414            .map(|key| {
415                serde_json::json!({
416                    "kty": "OKP",
417                    "use": "sig",
418                    "kid": key.key_id(),
419                    "x": b64(key.to_bytes()),
420                })
421            })
422            .collect();
423        serde_json::json!({ "keys": keys }).to_string()
424    }
425
426    /// Serve `body` as JSON to every connection until the test ends.
427    #[cfg(feature = "jwks")]
428    async fn serve_jwks(body: String) -> std::net::SocketAddr {
429        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
430        let addr = listener.local_addr().unwrap();
431        tokio::spawn(async move {
432            loop {
433                let (mut socket, _) = listener.accept().await.unwrap();
434                let body = body.clone();
435                tokio::spawn(async move {
436                    let mut buffer = [0u8; 1024];
437                    let _ = socket.read(&mut buffer).await;
438                    let response = format!(
439                        "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
440                        body.len(),
441                        body
442                    );
443                    let _ = socket.write_all(response.as_bytes()).await;
444                });
445            }
446        });
447        addr
448    }
449
450    fn signed_token(signing_key: SigningKey) -> String {
451        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
452            .with_scope("read")
453            .with_metering_key("meter-123")
454            .with_key_class(KeyClass::Publishable)
455            .build();
456        TokenSigner::new(signing_key, "test-issuer")
457            .sign(claims)
458            .unwrap()
459    }
460
461    /// A token "signed" for the all-zero public key, which is a small-order
462    /// point: with R = identity and s = 0 the lax Ed25519 equation holds for
463    /// about one message in four, so no private key is needed.
464    fn forge_low_order_token(kid: Option<&str>) -> String {
465        use ed25519_dalek::Verifier as _;
466
467        let zero_key = ed25519_dalek::VerifyingKey::from_bytes(&[0u8; 32]).unwrap();
468        let mut signature_bytes = [0u8; 64];
469        signature_bytes[0] = 1; // compressed identity point; s stays 0
470        let signature = ed25519_dalek::Signature::from_bytes(&signature_bytes);
471
472        let mut header = serde_json::json!({ "alg": "EdDSA", "typ": "JWT" });
473        if let Some(kid) = kid {
474            header["kid"] = kid.into();
475        }
476        let header_b64 = b64(header.to_string());
477
478        for attempt in 0..512 {
479            let claims = SessionClaims::builder(
480                "test-issuer",
481                format!("forged-subject-{attempt}"),
482                "test-audience",
483            )
484            .with_scope("admin")
485            .with_metering_key("forged")
486            .with_key_class(KeyClass::Secret)
487            .build();
488            let message = format!("{header_b64}.{}", b64(serde_json::to_vec(&claims).unwrap()));
489            if zero_key.verify(message.as_bytes(), &signature).is_ok() {
490                return format!("{message}.{}", b64(signature_bytes));
491            }
492        }
493        panic!("no low-order forgery found");
494    }
495
496    fn zero_key() -> VerifyingKey {
497        VerifyingKey::from_bytes(&[0u8; 32]).unwrap()
498    }
499
500    #[cfg(feature = "jwks")]
501    #[tokio::test]
502    async fn jwks_url_verify_uses_the_fetched_keys() {
503        let signing_key = SigningKey::generate();
504        let addr = serve_jwks(jwks_json(&[&signing_key.verifying_key()])).await;
505        let verifier = AsyncVerifier::with_jwks_url(
506            format!("http://{addr}/jwks"),
507            "test-issuer",
508            "test-audience",
509        );
510
511        let token = signed_token(signing_key);
512        let context = verifier.verify(&token, None, None).await.unwrap();
513        assert_eq!(context.subject, "test-subject");
514        let context = verifier
515            .verify_with_cache(&token, None, None)
516            .await
517            .unwrap();
518        assert_eq!(context.subject, "test-subject");
519    }
520
521    #[cfg(feature = "jwks")]
522    #[tokio::test]
523    async fn jwks_url_verify_fails_closed_when_the_jwks_is_unreachable() {
524        // Bind then drop, so nothing is listening on the port.
525        let addr = tokio::net::TcpListener::bind("127.0.0.1:0")
526            .await
527            .unwrap()
528            .local_addr()
529            .unwrap();
530        let verifier = AsyncVerifier::with_jwks_url(
531            format!("http://{addr}/jwks"),
532            "test-issuer",
533            "test-audience",
534        );
535
536        for token in [
537            signed_token(SigningKey::generate()),
538            forge_low_order_token(None),
539        ] {
540            assert!(verifier.verify(&token, None, None).await.is_err());
541            assert!(verifier
542                .verify_with_cache(&token, None, None)
543                .await
544                .is_err());
545        }
546    }
547
548    #[cfg(feature = "jwks")]
549    #[tokio::test]
550    async fn jwks_url_verify_rejects_unknown_keys_and_low_order_forgeries() {
551        let trusted = SigningKey::generate();
552        let trusted_key = trusted.verifying_key();
553        let addr = serve_jwks(jwks_json(&[&trusted_key])).await;
554        let verifier = AsyncVerifier::with_jwks_url(
555            format!("http://{addr}/jwks"),
556            "test-issuer",
557            "test-audience",
558        );
559
560        let untrusted = signed_token(SigningKey::generate());
561        let forged = forge_low_order_token(None);
562        let forged_with_trusted_kid = forge_low_order_token(Some(&trusted_key.key_id()));
563
564        for token in [&untrusted, &forged, &forged_with_trusted_kid] {
565            assert!(verifier.verify(token, None, None).await.is_err());
566            assert!(verifier.verify_with_cache(token, None, None).await.is_err());
567        }
568        assert!(matches!(
569            verifier.verify(&forged_with_trusted_kid, None, None).await,
570            Err(VerifyError::InvalidSignature)
571        ));
572    }
573
574    #[tokio::test]
575    async fn low_order_forgeries_are_rejected_by_every_verifier() {
576        let zero = zero_key();
577        let forged = forge_low_order_token(None);
578        let forged_with_kid = forge_low_order_token(Some(&zero.key_id()));
579
580        for token in [&forged, &forged_with_kid] {
581            assert!(matches!(
582                TokenVerifier::new(zero.clone(), "test-issuer", "test-audience")
583                    .verify(token, None, None),
584                Err(VerifyError::InvalidSignature)
585            ));
586            assert!(matches!(
587                SimpleVerifier::new(zero.clone(), "test-issuer", "test-audience")
588                    .verify(token, None, None),
589                Err(VerifyError::InvalidSignature)
590            ));
591
592            let static_verifier =
593                AsyncVerifier::with_static_key(zero.clone(), "test-issuer", "test-audience");
594            assert!(static_verifier.verify(token, None, None).await.is_err());
595        }
596
597        let jwks: crate::token::Jwks = serde_json::from_str(&jwks_json(&[&zero])).unwrap();
598        assert!(matches!(
599            JwksVerifier::new(jwks.clone(), "test-issuer", "test-audience").verify(
600                &forged_with_kid,
601                None,
602                None
603            ),
604            Err(VerifyError::InvalidSignature)
605        ));
606        let jwks_verifier = AsyncVerifier::with_jwks(jwks, "test-issuer", "test-audience");
607        assert!(matches!(
608            jwks_verifier.verify(&forged_with_kid, None, None).await,
609            Err(VerifyError::InvalidSignature)
610        ));
611    }
612}