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
36/// Build the rustls client config for the JWKS HTTPS client.
37///
38/// Uses the pure-Rust `ring` provider (installed per-config, not process-global)
39/// and bundles Mozilla's root store via `webpki-roots`, so the binary needs no
40/// system CA bundle and works in musl / scratch / distroless images.
41fn build_tls_config() -> rustls::ClientConfig {
42    let mut roots = rustls::RootCertStore::empty();
43    roots.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
44    rustls::ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
45        .with_safe_default_protocol_versions()
46        .expect("ring provider supports the default TLS protocol versions")
47        .with_root_certificates(roots)
48        .with_no_client_auth()
49}
50
51impl JwksCache {
52    /// Create a cache for `uri` (keys are loaded lazily on first lookup).
53    pub fn new(uri: String) -> Self {
54        let client = reqwest::Client::builder()
55            .timeout(JWKS_HTTP_TIMEOUT)
56            // Hand reqwest a fully preconfigured rustls backend rather than
57            // relying on a process-global default provider: no install ordering
58            // constraint, no global side effect, safe for library/test callers.
59            .tls_backend_preconfigured(build_tls_config())
60            .build()
61            .unwrap_or_default();
62        Self {
63            uri,
64            client,
65            keys: RwLock::new(HashMap::new()),
66            last_refresh: Mutex::new(None),
67        }
68    }
69
70    /// Resolve the verifying key for `kid`, refreshing from the JWKS endpoint
71    /// once (throttled) if it is not already cached.
72    pub async fn key_for(&self, kid: &str) -> Option<VerifyingKey> {
73        if let Some(k) = self.keys.read().await.get(kid).cloned() {
74            return Some(k);
75        }
76        if self.refresh().await.is_err() {
77            return None;
78        }
79        self.keys.read().await.get(kid).cloned()
80    }
81
82    /// Fetch the JWKS and replace the cache. Throttled by [`MIN_REFRESH_INTERVAL`]
83    /// unless the cache is still empty (first load).
84    async fn refresh(&self) -> Result<(), String> {
85        // Claim the refresh slot atomically: hold the lock across the throttle
86        // check and the timestamp update so concurrent callers cannot all pass.
87        {
88            let mut last = self.last_refresh.lock().await;
89            if let Some(t) = *last {
90                let empty = self.keys.read().await.is_empty();
91                if !empty && t.elapsed() < MIN_REFRESH_INTERVAL {
92                    return Err("refresh throttled".to_string());
93                }
94            }
95            *last = Some(Instant::now());
96        }
97
98        let set: JwkSet = self
99            .client
100            .get(&self.uri)
101            .send()
102            .await
103            .map_err(|e| format!("JWKS fetch failed: {e}"))?
104            .json()
105            .await
106            .map_err(|e| format!("JWKS decode failed: {e}"))?;
107
108        let new_keys = parse_jwks(&set);
109        *self.keys.write().await = new_keys;
110        Ok(())
111    }
112}
113
114/// Build the `kid → VerifyingKey` map from a JWK set, skipping keys without a
115/// `kid`, symmetric keys, or those that fail to convert.
116fn parse_jwks(set: &JwkSet) -> HashMap<String, VerifyingKey> {
117    let mut map = HashMap::new();
118    for jwk in &set.keys {
119        let Some(kid) = jwk.common.key_id.clone() else {
120            continue;
121        };
122        let Some(algorithm) = algorithm_for(jwk) else {
123            continue;
124        };
125        if let Ok(key) = DecodingKey::from_jwk(jwk) {
126            map.insert(
127                kid,
128                VerifyingKey {
129                    key: Arc::new(key),
130                    algorithm,
131                },
132            );
133        }
134    }
135    map
136}
137
138/// Pick the signature algorithm for a key.
139///
140/// The JWK's explicit `alg` is authoritative (so ES384 / RS512 / PS256 keys are
141/// not mis-pinned). Without it, fall back to the key type and EC curve.
142/// Symmetric keys (`OctetKey`) and unsupported variants are rejected.
143fn algorithm_for(jwk: &Jwk) -> Option<Algorithm> {
144    if let Some(alg) = jwk.common.key_algorithm.and_then(key_algorithm_to_alg) {
145        return Some(alg);
146    }
147    match &jwk.algorithm {
148        AlgorithmParameters::RSA(_) => Some(Algorithm::RS256),
149        AlgorithmParameters::EllipticCurve(ec) => match ec.curve {
150            EllipticCurve::P256 => Some(Algorithm::ES256),
151            EllipticCurve::P384 => Some(Algorithm::ES384),
152            // P-521 (ES512) is not supported by the verifier.
153            _ => None,
154        },
155        AlgorithmParameters::OctetKeyPair(_) => Some(Algorithm::EdDSA),
156        AlgorithmParameters::OctetKey(_) => None,
157    }
158}
159
160/// Map a JWK signature `alg` to a verifier algorithm, rejecting symmetric and
161/// encryption algorithms (only asymmetric signatures are usable from a JWKS).
162fn key_algorithm_to_alg(ka: KeyAlgorithm) -> Option<Algorithm> {
163    Some(match ka {
164        KeyAlgorithm::ES256 => Algorithm::ES256,
165        KeyAlgorithm::ES384 => Algorithm::ES384,
166        KeyAlgorithm::RS256 => Algorithm::RS256,
167        KeyAlgorithm::RS384 => Algorithm::RS384,
168        KeyAlgorithm::RS512 => Algorithm::RS512,
169        KeyAlgorithm::PS256 => Algorithm::PS256,
170        KeyAlgorithm::PS384 => Algorithm::PS384,
171        KeyAlgorithm::PS512 => Algorithm::PS512,
172        KeyAlgorithm::EdDSA => Algorithm::EdDSA,
173        _ => return None,
174    })
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn parse_jwks_keeps_asymmetric_keys_and_maps_algorithms() {
183        // A minimal RSA JWK with a kid (values are a real test key from the
184        // jsonwebtoken test vectors).
185        let set: JwkSet = serde_json::from_value(serde_json::json!({
186            "keys": [{
187                "kty": "RSA",
188                "kid": "rsa-1",
189                "use": "sig",
190                "n": "0vx7agoebGcQSuuPiLJXZptN9nndrQmbXEps2aiAFbWhM78LhWx4cbbfAAtVT86zwu1RK7aPFFxuhDR1L6tSoc_BJECPebWKRXjBZCiFV4n3oknjhMstn64tZ_2W-5JsGY4Hc5n9yBXArwl93lqt7_RN5w6Cf0h4QyQ5v-65YGjQR0_FDW2QvzqY368Qen-JS7-zw04o6sJ9qjp6lFm5_T4nzcCqRfMOgRA_g_S0d7e9k7B0v0vqHr0e1V_o-z0ow5dWpql8-zKj4hQp8sg_Pn8O0R5ZQS4t8hUE-3-r3ftt1YzQ",
191                "e": "AQAB"
192            }]
193        })).unwrap();
194        let keys = parse_jwks(&set);
195        assert!(keys.contains_key("rsa-1"));
196        assert_eq!(keys["rsa-1"].algorithm, Algorithm::RS256);
197    }
198
199    #[test]
200    fn algorithm_prefers_explicit_jwk_alg() {
201        // An EC key that explicitly declares ES384 must not be pinned to ES256.
202        let jwk: Jwk = serde_json::from_value(serde_json::json!({
203            "kty": "EC", "crv": "P-384", "alg": "ES384", "kid": "k",
204            "x": "AAAA", "y": "AAAA"
205        }))
206        .unwrap();
207        assert_eq!(algorithm_for(&jwk), Some(Algorithm::ES384));
208    }
209
210    #[test]
211    fn algorithm_falls_back_to_curve_not_es256() {
212        // No alg field → infer from the curve, not a blanket ES256.
213        let jwk: Jwk = serde_json::from_value(serde_json::json!({
214            "kty": "EC", "crv": "P-384", "kid": "k", "x": "AAAA", "y": "AAAA"
215        }))
216        .unwrap();
217        assert_eq!(algorithm_for(&jwk), Some(Algorithm::ES384));
218    }
219
220    #[test]
221    fn parse_jwks_skips_symmetric_and_keyless() {
222        let set: JwkSet = serde_json::from_value(serde_json::json!({
223            "keys": [
224                { "kty": "oct", "kid": "hmac", "k": "c2VjcmV0" },
225                { "kty": "RSA", "n": "0vx7ag", "e": "AQAB" }
226            ]
227        }))
228        .unwrap();
229        // Symmetric key rejected; RSA without a kid skipped.
230        assert!(parse_jwks(&set).is_empty());
231    }
232}
structured_proxy/auth/jwks.rs

structured_proxy/auth/
jwks.rs
