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: Detect key rotation for audit purposes
133        if self.detect_key_rotation(&jwks) {
134            tracing::warn!(
135                "OIDC key rotation detected: some previously cached keys no longer available"
136            );
137        }
138
139        // Find the key index first, then we can clone the key
140        let key_index =
141            jwks.keys.iter().position(|k| k.kid.as_deref() == Some(kid)).ok_or_else(|| {
142                tracing::debug!(kid = %kid, "Key not found in JWKS");
143                SecurityError::InvalidToken
144            })?;
145
146        // Clone the key before caching (keys are small, cloning is fine)
147        let key = jwks.keys[key_index].clone();
148
149        // Cache the JWKS
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        self.jwk_to_decoding_key(&key)
160    }
161
162    /// Fetch JWKS from the provider.
163    ///
164    /// # Errors
165    ///
166    /// Returns `SecurityError::SecurityConfigError` if the HTTP request fails or
167    /// the response cannot be parsed as a valid JWKS.
168    async fn fetch_jwks(&self) -> Result<Jwks> {
169        tracing::debug!(uri = %self.jwks_uri, "Fetching JWKS");
170
171        let response = self.http_client.get(&self.jwks_uri).send().await.map_err(|e| {
172            tracing::error!(error = %e, "Failed to fetch JWKS");
173            SecurityError::SecurityConfigError(format!("Failed to fetch JWKS: {e}"))
174        })?;
175
176        if !response.status().is_success() {
177            return Err(SecurityError::SecurityConfigError(format!(
178                "JWKS fetch failed with status: {}",
179                response.status()
180            )));
181        }
182
183        // Cap the response body before deserialising to prevent memory exhaustion
184        // from a malicious or compromised OIDC provider sending an oversized payload.
185        let body_bytes = response.bytes().await.map_err(|e| {
186            SecurityError::SecurityConfigError(format!("Failed to read JWKS response body: {e}"))
187        })?;
188
189        if body_bytes.len() > MAX_JWKS_RESPONSE_BYTES {
190            return Err(SecurityError::SecurityConfigError(format!(
191                "JWKS response body too large ({} bytes, max {MAX_JWKS_RESPONSE_BYTES})",
192                body_bytes.len()
193            )));
194        }
195
196        let jwks: Jwks = serde_json::from_slice(&body_bytes).map_err(|e| {
197            SecurityError::SecurityConfigError(format!("Invalid JWKS response: {e}"))
198        })?;
199
200        tracing::debug!(key_count = jwks.keys.len(), "JWKS fetched successfully");
201
202        Ok(jwks)
203    }
204
205    /// Find a key in the JWKS by key ID.
206    pub(super) fn find_key<'a>(&self, jwks: &'a Jwks, kid: &str) -> Option<&'a Jwk> {
207        jwks.keys.iter().find(|k| k.kid.as_deref() == Some(kid))
208    }
209
210    /// Detect if JWKS keys have been rotated (old keys removed).
211    ///
212    /// Compares current cached keys with newly fetched keys.
213    /// Returns true if any previously cached keys are missing from the new JWKS.
214    pub(super) fn detect_key_rotation(&self, new_jwks: &Jwks) -> bool {
215        let cache = self.jwks_cache.read();
216        if let Some(ref cached) = *cache {
217            // Get set of old key IDs
218            let old_kids: std::collections::HashSet<_> =
219                cached.jwks.keys.iter().filter_map(|k| k.kid.as_deref()).collect();
220
221            // Get set of new key IDs
222            let new_kids: std::collections::HashSet<_> =
223                new_jwks.keys.iter().filter_map(|k| k.kid.as_deref()).collect();
224
225            // Rotation detected if any old keys are missing
226            !old_kids.is_subset(&new_kids)
227        } else {
228            false
229        }
230    }
231
232    /// Convert a JWK to a jsonwebtoken `DecodingKey`.
233    ///
234    /// # Errors
235    ///
236    /// Returns `SecurityError::InvalidToken` if the key type is unsupported or
237    /// required RSA components (n, e) are missing.
238    pub(super) fn jwk_to_decoding_key(&self, jwk: &Jwk) -> Result<DecodingKey> {
239        match jwk.kty.as_str() {
240            "RSA" => {
241                let n = jwk.n.as_ref().ok_or(SecurityError::InvalidToken)?;
242                let e = jwk.e.as_ref().ok_or(SecurityError::InvalidToken)?;
243
244                DecodingKey::from_rsa_components(n, e).map_err(|e| {
245                    tracing::debug!(error = %e, "Failed to create RSA decoding key");
246                    SecurityError::InvalidToken
247                })
248            },
249            other => {
250                tracing::debug!(key_type = %other, "Unsupported key type");
251                Err(SecurityError::InvalidTokenAlgorithm {
252                    algorithm: other.to_string(),
253                })
254            },
255        }
256    }
257}