Skip to main content

fraiseql_core/security/oidc/
jwks.rs

1//! JWKS (JSON Web Key Set) types, cache, key selection, and fetch logic.
2//!
3//! The `impl OidcValidator` block here adds JWKS-specific methods to the
4//! validator type defined in [`super::token`].
5
6use std::time::{Duration, Instant};
7
8use jsonwebtoken::DecodingKey;
9use serde::Deserialize;
10
11use crate::security::{
12    errors::{Result, SecurityError},
13    oidc::token::OidcValidator,
14};
15
16/// Maximum byte length accepted from a JWKS endpoint response.
17///
18/// A legitimate JWKS document (a few RSA/EC public keys) is well under 64 `KiB`.
19/// A 1 `MiB` cap prevents a malicious or compromised OIDC provider from sending
20/// a response large enough to exhaust server memory.
21pub const MAX_JWKS_RESPONSE_BYTES: usize = 1024 * 1024; // 1 MiB
22
23// ============================================================================
24// OIDC Discovery Response
25// ============================================================================
26
27/// OIDC Discovery document (partial).
28///
29/// Contains the fields we need from `/.well-known/openid-configuration`.
30#[derive(Debug, Clone, Deserialize)]
31pub struct OidcDiscoveryDocument {
32    /// Issuer identifier
33    pub issuer: String,
34
35    /// JWKS URI for fetching public keys
36    pub jwks_uri: String,
37
38    /// Supported signing algorithms
39    #[serde(default)]
40    pub id_token_signing_alg_values_supported: Vec<String>,
41
42    /// Authorization endpoint (for reference)
43    #[serde(default)]
44    pub authorization_endpoint: Option<String>,
45
46    /// Token endpoint (for reference)
47    #[serde(default)]
48    pub token_endpoint: Option<String>,
49}
50
51// ============================================================================
52// JWKS Types
53// ============================================================================
54
55/// JSON Web Key Set.
56#[derive(Debug, Clone, Deserialize)]
57pub struct Jwks {
58    /// Array of JSON Web Keys
59    pub keys: Vec<Jwk>,
60}
61
62/// JSON Web Key.
63#[derive(Debug, Clone, Deserialize)]
64pub struct Jwk {
65    /// Key type (e.g., "RSA")
66    pub kty: String,
67
68    /// Key ID (used to match with JWT header)
69    pub kid: Option<String>,
70
71    /// Algorithm (e.g., "RS256")
72    #[serde(default)]
73    pub alg: Option<String>,
74
75    /// Intended use (e.g., "sig" for signature)
76    #[serde(rename = "use")]
77    pub key_use: Option<String>,
78
79    /// RSA modulus (base64url encoded)
80    pub n: Option<String>,
81
82    /// RSA exponent (base64url encoded)
83    pub e: Option<String>,
84
85    /// X.509 certificate chain
86    #[serde(default)]
87    pub x5c: Vec<String>,
88}
89
90/// Cached JWKS with expiration.
91#[derive(Debug)]
92pub struct CachedJwks {
93    pub(super) jwks:       Jwks,
94    pub(super) fetched_at: Instant,
95    pub(super) ttl:        Duration,
96}
97
98impl CachedJwks {
99    pub(super) fn is_expired(&self) -> bool {
100        self.fetched_at.elapsed() > self.ttl
101    }
102}
103
104// ============================================================================
105// OidcValidator — JWKS fetch, cache and key-resolution methods
106// ============================================================================
107
108impl OidcValidator {
109    /// Get the decoding key for a specific key ID.
110    ///
111    /// Checks the cache first; fetches fresh JWKS on miss or expiry.
112    ///
113    /// # Errors
114    ///
115    /// Returns `SecurityError::InvalidToken` if the key is not found or cannot be decoded.
116    pub(super) async fn get_decoding_key(&self, kid: &str) -> Result<DecodingKey> {
117        // Check cache first
118        {
119            let cache = self.jwks_cache.read();
120            if let Some(ref cached) = *cache {
121                if !cached.is_expired() {
122                    if let Some(key) = self.find_key(&cached.jwks, kid) {
123                        return self.jwk_to_decoding_key(key);
124                    }
125                }
126            }
127        }
128
129        // Fetch fresh JWKS
130        let jwks = self.fetch_jwks().await?;
131
132        // SECURITY (#361): a key missing from the fresh set means the IdP rotated it
133        // out — tokens signed by it must stop validating. Detect this against the
134        // currently-cached set BEFORE we replace the cache below.
135        if self.detect_key_rotation(&jwks) {
136            tracing::warn!(
137                "OIDC key rotation detected: previously cached keys are no longer published \
138                 by the provider; evicting them so tokens signed by rotated-out keys are rejected"
139            );
140        }
141
142        // Resolve the requested key from the freshly-fetched set first, so that a
143        // missing `kid` does not short-circuit the cache replacement below.
144        let key = jwks.keys.iter().find(|k| k.kid.as_deref() == Some(kid)).cloned();
145
146        // SECURITY (#361): ALWAYS replace the cache with the freshly-fetched JWKS,
147        // even when `kid` was not found. This evicts rotated-out keys so a token
148        // signed by a removed key cannot keep validating off a stale cache entry on
149        // a later cache hit.
150        {
151            let mut cache = self.jwks_cache.write();
152            *cache = Some(CachedJwks {
153                jwks,
154                fetched_at: Instant::now(),
155                ttl: Duration::from_secs(self.config.jwks_cache_ttl_secs),
156            });
157        }
158
159        let key = key.ok_or_else(|| {
160            tracing::debug!(kid = %kid, "Key not found in JWKS");
161            SecurityError::InvalidToken
162        })?;
163        self.jwk_to_decoding_key(&key)
164    }
165
166    /// Invalidate the cached JWKS so the next token validation refetches keys.
167    ///
168    /// Use this when an operator learns of an `IdP`-side key compromise or rotation
169    /// and wants to close the stolen-key replay window immediately, rather than
170    /// waiting up to `jwks_cache_ttl_secs` for the cache to expire. The next token
171    /// validation performs a fresh fetch from the provider.
172    pub fn invalidate_jwks_cache(&self) {
173        *self.jwks_cache.write() = None;
174        tracing::info!(
175            "JWKS cache invalidated; next token validation will refetch keys from the provider"
176        );
177    }
178
179    /// Force an immediate JWKS refetch, replacing the cache with the provider's
180    /// current key set.
181    ///
182    /// Returns the number of keys fetched. Backs the operator-facing
183    /// `/admin/v1/auth/refresh-jwks` endpoint, which closes the stolen-key replay
184    /// window on demand and confirms the refresh succeeded.
185    ///
186    /// # Errors
187    ///
188    /// Returns `SecurityError::SecurityConfigError` if the JWKS endpoint cannot be
189    /// reached or the response cannot be parsed.
190    pub async fn refresh_jwks(&self) -> Result<usize> {
191        let jwks = self.fetch_jwks().await?;
192        let key_count = jwks.keys.len();
193        {
194            let mut cache = self.jwks_cache.write();
195            *cache = Some(CachedJwks {
196                jwks,
197                fetched_at: Instant::now(),
198                ttl: Duration::from_secs(self.config.jwks_cache_ttl_secs),
199            });
200        }
201        tracing::info!(key_count, "JWKS cache force-refreshed from provider");
202        Ok(key_count)
203    }
204
205    /// Fetch JWKS from the provider.
206    ///
207    /// # Errors
208    ///
209    /// Returns `SecurityError::SecurityConfigError` if the HTTP request fails or
210    /// the response cannot be parsed as a valid JWKS.
211    async fn fetch_jwks(&self) -> Result<Jwks> {
212        tracing::debug!(uri = %self.jwks_uri, "Fetching JWKS");
213
214        let response = self.http_client.get(&self.jwks_uri).send().await.map_err(|e| {
215            tracing::error!(error = %e, "Failed to fetch JWKS");
216            SecurityError::SecurityConfigError(format!("Failed to fetch JWKS: {e}"))
217        })?;
218
219        if !response.status().is_success() {
220            return Err(SecurityError::SecurityConfigError(format!(
221                "JWKS fetch failed with status: {}",
222                response.status()
223            )));
224        }
225
226        // Cap the response body before deserialising to prevent memory exhaustion
227        // from a malicious or compromised OIDC provider sending an oversized payload.
228        let body_bytes = response.bytes().await.map_err(|e| {
229            SecurityError::SecurityConfigError(format!("Failed to read JWKS response body: {e}"))
230        })?;
231
232        if body_bytes.len() > MAX_JWKS_RESPONSE_BYTES {
233            return Err(SecurityError::SecurityConfigError(format!(
234                "JWKS response body too large ({} bytes, max {MAX_JWKS_RESPONSE_BYTES})",
235                body_bytes.len()
236            )));
237        }
238
239        let jwks: Jwks = serde_json::from_slice(&body_bytes).map_err(|e| {
240            SecurityError::SecurityConfigError(format!("Invalid JWKS response: {e}"))
241        })?;
242
243        tracing::debug!(key_count = jwks.keys.len(), "JWKS fetched successfully");
244
245        Ok(jwks)
246    }
247
248    /// Find a key in the JWKS by key ID.
249    pub(super) fn find_key<'a>(&self, jwks: &'a Jwks, kid: &str) -> Option<&'a Jwk> {
250        jwks.keys.iter().find(|k| k.kid.as_deref() == Some(kid))
251    }
252
253    /// Detect if JWKS keys have been rotated (old keys removed).
254    ///
255    /// Compares current cached keys with newly fetched keys.
256    /// Returns true if any previously cached keys are missing from the new JWKS.
257    pub(super) fn detect_key_rotation(&self, new_jwks: &Jwks) -> bool {
258        let cache = self.jwks_cache.read();
259        if let Some(ref cached) = *cache {
260            // Get set of old key IDs
261            let old_kids: std::collections::HashSet<_> =
262                cached.jwks.keys.iter().filter_map(|k| k.kid.as_deref()).collect();
263
264            // Get set of new key IDs
265            let new_kids: std::collections::HashSet<_> =
266                new_jwks.keys.iter().filter_map(|k| k.kid.as_deref()).collect();
267
268            // Rotation detected if any old keys are missing
269            !old_kids.is_subset(&new_kids)
270        } else {
271            false
272        }
273    }
274
275    /// Convert a JWK to a jsonwebtoken `DecodingKey`.
276    ///
277    /// # Errors
278    ///
279    /// Returns `SecurityError::InvalidToken` if the key type is unsupported or
280    /// required RSA components (n, e) are missing.
281    pub(super) fn jwk_to_decoding_key(&self, jwk: &Jwk) -> Result<DecodingKey> {
282        match jwk.kty.as_str() {
283            "RSA" => {
284                let n = jwk.n.as_ref().ok_or(SecurityError::InvalidToken)?;
285                let e = jwk.e.as_ref().ok_or(SecurityError::InvalidToken)?;
286
287                DecodingKey::from_rsa_components(n, e).map_err(|e| {
288                    tracing::debug!(error = %e, "Failed to create RSA decoding key");
289                    SecurityError::InvalidToken
290                })
291            },
292            other => {
293                tracing::debug!(key_type = %other, "Unsupported key type");
294                Err(SecurityError::InvalidTokenAlgorithm {
295                    algorithm: other.to_string(),
296                })
297            },
298        }
299    }
300}