structured_proxy/auth/
jwks.rs

1//! JWKS fetching and key cache.
2//!
3//! Keys are fetched from the configured JWKS URI and cached by `kid`. An unknown
4//! `kid` triggers a refresh (throttled), which is how key rotation is picked up.
5
6use std::collections::HashMap;
7use std::sync::Arc;
8use std::time::{Duration, Instant};
9
10use jsonwebtoken::jwk::{AlgorithmParameters, EllipticCurve, Jwk, JwkSet, KeyAlgorithm};
11use jsonwebtoken::{Algorithm, DecodingKey};
12use tokio::sync::{Mutex, RwLock};
13
14/// A decoding key plus the signature algorithm it is valid for.
15#[derive(Clone)]
16pub struct VerifyingKey {
17    pub key: Arc<DecodingKey>,
18    pub algorithm: Algorithm,
19}
20
21/// Fetches and caches JWKS keys by `kid`.
22pub struct JwksCache {
23    uri: String,
24    client: reqwest::Client,
25    keys: RwLock<HashMap<String, VerifyingKey>>,
26    last_refresh: Mutex<Option<Instant>>,
27}
28
29/// Minimum spacing between refreshes triggered by an unknown `kid`, so a flood
30/// of bogus `kid`s cannot hammer the JWKS endpoint.
31const MIN_REFRESH_INTERVAL: Duration = Duration::from_secs(60);
32
33/// Bound the worst-case latency of a slow/stalled JWKS endpoint.
34const JWKS_HTTP_TIMEOUT: Duration = Duration::from_secs(5);
35
36impl JwksCache {
37    /// Create a cache for `uri` (keys are loaded lazily on first lookup).
38    pub fn new(uri: String) -> Self {
39        let client = reqwest::Client::builder()
40            .timeout(JWKS_HTTP_TIMEOUT)
41            // Hand reqwest a fully preconfigured rustls backend rather than
42            // relying on a process-global default provider: no install ordering
43            // constraint, no global side effect, safe for library/test callers.
44            .tls_backend_preconfigured(crate::tls::client_config())
45            .build()
46            .unwrap_or_default();
47        Self {
48            uri,
49            client,
50            keys: RwLock::new(HashMap::new()),
51            last_refresh: Mutex::new(None),
52        }
53    }
54
55    /// Resolve the verifying key for `kid`, refreshing from the JWKS endpoint
56    /// once (throttled) if it is not already cached.
57    pub async fn key_for(&self, kid: &str) -> Option<VerifyingKey> {
58        if let Some(k) = self.keys.read().await.get(kid).cloned() {
59            return Some(k);
60        }
61        if self.refresh().await.is_err() {
62            return None;
63        }
64        self.keys.read().await.get(kid).cloned()
65    }
66
67    /// Fetch the JWKS and replace the cache. Throttled by [`MIN_REFRESH_INTERVAL`]
68    /// unless the cache is still empty (first load).
69    async fn refresh(&self) -> Result<(), String> {
70        // Claim the refresh slot atomically: hold the lock across the throttle
71        // check and the timestamp update so concurrent callers cannot all pass.
72        {
73            let mut last = self.last_refresh.lock().await;
74            if let Some(t) = *last {
75                let empty = self.keys.read().await.is_empty();
76                if !empty && t.elapsed() < MIN_REFRESH_INTERVAL {
77                    return Err("refresh throttled".to_string());
78                }
79            }
80            *last = Some(Instant::now());
81        }
82
83        let set: JwkSet = self
84            .client
85            .get(&self.uri)
86            .send()
87            .await
88            .map_err(|e| format!("JWKS fetch failed: {e}"))?
89            .json()
90            .await
91            .map_err(|e| format!("JWKS decode failed: {e}"))?;
92
93        let new_keys = parse_jwks(&set);
94        *self.keys.write().await = new_keys;
95        Ok(())
96    }
97}
98
99/// Build the `kid → VerifyingKey` map from a JWK set, skipping keys without a
100/// `kid`, symmetric keys, or those that fail to convert.
101fn parse_jwks(set: &JwkSet) -> HashMap<String, VerifyingKey> {
102    let mut map = HashMap::new();
103    for jwk in &set.keys {
104        let Some(kid) = jwk.common.key_id.clone() else {
105            continue;
106        };
107        let Some(algorithm) = algorithm_for(jwk) else {
108            continue;
109        };
110        if let Ok(key) = DecodingKey::from_jwk(jwk) {
111            map.insert(
112                kid,
113                VerifyingKey {
114                    key: Arc::new(key),
115                    algorithm,
116                },
117            );
118        }
119    }
120    map
121}
122
123/// Pick the signature algorithm for a key.
124///
125/// The JWK's explicit `alg` is authoritative (so ES384 / RS512 / PS256 keys are
126/// not mis-pinned). Without it, fall back to the key type and EC curve.
127/// Symmetric keys (`OctetKey`) and unsupported variants are rejected.
128fn algorithm_for(jwk: &Jwk) -> Option<Algorithm> {
129    if let Some(alg) = jwk.common.key_algorithm.and_then(key_algorithm_to_alg) {
130        return Some(alg);
131    }
132    match &jwk.algorithm {
133        AlgorithmParameters::RSA(_) => Some(Algorithm::RS256),
134        AlgorithmParameters::EllipticCurve(ec) => match ec.curve {
135            EllipticCurve::P256 => Some(Algorithm::ES256),
136            EllipticCurve::P384 => Some(Algorithm::ES384),
137            // P-521 (ES512) is not supported by the verifier.
138            _ => None,
139        },
140        AlgorithmParameters::OctetKeyPair(_) => Some(Algorithm::EdDSA),
141        AlgorithmParameters::OctetKey(_) => None,
142        // The enum is non-exhaustive: key types added upstream are not
143        // verifiable here until they are mapped explicitly.
144        _ => None,
145    }
146}
147
148/// Map a JWK signature `alg` to a verifier algorithm, rejecting symmetric and
149/// encryption algorithms (only asymmetric signatures are usable from a JWKS).
150fn key_algorithm_to_alg(ka: KeyAlgorithm) -> Option<Algorithm> {
151    Some(match ka {
152        KeyAlgorithm::ES256 => Algorithm::ES256,
153        KeyAlgorithm::ES384 => Algorithm::ES384,
154        KeyAlgorithm::RS256 => Algorithm::RS256,
155        KeyAlgorithm::RS384 => Algorithm::RS384,
156        KeyAlgorithm::RS512 => Algorithm::RS512,
157        KeyAlgorithm::PS256 => Algorithm::PS256,
158        KeyAlgorithm::PS384 => Algorithm::PS384,
159        KeyAlgorithm::PS512 => Algorithm::PS512,
160        KeyAlgorithm::EdDSA => Algorithm::EdDSA,
161        _ => return None,
162    })
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168
169    #[test]
170    fn parse_jwks_keeps_asymmetric_keys_and_maps_algorithms() {
171        // A minimal RSA JWK with a kid (values are a real test key from the
172        // jsonwebtoken test vectors).
173        let set: JwkSet = serde_json::from_value(serde_json::json!({
174            "keys": [{
175                "kty": "RSA",
176                "kid": "rsa-1",
177                "use": "sig",
178                "n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368Qen-JS7-zw04o6sJ9qjp6lFm5_T4nzcCqRfMOgRA_g_S0d7e9k7B0v0vqHr0e1V_o-z0ow5dWpql8-zKj4hQp8sg_Pn8O0R5ZQS4t8hUE-3-r3ftt1YzQ",
179                "e": "AQAB"
180            }]
181        })).unwrap();
182        let keys = parse_jwks(&set);
183        assert!(keys.contains_key("rsa-1"));
184        assert_eq!(keys["rsa-1"].algorithm, Algorithm::RS256);
185    }
186
187    #[test]
188    fn algorithm_prefers_explicit_jwk_alg() {
189        // An EC key that explicitly declares ES384 must not be pinned to ES256.
190        let jwk: Jwk = serde_json::from_value(serde_json::json!({
191            "kty": "EC", "crv": "P-384", "alg": "ES384", "kid": "k",
192            "x": "AAAA", "y": "AAAA"
193        }))
194        .unwrap();
195        assert_eq!(algorithm_for(&jwk), Some(Algorithm::ES384));
196    }
197
198    #[test]
199    fn algorithm_falls_back_to_curve_not_es256() {
200        // No alg field → infer from the curve, not a blanket ES256.
201        let jwk: Jwk = serde_json::from_value(serde_json::json!({
202            "kty": "EC", "crv": "P-384", "kid": "k", "x": "AAAA", "y": "AAAA"
203        }))
204        .unwrap();
205        assert_eq!(algorithm_for(&jwk), Some(Algorithm::ES384));
206    }
207
208    #[test]
209    fn parse_jwks_skips_symmetric_and_keyless() {
210        let set: JwkSet = serde_json::from_value(serde_json::json!({
211            "keys": [
212                { "kty": "oct", "kid": "hmac", "k": "c2VjcmV0" },
213                { "kty": "RSA", "n": "0vx7ag", "e": "AQAB" }
214            ]
215        }))
216        .unwrap();
217        // Symmetric key rejected; RSA without a kid skipped.
218        assert!(parse_jwks(&set).is_empty());
219    }
220}
structured_proxy/auth/jwks.rs

structured_proxy/auth/
jwks.rs
