pep/oidc/
resource_server.rs

1//! Resource server functionality for JWT validation and API protection
2
3use std::{collections::HashMap, sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}};
4use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
5use jsonwebtoken::jwk::JwkSet;
6use reqwest::Client;
7use tokio::sync::RwLock;
8use tracing;
9
10use crate::error::{PepError, Result};
11use super::types::{JwtClaims, OidcDiscoveryDocument, CachedJwks, CachedDiscoveryRaw, JwtValidationOptions};
12
13/// Cache entry for userinfo endpoint responses
14#[derive(Clone)]
15pub struct CachedUserInfo {
16    /// The userinfo claims
17    pub claims: serde_json::Map<String, serde_json::Value>,
18    /// When the entry was cached
19    pub cached_at: SystemTime,
20    /// TTL in seconds (derived from token expiry)
21    pub ttl_secs: u64,
22}
23
24impl CachedUserInfo {
25    /// Check if this cache entry has expired
26    pub fn is_expired(&self) -> bool {
27        self.cached_at.elapsed().unwrap_or(Duration::from_secs(self.ttl_secs + 1)) >= Duration::from_secs(self.ttl_secs)
28    }
29}
30
31/// In-memory cache for OIDC userinfo responses.
32///
33/// Keyed by JWT `jti` (token ID) when available, falling back to `sub` (subject).
34/// Entries auto-expire based on the remaining token lifetime.
35#[derive(Clone)]
36pub struct UserInfoCache {
37    inner: Arc<RwLock<HashMap<String, CachedUserInfo>>>,
38}
39
40impl UserInfoCache {
41    /// Create a new empty userinfo cache
42    pub fn new() -> Self {
43        Self {
44            inner: Arc::new(RwLock::new(HashMap::new())),
45        }
46    }
47
48    /// Get cached userinfo if present and not expired
49    pub async fn get(&self, key: &str) -> Option<serde_json::Map<String, serde_json::Value>> {
50        let cache = self.inner.read().await;
51        cache.get(key).and_then(|entry| {
52            if entry.is_expired() {
53                None
54            } else {
55                Some(entry.claims.clone())
56            }
57        })
58    }
59
60    /// Store userinfo claims with a TTL derived from token expiry
61    pub async fn insert(
62        &self,
63        key: String,
64        claims: serde_json::Map<String, serde_json::Value>,
65        token_exp: i64,
66    ) {
67        let now = SystemTime::now()
68            .duration_since(UNIX_EPOCH)
69            .unwrap_or(Duration::from_secs(0))
70            .as_secs() as i64;
71
72        // TTL = remaining token lifetime, with a safety margin of 30 seconds
73        let remaining = if token_exp > now {
74            (token_exp - now) as u64
75        } else {
76            0
77        };
78        let ttl_secs = remaining.saturating_sub(30).max(30); // at least 30s, at most token remaining - 30s
79
80        let entry = CachedUserInfo {
81            claims,
82            cached_at: SystemTime::now(),
83            ttl_secs,
84        };
85
86        let mut cache = self.inner.write().await;
87        cache.insert(key, entry);
88    }
89
90    /// Purge expired entries (call periodically to free memory)
91    pub async fn purge_expired(&self) {
92        let mut cache = self.inner.write().await;
93        cache.retain(|_, entry| !entry.is_expired());
94    }
95}
96
97impl Default for UserInfoCache {
98    fn default() -> Self {
99        Self::new()
100    }
101}
102
103/// Map a JWK's algorithm parameters to a `jsonwebtoken::Algorithm`
104pub fn jwk_algorithm_to_algorithm(jwk: &jsonwebtoken::jwk::Jwk) -> Result<Algorithm> {
105    match &jwk.algorithm {
106        jsonwebtoken::jwk::AlgorithmParameters::RSA(_) => Ok(Algorithm::RS256),
107        jsonwebtoken::jwk::AlgorithmParameters::EllipticCurve(params) => match &params.curve {
108            jsonwebtoken::jwk::EllipticCurve::P256 => Ok(Algorithm::ES256),
109            jsonwebtoken::jwk::EllipticCurve::P384 => Ok(Algorithm::ES384),
110            other => Err(PepError::BadRequest(format!("Unsupported elliptic curve for JWK: {:?}", other))),
111        },
112        jsonwebtoken::jwk::AlgorithmParameters::OctetKey(_) => Err(PepError::BadRequest("HMAC keys not supported for OIDC verification".to_string())),
113        jsonwebtoken::jwk::AlgorithmParameters::OctetKeyPair(_) => Ok(Algorithm::EdDSA),
114    }
115}
116
117/// Resource server client for JWT validation
118#[derive(Clone)]
119pub struct ResourceServerClient {
120    /// HTTP client
121    pub http_client: Client,
122    /// JWKS cache
123    pub jwks_cache: Arc<RwLock<HashMap<String, CachedJwks>>>,
124    /// Discovery document cache (parsed)
125    pub discovery_cache: Arc<RwLock<HashMap<String, (OidcDiscoveryDocument, SystemTime)>>>,
126    /// Discovery document cache (raw JSON for proxying)
127    pub discovery_cache_raw: Arc<RwLock<HashMap<String, CachedDiscoveryRaw>>>,
128    /// Userinfo response cache (keyed by jti/sub, TTL = token remaining lifetime)
129    pub userinfo_cache: Arc<UserInfoCache>,
130}
131
132impl ResourceServerClient {
133    /// Create a new resource server client
134    pub fn new() -> Self {
135        Self {
136            http_client: Client::new(),
137            jwks_cache: Arc::new(RwLock::new(HashMap::new())),
138            discovery_cache: Arc::new(RwLock::new(HashMap::new())),
139            discovery_cache_raw: Arc::new(RwLock::new(HashMap::new())),
140            userinfo_cache: Arc::new(UserInfoCache::new()),
141        }
142    }
143
144    /// Fetch OIDC discovery document with caching
145    pub async fn get_discovery_document(&self, issuer_url: &str) -> Result<OidcDiscoveryDocument> {
146        // Check cache first
147        {
148            let cache = self.discovery_cache.read().await;
149            if let Some((doc, fetched_at)) = cache.get(issuer_url) {
150                // Cache for 1 hour
151                if fetched_at.elapsed().unwrap_or(Duration::from_secs(3600)) < Duration::from_secs(3600) {
152                    return Ok(doc.clone());
153                }
154            }
155        }
156
157        // Fetch discovery document
158        let discovery_url = format!("{}/.well-known/openid-configuration", issuer_url.trim_end_matches('/'));
159        tracing::debug!("Fetching OIDC discovery document from: {}", discovery_url);
160
161        let response = self.http_client
162            .get(&discovery_url)
163            .header("Accept", "application/json")
164            .send()
165            .await
166            .map_err(|e| PepError::OidcDiscovery(format!("Failed to fetch discovery document: {}", e)))?;
167
168        if !response.status().is_success() {
169            return Err(PepError::OidcDiscovery(format!("Discovery document fetch failed with status: {}", response.status())));
170        }
171
172        let discovery_doc: OidcDiscoveryDocument = response
173            .json()
174            .await
175            .map_err(|e| PepError::OidcDiscovery(format!("Failed to parse discovery document: {}", e)))?;
176
177        // Cache the document
178        {
179            let mut cache = self.discovery_cache.write().await;
180            cache.insert(issuer_url.to_string(), (discovery_doc.clone(), SystemTime::now()));
181        }
182
183        Ok(discovery_doc)
184    }
185
186    /// Fetch OIDC discovery document as raw JSON with caching
187    pub async fn get_discovery_document_raw(&self, issuer_url: &str) -> Result<String> {
188        let cache_duration = Duration::from_secs(3600);
189        
190        {
191            let cache = self.discovery_cache_raw.read().await;
192            if let Some(cached) = cache.get(issuer_url) {
193                if cached.fetched_at.elapsed().unwrap_or(cache_duration) < cache_duration {
194                    return Ok(cached.raw_json.clone());
195                }
196            }
197        }
198
199        let discovery_url = format!("{}/.well-known/openid-configuration", issuer_url.trim_end_matches('/'));
200        tracing::debug!("Fetching raw OIDC discovery document from: {}", discovery_url);
201
202        let response = self.http_client
203            .get(&discovery_url)
204            .header("Accept", "application/json")
205            .send()
206            .await
207            .map_err(|e| PepError::OidcDiscovery(format!("Failed to fetch discovery document: {}", e)))?;
208
209        if !response.status().is_success() {
210            return Err(PepError::OidcDiscovery(format!("Discovery document fetch failed with status: {}", response.status())));
211        }
212
213        let raw_json = response
214            .text()
215            .await
216            .map_err(|e| PepError::OidcDiscovery(format!("Failed to read discovery document: {}", e)))?;
217
218        let cached = CachedDiscoveryRaw {
219            raw_json: raw_json.clone(),
220            fetched_at: SystemTime::now(),
221            cache_duration,
222        };
223        {
224            let mut cache = self.discovery_cache_raw.write().await;
225            cache.insert(issuer_url.to_string(), cached);
226        }
227
228        Ok(raw_json)
229    }
230
231    /// Fetch JWKS with caching
232    pub async fn get_jwks(&self, jwks_uri: &str) -> Result<HashMap<String, (DecodingKey, Algorithm)>> {
233        // Check cache first
234        {
235            let cache = self.jwks_cache.read().await;
236            if let Some(cached) = cache.get(jwks_uri) {
237                // Cache for 1 hour
238                if cached.fetched_at.elapsed().unwrap_or(cached.cache_duration) < cached.cache_duration {
239                    return Ok(cached.keys.clone());
240                }
241            }
242        }
243
244        // Fetch JWKS
245        tracing::debug!("Fetching JWKS from: {}", jwks_uri);
246
247        let response = self.http_client
248            .get(jwks_uri)
249            .header("Accept", "application/json")
250            .send()
251            .await
252            .map_err(|e| PepError::JwksFetch(format!("Failed to fetch JWKS: {}", e)))?;
253
254        if !response.status().is_success() {
255            return Err(PepError::JwksFetch(format!("JWKS fetch failed with status: {}", response.status())));
256        }
257
258        let jwks_text = response
259            .text()
260            .await
261            .map_err(|e| PepError::JwksFetch(format!("Failed to read JWKS response: {}", e)))?;
262
263        let jwk_set: JwkSet = serde_json::from_str(&jwks_text)
264            .map_err(|e| PepError::JwksFetch(format!("Failed to parse JWKS: {}", e)))?;
265
266        // Parse the keys
267        let mut keys = HashMap::new();
268        for jwk in jwk_set.keys {
269            if let Some(kid) = &jwk.common.key_id {
270                match DecodingKey::from_jwk(&jwk) {
271                    Ok(decoding_key) => {
272                        match jwk_algorithm_to_algorithm(&jwk) {
273                            Ok(algorithm) => {
274                                keys.insert(kid.clone(), (decoding_key, algorithm));
275                                tracing::debug!("Successfully parsed key {}: algorithm={:?}", kid, algorithm);
276                            }
277                            Err(e) => {
278                                tracing::warn!("Unsupported algorithm for kid {}: {}", kid, e);
279                                continue;
280                            }
281                        }
282                    }
283                    Err(err) => {
284                        tracing::warn!("Failed to create decoding key for kid {}: {}", kid, err);
285                    }
286                }
287            } else {
288                tracing::warn!("JWK missing kid field, skipping");
289            }
290        }
291
292        // Cache the keys
293        let cached = CachedJwks {
294            keys: keys.clone(),
295            fetched_at: SystemTime::now(),
296            cache_duration: Duration::from_secs(3600), // 1 hour
297        };
298        {
299            let mut cache = self.jwks_cache.write().await;
300            cache.insert(jwks_uri.to_string(), cached);
301        }
302
303        Ok(keys)
304    }
305
306    /// Validate JWT token with custom validation options
307    pub async fn validate_jwt_with_options(
308        &self,
309        token: &str,
310        issuer_url: &str,
311        client_id: &str,
312        options: &JwtValidationOptions,
313    ) -> Result<JwtClaims> {
314        // Decode header to get kid and algorithm
315        let header = decode_header(token)
316            .map_err(|e| PepError::JwtValidation(format!("Invalid JWT header: {}", e)))?;
317
318        let kid = header.kid
319            .ok_or_else(|| PepError::JwtValidation("JWT missing kid in header".to_string()))?;
320
321        // Get discovery document
322        let discovery_doc = self.get_discovery_document(issuer_url).await?;
323
324        // Get JWKS
325        let keys = self.get_jwks(&discovery_doc.jwks_uri).await?;
326
327        // Find the key for this kid
328        let (decoding_key, key_algorithm) = keys.get(&kid)
329            .ok_or_else(|| PepError::JwtValidation(format!("No key found for kid: {}", kid)))?;
330
331        // Determine which algorithm to use:
332        let algorithm = {
333            let jwt_alg = header.alg;
334            if jwt_alg == *key_algorithm {
335                jwt_alg
336            } else {
337                tracing::warn!(
338                    "JWT header algorithm ({:?}) doesn't match key algorithm ({:?}) for kid {}. Using key algorithm.",
339                    jwt_alg, key_algorithm, kid
340                );
341                *key_algorithm
342            }
343        };
344
345        tracing::debug!("Validating JWT with kid: {}, algorithm: {:?}", kid, algorithm);
346
347        // Set up validation
348        let mut validation = Validation::new(algorithm);
349        // Allow 60 seconds of clock skew between issuer, proxy, and resource server.
350        // Without this, sub-second timing differences between Torpi (proxy) and
351        // Trustee (resource server) can cause ExpiredSignature on valid tokens.
352        validation.leeway = 60;
353
354        // Configure issuer validation
355        if options.skip_issuer_validation {
356            tracing::debug!("Skipping issuer validation as configured");
357        } else {
358            validation.set_issuer(&[issuer_url]);
359        }
360
361        // Configure audience validation
362        if options.skip_audience_validation {
363            tracing::debug!("Skipping audience validation as configured");
364            validation.validate_aud = false;
365        } else {
366            let audience = options.expected_audience.as_deref().unwrap_or(client_id);
367            tracing::debug!("Validating audience against: {}", audience);
368            validation.set_audience(&[audience]);
369        }
370
371        // Decode and validate the token
372        let token_data = decode::<JwtClaims>(token, decoding_key, &validation)
373            .map_err(|e| PepError::JwtValidation(format!("JWT validation failed: {}", e)))?;
374
375        Ok(token_data.claims)
376    }
377
378    /// Validate JWT token with default options
379    pub async fn validate_jwt(&self, token: &str, issuer_url: &str, client_id: &str) -> Result<JwtClaims> {
380        self.validate_jwt_with_options(token, issuer_url, client_id, &JwtValidationOptions::default()).await
381    }
382
383    /// Adaptive claims enrichment: fill missing `groups` / `role` from the OIDC `/userinfo` endpoint.
384    ///
385    /// **IdP-agnostic design:**
386    /// 1. If `claims.extra` already contains `groups` or `role`, use them directly (zero cost).
387    /// 2. Otherwise, call the `/userinfo` endpoint with the access token as Bearer.
388    /// 3. Merge `groups` and `role` from userinfo into `claims.extra`.
389    /// 4. Cache the userinfo response by `jti` (or `sub` fallback) until the token expires.
390    ///
391    /// This works with Kanidm (no groups in AT), Keycloak/Auth0 (groups in AT → fast path),
392    /// Okta, Azure AD, Google, and any other OIDC-compliant provider.
393    ///
394    /// # Arguments
395    ///
396    /// * `claims` - The JWT claims returned by `validate_jwt_*`. Mutated in place.
397    /// * `token` - The raw access token string (used as Bearer for the userinfo call).
398    /// * `issuer_url` - The OIDC issuer URL (used to derive the userinfo endpoint).
399    /// * `userinfo_url_override` - Optional explicit userinfo URL. If `None`, the URL is
400    ///   derived from the discovery document's `userinfo_endpoint`, falling back to
401    ///   `{issuer_url}/userinfo`.
402    pub async fn enrich_claims_with_userinfo(
403        &self,
404        claims: &mut JwtClaims,
405        token: &str,
406        issuer_url: &str,
407        userinfo_url_override: Option<&str>,
408    ) -> Result<()> {
409        // Fast path: claims already have groups — sufficient for authorization.
410        // `role` is an application-level concept not issued by standard OIDC providers,
411        // so requiring it here would cause unnecessary userinfo calls on every request.
412        let has_groups = claims.extra.contains_key("groups");
413
414        if has_groups {
415            tracing::debug!("Claims already contain groups — skipping userinfo enrichment");
416            return Ok(());
417        }
418
419        tracing::debug!("Claims missing groups — attempting userinfo enrichment");
420
421        // Build cache key from jti (if present in extra) or sub
422        let cache_key = claims
423            .extra
424            .get("jti")
425            .and_then(|v| v.as_str())
426            .map(|s| s.to_string())
427            .unwrap_or_else(|| claims.sub.clone());
428
429        // Check userinfo cache first
430        if let Some(cached_claims) = self.userinfo_cache.get(&cache_key).await {
431            tracing::debug!(cache_key = %cache_key, "Using cached userinfo for claims enrichment");
432            merge_userinfo_into_claims(claims, &cached_claims);
433            return Ok(());
434        }
435
436        // Resolve the userinfo endpoint URL
437        let userinfo_url = match userinfo_url_override {
438            Some(url) => url.to_string(),
439            None => {
440                // Try discovery document first (has the canonical userinfo_endpoint)
441                match self.get_discovery_document(issuer_url).await {
442                    Ok(doc) if doc.userinfo_endpoint.is_some() => {
443                        doc.userinfo_endpoint.unwrap()
444                    }
445                    Ok(_) => {
446                        // Fallback: derive from issuer URL
447                        format!("{}/userinfo", issuer_url.trim_end_matches('/'))
448                    }
449                    Err(e) => {
450                        tracing::warn!("Failed to fetch discovery for userinfo URL: {}. Deriving from issuer.", e);
451                        format!("{}/userinfo", issuer_url.trim_end_matches('/'))
452                    }
453                }
454            }
455        };
456
457        tracing::debug!(userinfo_url = %userinfo_url, "Calling userinfo endpoint for claims enrichment");
458
459        // Call the userinfo endpoint with the access token as Bearer
460        let response = self.http_client
461            .get(&userinfo_url)
462            .header("Authorization", format!("Bearer {}", token))
463            .header("Accept", "application/json")
464            .send()
465            .await
466            .map_err(|e| PepError::Userinfo(format!("Userinfo request failed: {}", e)))?;
467
468        if !response.status().is_success() {
469            let status = response.status();
470            let body = response.text().await.unwrap_or_default();
471            tracing::warn!(
472                userinfo_url = %userinfo_url, status = %status,
473                "Userinfo endpoint returned error: {}", body
474            );
475            // Non-fatal: enrichment failed but JWT is still valid.
476            // Return Ok so the request proceeds with whatever claims we have.
477            return Ok(());
478        }
479
480        let userinfo: serde_json::Map<String, serde_json::Value> = response
481            .json()
482            .await
483            .map_err(|e| PepError::Userinfo(format!("Failed to parse userinfo response: {}", e)))?;
484
485        tracing::debug!(
486            userinfo_keys = ?userinfo.keys().collect::<Vec<_>>(),
487            "Userinfo response received"
488        );
489
490        // Cache the userinfo response (TTL = remaining token lifetime)
491        self.userinfo_cache.insert(cache_key.clone(), userinfo.clone(), claims.exp).await;
492
493        // Merge userinfo claims into JWT claims
494        merge_userinfo_into_claims(claims, &userinfo);
495
496        Ok(())
497    }
498}
499
500/// Merge select fields from the userinfo endpoint response into JWT claims `extra`.
501///
502/// Only merges fields that are relevant for authorization decisions (`groups`, `role`)
503/// and are not already present in `claims.extra`. This avoids overwriting values that
504/// the IdP may have already put into the access token.
505fn merge_userinfo_into_claims(
506    claims: &mut JwtClaims,
507    userinfo: &serde_json::Map<String, serde_json::Value>,
508) {
509    let fields_to_merge = ["groups", "role"];
510    for field in &fields_to_merge {
511        if !claims.extra.contains_key(*field) {
512            if let Some(value) = userinfo.get(*field) {
513                claims.extra.insert(field.to_string(), value.clone());
514                tracing::debug!(
515                    field,
516                    value = %value,
517                    "Merged field from userinfo into claims"
518                );
519            }
520        }
521    }
522}
523
524impl Default for ResourceServerClient {
525    fn default() -> Self {
526        Self::new()
527    }
528}
pep/oidc/resource_server.rs

pep/oidc/
resource_server.rs
