modkit_auth/providers/
jwks.rs

1use crate::{claims_error::ClaimsError, traits::KeyProvider};
2use arc_swap::ArcSwap;
3use async_trait::async_trait;
4use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
5use jsonwebtoken::{DecodingKey, Header, Validation, decode, decode_header};
6use serde::Deserialize;
7use serde_json::Value;
8use std::collections::{HashMap, HashSet};
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::sync::RwLock;
12use tokio::time::Instant;
13use tokio_util::sync::CancellationToken;
14
15#[derive(Debug, Clone, Deserialize)]
16struct Jwk {
17    kid: String,
18    kty: String,
19    #[serde(rename = "use")]
20    #[allow(dead_code)]
21    use_: Option<String>,
22    n: String,
23    e: String,
24    #[allow(dead_code)]
25    alg: Option<String>,
26}
27
28#[derive(Debug, Clone, Deserialize)]
29struct JwksResponse {
30    keys: Vec<Jwk>,
31}
32
33/// Handler for non-string custom JWT header fields; return `Some` to keep as string, or `None` to drop.
34type HeaderExtrasHandler = dyn Fn(&str, &Value) -> Option<String> + Send + Sync;
35
36/// Standard JWT header field names from RFC 7515 (JWS), RFC 7516 (JWE),
37/// RFC 7518 (JWA), RFC 7797 (b64), and RFC 8555 (ACME).
38const STANDARD_HEADER_FIELDS: &[&str] = &[
39    "typ", "alg", "cty", "jku", "jwk", "kid", "x5u", "x5c", "x5t", "x5t#S256", "crit", "enc",
40    "zip", "url", "nonce", "epk", "apu", "apv", "iv", "tag", "p2s", "p2c", "b64",
41];
42
43/// JWKS-based key provider with lock-free reads
44///
45/// Uses `ArcSwap` for lock-free key lookups and background refresh with exponential backoff.
46#[must_use]
47pub struct JwksKeyProvider {
48    /// JWKS endpoint URL
49    jwks_uri: String,
50
51    /// Keys stored in `ArcSwap` for lock-free reads
52    keys: Arc<ArcSwap<HashMap<String, DecodingKey>>>,
53
54    /// Last refresh time and error tracking for backoff
55    refresh_state: Arc<RwLock<RefreshState>>,
56
57    /// Shared HTTP client for JWKS fetches (pooled connections)
58    /// `HttpClient` is `Clone + Send + Sync`, no external locking needed.
59    client: modkit_http::HttpClient,
60
61    /// Refresh interval (default: 5 minutes)
62    refresh_interval: Duration,
63
64    /// Maximum backoff duration (default: 1 hour)
65    max_backoff: Duration,
66
67    /// Cooldown for on-demand refresh (default: 60 seconds)
68    on_demand_refresh_cooldown: Duration,
69
70    /// Optional handler for non-string custom JWT header fields.
71    /// Called for each non-standard field whose value is not a JSON string.
72    /// Return `Some(s)` to keep, `None` to drop.
73    header_extras_handler: Option<Arc<HeaderExtrasHandler>>,
74}
75
76#[derive(Debug, Default)]
77struct RefreshState {
78    last_refresh: Option<Instant>,
79    last_on_demand_refresh: Option<Instant>,
80    consecutive_failures: u32,
81    last_error: Option<String>,
82    failed_kids: HashSet<String>,
83}
84
85impl JwksKeyProvider {
86    /// Create a new JWKS key provider
87    ///
88    /// # Errors
89    /// Returns error if HTTP client initialization fails (e.g., TLS setup)
90    pub fn new(jwks_uri: impl Into<String>) -> Result<Self, modkit_http::HttpError> {
91        Self::with_http_timeout(jwks_uri, Duration::from_secs(10))
92    }
93
94    /// Create a new JWKS key provider with custom HTTP timeout
95    ///
96    /// # Errors
97    /// Returns error if HTTP client initialization fails (e.g., TLS setup)
98    pub fn with_http_timeout(
99        jwks_uri: impl Into<String>,
100        timeout: Duration,
101    ) -> Result<Self, modkit_http::HttpError> {
102        let client = modkit_http::HttpClient::builder()
103            .timeout(timeout)
104            .retry(None) // JWKS provider handles its own retry logic
105            .build()?;
106
107        Ok(Self {
108            jwks_uri: jwks_uri.into(),
109            keys: Arc::new(ArcSwap::from_pointee(HashMap::new())),
110            refresh_state: Arc::new(RwLock::new(RefreshState::default())),
111            client,
112            refresh_interval: Duration::from_secs(300), // 5 minutes
113            max_backoff: Duration::from_secs(3600),     // 1 hour
114            on_demand_refresh_cooldown: Duration::from_secs(60), // 1 minute
115            header_extras_handler: None,
116        })
117    }
118
119    /// Create a new JWKS key provider (alias for new, kept for compatibility)
120    ///
121    /// # Errors
122    /// Returns error if HTTP client initialization fails (e.g., TLS setup)
123    pub fn try_new(jwks_uri: impl Into<String>) -> Result<Self, modkit_http::HttpError> {
124        Self::new(jwks_uri)
125    }
126
127    /// Create with custom refresh interval
128    pub fn with_refresh_interval(mut self, interval: Duration) -> Self {
129        self.refresh_interval = interval;
130        self
131    }
132
133    /// Create with custom max backoff
134    pub fn with_max_backoff(mut self, max_backoff: Duration) -> Self {
135        self.max_backoff = max_backoff;
136        self
137    }
138
139    /// Create with custom on-demand refresh cooldown
140    pub fn with_on_demand_refresh_cooldown(mut self, cooldown: Duration) -> Self {
141        self.on_demand_refresh_cooldown = cooldown;
142        self
143    }
144
145    /// Stringify all non-string custom JWT header fields.
146    ///
147    /// Convenience wrapper around [`with_header_extras_handler`](Self::with_header_extras_handler)
148    /// that converts every non-string value to its JSON representation
149    /// (e.g. `123` → `"123"`, `true` → `"true"`, `[1,2]` → `"[1,2]"`).
150    pub fn with_header_extras_stringified(self) -> Self {
151        self.with_header_extras_handler(|_, v| Some(v.to_string()))
152    }
153
154    /// Set a handler for non-string custom JWT header fields.
155    ///
156    /// `jsonwebtoken::Header::extras` is `HashMap<String, String>` and rejects
157    /// non-string values. This callback is invoked for each such field.
158    /// Return `Some(s)` to keep, `None` to drop.
159    /// Without a handler, upstream `decode_header` is used as-is.
160    pub fn with_header_extras_handler(
161        mut self,
162        handler: impl Fn(&str, &Value) -> Option<String> + Send + Sync + 'static,
163    ) -> Self {
164        self.header_extras_handler = Some(Arc::new(handler));
165        self
166    }
167
168    /// Fetch JWKS from the endpoint
169    async fn fetch_jwks(&self) -> Result<HashMap<String, DecodingKey>, ClaimsError> {
170        // HttpClient is Clone + Send + Sync, no locking needed
171        let jwks: JwksResponse = self
172            .client
173            .get(&self.jwks_uri)
174            .send()
175            .await
176            .map_err(|e| map_http_error(&e))?
177            .json()
178            .await
179            .map_err(|e| map_http_error(&e))?;
180
181        let mut keys = HashMap::new();
182        for jwk in jwks.keys {
183            if jwk.kty == "RSA" {
184                let key = DecodingKey::from_rsa_components(&jwk.n, &jwk.e)
185                    .map_err(|e| ClaimsError::JwksFetchFailed(format!("Invalid RSA key: {e}")))?;
186                keys.insert(jwk.kid, key);
187            }
188        }
189
190        if keys.is_empty() {
191            return Err(ClaimsError::JwksFetchFailed(
192                "No valid RSA keys found in JWKS".into(),
193            ));
194        }
195
196        Ok(keys)
197    }
198
199    /// Calculate backoff duration based on consecutive failures
200    fn calculate_backoff(&self, failures: u32) -> Duration {
201        let base = Duration::from_secs(60); // 1 minute base
202        let exponential = base * 2u32.pow(failures.min(10)); // Cap at 2^10
203        exponential.min(self.max_backoff)
204    }
205
206    /// Check if refresh is needed based on interval and backoff
207    async fn should_refresh(&self) -> bool {
208        let state = self.refresh_state.read().await;
209
210        match state.last_refresh {
211            None => true, // Never refreshed
212            Some(last) => {
213                let elapsed = last.elapsed();
214                if state.consecutive_failures == 0 {
215                    // Normal refresh interval
216                    elapsed >= self.refresh_interval
217                } else {
218                    // Exponential backoff
219                    elapsed >= self.calculate_backoff(state.consecutive_failures)
220                }
221            }
222        }
223    }
224
225    /// Perform key refresh with error tracking
226    async fn perform_refresh(&self) -> Result<(), ClaimsError> {
227        match self.fetch_jwks().await {
228            Ok(new_keys) => {
229                // Update keys atomically
230                self.keys.store(Arc::new(new_keys));
231
232                // Update refresh state
233                let mut state = self.refresh_state.write().await;
234                state.last_refresh = Some(Instant::now());
235                state.consecutive_failures = 0;
236                state.last_error = None;
237
238                Ok(())
239            }
240            Err(e) => {
241                // Update failure state
242                let mut state = self.refresh_state.write().await;
243                state.last_refresh = Some(Instant::now());
244                state.consecutive_failures += 1;
245                state.last_error = Some(e.to_string());
246
247                Err(e)
248            }
249        }
250    }
251
252    /// Check if a key exists in the cache
253    fn key_exists(&self, kid: &str) -> bool {
254        let keys = self.keys.load();
255        keys.contains_key(kid)
256    }
257
258    /// Check if we're in cooldown period and handle throttling logic
259    async fn check_refresh_throttle(&self, kid: &str) -> Result<(), ClaimsError> {
260        let state = self.refresh_state.read().await;
261        if let Some(last_on_demand) = state.last_on_demand_refresh {
262            let elapsed = last_on_demand.elapsed();
263            if elapsed < self.on_demand_refresh_cooldown {
264                let remaining = self.on_demand_refresh_cooldown.saturating_sub(elapsed);
265                tracing::debug!(
266                    kid = kid,
267                    remaining_secs = remaining.as_secs(),
268                    "On-demand JWKS refresh throttled (cooldown active)"
269                );
270
271                // Check if this kid has failed before
272                if state.failed_kids.contains(kid) {
273                    tracing::warn!(
274                        kid = kid,
275                        "Unknown kid repeatedly requested despite recent refresh attempts"
276                    );
277                }
278
279                return Err(ClaimsError::UnknownKeyId(kid.to_owned()));
280            }
281        }
282        Ok(())
283    }
284
285    /// Update state after successful refresh and check if kid is now available
286    async fn handle_refresh_success(&self, kid: &str) -> Result<(), ClaimsError> {
287        let mut state = self.refresh_state.write().await;
288        state.last_on_demand_refresh = Some(Instant::now());
289
290        // Check if the kid now exists
291        if self.key_exists(kid) {
292            // Kid found - remove from failed list if present
293            state.failed_kids.remove(kid);
294        } else {
295            // Kid still not found after refresh - track it
296            state.failed_kids.insert(kid.to_owned());
297            tracing::warn!(
298                kid = kid,
299                "Kid still not found after on-demand JWKS refresh"
300            );
301        }
302
303        Ok(())
304    }
305
306    /// Update state after failed refresh
307    async fn handle_refresh_failure(&self, kid: &str, error: ClaimsError) -> ClaimsError {
308        let mut state = self.refresh_state.write().await;
309        state.last_on_demand_refresh = Some(Instant::now());
310        state.failed_kids.insert(kid.to_owned());
311        error
312    }
313
314    /// Try to refresh keys if unknown kid is encountered
315    /// Implements throttling to prevent excessive refreshes
316    async fn on_demand_refresh(&self, kid: &str) -> Result<(), ClaimsError> {
317        // Check if key exists
318        if self.key_exists(kid) {
319            return Ok(());
320        }
321
322        // Check if we're in cooldown period
323        self.check_refresh_throttle(kid).await?;
324
325        // Attempt refresh and track the kid if it fails
326        tracing::info!(
327            kid = kid,
328            "Performing on-demand JWKS refresh for unknown kid"
329        );
330
331        match self.perform_refresh().await {
332            Ok(()) => self.handle_refresh_success(kid).await,
333            Err(e) => Err(self.handle_refresh_failure(kid, e).await),
334        }
335    }
336
337    /// Get a key by kid (lock-free read)
338    fn get_key(&self, kid: &str) -> Option<DecodingKey> {
339        let keys = self.keys.load();
340        keys.get(kid).cloned()
341    }
342
343    /// Validate JWT and decode into header + raw claims
344    fn validate_token(
345        token: &str,
346        key: &DecodingKey,
347        header: &Header,
348    ) -> Result<Value, ClaimsError> {
349        let mut validation = Validation::new(header.alg);
350
351        // Disable all built-in validations - we'll do them separately
352        validation.validate_exp = false;
353        validation.validate_nbf = false;
354        validation.validate_aud = false;
355
356        // Don't require any standard claims
357        let empty_claims: &[&str] = &[];
358        validation.set_required_spec_claims(empty_claims);
359
360        let token_data = decode::<Value>(token, key, &validation)
361            .map_err(|e| ClaimsError::DecodeFailed(format!("JWT validation failed: {e}")))?;
362
363        Ok(token_data.claims)
364    }
365}
366
367#[async_trait]
368impl KeyProvider for JwksKeyProvider {
369    fn name(&self) -> &'static str {
370        "jwks"
371    }
372
373    async fn validate_and_decode(&self, token: &str) -> Result<(Header, Value), ClaimsError> {
374        // Strip "Bearer " prefix if present
375        let token = token.trim_start_matches("Bearer ").trim();
376
377        // Decode header to get kid and algorithm
378        let header = match &self.header_extras_handler {
379            Some(handler) => decode_header_with_handler(token, handler.as_ref()),
380            None => decode_header(token),
381        }
382        .map_err(|e| ClaimsError::DecodeFailed(format!("Invalid JWT header: {e}")))?;
383
384        let kid = header
385            .kid
386            .as_ref()
387            .ok_or_else(|| ClaimsError::DecodeFailed("Missing kid in JWT header".into()))?;
388
389        // Try to get key from cache
390        let key = if let Some(k) = self.get_key(kid) {
391            k
392        } else {
393            // Key not in cache, try on-demand refresh
394            self.on_demand_refresh(kid).await?;
395
396            // Try again after refresh
397            self.get_key(kid)
398                .ok_or_else(|| ClaimsError::UnknownKeyId(kid.clone()))?
399        };
400
401        // Validate signature and decode claims
402        let claims = Self::validate_token(token, &key, &header)?;
403
404        Ok((header, claims))
405    }
406
407    async fn refresh_keys(&self) -> Result<(), ClaimsError> {
408        if self.should_refresh().await {
409            self.perform_refresh().await
410        } else {
411            Ok(())
412        }
413    }
414}
415
416/// Background task to periodically refresh JWKS
417///
418/// This task will run until the `cancellation_token` is cancelled, enabling
419/// graceful shutdown per `ModKit` patterns. Without cancellation support, this
420/// task would run indefinitely and potentially cause process hang on shutdown.
421///
422/// # Example
423///
424/// ```ignore
425/// use tokio_util::sync::CancellationToken;
426/// use std::sync::Arc;
427///
428/// let provider = Arc::new(JwksKeyProvider::new("https://issuer/.well-known/jwks.json")?);
429/// let cancel_token = CancellationToken::new();
430///
431/// // Spawn the refresh task
432/// let task_handle = tokio::spawn(run_jwks_refresh_task(provider.clone(), cancel_token.clone()));
433///
434/// // On shutdown:
435/// cancel_token.cancel();
436/// task_handle.await?;
437/// ```
438pub async fn run_jwks_refresh_task(
439    provider: Arc<JwksKeyProvider>,
440    cancellation_token: CancellationToken,
441) {
442    let mut interval = tokio::time::interval(Duration::from_secs(60)); // Check every minute
443
444    loop {
445        tokio::select! {
446            () = cancellation_token.cancelled() => {
447                tracing::info!("JWKS refresh task shutting down");
448                break;
449            }
450            _ = interval.tick() => {
451                if let Err(e) = provider.refresh_keys().await {
452                    tracing::warn!("JWKS refresh failed: {}", e);
453                }
454            }
455        }
456    }
457}
458
459/// Decode a JWT header, routing non-string custom fields through `handler`.
460///
461/// Returns `Some(s)` to keep the field, `None` to drop it.
462fn decode_header_with_handler(
463    token: &str,
464    handler: &dyn Fn(&str, &Value) -> Option<String>,
465) -> Result<Header, jsonwebtoken::errors::Error> {
466    let header_b64 = token
467        .split('.')
468        .next()
469        .ok_or(jsonwebtoken::errors::ErrorKind::InvalidToken)?;
470
471    let header_bytes = URL_SAFE_NO_PAD
472        .decode(header_b64.trim_end_matches('='))
473        .map_err(jsonwebtoken::errors::ErrorKind::Base64)?;
474
475    let mut json: serde_json::Map<String, Value> = serde_json::from_slice(&header_bytes)?;
476
477    json.retain(|key, value| {
478        if STANDARD_HEADER_FIELDS.contains(&key.as_str()) || value.is_string() {
479            return true;
480        }
481        match handler(key, value) {
482            Some(s) => {
483                *value = Value::String(s);
484                true
485            }
486            None => false,
487        }
488    });
489
490    Ok(serde_json::from_value(Value::Object(json))?)
491}
492
493/// Map `HttpError` variants to appropriate `ClaimsError` messages
494fn map_http_error(e: &modkit_http::HttpError) -> ClaimsError {
495    ClaimsError::JwksFetchFailed(crate::http_error::format_http_error(e, "JWKS"))
496}
497
498#[cfg(test)]
499#[cfg_attr(coverage_nightly, coverage(off))]
500mod tests {
501    use super::*;
502    use httpmock::prelude::*;
503
504    /// Create a test provider with insecure HTTP allowed (for httpmock) and no retries
505    fn test_provider_with_http(uri: &str) -> JwksKeyProvider {
506        let client = modkit_http::HttpClient::builder()
507            .timeout(Duration::from_secs(5))
508            .retry(None)
509            .build()
510            .expect("failed to create test HTTP client");
511
512        JwksKeyProvider {
513            jwks_uri: uri.to_owned(),
514            keys: Arc::new(ArcSwap::from_pointee(HashMap::new())),
515            refresh_state: Arc::new(RwLock::new(RefreshState::default())),
516            client,
517            refresh_interval: Duration::from_secs(300),
518            max_backoff: Duration::from_secs(3600),
519            on_demand_refresh_cooldown: Duration::from_secs(60),
520            header_extras_handler: None,
521        }
522    }
523
524    /// Create a basic test provider (HTTPS only, for non-network tests)
525    fn test_provider(uri: &str) -> JwksKeyProvider {
526        JwksKeyProvider::new(uri).expect("failed to create test provider")
527    }
528
529    /// Valid JWKS JSON response with a single RSA key
530    fn valid_jwks_json() -> &'static str {
531        r#"{
532            "keys": [{
533                "kty": "RSA",
534                "kid": "test-key-1",
535                "use": "sig",
536                "n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368QQMicAtaSqzs8KJZgnYb9c7d0zgdAZHzu6qMQvRL5hajrn1n91CbOpbISD08qNLyrdkt-bFTWhAI4vMQFh6WeZu0fM4lFd2NcRwr3XPksINHaQ-G_xBniIqbw0Ls1jF44-csFCur-kEgU8awapJzKnqDKgw",
537                "e": "AQAB",
538                "alg": "RS256"
539            }]
540        }"#
541    }
542
543    #[tokio::test]
544    async fn test_calculate_backoff() {
545        let provider = test_provider("https://example.com/jwks");
546
547        assert_eq!(provider.calculate_backoff(0), Duration::from_secs(60));
548        assert_eq!(provider.calculate_backoff(1), Duration::from_secs(120));
549        assert_eq!(provider.calculate_backoff(2), Duration::from_secs(240));
550        assert_eq!(provider.calculate_backoff(3), Duration::from_secs(480));
551
552        // Should cap at max_backoff
553        assert_eq!(provider.calculate_backoff(100), provider.max_backoff);
554    }
555
556    #[tokio::test]
557    async fn test_should_refresh_on_first_call() {
558        let provider = test_provider("https://example.com/jwks");
559        assert!(provider.should_refresh().await);
560    }
561
562    #[tokio::test]
563    async fn test_key_storage() {
564        let provider = test_provider("https://example.com/jwks");
565
566        // Initially empty
567        assert!(provider.get_key("test-kid").is_none());
568
569        // Store a dummy key
570        let mut keys = HashMap::new();
571        keys.insert("test-kid".to_owned(), DecodingKey::from_secret(b"secret"));
572        provider.keys.store(Arc::new(keys));
573
574        // Should be retrievable
575        assert!(provider.get_key("test-kid").is_some());
576    }
577
578    #[tokio::test]
579    async fn test_on_demand_refresh_returns_ok_when_key_exists() {
580        let provider = test_provider("https://example.com/jwks");
581
582        // Pre-populate with a key
583        let mut keys = HashMap::new();
584        keys.insert(
585            "existing-kid".to_owned(),
586            DecodingKey::from_secret(b"secret"),
587        );
588        provider.keys.store(Arc::new(keys));
589
590        // Should return Ok immediately without any refresh
591        let result = provider.on_demand_refresh("existing-kid").await;
592        assert!(result.is_ok());
593    }
594
595    #[tokio::test]
596    async fn test_try_new_returns_result() {
597        // Valid URL should work
598        let result = JwksKeyProvider::try_new("https://example.com/jwks");
599        assert!(result.is_ok());
600    }
601
602    // ==================== httpmock-based tests ====================
603
604    #[tokio::test]
605    async fn test_fetch_jwks_success_with_valid_json() {
606        let server = MockServer::start();
607
608        let mock = server.mock(|when, then| {
609            when.method(GET).path("/jwks");
610            then.status(200)
611                .header("content-type", "application/json")
612                .body(valid_jwks_json());
613        });
614
615        let jwks_url = server.url("/jwks");
616        let provider = test_provider_with_http(&jwks_url);
617
618        let result = provider.perform_refresh().await;
619        assert!(result.is_ok(), "Expected success, got: {result:?}");
620
621        // Verify key was stored
622        assert!(
623            provider.get_key("test-key-1").is_some(),
624            "Expected key 'test-key-1' to be stored"
625        );
626
627        mock.assert();
628    }
629
630    #[tokio::test]
631    async fn test_fetch_jwks_http_404_error_mapping() {
632        let server = MockServer::start();
633
634        let mock = server.mock(|when, then| {
635            when.method(GET).path("/jwks");
636            then.status(404).body("Not Found");
637        });
638
639        let jwks_url = server.url("/jwks");
640        let provider = test_provider_with_http(&jwks_url);
641
642        let result = provider.perform_refresh().await;
643        assert!(result.is_err());
644
645        let err = result.unwrap_err();
646        let err_msg = err.to_string();
647        assert!(
648            err_msg.contains("JWKS HTTP 404"),
649            "Expected error to contain 'JWKS HTTP 404', got: {err_msg}"
650        );
651        // Must NOT say "parse"
652        assert!(
653            !err_msg.to_lowercase().contains("parse"),
654            "HTTP status error should not mention 'parse', got: {err_msg}"
655        );
656
657        mock.assert();
658    }
659
660    #[tokio::test]
661    async fn test_fetch_jwks_http_500_error_mapping() {
662        let server = MockServer::start();
663
664        let mock = server.mock(|when, then| {
665            when.method(GET).path("/jwks");
666            then.status(500).body("Internal Server Error");
667        });
668
669        let jwks_url = server.url("/jwks");
670        let provider = test_provider_with_http(&jwks_url);
671
672        let result = provider.perform_refresh().await;
673        assert!(result.is_err());
674
675        let err = result.unwrap_err();
676        let err_msg = err.to_string();
677        assert!(
678            err_msg.contains("JWKS HTTP 500"),
679            "Expected error to contain 'JWKS HTTP 500', got: {err_msg}"
680        );
681
682        mock.assert();
683    }
684
685    #[tokio::test]
686    async fn test_fetch_jwks_invalid_json_error_mapping() {
687        let server = MockServer::start();
688
689        let mock = server.mock(|when, then| {
690            when.method(GET).path("/jwks");
691            then.status(200)
692                .header("content-type", "application/json")
693                .body("this is not valid json");
694        });
695
696        let jwks_url = server.url("/jwks");
697        let provider = test_provider_with_http(&jwks_url);
698
699        let result = provider.perform_refresh().await;
700        assert!(result.is_err());
701
702        let err = result.unwrap_err();
703        let err_msg = err.to_string();
704        assert!(
705            err_msg.contains("JWKS JSON parse failed"),
706            "Expected error to contain 'JWKS JSON parse failed', got: {err_msg}"
707        );
708
709        mock.assert();
710    }
711
712    #[tokio::test]
713    async fn test_fetch_jwks_empty_keys_error() {
714        let server = MockServer::start();
715
716        let mock = server.mock(|when, then| {
717            when.method(GET).path("/jwks");
718            then.status(200)
719                .header("content-type", "application/json")
720                .body(r#"{"keys": []}"#);
721        });
722
723        let jwks_url = server.url("/jwks");
724        let provider = test_provider_with_http(&jwks_url);
725
726        let result = provider.perform_refresh().await;
727        assert!(result.is_err());
728
729        let err = result.unwrap_err();
730        let err_msg = err.to_string();
731        assert!(
732            err_msg.contains("No valid RSA keys"),
733            "Expected error about no RSA keys, got: {err_msg}"
734        );
735
736        mock.assert();
737    }
738
739    #[tokio::test]
740    async fn test_on_demand_refresh_respects_cooldown() {
741        let server = MockServer::start();
742
743        // First request will return 404
744        let mock = server.mock(|when, then| {
745            when.method(GET).path("/jwks");
746            then.status(404).body("Not Found");
747        });
748
749        let jwks_url = server.url("/jwks");
750        let provider = test_provider_with_http(&jwks_url)
751            .with_on_demand_refresh_cooldown(Duration::from_secs(60));
752
753        // First attempt - should try to refresh and fail
754        let result1 = provider.on_demand_refresh("test-kid").await;
755        assert!(result1.is_err());
756
757        // Immediate second attempt - should be throttled (no network call)
758        let result2 = provider.on_demand_refresh("test-kid").await;
759        assert!(result2.is_err());
760
761        // Should return UnknownKeyId due to cooldown
762        match result2.unwrap_err() {
763            ClaimsError::UnknownKeyId(_) => {}
764            other => panic!("Expected UnknownKeyId during cooldown, got: {other:?}"),
765        }
766
767        // Only one request should have been made (first attempt)
768        mock.assert_calls(1);
769    }
770
771    #[tokio::test]
772    async fn test_on_demand_refresh_tracks_failed_kids() {
773        let server = MockServer::start();
774
775        server.mock(|when, then| {
776            when.method(GET).path("/jwks");
777            then.status(404).body("Not Found");
778        });
779
780        let jwks_url = server.url("/jwks");
781        let provider = test_provider_with_http(&jwks_url)
782            .with_on_demand_refresh_cooldown(Duration::from_millis(100));
783
784        // Attempt refresh - will fail and track the kid
785        let result = provider.on_demand_refresh("failed-kid").await;
786        assert!(result.is_err());
787
788        // Check that failed_kids contains the kid
789        let state = provider.refresh_state.read().await;
790        assert!(state.failed_kids.contains("failed-kid"));
791    }
792
793    #[tokio::test]
794    async fn test_perform_refresh_updates_state_on_failure() {
795        let server = MockServer::start();
796
797        server.mock(|when, then| {
798            when.method(GET).path("/jwks");
799            then.status(500).body("Server Error");
800        });
801
802        let jwks_url = server.url("/jwks");
803        let provider = test_provider_with_http(&jwks_url);
804
805        // Mark as previously failed
806        {
807            let mut state = provider.refresh_state.write().await;
808            state.consecutive_failures = 3;
809            state.last_error = Some("Previous error".to_owned());
810        }
811
812        // This will fail
813        _ = provider.perform_refresh().await;
814
815        // Check that consecutive_failures increased
816        let state = provider.refresh_state.read().await;
817        assert_eq!(state.consecutive_failures, 4);
818        assert!(state.last_error.is_some());
819    }
820
821    #[tokio::test]
822    async fn test_perform_refresh_resets_state_on_success() {
823        let server = MockServer::start();
824
825        server.mock(|when, then| {
826            when.method(GET).path("/jwks");
827            then.status(200)
828                .header("content-type", "application/json")
829                .body(valid_jwks_json());
830        });
831
832        let jwks_url = server.url("/jwks");
833        let provider = test_provider_with_http(&jwks_url);
834
835        // Mark as previously failed
836        {
837            let mut state = provider.refresh_state.write().await;
838            state.consecutive_failures = 5;
839            state.last_error = Some("Previous error".to_owned());
840        }
841
842        // This should succeed
843        let result = provider.perform_refresh().await;
844        assert!(result.is_ok());
845
846        // Check that state was reset
847        let state = provider.refresh_state.read().await;
848        assert_eq!(state.consecutive_failures, 0);
849        assert!(state.last_error.is_none());
850    }
851
852    #[tokio::test]
853    async fn test_validate_and_decode_with_missing_kid() {
854        let server = MockServer::start();
855
856        // Return valid JWKS but without the requested kid
857        server.mock(|when, then| {
858            when.method(GET).path("/jwks");
859            then.status(200)
860                .header("content-type", "application/json")
861                .body(valid_jwks_json());
862        });
863
864        let jwks_url = server.url("/jwks");
865        let provider = test_provider_with_http(&jwks_url)
866            .with_on_demand_refresh_cooldown(Duration::from_millis(100));
867
868        // Create a minimal JWT with a kid that doesn't exist in JWKS
869        // Header: {"alg":"RS256","kid":"nonexistent-kid"}
870        let token = "eyJhbGciOiJSUzI1NiIsImtpZCI6Im5vbmV4aXN0ZW50LWtpZCJ9.\
871                     eyJzdWIiOiIxMjM0NTY3ODkwIn0.invalid";
872
873        // Should attempt on-demand refresh but kid still won't exist
874        let result = provider.validate_and_decode(token).await;
875        assert!(result.is_err());
876
877        match result.unwrap_err() {
878            ClaimsError::UnknownKeyId(kid) => {
879                assert_eq!(kid, "nonexistent-kid");
880            }
881            other => panic!("Expected UnknownKeyId, got: {other:?}"),
882        }
883    }
884
885    #[test]
886    fn test_decode_header_with_handler_coerces_non_string_extras() {
887        use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
888
889        // Header with non-standard fields: integer, string, and array
890        let header_json = r#"{"alg":"RS256","eap":1,"iri":"some-string-id","irn":["role_a"],"kid":"kid-1","typ":"at+jwt"}"#;
891        let header_b64 = URL_SAFE_NO_PAD.encode(header_json.as_bytes());
892        let payload_b64 = URL_SAFE_NO_PAD.encode(b"{}");
893        let token = format!("{header_b64}.{payload_b64}.fake");
894
895        let header = decode_header_with_handler(&token, &|_key, value| Some(value.to_string()))
896            .expect("should handle non-standard header fields");
897
898        assert_eq!(header.alg, jsonwebtoken::Algorithm::RS256);
899        assert_eq!(header.kid.as_deref(), Some("kid-1"));
900        assert_eq!(header.typ.as_deref(), Some("at+jwt"));
901
902        // Non-string extras coerced to JSON text
903        assert_eq!(header.extras.get("eap").map(String::as_str), Some("1"));
904        assert_eq!(
905            header.extras.get("irn").map(String::as_str),
906            Some(r#"["role_a"]"#)
907        );
908        // String extras preserved as-is
909        assert_eq!(
910            header.extras.get("iri").map(String::as_str),
911            Some("some-string-id")
912        );
913    }
914
915    #[test]
916    fn test_decode_header_with_handler_can_drop_fields() {
917        use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
918
919        let header_json = r#"{"alg":"RS256","eap":1,"iri":"keep-me","kid":"kid-1","typ":"JWT"}"#;
920        let header_b64 = URL_SAFE_NO_PAD.encode(header_json.as_bytes());
921        let token = format!("{header_b64}.e30.fake");
922
923        let header = decode_header_with_handler(&token, &|_key, _value| None)
924            .expect("should succeed when handler drops non-string fields");
925
926        assert_eq!(header.alg, jsonwebtoken::Algorithm::RS256);
927        assert!(!header.extras.contains_key("eap"));
928        assert_eq!(
929            header.extras.get("iri").map(String::as_str),
930            Some("keep-me")
931        );
932    }
933
934    #[tokio::test]
935    async fn test_with_header_extras_stringified_coerces_non_string_extras() {
936        let server = MockServer::start();
937
938        server.mock(|when, then| {
939            when.method(GET).path("/jwks");
940            then.status(200)
941                .header("content-type", "application/json")
942                .body(valid_jwks_json());
943        });
944
945        let jwks_url = server.url("/jwks");
946        let provider = test_provider_with_http(&jwks_url).with_header_extras_stringified();
947
948        // Header with non-string extras: integer and array
949        let header_json =
950            r#"{"alg":"RS256","kid":"test-key-1","typ":"JWT","eap":1,"irn":["role_a"]}"#;
951        let header_b64 = URL_SAFE_NO_PAD.encode(header_json.as_bytes());
952        let payload_b64 = URL_SAFE_NO_PAD.encode(b"{}");
953        let token = format!("{header_b64}.{payload_b64}.AAAA");
954
955        let result = provider.validate_and_decode(&token).await;
956
957        // The handler lets header decode succeed; error must come from signature
958        // validation, not from header parsing.
959        let err = result.expect_err("fake signature should fail validation");
960        match &err {
961            ClaimsError::DecodeFailed(msg) => {
962                assert!(
963                    msg.contains("JWT validation failed"),
964                    "Expected signature-validation error, got: {msg}"
965                );
966            }
967            other => panic!("Expected DecodeFailed, got: {other:?}"),
968        }
969    }
970
971    #[tokio::test]
972    async fn test_validate_and_decode_uses_header_extras_handler() {
973        let server = MockServer::start();
974
975        server.mock(|when, then| {
976            when.method(GET).path("/jwks");
977            then.status(200)
978                .header("content-type", "application/json")
979                .body(valid_jwks_json());
980        });
981
982        let jwks_url = server.url("/jwks");
983        let provider = test_provider_with_http(&jwks_url)
984            .with_header_extras_handler(|_key, value| Some(value.to_string()));
985
986        // Header with a non-string extra ("eap":1) that would reject without handler
987        let header_json = r#"{"alg":"RS256","kid":"test-key-1","typ":"JWT","eap":1}"#;
988        let header_b64 = URL_SAFE_NO_PAD.encode(header_json.as_bytes());
989        let payload_b64 = URL_SAFE_NO_PAD.encode(b"{}");
990        let token = format!("{header_b64}.{payload_b64}.AAAA");
991
992        let result = provider.validate_and_decode(&token).await;
993
994        // Handler lets header decode succeed → error must come from signature
995        // validation, not from header parsing.
996        let err = result.expect_err("fake signature should fail validation");
997        match &err {
998            ClaimsError::DecodeFailed(msg) => {
999                assert!(
1000                    msg.contains("JWT validation failed"),
1001                    "Expected signature-validation error, got: {msg}"
1002                );
1003            }
1004            other => panic!("Expected DecodeFailed, got: {other:?}"),
1005        }
1006    }
1007
1008    #[test]
1009    fn test_decode_header_without_handler_rejects_non_string_extras() {
1010        use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
1011
1012        let header_json = r#"{"alg":"RS256","eap":1,"kid":"kid-1","typ":"JWT"}"#;
1013        let header_b64 = URL_SAFE_NO_PAD.encode(header_json.as_bytes());
1014        let token = format!("{header_b64}.e30.fake");
1015
1016        let result = decode_header(&token);
1017        assert!(result.is_err());
1018        let err = result.unwrap_err().to_string();
1019        assert!(
1020            err.contains("invalid type: integer"),
1021            "expected type error, got: {err}"
1022        );
1023    }
1024}
modkit_auth/providers/jwks.rs

modkit_auth/providers/
jwks.rs
