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, PoisonError, RwLock};
7use std::time::{Duration, Instant};
8
9/// Default minimum time between two JWKS fetches, whatever triggers them.
10/// See [`AsyncVerifier::with_min_refresh_interval`].
11pub const DEFAULT_MIN_JWKS_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
12
13/// Upper bound on a single JWKS fetch.
14#[cfg(feature = "jwks")]
15const JWKS_FETCH_TIMEOUT: Duration = Duration::from_secs(10);
16
17// The cache is only filled with the `jwks` feature.
18/// Keys from the last successful JWKS fetch.
19#[cfg_attr(not(feature = "jwks"), allow(dead_code))]
20struct CachedJwks {
21    verifier: Arc<JwksVerifier>,
22    fetched_at: Instant,
23}
24
25#[derive(Default)]
26struct JwksCacheState {
27    /// Last keys fetched successfully. A failed fetch leaves them in place.
28    current: Option<Arc<CachedJwks>>,
29    /// End of the last fetch attempt, successful or not.
30    last_attempt: Option<Instant>,
31}
32
33/// JWKS cache shared by every request on one verifier.
34///
35/// Fetches are single-flight, and a new one starts no sooner than the minimum
36/// refresh interval after the previous one finished, whether it is triggered
37/// by expiry, an unknown `kid`, or a previous failure. Requests never queue
38/// behind a fetch unless there are no keys at all yet.
39#[derive(Default)]
40#[cfg_attr(not(feature = "jwks"), allow(dead_code))]
41struct JwksCache {
42    state: RwLock<JwksCacheState>,
43    /// Held while a fetch is in flight.
44    fetch: tokio::sync::Mutex<()>,
45}
46
47/// Records a fetch attempt when dropped, however the fetch ended.
48#[cfg(feature = "jwks")]
49struct MarkAttemptOnDrop<'a>(&'a JwksCache);
50
51#[cfg(feature = "jwks")]
52impl Drop for MarkAttemptOnDrop<'_> {
53    fn drop(&mut self) {
54        self.0.mark_attempt();
55    }
56}
57
58#[cfg_attr(not(feature = "jwks"), allow(dead_code))]
59impl JwksCache {
60    fn current(&self) -> Option<Arc<CachedJwks>> {
61        self.state
62            .read()
63            .unwrap_or_else(PoisonError::into_inner)
64            .current
65            .clone()
66    }
67
68    fn refresh_due(&self, min_interval: Duration) -> bool {
69        self.state
70            .read()
71            .unwrap_or_else(PoisonError::into_inner)
72            .last_attempt
73            .is_none_or(|at| at.elapsed() >= min_interval)
74    }
75
76    fn mark_attempt(&self) {
77        self.state
78            .write()
79            .unwrap_or_else(PoisonError::into_inner)
80            .last_attempt = Some(Instant::now());
81    }
82
83    fn store(&self, verifier: JwksVerifier) {
84        self.state
85            .write()
86            .unwrap_or_else(PoisonError::into_inner)
87            .current = Some(Arc::new(CachedJwks {
88            verifier: Arc::new(verifier),
89            fetched_at: Instant::now(),
90        }));
91    }
92}
93
94/// Async verifier with JWKS caching support
95pub struct AsyncVerifier {
96    inner: VerifierInner,
97    jwks_url: Option<String>,
98    cache_duration: Duration,
99    min_refresh_interval: Duration,
100    jwks_cache: Arc<JwksCache>,
101    /// Issuer for JWKS-based verification
102    issuer: String,
103    /// Audience for JWKS-based verification
104    audiences: crate::AudienceSet,
105    require_origin: bool,
106    claim_seal_keys: Option<ClaimSealKeys>,
107}
108
109enum VerifierInner {
110    Static(TokenVerifier),
111    Jwks(JwksVerifier),
112    /// Keys are fetched from `jwks_url`; there is no local key to fall back to.
113    #[cfg(feature = "jwks")]
114    Remote,
115}
116
117impl AsyncVerifier {
118    /// Create a verifier with a static key
119    pub fn with_static_key(
120        key: VerifyingKey,
121        issuer: impl Into<String>,
122        audience: impl Into<String>,
123    ) -> Self {
124        let issuer_str = issuer.into();
125        let audience_str = audience.into();
126        Self {
127            inner: VerifierInner::Static(TokenVerifier::new(
128                key,
129                issuer_str.clone(),
130                audience_str.clone(),
131            )),
132            jwks_url: None,
133            cache_duration: Duration::from_secs(3600), // 1 hour default
134            min_refresh_interval: DEFAULT_MIN_JWKS_REFRESH_INTERVAL,
135            jwks_cache: Arc::default(),
136            issuer: issuer_str,
137            audiences: crate::AudienceSet::single(audience_str),
138            require_origin: false,
139            claim_seal_keys: None,
140        }
141    }
142
143    /// Create a verifier with JWKS
144    pub fn with_jwks(
145        jwks: crate::token::Jwks,
146        issuer: impl Into<String>,
147        audience: impl Into<String>,
148    ) -> Self {
149        let issuer_str = issuer.into();
150        let audience_str = audience.into();
151        Self {
152            inner: VerifierInner::Jwks(JwksVerifier::new(
153                jwks,
154                issuer_str.clone(),
155                audience_str.clone(),
156            )),
157            jwks_url: None,
158            cache_duration: Duration::from_secs(3600),
159            min_refresh_interval: DEFAULT_MIN_JWKS_REFRESH_INTERVAL,
160            jwks_cache: Arc::default(),
161            issuer: issuer_str,
162            audiences: crate::AudienceSet::single(audience_str),
163            require_origin: false,
164            claim_seal_keys: None,
165        }
166    }
167
168    /// Create a verifier that fetches JWKS from a URL
169    #[cfg(feature = "jwks")]
170    pub fn with_jwks_url(
171        url: impl Into<String>,
172        issuer: impl Into<String>,
173        audience: impl Into<String>,
174    ) -> Self {
175        Self {
176            inner: VerifierInner::Remote,
177            jwks_url: Some(url.into()),
178            issuer: issuer.into(),
179            audiences: crate::AudienceSet::single(audience.into()),
180            cache_duration: Duration::from_secs(3600),
181            min_refresh_interval: DEFAULT_MIN_JWKS_REFRESH_INTERVAL,
182            jwks_cache: Arc::default(),
183            require_origin: false,
184            claim_seal_keys: None,
185        }
186    }
187
188    /// Require origin validation on verified tokens.
189    pub fn with_origin_validation(mut self) -> Self {
190        self.require_origin = true;
191        self.inner = match self.inner {
192            VerifierInner::Static(verifier) => {
193                VerifierInner::Static(verifier.with_origin_validation())
194            }
195            VerifierInner::Jwks(verifier) => VerifierInner::Jwks(verifier.with_origin_validation()),
196            #[cfg(feature = "jwks")]
197            VerifierInner::Remote => VerifierInner::Remote,
198        };
199        self.jwks_cache = Arc::default();
200        self
201    }
202
203    /// Open sealed claims with these keys. Without them, a token that carries
204    /// sealed claims is refused.
205    pub fn with_claim_seal_keys(mut self, keys: ClaimSealKeys) -> Self {
206        self.inner = match self.inner {
207            VerifierInner::Static(verifier) => {
208                VerifierInner::Static(verifier.with_claim_seal_keys(keys.clone()))
209            }
210            VerifierInner::Jwks(verifier) => {
211                VerifierInner::Jwks(verifier.with_claim_seal_keys(keys.clone()))
212            }
213            #[cfg(feature = "jwks")]
214            VerifierInner::Remote => VerifierInner::Remote,
215        };
216        self.claim_seal_keys = Some(keys);
217        // A verifier cached before this call would still open with the old keys.
218        self.jwks_cache = Arc::default();
219        self
220    }
221
222    /// Set how long fetched JWKS keys count as fresh. Once they are stale,
223    /// requests try to refetch them (subject to the minimum refresh interval)
224    /// and keep verifying against them until a fetch succeeds.
225    pub fn with_cache_duration(mut self, duration: Duration) -> Self {
226        self.cache_duration = duration;
227        self
228    }
229
230    /// Set the minimum time from the end of one JWKS fetch to the start of
231    /// the next, whatever triggers it: expiry, an unknown `kid`, or a failed
232    /// fetch. A key published at the JWKS URL is picked up within about this
233    /// interval. Defaults to [`DEFAULT_MIN_JWKS_REFRESH_INTERVAL`].
234    pub fn with_min_refresh_interval(mut self, interval: Duration) -> Self {
235        self.min_refresh_interval = interval;
236        self
237    }
238
239    /// Verify a token with automatic JWKS fetching and caching
240    #[cfg(feature = "jwks")]
241    pub async fn verify(
242        &self,
243        token: &str,
244        expected_origin: Option<&str>,
245        expected_client_ip: Option<&str>,
246    ) -> Result<AuthContext, VerifyError> {
247        match &self.inner {
248            VerifierInner::Static(verifier) => {
249                verifier.verify(token, expected_origin, expected_client_ip)
250            }
251            VerifierInner::Jwks(verifier) => {
252                verifier.verify(token, expected_origin, expected_client_ip)
253            }
254            // A URL-backed verifier only trusts keys fetched from its JWKS URL.
255            VerifierInner::Remote => {
256                self.verify_with_cache(token, expected_origin, expected_client_ip)
257                    .await
258            }
259        }
260    }
261
262    /// Verify a token (non-JWKS version)
263    #[cfg(not(feature = "jwks"))]
264    pub fn verify(
265        &self,
266        token: &str,
267        expected_origin: Option<&str>,
268        expected_client_ip: Option<&str>,
269    ) -> Result<AuthContext, VerifyError> {
270        match &self.inner {
271            VerifierInner::Static(verifier) => {
272                verifier.verify(token, expected_origin, expected_client_ip)
273            }
274            VerifierInner::Jwks(verifier) => {
275                verifier.verify(token, expected_origin, expected_client_ip)
276            }
277        }
278    }
279
280    /// Fetch the JWKS from the configured URL now, ignoring the minimum
281    /// refresh interval. Waits for a fetch already in flight. On failure the
282    /// previously fetched keys stay in use.
283    #[cfg(feature = "jwks")]
284    pub async fn refresh_cache(&self) -> Result<(), VerifyError> {
285        let Some(jwks_url) = self.jwks_url.as_deref() else {
286            return Ok(());
287        };
288        let _fetch = self.jwks_cache.fetch.lock().await;
289        self.fetch_and_store(jwks_url).await
290    }
291
292    /// Fetch the JWKS unless one finished within the minimum refresh interval
293    /// or one is in flight. Failures are not returned: callers fall
294    /// back to whatever keys are cached.
295    #[cfg(feature = "jwks")]
296    async fn refresh_if_due(&self) {
297        let Some(jwks_url) = self.jwks_url.as_deref() else {
298            return;
299        };
300        let _fetch = if self.jwks_cache.current().is_some() {
301            if !self.jwks_cache.refresh_due(self.min_refresh_interval) {
302                return;
303            }
304            // Another request is fetching: use the cached keys rather than
305            // holding this request.
306            match self.jwks_cache.fetch.try_lock() {
307                Ok(guard) => guard,
308                Err(_) => return,
309            }
310        } else {
311            // No keys yet: wait for a fetch in flight instead of failing.
312            self.jwks_cache.fetch.lock().await
313        };
314        if !self.jwks_cache.refresh_due(self.min_refresh_interval) {
315            return;
316        }
317        if let Err(error) = self.fetch_and_store(jwks_url).await {
318            tracing::warn!(%error, "JWKS refresh failed; keeping the cached keys");
319        }
320    }
321
322    /// Fetch and cache the JWKS. The caller holds the fetch lock.
323    #[cfg(feature = "jwks")]
324    async fn fetch_and_store(&self, jwks_url: &str) -> Result<(), VerifyError> {
325        // The interval runs from the end of a fetch, failed or not. Marking on
326        // drop also covers a fetch cancelled with its request (for example by
327        // a handshake deadline), which would otherwise let the next request
328        // fetch at once.
329        let _mark = MarkAttemptOnDrop(&self.jwks_cache);
330        let jwks = tokio::time::timeout(
331            JWKS_FETCH_TIMEOUT,
332            crate::token::JwksVerifier::fetch_jwks(jwks_url),
333        )
334        .await
335        .map_err(|_| VerifyError::InvalidFormat("Failed to fetch JWKS: timed out".to_string()))?
336        .map_err(|e| VerifyError::InvalidFormat(format!("Failed to fetch JWKS: {}", e)))?;
337
338        let verifier = JwksVerifier::with_audience_set(jwks, &self.issuer, self.audiences.clone());
339        let verifier = if self.require_origin {
340            verifier.with_origin_validation()
341        } else {
342            verifier
343        };
344        let verifier = match &self.claim_seal_keys {
345            Some(keys) => verifier.with_claim_seal_keys(keys.clone()),
346            None => verifier,
347        };
348        self.jwks_cache.store(verifier);
349        Ok(())
350    }
351
352    /// Verify a token against the JWKS fetched from the configured URL.
353    ///
354    /// Fresh keys are used directly. An unknown `kid`, stale keys, or no keys
355    /// trigger a refetch, but at most one fetch runs at a time and the next
356    /// waits out the minimum refresh interval, so tokens with made-up key IDs
357    /// cannot drive outbound requests. If the fetch fails or is skipped, the
358    /// last successfully fetched keys are used.
359    #[cfg(feature = "jwks")]
360    pub async fn verify_with_cache(
361        &self,
362        token: &str,
363        expected_origin: Option<&str>,
364        expected_client_ip: Option<&str>,
365    ) -> Result<AuthContext, VerifyError> {
366        if let Some(cached) = self.jwks_cache.current() {
367            if cached.fetched_at.elapsed() < self.cache_duration {
368                match cached
369                    .verifier
370                    .verify(token, expected_origin, expected_client_ip)
371                {
372                    Err(VerifyError::KeyNotFound(_)) => {}
373                    result => return result,
374                }
375            }
376        }
377
378        self.refresh_if_due().await;
379
380        if let Some(cached) = self.jwks_cache.current() {
381            return cached
382                .verifier
383                .verify(token, expected_origin, expected_client_ip);
384        }
385
386        // No keys fetched yet: verifiers with a local key fall back to it; a
387        // URL-backed verifier has none and must fail closed.
388        match &self.inner {
389            VerifierInner::Static(verifier) => {
390                verifier.verify(token, expected_origin, expected_client_ip)
391            }
392            VerifierInner::Jwks(verifier) => {
393                verifier.verify(token, expected_origin, expected_client_ip)
394            }
395            VerifierInner::Remote => {
396                Err(VerifyError::InvalidFormat("JWKS unavailable".to_string()))
397            }
398        }
399    }
400}
401
402/// Simple synchronous verifier for use in non-async contexts
403pub struct SimpleVerifier {
404    inner: TokenVerifier,
405}
406
407impl SimpleVerifier {
408    /// Create a new simple verifier
409    pub fn new(key: VerifyingKey, issuer: impl Into<String>, audience: impl Into<String>) -> Self {
410        Self {
411            inner: TokenVerifier::new(key, issuer, audience),
412        }
413    }
414
415    /// Verify a token synchronously
416    pub fn verify(
417        &self,
418        token: &str,
419        expected_origin: Option<&str>,
420        expected_client_ip: Option<&str>,
421    ) -> Result<AuthContext, VerifyError> {
422        self.inner
423            .verify(token, expected_origin, expected_client_ip)
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use super::*;
430    use crate::claims::{KeyClass, SessionClaims};
431    use crate::keys::SigningKey;
432    use crate::token::TokenSigner;
433    use base64::Engine;
434
435    #[cfg(feature = "jwks")]
436    use tokio::io::{AsyncReadExt, AsyncWriteExt};
437
438    #[tokio::test]
439    async fn test_async_verifier_with_static_key() {
440        let signing_key = SigningKey::generate();
441        let verifying_key = signing_key.verifying_key();
442
443        let signer = TokenSigner::new(signing_key, "test-issuer");
444        let verifier =
445            AsyncVerifier::with_static_key(verifying_key, "test-issuer", "test-audience");
446
447        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
448            .with_scope("read")
449            .with_metering_key("meter-123")
450            .with_key_class(KeyClass::Publishable)
451            .build();
452
453        let token = signer.sign(claims).unwrap();
454        let context = verifier.verify(&token, None, None).await.unwrap();
455
456        assert_eq!(context.subject, "test-subject");
457    }
458
459    #[test]
460    fn test_simple_verifier() {
461        let signing_key = SigningKey::generate();
462        let verifying_key = signing_key.verifying_key();
463
464        let signer = TokenSigner::new(signing_key, "test-issuer");
465        let verifier = SimpleVerifier::new(verifying_key, "test-issuer", "test-audience");
466
467        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
468            .with_scope("read")
469            .with_metering_key("meter-123")
470            .with_key_class(KeyClass::Publishable)
471            .build();
472
473        let token = signer.sign(claims).unwrap();
474        let context = verifier.verify(&token, None, None).unwrap();
475
476        assert_eq!(context.subject, "test-subject");
477        assert_eq!(context.metering_key, "meter-123");
478    }
479
480    fn b64(bytes: impl AsRef<[u8]>) -> String {
481        base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
482    }
483
484    fn jwks_json(keys: &[&VerifyingKey]) -> String {
485        let keys: Vec<_> = keys
486            .iter()
487            .map(|key| {
488                serde_json::json!({
489                    "kty": "OKP",
490                    "use": "sig",
491                    "kid": key.key_id(),
492                    "x": b64(key.to_bytes()),
493                })
494            })
495            .collect();
496        serde_json::json!({ "keys": keys }).to_string()
497    }
498
499    /// Serve `body` as JSON to every connection until the test ends.
500    #[cfg(feature = "jwks")]
501    async fn serve_jwks(body: String) -> std::net::SocketAddr {
502        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
503        let addr = listener.local_addr().unwrap();
504        tokio::spawn(async move {
505            loop {
506                let (mut socket, _) = listener.accept().await.unwrap();
507                let body = body.clone();
508                tokio::spawn(async move {
509                    let mut buffer = [0u8; 1024];
510                    let _ = socket.read(&mut buffer).await;
511                    let response = format!(
512                        "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
513                        body.len(),
514                        body
515                    );
516                    let _ = socket.write_all(response.as_bytes()).await;
517                });
518            }
519        });
520        addr
521    }
522
523    fn signed_token(signing_key: SigningKey) -> String {
524        let claims = SessionClaims::builder("test-issuer", "test-subject", "test-audience")
525            .with_scope("read")
526            .with_metering_key("meter-123")
527            .with_key_class(KeyClass::Publishable)
528            .build();
529        TokenSigner::new(signing_key, "test-issuer")
530            .sign(claims)
531            .unwrap()
532    }
533
534    /// A token "signed" for the all-zero public key, which is a small-order
535    /// point: with R = identity and s = 0 the lax Ed25519 equation holds for
536    /// about one message in four, so no private key is needed.
537    fn forge_low_order_token(kid: Option<&str>) -> String {
538        use ed25519_dalek::Verifier as _;
539
540        let zero_key = ed25519_dalek::VerifyingKey::from_bytes(&[0u8; 32]).unwrap();
541        let mut signature_bytes = [0u8; 64];
542        signature_bytes[0] = 1; // compressed identity point; s stays 0
543        let signature = ed25519_dalek::Signature::from_bytes(&signature_bytes);
544
545        let mut header = serde_json::json!({ "alg": "EdDSA", "typ": "JWT" });
546        if let Some(kid) = kid {
547            header["kid"] = kid.into();
548        }
549        let header_b64 = b64(header.to_string());
550
551        for attempt in 0..512 {
552            let claims = SessionClaims::builder(
553                "test-issuer",
554                format!("forged-subject-{attempt}"),
555                "test-audience",
556            )
557            .with_scope("admin")
558            .with_metering_key("forged")
559            .with_key_class(KeyClass::Secret)
560            .build();
561            let message = format!("{header_b64}.{}", b64(serde_json::to_vec(&claims).unwrap()));
562            if zero_key.verify(message.as_bytes(), &signature).is_ok() {
563                return format!("{message}.{}", b64(signature_bytes));
564            }
565        }
566        panic!("no low-order forgery found");
567    }
568
569    fn zero_key() -> VerifyingKey {
570        VerifyingKey::from_bytes(&[0u8; 32]).unwrap()
571    }
572
573    #[cfg(feature = "jwks")]
574    #[tokio::test]
575    async fn jwks_url_verify_uses_the_fetched_keys() {
576        let signing_key = SigningKey::generate();
577        let addr = serve_jwks(jwks_json(&[&signing_key.verifying_key()])).await;
578        let verifier = AsyncVerifier::with_jwks_url(
579            format!("http://{addr}/jwks"),
580            "test-issuer",
581            "test-audience",
582        );
583
584        let token = signed_token(signing_key);
585        let context = verifier.verify(&token, None, None).await.unwrap();
586        assert_eq!(context.subject, "test-subject");
587        let context = verifier
588            .verify_with_cache(&token, None, None)
589            .await
590            .unwrap();
591        assert_eq!(context.subject, "test-subject");
592    }
593
594    #[cfg(feature = "jwks")]
595    #[tokio::test]
596    async fn jwks_url_verify_fails_closed_when_the_jwks_is_unreachable() {
597        // Bind then drop, so nothing is listening on the port.
598        let addr = tokio::net::TcpListener::bind("127.0.0.1:0")
599            .await
600            .unwrap()
601            .local_addr()
602            .unwrap();
603        let verifier = AsyncVerifier::with_jwks_url(
604            format!("http://{addr}/jwks"),
605            "test-issuer",
606            "test-audience",
607        );
608
609        for token in [
610            signed_token(SigningKey::generate()),
611            forge_low_order_token(None),
612        ] {
613            assert!(verifier.verify(&token, None, None).await.is_err());
614            assert!(verifier
615                .verify_with_cache(&token, None, None)
616                .await
617                .is_err());
618        }
619    }
620
621    #[cfg(feature = "jwks")]
622    #[tokio::test]
623    async fn jwks_url_verify_rejects_unknown_keys_and_low_order_forgeries() {
624        let trusted = SigningKey::generate();
625        let trusted_key = trusted.verifying_key();
626        let addr = serve_jwks(jwks_json(&[&trusted_key])).await;
627        let verifier = AsyncVerifier::with_jwks_url(
628            format!("http://{addr}/jwks"),
629            "test-issuer",
630            "test-audience",
631        );
632
633        let untrusted = signed_token(SigningKey::generate());
634        let forged = forge_low_order_token(None);
635        let forged_with_trusted_kid = forge_low_order_token(Some(&trusted_key.key_id()));
636
637        for token in [&untrusted, &forged, &forged_with_trusted_kid] {
638            assert!(verifier.verify(token, None, None).await.is_err());
639            assert!(verifier.verify_with_cache(token, None, None).await.is_err());
640        }
641        assert!(matches!(
642            verifier.verify(&forged_with_trusted_kid, None, None).await,
643            Err(VerifyError::InvalidSignature)
644        ));
645    }
646
647    /// A JWKS endpoint that counts requests and can be switched to new keys
648    /// or to failing with a 500.
649    #[cfg(feature = "jwks")]
650    struct CountingJwks {
651        addr: std::net::SocketAddr,
652        hits: Arc<std::sync::atomic::AtomicUsize>,
653        body: Arc<std::sync::Mutex<Option<String>>>,
654    }
655
656    #[cfg(feature = "jwks")]
657    impl CountingJwks {
658        /// Serve `body`, answering each request after `delay`.
659        async fn start(body: String, delay: Duration) -> Self {
660            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
661            let addr = listener.local_addr().unwrap();
662            let hits = Arc::new(std::sync::atomic::AtomicUsize::new(0));
663            let body = Arc::new(std::sync::Mutex::new(Some(body)));
664            let (served_hits, served_body) = (hits.clone(), body.clone());
665            tokio::spawn(async move {
666                loop {
667                    let (mut socket, _) = listener.accept().await.unwrap();
668                    let (hits, body) = (served_hits.clone(), served_body.clone());
669                    tokio::spawn(async move {
670                        let mut buffer = [0u8; 1024];
671                        let _ = socket.read(&mut buffer).await;
672                        hits.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
673                        tokio::time::sleep(delay).await;
674                        let (status, body) = match body.lock().unwrap().clone() {
675                            Some(body) => ("200 OK", body),
676                            None => ("500 Internal Server Error", "unavailable".to_string()),
677                        };
678                        let response = format!(
679                            "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
680                            body.len(),
681                            body
682                        );
683                        let _ = socket.write_all(response.as_bytes()).await;
684                    });
685                }
686            });
687            Self { addr, hits, body }
688        }
689
690        fn url(&self) -> String {
691            format!("http://{}/jwks", self.addr)
692        }
693
694        fn hits(&self) -> usize {
695            self.hits.load(std::sync::atomic::Ordering::SeqCst)
696        }
697
698        fn serve(&self, body: Option<String>) {
699            *self.body.lock().unwrap() = body;
700        }
701    }
702
703    #[cfg(feature = "jwks")]
704    fn copy_key(key: &SigningKey) -> SigningKey {
705        SigningKey::from_bytes(&key.to_bytes())
706    }
707
708    #[cfg(feature = "jwks")]
709    async fn verify_concurrently(
710        verifier: &Arc<AsyncVerifier>,
711        tokens: Vec<String>,
712    ) -> Vec<Result<AuthContext, VerifyError>> {
713        let mut tasks = tokio::task::JoinSet::new();
714        for token in tokens {
715            let verifier = verifier.clone();
716            tasks.spawn(async move { verifier.verify(&token, None, None).await });
717        }
718        tasks.join_all().await
719    }
720
721    #[cfg(feature = "jwks")]
722    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
723    async fn unknown_kids_trigger_at_most_one_fetch_per_interval() {
724        let trusted = SigningKey::generate();
725        let server = CountingJwks::start(
726            jwks_json(&[&trusted.verifying_key()]),
727            Duration::from_millis(50),
728        )
729        .await;
730        let verifier = Arc::new(AsyncVerifier::with_jwks_url(
731            server.url(),
732            "test-issuer",
733            "test-audience",
734        ));
735        let random_kid_tokens = || -> Vec<String> {
736            (0..64)
737                .map(|_| signed_token(SigningKey::generate()))
738                .collect()
739        };
740
741        // Cold cache: every request misses, one fetch serves them all.
742        let results = verify_concurrently(&verifier, random_kid_tokens()).await;
743        assert!(results
744            .iter()
745            .all(|result| matches!(result, Err(VerifyError::KeyNotFound(_)))));
746        assert_eq!(server.hits(), 1);
747
748        // Warm cache: more unknown kids within the interval fetch nothing.
749        let results = verify_concurrently(&verifier, random_kid_tokens()).await;
750        assert!(results
751            .iter()
752            .all(|result| matches!(result, Err(VerifyError::KeyNotFound(_)))));
753        assert_eq!(server.hits(), 1);
754
755        let context = verifier
756            .verify(&signed_token(trusted), None, None)
757            .await
758            .unwrap();
759        assert_eq!(context.subject, "test-subject");
760        assert_eq!(server.hits(), 1);
761    }
762
763    #[cfg(feature = "jwks")]
764    async fn wait_for_hits(server: &CountingJwks, hits: usize) {
765        tokio::time::timeout(Duration::from_secs(10), async {
766            while server.hits() < hits {
767                tokio::time::sleep(Duration::from_millis(10)).await;
768            }
769        })
770        .await
771        .expect("JWKS request never arrived");
772    }
773
774    #[cfg(feature = "jwks")]
775    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
776    async fn cached_tokens_do_not_wait_behind_a_due_refresh() {
777        let trusted = SigningKey::generate();
778        let fetch_time = Duration::from_secs(2);
779        let server = CountingJwks::start(jwks_json(&[&trusted.verifying_key()]), fetch_time).await;
780        let interval = Duration::from_millis(100);
781        // Keys go stale at once, so every request after the interval is due
782        // to refresh.
783        let verifier = Arc::new(
784            AsyncVerifier::with_jwks_url(server.url(), "test-issuer", "test-audience")
785                .with_cache_duration(Duration::ZERO)
786                .with_min_refresh_interval(interval),
787        );
788        verifier.refresh_cache().await.unwrap();
789        tokio::time::sleep(interval + Duration::from_millis(50)).await;
790
791        // One request starts the due refresh and holds it open.
792        let refreshing = {
793            let verifier = verifier.clone();
794            let token = signed_token(copy_key(&trusted));
795            tokio::spawn(async move { verifier.verify(&token, None, None).await })
796        };
797        wait_for_hits(&server, 2).await;
798
799        // Requests arriving meanwhile use the cached keys at once and start
800        // no fetch of their own.
801        let tokens = (0..16).map(|_| signed_token(copy_key(&trusted))).collect();
802        let results = tokio::time::timeout(fetch_time / 2, verify_concurrently(&verifier, tokens))
803            .await
804            .expect("requests waited behind the in-flight fetch");
805        assert!(results.iter().all(Result::is_ok));
806        assert!(!refreshing.is_finished());
807        assert_eq!(server.hits(), 2);
808
809        refreshing.await.unwrap().unwrap();
810        assert_eq!(server.hits(), 2);
811    }
812
813    #[cfg(feature = "jwks")]
814    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
815    async fn cancelled_fetch_still_counts_against_the_interval() {
816        let server = CountingJwks::start(
817            jwks_json(&[&SigningKey::generate().verifying_key()]),
818            Duration::from_secs(2),
819        )
820        .await;
821        let verifier = Arc::new(AsyncVerifier::with_jwks_url(
822            server.url(),
823            "test-issuer",
824            "test-audience",
825        ));
826
827        // A request is dropped mid-fetch, as a handshake deadline would.
828        let cancelled = {
829            let verifier = verifier.clone();
830            let token = signed_token(SigningKey::generate());
831            tokio::spawn(async move { verifier.verify(&token, None, None).await })
832        };
833        wait_for_hits(&server, 1).await;
834        cancelled.abort();
835        assert!(cancelled.await.unwrap_err().is_cancelled());
836
837        // The next unknown kid does not fetch again inside the interval.
838        assert!(verifier
839            .verify(&signed_token(SigningKey::generate()), None, None)
840            .await
841            .is_err());
842        tokio::time::sleep(Duration::from_millis(100)).await;
843        assert_eq!(server.hits(), 1);
844    }
845
846    #[cfg(feature = "jwks")]
847    #[tokio::test]
848    async fn rotated_key_is_picked_up_after_the_interval() {
849        let old_key = SigningKey::generate();
850        let new_key = SigningKey::generate();
851        let new_token = signed_token(copy_key(&new_key));
852        let server =
853            CountingJwks::start(jwks_json(&[&old_key.verifying_key()]), Duration::ZERO).await;
854        let interval = Duration::from_millis(300);
855        let verifier = AsyncVerifier::with_jwks_url(server.url(), "test-issuer", "test-audience")
856            .with_min_refresh_interval(interval);
857
858        verifier
859            .verify(&signed_token(copy_key(&old_key)), None, None)
860            .await
861            .unwrap();
862        assert_eq!(server.hits(), 1);
863
864        server.serve(Some(jwks_json(&[
865            &old_key.verifying_key(),
866            &new_key.verifying_key(),
867        ])));
868        assert!(matches!(
869            verifier.verify(&new_token, None, None).await,
870            Err(VerifyError::KeyNotFound(_))
871        ));
872        assert_eq!(server.hits(), 1, "no fetch inside the interval");
873
874        tokio::time::sleep(interval + Duration::from_millis(50)).await;
875        let context = verifier.verify(&new_token, None, None).await.unwrap();
876        assert_eq!(context.subject, "test-subject");
877        assert_eq!(server.hits(), 2);
878        verifier
879            .verify(&signed_token(old_key), None, None)
880            .await
881            .unwrap();
882        assert_eq!(server.hits(), 2);
883    }
884
885    #[cfg(feature = "jwks")]
886    #[tokio::test]
887    async fn failed_refresh_keeps_serving_cached_keys() {
888        let trusted = SigningKey::generate();
889        let server =
890            CountingJwks::start(jwks_json(&[&trusted.verifying_key()]), Duration::ZERO).await;
891        let interval = Duration::from_millis(300);
892        // Keys go stale at once, so every request after the interval refetches.
893        let verifier = AsyncVerifier::with_jwks_url(server.url(), "test-issuer", "test-audience")
894            .with_cache_duration(Duration::ZERO)
895            .with_min_refresh_interval(interval);
896
897        verifier
898            .verify(&signed_token(copy_key(&trusted)), None, None)
899            .await
900            .unwrap();
901        assert_eq!(server.hits(), 1);
902
903        server.serve(None);
904        tokio::time::sleep(interval + Duration::from_millis(50)).await;
905        let context = verifier
906            .verify(&signed_token(copy_key(&trusted)), None, None)
907            .await
908            .unwrap();
909        assert_eq!(context.subject, "test-subject");
910        assert_eq!(server.hits(), 2, "the failed fetch was attempted");
911
912        // Failures are throttled like successes.
913        assert!(matches!(
914            verifier
915                .verify(&signed_token(SigningKey::generate()), None, None)
916                .await,
917            Err(VerifyError::KeyNotFound(_))
918        ));
919        verifier
920            .verify(&signed_token(copy_key(&trusted)), None, None)
921            .await
922            .unwrap();
923        assert_eq!(server.hits(), 2);
924
925        // An explicit refresh reports the failure and still keeps the keys.
926        assert!(verifier.refresh_cache().await.is_err());
927        assert_eq!(server.hits(), 3);
928        verifier
929            .verify(&signed_token(trusted), None, None)
930            .await
931            .unwrap();
932    }
933
934    #[tokio::test]
935    async fn low_order_forgeries_are_rejected_by_every_verifier() {
936        let zero = zero_key();
937        let forged = forge_low_order_token(None);
938        let forged_with_kid = forge_low_order_token(Some(&zero.key_id()));
939
940        for token in [&forged, &forged_with_kid] {
941            assert!(matches!(
942                TokenVerifier::new(zero.clone(), "test-issuer", "test-audience")
943                    .verify(token, None, None),
944                Err(VerifyError::InvalidSignature)
945            ));
946            assert!(matches!(
947                SimpleVerifier::new(zero.clone(), "test-issuer", "test-audience")
948                    .verify(token, None, None),
949                Err(VerifyError::InvalidSignature)
950            ));
951
952            let static_verifier =
953                AsyncVerifier::with_static_key(zero.clone(), "test-issuer", "test-audience");
954            assert!(static_verifier.verify(token, None, None).await.is_err());
955        }
956
957        let jwks: crate::token::Jwks = serde_json::from_str(&jwks_json(&[&zero])).unwrap();
958        assert!(matches!(
959            JwksVerifier::new(jwks.clone(), "test-issuer", "test-audience").verify(
960                &forged_with_kid,
961                None,
962                None
963            ),
964            Err(VerifyError::InvalidSignature)
965        ));
966        let jwks_verifier = AsyncVerifier::with_jwks(jwks, "test-issuer", "test-audience");
967        assert!(matches!(
968            jwks_verifier.verify(&forged_with_kid, None, None).await,
969            Err(VerifyError::InvalidSignature)
970        ));
971    }
972}