Skip to main content

rws_server/sso/
jwks.rs

1//! JSON Web Key Set (JWKS) cache and JWT verification.
2//!
3//! [`JwksCache`] fetches and caches public keys from a JWKS endpoint, then
4//! verifies RS256 and ES256 JWTs with those keys.  It re-fetches on the first
5//! verification failure to handle key rotation.
6//!
7//! [`OidcClaims`] holds the standard OIDC claims extracted from a verified
8//! id_token payload.
9
10use std::sync::Mutex;
11
12use p256::ecdsa::signature::Verifier as EcVerifier;
13use p256::ecdsa::{Signature as EcSignature, VerifyingKey as EcVerifyingKey};
14use crate::rsa_key::RsaPublicKey;
15use serde::{Deserialize, Serialize};
16
17use crate::http_client::Client;
18
19use super::pkce::base64url_decode;
20use super::SsoError;
21
22/// Standard OIDC claims extracted from an id_token or userinfo response.
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct OidcClaims {
25    pub sub:            String,
26    pub iss:            String,
27    pub aud:            Vec<String>,
28    pub exp:            u64,
29    pub iat:            u64,
30    pub nonce:          Option<String>,
31    pub email:          Option<String>,
32    pub email_verified: Option<bool>,
33    pub name:           Option<String>,
34    pub given_name:     Option<String>,
35    pub family_name:    Option<String>,
36    pub picture:        Option<String>,
37    pub locale:         Option<String>,
38}
39
40/// Options for JWT verification.
41pub struct VerifyOptions<'a> {
42    /// Expected `aud` claim value — typically the OAuth client ID.
43    pub audience:    &'a str,
44    /// Expected `iss` claim value.
45    pub issuer:      &'a str,
46    /// Clock skew tolerance in seconds.
47    pub leeway_secs: u64,
48}
49
50// ── internal key cache ────────────────────────────────────────────────────────
51
52enum CachedKey {
53    Rsa { n: Vec<u8>, e: Vec<u8> },
54    Ec  { x: Vec<u8>, y: Vec<u8> },
55}
56
57struct JwkEntry {
58    kid: Option<String>,
59    key: CachedKey,
60}
61
62/// Thread-safe cache of public keys from a JWKS endpoint.
63pub struct JwksCache {
64    uri:  String,
65    keys: Mutex<Vec<JwkEntry>>,
66}
67
68impl JwksCache {
69    /// Create a new (empty) cache for the given JWKS URI.
70    pub fn new(uri: &str) -> Self {
71        JwksCache { uri: uri.to_string(), keys: Mutex::new(Vec::new()) }
72    }
73
74    /// Fetch keys from the JWKS URI and replace the cache.
75    pub fn fetch(&self) -> Result<(), SsoError> {
76        let resp = Client::new()
77            .get(&self.uri)
78            .timeout_ms(10_000)
79            .send()
80            .map_err(|e| SsoError(format!("JWKS fetch failed: {e}")))?;
81        if !resp.is_success() {
82            return Err(SsoError(format!("JWKS returned {}", resp.status())));
83        }
84        let body = resp.text().map_err(|e| SsoError(e.to_string()))?;
85        let entries = parse_jwks(&body)?;
86        *self.keys.lock().unwrap() = entries;
87        Ok(())
88    }
89
90    /// Verify a JWT and return the parsed claims.
91    ///
92    /// Fetches keys on first call or on signature failure (key rotation).
93    pub fn verify_jwt(
94        &self,
95        token: &str,
96        opts: &VerifyOptions<'_>,
97    ) -> Result<OidcClaims, SsoError> {
98        let parts: Vec<&str> = token.splitn(3, '.').collect();
99        if parts.len() != 3 {
100            return Err(SsoError("JWT must have 3 parts".into()));
101        }
102
103        let header_bytes = base64url_decode(parts[0])?;
104        let header_json = std::str::from_utf8(&header_bytes)
105            .map_err(|_| SsoError("JWT header is not UTF-8".into()))?;
106
107        let alg = json_str(header_json, "alg").unwrap_or_default();
108        let kid = json_str(header_json, "kid");
109
110        let message = format!("{}.{}", parts[0], parts[1]).into_bytes();
111        let signature = base64url_decode(parts[2])?;
112
113        // Lazy-load keys on first use
114        {
115            let keys = self.keys.lock().unwrap();
116            if keys.is_empty() {
117                drop(keys);
118                self.fetch()?;
119            }
120        }
121
122        let verified =
123            self.try_verify(&message, &signature, &alg, kid.as_deref())?;
124        if !verified {
125            // Try re-fetching in case of key rotation
126            self.fetch()?;
127            if !self.try_verify(&message, &signature, &alg, kid.as_deref())? {
128                return Err(SsoError("JWT signature verification failed".into()));
129            }
130        }
131
132        let payload_bytes = base64url_decode(parts[1])?;
133        let payload_json = std::str::from_utf8(&payload_bytes)
134            .map_err(|_| SsoError("JWT payload is not UTF-8".into()))?;
135
136        let claims = parse_claims(payload_json)?;
137
138        // Validate standard claims
139        let now = unix_now();
140        if claims.exp + opts.leeway_secs < now {
141            return Err(SsoError(format!(
142                "JWT expired (exp={}, now={})",
143                claims.exp, now
144            )));
145        }
146        if claims.iat > now + opts.leeway_secs {
147            return Err(SsoError("JWT issued in the future".into()));
148        }
149        if claims.iss != opts.issuer {
150            return Err(SsoError(format!(
151                "JWT issuer mismatch: expected {}, got {}",
152                opts.issuer, claims.iss
153            )));
154        }
155        if !claims.aud.iter().any(|a| a == opts.audience) {
156            return Err(SsoError(format!(
157                "JWT audience does not include {}",
158                opts.audience
159            )));
160        }
161
162        Ok(claims)
163    }
164
165    #[allow(deprecated)]
166    fn try_verify(
167        &self,
168        message: &[u8],
169        signature: &[u8],
170        alg: &str,
171        kid: Option<&str>,
172    ) -> Result<bool, SsoError> {
173        let keys = self.keys.lock().unwrap();
174        let candidates: Vec<&JwkEntry> = if let Some(kid_val) = kid {
175            keys.iter()
176                .filter(|k| k.kid.as_deref() == Some(kid_val))
177                .collect()
178        } else {
179            keys.iter().collect()
180        };
181
182        for entry in candidates {
183            match (&entry.key, alg) {
184                (CachedKey::Rsa { n, e }, "RS256") => {
185                    // A key we can't use (e.g. under 2048 bits) is skipped,
186                    // not fatal: another key in the set may match.
187                    let Ok(pub_key) = RsaPublicKey::from_components(n, e) else { continue };
188                    if pub_key.verify_pkcs1v15_sha256(message, signature) {
189                        return Ok(true);
190                    }
191                }
192                (CachedKey::Ec { x, y }, "ES256") => {
193                    if x.len() != 32 || y.len() != 32 {
194                        continue;
195                    }
196                    // Uncompressed SEC1 point: 0x04 || x || y.
197                    let mut sec1 = Vec::with_capacity(65);
198                    sec1.push(0x04);
199                    sec1.extend_from_slice(x);
200                    sec1.extend_from_slice(y);
201                    if let Ok(vk) = EcVerifyingKey::from_sec1_bytes(&sec1) {
202                        // JWT ES256 signatures are r||s (64 bytes), not DER
203                        if signature.len() != 64 {
204                            return Err(SsoError("ES256 sig wrong length".into()));
205                        }
206                        if let Ok(sig) = EcSignature::from_slice(signature) {
207                            if vk.verify(message, &sig).is_ok() {
208                                return Ok(true);
209                            }
210                        }
211                    }
212                }
213                _ => {}
214            }
215        }
216        Ok(false)
217    }
218}
219
220// ── JWKS JSON parsing ─────────────────────────────────────────────────────────
221
222fn parse_jwks(json: &str) -> Result<Vec<JwkEntry>, SsoError> {
223    // Find the start of the "keys" array
224    let keys_start = json
225        .find("\"keys\"")
226        .and_then(|p| json[p..].find('[').map(|q| p + q + 1))
227        .ok_or_else(|| SsoError("JWKS missing 'keys' array".into()))?;
228
229    let arr = &json[keys_start..];
230    let mut entries = Vec::new();
231    let mut depth = 0i32;
232    let mut obj_start = None;
233
234    for (i, ch) in arr.char_indices() {
235        match ch {
236            '{' => {
237                if depth == 0 {
238                    obj_start = Some(i);
239                }
240                depth += 1;
241            }
242            '}' => {
243                depth -= 1;
244                if depth == 0 {
245                    if let Some(start) = obj_start {
246                        let obj = &arr[start..=i];
247                        if let Some(entry) = parse_jwk_object(obj) {
248                            entries.push(entry);
249                        }
250                        obj_start = None;
251                    }
252                }
253            }
254            _ => {}
255        }
256    }
257    Ok(entries)
258}
259
260fn parse_jwk_object(obj: &str) -> Option<JwkEntry> {
261    let kty = json_str(obj, "kty")?;
262    let kid = json_str(obj, "kid");
263    match kty.as_str() {
264        "RSA" => {
265            let n = base64url_decode(&json_str(obj, "n")?).ok()?;
266            let e = base64url_decode(&json_str(obj, "e")?).ok()?;
267            Some(JwkEntry { kid, key: CachedKey::Rsa { n, e } })
268        }
269        "EC" => {
270            let crv = json_str(obj, "crv")?;
271            if crv != "P-256" {
272                return None;
273            }
274            let x = base64url_decode(&json_str(obj, "x")?).ok()?;
275            let y = base64url_decode(&json_str(obj, "y")?).ok()?;
276            Some(JwkEntry { kid, key: CachedKey::Ec { x, y } })
277        }
278        _ => None,
279    }
280}
281
282// ── claims parsing ────────────────────────────────────────────────────────────
283
284fn parse_claims(json: &str) -> Result<OidcClaims, SsoError> {
285    let sub = json_str(json, "sub")
286        .ok_or_else(|| SsoError("JWT missing 'sub' claim".into()))?;
287    let iss = json_str(json, "iss").unwrap_or_default();
288    let exp = json_u64(json, "exp")
289        .ok_or_else(|| SsoError("JWT missing 'exp' claim".into()))?;
290    let iat = json_u64(json, "iat").unwrap_or(0);
291    let nonce = json_str(json, "nonce");
292    let email = json_str(json, "email");
293    let email_verified = json_bool(json, "email_verified");
294    let name = json_str(json, "name");
295    let given_name = json_str(json, "given_name");
296    let family_name = json_str(json, "family_name");
297    let picture = json_str(json, "picture");
298    let locale = json_str(json, "locale");
299    let aud = parse_aud(json);
300
301    Ok(OidcClaims {
302        sub,
303        iss,
304        aud,
305        exp,
306        iat,
307        nonce,
308        email,
309        email_verified,
310        name,
311        given_name,
312        family_name,
313        picture,
314        locale,
315    })
316}
317
318fn parse_aud(json: &str) -> Vec<String> {
319    let needle = "\"aud\"";
320    let Some(pos) = json.find(needle) else {
321        return vec![];
322    };
323    let rest =
324        json[pos + needle.len()..].trim_start_matches(|c: char| c.is_whitespace() || c == ':');
325    if rest.starts_with('[') {
326        // Array form
327        let end = rest.find(']').unwrap_or(rest.len());
328        rest[1..end]
329            .split(',')
330            .filter_map(|s| {
331                let s = s.trim();
332                if s.starts_with('"') && s.ends_with('"') {
333                    Some(s[1..s.len() - 1].to_string())
334                } else {
335                    None
336                }
337            })
338            .collect()
339    } else if rest.starts_with('"') {
340        let inner = &rest[1..];
341        let end = inner.find('"').unwrap_or(inner.len());
342        vec![inner[..end].to_string()]
343    } else {
344        vec![]
345    }
346}
347
348// ── minimal JSON field extractors ─────────────────────────────────────────────
349
350/// Extract a JSON string field by key (handles simple escape sequences).
351#[doc(hidden)]
352pub fn json_str(json: &str, key: &str) -> Option<String> {
353    let needle = format!("\"{key}\"");
354    let start = json.find(&needle)? + needle.len();
355    let rest = json[start..].trim_start_matches(|c: char| c.is_whitespace() || c == ':');
356    if rest.starts_with('"') {
357        let inner = &rest[1..];
358        let mut escaped = false;
359        let mut out = String::new();
360        for ch in inner.chars() {
361            if escaped {
362                out.push(ch);
363                escaped = false;
364                continue;
365            }
366            if ch == '\\' {
367                escaped = true;
368                continue;
369            }
370            if ch == '"' {
371                return Some(out);
372            }
373            out.push(ch);
374        }
375        None
376    } else {
377        None
378    }
379}
380
381fn json_u64(json: &str, key: &str) -> Option<u64> {
382    let needle = format!("\"{key}\"");
383    let start = json.find(&needle)? + needle.len();
384    let rest = json[start..].trim_start_matches(|c: char| c.is_whitespace() || c == ':');
385    let end = rest.find(|c: char| !c.is_ascii_digit()).unwrap_or(rest.len());
386    rest[..end].parse().ok()
387}
388
389fn json_bool(json: &str, key: &str) -> Option<bool> {
390    let needle = format!("\"{key}\"");
391    let start = json.find(&needle)? + needle.len();
392    let rest = json[start..].trim_start_matches(|c: char| c.is_whitespace() || c == ':');
393    if rest.starts_with("true") {
394        Some(true)
395    } else if rest.starts_with("false") {
396        Some(false)
397    } else {
398        None
399    }
400}
401
402fn unix_now() -> u64 {
403    std::time::SystemTime::now()
404        .duration_since(std::time::UNIX_EPOCH)
405        .map(|d| d.as_secs())
406        .unwrap_or(0)
407}