klieo-auth-oauth 3.5.0

//! JWKS cache. Fetched at builder construction; refreshed on
//! signature-verification failure (one-shot rotation handling).
//!
//! The cache holds pre-parsed [`DecodingKey`] values keyed by `kid` so
//! the hot verify path never re-parses JWK JSON. `refresh` swaps the
//! map under a write lock; concurrent verifies may briefly see the
//! pre-refresh snapshot, which is harmless because a stale key still
//! refers to a previously-valid issuer key.

use crate::error::OAuthError;
use jsonwebtoken::{jwk::JwkSet, DecodingKey};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// JWKS cache: maps `kid` → pre-parsed [`DecodingKey`] with one-shot
/// refresh on verification failure.
#[derive(Clone)]
pub(crate) struct Cache {
    inner: Arc<RwLock<HashMap<String, DecodingKey>>>,
    http: reqwest::Client,
    jwks_uri: String,
}

impl Cache {
    /// Build an empty cache with a stub URI — for unit tests that do
    /// not exercise the JWKS fetch path.
    #[cfg(test)]
    pub(crate) fn empty() -> Self {
        Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
            http: reqwest::Client::new(),
            jwks_uri: "http://localhost/not-used".into(),
        }
    }

    /// Build a cache + fetch JWKS once. Returns an error if the
    /// endpoint is unreachable or returns a malformed JWK set.
    pub(crate) async fn fetch(http: reqwest::Client, jwks_uri: String) -> Result<Self, OAuthError> {
        let cache = Self {
            inner: Arc::new(RwLock::new(HashMap::new())),
            http,
            jwks_uri,
        };
        cache.refresh().await?;
        Ok(cache)
    }

    /// Re-fetch JWKS and replace the cache contents. Called on
    /// signature-verification failure (one-shot rotation handling).
    pub(crate) async fn refresh(&self) -> Result<(), OAuthError> {
        let resp = self
            .http
            .get(&self.jwks_uri)
            .send()
            .await
            .map_err(|e| OAuthError::JwksFetch(e.to_string()))?;
        if !resp.status().is_success() {
            return Err(OAuthError::JwksFetch(format!("status {}", resp.status())));
        }
        let set: JwkSet = resp
            .json()
            .await
            .map_err(|e| OAuthError::JwksFetch(e.to_string()))?;
        let next = parse_jwk_set(&set)?;
        let mut guard = self.inner.write().await;
        *guard = next;
        Ok(())
    }

    /// Borrow the decoding key for `kid` if cached.
    pub(crate) async fn get(&self, kid: &str) -> Option<DecodingKey> {
        self.inner.read().await.get(kid).cloned()
    }
}

/// Parse a JwkSet into a kid→DecodingKey map. JWKs without a `kid` are
/// skipped (per RFC 7517 §4.5 `kid` is optional, but multi-key sets
/// without `kid` cannot be addressed by JWT headers).
fn parse_jwk_set(set: &JwkSet) -> Result<HashMap<String, DecodingKey>, OAuthError> {
    let mut next: HashMap<String, DecodingKey> = HashMap::new();
    for jwk in &set.keys {
        let Some(kid) = jwk.common.key_id.as_ref() else {
            continue;
        };
        let dk = DecodingKey::from_jwk(jwk)
            .map_err(|e| OAuthError::JwksFetch(format!("from_jwk: {e}")))?;
        next.insert(kid.clone(), dk);
    }
    Ok(next)
}