1use std::collections::HashMap;
42use std::sync::Arc;
43use std::time::{Duration, Instant};
44
45use jsonwebtoken::jwk::{Jwk, JwkSet};
46use jsonwebtoken::{Algorithm, DecodingKey};
47use parking_lot::RwLock;
48use serde::Deserialize;
49use tokio::sync::Mutex as AsyncMutex;
50
51use crate::auth::error::AuthError;
52
53const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
55
56#[derive(Debug, Deserialize)]
62struct OidcDiscoveryDocument {
63 jwks_uri: String,
65 issuer: String,
67}
68
69fn tls_config() -> Result<rustls::ClientConfig, rustls::Error> {
87 use rustls_platform_verifier::BuilderVerifierExt;
88
89 let provider = rustls::crypto::CryptoProvider::get_default()
90 .cloned()
91 .unwrap_or_else(|| std::sync::Arc::new(rustls::crypto::aws_lc_rs::default_provider()));
92
93 Ok(rustls::ClientConfig::builder_with_provider(provider)
94 .with_safe_default_protocol_versions()?
95 .with_platform_verifier()?
96 .with_no_client_auth())
97}
98
99fn http_client() -> Result<reqwest::Client, AuthError> {
112 let tls = tls_config()
113 .map_err(|e| AuthError::Config(format!("failed to build TLS configuration: {e}")))?;
114 reqwest::Client::builder()
115 .timeout(HTTP_TIMEOUT)
116 .use_preconfigured_tls(tls)
117 .build()
118 .map_err(|e| AuthError::Config(format!("failed to build HTTP client: {e}")))
119}
120
121pub async fn discover_jwks_url(issuer: &str) -> Result<String, AuthError> {
126 let url = format!(
127 "{}/.well-known/openid-configuration",
128 issuer.trim_end_matches('/')
129 );
130
131 let client = http_client()?;
132
133 let resp =
134 client.get(&url).send().await.map_err(|e| {
135 AuthError::Config(format!("OIDC discovery request failed for {url}: {e}"))
136 })?;
137
138 if !resp.status().is_success() {
139 return Err(AuthError::Config(format!(
140 "OIDC discovery returned HTTP {}",
141 resp.status()
142 )));
143 }
144
145 let doc: OidcDiscoveryDocument = resp
146 .json()
147 .await
148 .map_err(|e| AuthError::Config(format!("failed to parse OIDC discovery document: {e}")))?;
149
150 let expected = issuer.trim_end_matches('/');
153 let actual = doc.issuer.trim_end_matches('/');
154 if expected != actual {
155 return Err(AuthError::Config(format!(
156 "OIDC issuer mismatch: expected {expected}, document reports {actual}"
157 )));
158 }
159
160 Ok(doc.jwks_uri)
161}
162
163pub async fn fetch_jwks(url: &str) -> Result<JwkSet, AuthError> {
169 let client = http_client()?;
170
171 let resp = client
172 .get(url)
173 .send()
174 .await
175 .map_err(|e| AuthError::Config(format!("JWKS fetch failed for {url}: {e}")))?;
176
177 if !resp.status().is_success() {
178 return Err(AuthError::Config(format!(
179 "JWKS endpoint returned HTTP {}",
180 resp.status()
181 )));
182 }
183
184 resp.json::<JwkSet>()
185 .await
186 .map_err(|e| AuthError::Config(format!("failed to parse JWKS: {e}")))
187}
188
189fn decoding_key_from_jwk(jwk: &Jwk) -> Result<DecodingKey, AuthError> {
191 DecodingKey::from_jwk(jwk)
192 .map_err(|e| AuthError::Config(format!("failed to create decoding key from JWK: {e}")))
193}
194
195fn algorithm_for_jwk(jwk: &Jwk) -> Algorithm {
197 jwk.common
198 .key_algorithm
199 .map_or(Algorithm::RS256, |ka| match ka {
200 jsonwebtoken::jwk::KeyAlgorithm::RS256 => Algorithm::RS256,
201 jsonwebtoken::jwk::KeyAlgorithm::RS384 => Algorithm::RS384,
202 jsonwebtoken::jwk::KeyAlgorithm::RS512 => Algorithm::RS512,
203 jsonwebtoken::jwk::KeyAlgorithm::ES256 => Algorithm::ES256,
204 jsonwebtoken::jwk::KeyAlgorithm::ES384 => Algorithm::ES384,
205 _ => Algorithm::RS256,
206 })
207}
208
209#[derive(Clone)]
215pub struct CachedKey {
216 pub key: DecodingKey,
218 pub algorithm: Algorithm,
220}
221
222impl std::fmt::Debug for CachedKey {
223 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
224 f.debug_struct("CachedKey")
225 .field("algorithm", &self.algorithm)
226 .finish_non_exhaustive()
227 }
228}
229
230struct CacheInner {
236 keys: HashMap<String, CachedKey>,
238 fetched_at: Option<Instant>,
240}
241
242pub struct JwksCache {
249 jwks_url: String,
251 ttl: Duration,
253 max_stale: Duration,
255 inner: Arc<RwLock<CacheInner>>,
257 refresh_mutex: Arc<AsyncMutex<()>>,
259}
260
261impl std::fmt::Debug for JwksCache {
262 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263 f.debug_struct("JwksCache")
264 .field("jwks_url", &self.jwks_url)
265 .field("ttl", &self.ttl)
266 .field("cached_keys", &self.inner.read().keys.len())
267 .finish()
268 }
269}
270
271impl JwksCache {
272 #[must_use]
278 pub fn new(jwks_url: String, ttl: Duration) -> Self {
279 Self {
280 jwks_url,
281 ttl,
282 max_stale: Duration::from_secs(24 * 3600), inner: Arc::new(RwLock::new(CacheInner {
284 keys: HashMap::new(),
285 fetched_at: None,
286 })),
287 refresh_mutex: Arc::new(AsyncMutex::new(())),
288 }
289 }
290
291 pub async fn from_oidc_issuer(issuer: &str, ttl: Duration) -> Result<Self, AuthError> {
294 let jwks_url = discover_jwks_url(issuer).await?;
295 let cache = Self::new(jwks_url, ttl);
296 cache.refresh().await?;
297 Ok(cache)
298 }
299
300 #[must_use]
302 pub fn from_keys(keys: HashMap<String, CachedKey>) -> Self {
303 Self {
304 jwks_url: String::new(),
305 ttl: Duration::from_secs(3600),
306 max_stale: Duration::from_secs(24 * 3600),
307 inner: Arc::new(RwLock::new(CacheInner {
308 keys,
309 fetched_at: Some(Instant::now()),
310 })),
311 refresh_mutex: Arc::new(AsyncMutex::new(())),
312 }
313 }
314
315 #[must_use]
317 pub fn is_expired(&self) -> bool {
318 let inner = self.inner.read();
319 match inner.fetched_at {
320 None => true,
321 Some(t) => t.elapsed() > self.ttl,
322 }
323 }
324
325 #[must_use]
327 pub fn len(&self) -> usize {
328 self.inner.read().keys.len()
329 }
330
331 #[must_use]
333 pub fn is_empty(&self) -> bool {
334 self.inner.read().keys.is_empty()
335 }
336
337 pub async fn refresh(&self) -> Result<(), AuthError> {
343 if self.jwks_url.is_empty() {
344 return Err(AuthError::Config(
345 "cannot refresh: no JWKS URL configured (cache was created via from_keys)".into(),
346 ));
347 }
348
349 let jwk_set = match fetch_jwks(&self.jwks_url).await {
350 Ok(set) => set,
351 Err(e) => {
352 let inner = self.inner.read();
355 if !inner.keys.is_empty() {
356 let too_old = inner
357 .fetched_at
358 .is_none_or(|t| t.elapsed() > self.max_stale);
359 if too_old {
360 tracing::error!(
361 url = %self.jwks_url,
362 error = %e,
363 max_stale_secs = self.max_stale.as_secs(),
364 "JWKS refresh failed and stale keys exceeded max_stale — rejecting"
365 );
366 return Err(e);
367 }
368 tracing::warn!(
369 url = %self.jwks_url,
370 error = %e,
371 "JWKS refresh failed, serving stale keys"
372 );
373 return Ok(());
374 }
375 return Err(e);
376 }
377 };
378
379 let mut new_keys = HashMap::new();
380 for jwk in &jwk_set.keys {
381 let kid = jwk
382 .common
383 .key_id
384 .clone()
385 .unwrap_or_else(|| "default".to_string());
386 let dk = decoding_key_from_jwk(jwk)?;
387 let alg = algorithm_for_jwk(jwk);
388 new_keys.insert(
389 kid,
390 CachedKey {
391 key: dk,
392 algorithm: alg,
393 },
394 );
395 }
396
397 let mut inner = self.inner.write();
398 inner.keys = new_keys;
399 inner.fetched_at = Some(Instant::now());
400
401 tracing::debug!(
402 keys = inner.keys.len(),
403 url = %self.jwks_url,
404 "JWKS cache refreshed"
405 );
406
407 Ok(())
408 }
409
410 async fn ensure_fresh(&self) -> Result<(), AuthError> {
413 if !self.is_expired() {
414 return Ok(());
415 }
416 let _guard = self.refresh_mutex.lock().await;
418 if !self.is_expired() {
421 return Ok(());
422 }
423 self.refresh().await
424 }
425
426 pub async fn get_key(&self, kid: &str) -> Result<CachedKey, AuthError> {
431 self.ensure_fresh().await?;
432
433 let inner = self.inner.read();
434 inner
435 .keys
436 .get(kid)
437 .cloned()
438 .ok_or_else(|| AuthError::Config(format!("no JWKS key with kid={kid}")))
439 }
440
441 pub async fn get_all_keys(&self) -> Result<Vec<CachedKey>, AuthError> {
446 self.ensure_fresh().await?;
447
448 let inner = self.inner.read();
449 Ok(inner.keys.values().cloned().collect())
450 }
451
452 #[must_use]
454 pub fn jwks_url(&self) -> &str {
455 &self.jwks_url
456 }
457}
458
459#[cfg(test)]
464mod tests {
465 use super::*;
466
467 #[test]
468 fn cache_from_keys_not_expired() {
469 let mut keys = HashMap::new();
470 keys.insert(
471 "kid-1".to_string(),
472 CachedKey {
473 key: DecodingKey::from_secret(b"test"),
474 algorithm: Algorithm::HS256,
475 },
476 );
477 let cache = JwksCache::from_keys(keys);
478
479 assert!(!cache.is_expired());
480 assert_eq!(cache.len(), 1);
481 assert!(!cache.is_empty());
482 }
483
484 #[test]
485 fn empty_cache_is_expired() {
486 let cache = JwksCache::new("https://example.com/jwks".into(), Duration::from_secs(60));
487 assert!(cache.is_expired());
488 assert!(cache.is_empty());
489 }
490
491 #[test]
492 fn cache_expires_after_ttl() {
493 let cache = JwksCache {
494 jwks_url: String::new(),
495 ttl: Duration::from_millis(0), max_stale: Duration::from_secs(24 * 3600),
497 inner: Arc::new(RwLock::new(CacheInner {
498 keys: HashMap::new(),
499 fetched_at: Some(Instant::now().checked_sub(Duration::from_secs(1)).unwrap()),
500 })),
501 refresh_mutex: Arc::new(AsyncMutex::new(())),
502 };
503 assert!(cache.is_expired());
504 }
505
506 #[tokio::test]
507 async fn get_key_from_prepopulated_cache() {
508 let mut keys = HashMap::new();
509 keys.insert(
510 "my-kid".to_string(),
511 CachedKey {
512 key: DecodingKey::from_secret(b"secret"),
513 algorithm: Algorithm::HS256,
514 },
515 );
516 let cache = JwksCache::from_keys(keys);
517
518 let k = cache.get_key("my-kid").await.unwrap();
519 assert_eq!(k.algorithm, Algorithm::HS256);
520 }
521
522 #[tokio::test]
523 async fn get_key_missing_kid_returns_error() {
524 let cache = JwksCache::from_keys(HashMap::new());
525 let err = cache.get_key("nonexistent").await.unwrap_err();
526 assert!(matches!(err, AuthError::Config(_)));
527 }
528
529 #[tokio::test]
530 async fn get_all_keys_from_prepopulated_cache() {
531 let mut keys = HashMap::new();
532 keys.insert(
533 "kid-a".to_string(),
534 CachedKey {
535 key: DecodingKey::from_secret(b"a"),
536 algorithm: Algorithm::RS256,
537 },
538 );
539 keys.insert(
540 "kid-b".to_string(),
541 CachedKey {
542 key: DecodingKey::from_secret(b"b"),
543 algorithm: Algorithm::RS256,
544 },
545 );
546 let cache = JwksCache::from_keys(keys);
547
548 let all = cache.get_all_keys().await.unwrap();
549 assert_eq!(all.len(), 2);
550 }
551
552 #[test]
553 fn discover_url_construction() {
554 let issuer = "https://accounts.google.com";
556 let url = format!(
557 "{}/.well-known/openid-configuration",
558 issuer.trim_end_matches('/')
559 );
560 assert_eq!(
561 url,
562 "https://accounts.google.com/.well-known/openid-configuration"
563 );
564
565 let issuer = "https://example.com/";
567 let url = format!(
568 "{}/.well-known/openid-configuration",
569 issuer.trim_end_matches('/')
570 );
571 assert_eq!(url, "https://example.com/.well-known/openid-configuration");
572 }
573
574 #[test]
575 fn algorithm_for_jwk_defaults_to_rs256() {
576 let jwk: Jwk = serde_json::from_value(serde_json::json!({
577 "kty": "RSA",
578 "n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368QQMicAtaSqzs8KJZgnYb9c7d0zgdAZHzu6qMQvRL5hajrn1n91CbOpbISD08qNLyrdkt-bFTWhAI4vMQFh6WeZu0fM4lFd2NcRwr3XPksINHaQ-G_xBniIqbw0Ls1jF44-csFCur-kEgU8awapJzKnqDKgw",
579 "e": "AQAB",
580 "kid": "test-kid"
581 }))
582 .unwrap();
583
584 let alg = algorithm_for_jwk(&jwk);
585 assert_eq!(alg, Algorithm::RS256);
586 }
587
588 #[test]
589 fn decoding_key_from_valid_jwk() {
590 let jwk: Jwk = serde_json::from_value(serde_json::json!({
591 "kty": "RSA",
592 "n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368QQMicAtaSqzs8KJZgnYb9c7d0zgdAZHzu6qMQvRL5hajrn1n91CbOpbISD08qNLyrdkt-bFTWhAI4vMQFh6WeZu0fM4lFd2NcRwr3XPksINHaQ-G_xBniIqbw0Ls1jF44-csFCur-kEgU8awapJzKnqDKgw",
593 "e": "AQAB",
594 "kid": "test-kid"
595 }))
596 .unwrap();
597
598 let result = decoding_key_from_jwk(&jwk);
599 assert!(result.is_ok());
600 }
601}