1use std::collections::HashMap;
32use std::path::{Path, PathBuf};
33use std::sync::atomic::{AtomicU64, Ordering};
34use std::sync::{Arc, Mutex, RwLock};
35use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
36
37use base64::engine::general_purpose::URL_SAFE_NO_PAD;
38use base64::Engine as _;
39use serde::{Deserialize, Serialize};
40
41use cheers_core::KeyRole;
42
43#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49pub struct JwkKey {
50 pub kty: String,
51 #[serde(default)]
52 pub crv: Option<String>,
53 #[serde(default)]
54 pub x: Option<String>,
55 #[serde(default)]
56 pub kid: Option<String>,
57 #[serde(default, rename = "use", skip_serializing_if = "Option::is_none")]
58 pub use_: Option<String>,
59 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub alg: Option<String>,
61 pub principal: String,
63 pub role: KeyRole,
65}
66
67impl JwkKey {
68 pub fn ed25519(
70 kid: impl Into<String>,
71 public_key: &[u8; 32],
72 principal: impl Into<String>,
73 role: KeyRole,
74 ) -> Self {
75 Self {
76 kty: "OKP".into(),
77 crv: Some("Ed25519".into()),
78 x: Some(URL_SAFE_NO_PAD.encode(public_key)),
79 kid: Some(kid.into()),
80 use_: Some("sig".into()),
81 alg: Some("EdDSA".into()),
82 principal: principal.into(),
83 role,
84 }
85 }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
90pub struct JwksDoc {
91 pub keys: Vec<JwkKey>,
92}
93
94#[derive(Debug, Clone, PartialEq, Eq)]
96pub struct KeyEntry {
97 pub public_key: [u8; 32],
98 pub principal: String,
99 pub role: KeyRole,
100}
101
102#[derive(Debug, Clone)]
104pub struct KeySet {
105 keys: HashMap<String, KeyEntry>,
106 doc: JwksDoc,
108}
109
110impl KeySet {
111 pub fn from_doc(doc: JwksDoc) -> Result<Self, JwksError> {
114 let mut keys = HashMap::with_capacity(doc.keys.len());
115 for k in &doc.keys {
116 if k.kty != "OKP" || k.crv.as_deref() != Some("Ed25519") {
117 continue;
118 }
119 let kid = k.kid.clone().ok_or_else(|| JwksError::BadKey {
120 kid: "<missing>".into(),
121 reason: "OKP/Ed25519 entry missing kid".into(),
122 })?;
123 let x = k.x.as_ref().ok_or_else(|| JwksError::BadKey {
124 kid: kid.clone(),
125 reason: "OKP/Ed25519 entry missing x".into(),
126 })?;
127 let bytes = URL_SAFE_NO_PAD.decode(x).map_err(|e| JwksError::BadKey {
128 kid: kid.clone(),
129 reason: format!("base64url decode: {e}"),
130 })?;
131 let public_key: [u8; 32] =
132 bytes.try_into().map_err(|v: Vec<u8>| JwksError::BadKey {
133 kid: kid.clone(),
134 reason: format!("expected 32 bytes, got {}", v.len()),
135 })?;
136 keys.insert(
137 kid,
138 KeyEntry {
139 public_key,
140 principal: k.principal.clone(),
141 role: k.role,
142 },
143 );
144 }
145 if keys.is_empty() {
146 return Err(JwksError::Parse("JWKS contained no OKP/Ed25519 keys".into()));
147 }
148 Ok(Self { keys, doc })
149 }
150
151 pub fn get(&self, kid: &str) -> Option<&KeyEntry> {
152 self.keys.get(kid)
153 }
154
155 pub fn len(&self) -> usize {
156 self.keys.len()
157 }
158
159 pub fn is_empty(&self) -> bool {
160 self.keys.is_empty()
161 }
162
163 pub fn doc(&self) -> &JwksDoc {
164 &self.doc
165 }
166}
167
168#[derive(Debug, thiserror::Error)]
170pub enum JwksError {
171 #[error(
172 "insecure cheers issuer {issuer:?}: the JWKS must be fetched over \
173 https:// (plaintext http:// is allowed only for loopback hosts)"
174 )]
175 InsecureIssuer { issuer: String },
176
177 #[error("JWKS source unreachable at boot and no cache present at {cache_path:?}: {source}")]
178 BootFetchFatal {
179 cache_path: PathBuf,
180 #[source]
181 source: Box<JwksError>,
182 },
183
184 #[error("JWKS fetch failed: {0}")]
185 Fetch(String),
186
187 #[error("JWKS parse failed: {0}")]
188 Parse(String),
189
190 #[error("JWK {kid:?} has invalid Ed25519 public key: {reason}")]
191 BadKey { kid: String, reason: String },
192
193 #[error("JWKS cache I/O failure at {path:?}: {source}")]
194 Io {
195 path: PathBuf,
196 #[source]
197 source: std::io::Error,
198 },
199
200 #[error("JWKS cache serialize failed: {0}")]
201 Serialize(#[from] serde_json::Error),
202}
203
204#[async_trait::async_trait]
206pub trait JwksSource: Send + Sync + std::fmt::Debug {
207 async fn fetch(&self) -> Result<JwksDoc, JwksError>;
208}
209
210#[cfg(feature = "jwks-http")]
212#[derive(Debug, Clone)]
213pub struct HttpJwksSource {
214 http: reqwest::Client,
215 url: String,
216}
217
218#[cfg(feature = "jwks-http")]
219impl HttpJwksSource {
220 pub fn new(url: impl Into<String>) -> Result<Self, JwksError> {
223 let http = reqwest::Client::builder()
224 .timeout(Duration::from_secs(10))
225 .build()
226 .map_err(|e| JwksError::Fetch(e.to_string()))?;
227 Ok(Self {
228 http,
229 url: url.into(),
230 })
231 }
232
233 pub fn url(&self) -> &str {
234 &self.url
235 }
236}
237
238#[cfg(feature = "jwks-http")]
239#[async_trait::async_trait]
240impl JwksSource for HttpJwksSource {
241 async fn fetch(&self) -> Result<JwksDoc, JwksError> {
242 let fetch = |e: reqwest::Error| JwksError::Fetch(e.to_string());
243 self.http
244 .get(&self.url)
245 .send()
246 .await
247 .map_err(fetch)?
248 .error_for_status()
249 .map_err(fetch)?
250 .json::<JwksDoc>()
251 .await
252 .map_err(fetch)
253 }
254}
255
256#[derive(Debug, Clone)]
258pub struct JwksCacheConfig {
259 pub cache_path: PathBuf,
261 pub refresh_interval: Duration,
264 pub kid_miss_rate_limit: Duration,
266 pub serve_stale_on_failure: bool,
268}
269
270impl JwksCacheConfig {
271 pub fn new(cache_path: impl Into<PathBuf>) -> Self {
272 Self {
273 cache_path: cache_path.into(),
274 refresh_interval: Duration::from_secs(3600),
275 kid_miss_rate_limit: Duration::from_secs(1),
276 serve_stale_on_failure: true,
277 }
278 }
279}
280
281#[derive(Debug)]
283pub struct JwksCache {
284 source: Box<dyn JwksSource>,
285 config: JwksCacheConfig,
286 current: RwLock<Current>,
287 kid_miss_last: Mutex<Option<Instant>>,
289}
290
291#[derive(Debug, Clone)]
292struct Current {
293 keys: Arc<KeySet>,
294 last_refresh: SystemTime,
295}
296
297impl JwksCache {
298 pub async fn boot(
300 source: Box<dyn JwksSource>,
301 config: JwksCacheConfig,
302 ) -> Result<Self, JwksError> {
303 let current = match load_from_disk(&config.cache_path)? {
304 Some((keys, last_refresh)) => {
305 let stale = last_refresh
306 .elapsed()
307 .map(|age| age > config.refresh_interval)
308 .unwrap_or(true);
309 if !stale {
310 Current { keys: Arc::new(keys), last_refresh }
311 } else {
312 match fetch_and_persist(source.as_ref(), &config.cache_path).await {
313 Ok(c) => c,
314 Err(e) if config.serve_stale_on_failure => {
315 tracing::warn!(error = ?e, "JWKS source unreachable at boot; serving stale cache");
316 Current { keys: Arc::new(keys), last_refresh }
317 }
318 Err(e) => return Err(e),
319 }
320 }
321 }
322 None => fetch_and_persist(source.as_ref(), &config.cache_path)
323 .await
324 .map_err(|e| JwksError::BootFetchFatal {
325 cache_path: config.cache_path.clone(),
326 source: Box::new(e),
327 })?,
328 };
329 Ok(Self {
330 source,
331 config,
332 current: RwLock::new(current),
333 kid_miss_last: Mutex::new(None),
334 })
335 }
336
337 pub fn config(&self) -> &JwksCacheConfig {
338 &self.config
339 }
340
341 pub fn keys(&self) -> Arc<KeySet> {
343 self.current.read().expect("jwks lock").keys.clone()
344 }
345
346 pub fn last_refresh(&self) -> SystemTime {
348 self.current.read().expect("jwks lock").last_refresh
349 }
350
351 pub async fn refresh(&self) -> Result<(), JwksError> {
354 let next = fetch_and_persist(self.source.as_ref(), &self.config.cache_path).await?;
355 *self.current.write().expect("jwks lock") = next;
356 Ok(())
357 }
358
359 pub async fn keys_for(&self, kid: &str) -> Arc<KeySet> {
362 let keys = self.keys();
363 if keys.get(kid).is_some() {
364 return keys;
365 }
366 self.try_kid_miss_refresh().await;
367 self.keys()
368 }
369
370 async fn try_kid_miss_refresh(&self) {
371 {
372 let mut last = self.kid_miss_last.lock().expect("kid-miss lock");
373 let now = Instant::now();
374 if let Some(prev) = *last {
375 if now.duration_since(prev) < self.config.kid_miss_rate_limit {
376 return;
377 }
378 }
379 *last = Some(now);
380 }
381 if let Err(e) = self.refresh().await {
382 tracing::warn!(error = ?e, "kid-miss JWKS refresh failed");
383 }
384 }
385
386 #[cfg(test)]
387 fn kid_miss_last(&self) -> Option<Instant> {
388 *self.kid_miss_last.lock().unwrap()
389 }
390}
391
392async fn fetch_and_persist(source: &dyn JwksSource, path: &Path) -> Result<Current, JwksError> {
393 let keys = KeySet::from_doc(source.fetch().await?)?;
394 let last_refresh = SystemTime::now();
395 write_atomic(path, &keys, last_refresh)?;
396 Ok(Current {
397 keys: Arc::new(keys),
398 last_refresh,
399 })
400}
401
402#[derive(Debug, Serialize, Deserialize)]
407struct StoredCache {
408 format: u32,
409 doc: JwksDoc,
410 #[serde(default, skip_serializing_if = "Option::is_none")]
411 last_refresh_secs: Option<u64>,
412}
413
414const CACHE_FORMAT: u32 = 2;
415
416#[derive(Deserialize)]
417struct CacheHeader {
418 format: Option<u32>,
419}
420
421pub fn load_from_disk(path: &Path) -> Result<Option<(KeySet, SystemTime)>, JwksError> {
424 let io = |source| JwksError::Io {
425 path: path.to_path_buf(),
426 source,
427 };
428 let bytes = match std::fs::read(path) {
429 Ok(b) => b,
430 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
431 Err(e) => return Err(io(e)),
432 };
433 let header: CacheHeader = serde_json::from_slice(&bytes)?;
434 if header.format != Some(CACHE_FORMAT) {
435 tracing::warn!(path = %path.display(), found = ?header.format, expected = CACHE_FORMAT,
436 "discarding JWKS cache in an old format; will refetch");
437 std::fs::remove_file(path).map_err(io)?;
438 return Ok(None);
439 }
440 let stored: StoredCache = serde_json::from_slice(&bytes)?;
441 let keys = KeySet::from_doc(stored.doc)?;
442 let last_refresh = stored
443 .last_refresh_secs
444 .and_then(|s| UNIX_EPOCH.checked_add(Duration::from_secs(s)))
445 .unwrap_or(UNIX_EPOCH);
446 Ok(Some((keys, last_refresh)))
447}
448
449pub fn write_atomic(path: &Path, keys: &KeySet, last_refresh: SystemTime) -> Result<(), JwksError> {
456 let io = |p: &Path| {
457 let p = p.to_path_buf();
458 move |source| JwksError::Io { path: p, source }
459 };
460 if let Some(parent) = path.parent().filter(|p| !p.as_os_str().is_empty()) {
461 std::fs::create_dir_all(parent).map_err(io(parent))?;
462 }
463 let stored = StoredCache {
464 format: CACHE_FORMAT,
465 doc: keys.doc.clone(),
466 last_refresh_secs: last_refresh.duration_since(UNIX_EPOCH).ok().map(|d| d.as_secs()),
467 };
468 let body = serde_json::to_vec_pretty(&stored)?;
469 let tmp = staging_path(path);
470 std::fs::write(&tmp, &body).map_err(io(&tmp))?;
471 if let Err(source) = std::fs::rename(&tmp, path) {
472 let _ = std::fs::remove_file(&tmp);
473 return Err(JwksError::Io {
474 path: path.to_path_buf(),
475 source,
476 });
477 }
478 Ok(())
479}
480
481fn staging_path(final_path: &Path) -> PathBuf {
482 static SEQ: AtomicU64 = AtomicU64::new(0);
483 let seq = SEQ.fetch_add(1, Ordering::Relaxed);
484 let mut name = final_path.file_name().unwrap_or_default().to_os_string();
485 name.push(format!(".tmp.{}.{seq}", std::process::id()));
486 match final_path.parent() {
487 Some(parent) if !parent.as_os_str().is_empty() => parent.join(name),
488 _ => PathBuf::from(name),
489 }
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495 use serde_json::json;
496 use std::sync::atomic::AtomicUsize;
497
498 fn doc_with(kid: &str, pubkey: &[u8; 32]) -> JwksDoc {
499 JwksDoc {
500 keys: vec![JwkKey::ed25519(kid, pubkey, "https://cheers.test", KeyRole::Issuer)],
501 }
502 }
503
504 #[derive(Debug)]
505 struct CountingSource {
506 doc: Option<JwksDoc>,
507 calls: Arc<AtomicUsize>,
508 }
509
510 #[async_trait::async_trait]
511 impl JwksSource for CountingSource {
512 async fn fetch(&self) -> Result<JwksDoc, JwksError> {
513 self.calls.fetch_add(1, Ordering::SeqCst);
514 self.doc.clone().ok_or_else(|| JwksError::Fetch("unreachable".into()))
515 }
516 }
517
518 fn source(doc: Option<JwksDoc>) -> Box<CountingSource> {
519 counted(doc).0
520 }
521
522 fn counted(doc: Option<JwksDoc>) -> (Box<CountingSource>, Arc<AtomicUsize>) {
523 let calls = Arc::new(AtomicUsize::new(0));
524 (Box::new(CountingSource { doc, calls: calls.clone() }), calls)
525 }
526
527 #[test]
528 fn parses_okp_ed25519_skips_unknown_kty() {
529 let known = [7u8; 32];
530 let x = URL_SAFE_NO_PAD.encode(known);
531 let doc: JwksDoc = serde_json::from_value(json!({"keys": [
532 {"kty": "RSA", "n": "...", "e": "AQAB", "kid": "rsa-1", "principal": "i", "role": "issuer"},
533 {"kty": "OKP", "crv": "Ed25519", "x": x, "kid": "ed-1", "principal": "svc:a", "role": "self-signer"},
534 {"kty": "OKP", "crv": "X25519", "x": "AA", "kid": "x-skip", "principal": "i", "role": "issuer"},
535 ]}))
536 .unwrap();
537 let set = KeySet::from_doc(doc).unwrap();
538 assert_eq!(set.len(), 1);
539 let e = set.get("ed-1").unwrap();
540 assert_eq!((e.public_key, e.principal.as_str(), e.role), (known, "svc:a", KeyRole::SelfSigner));
541 assert!(set.get("rsa-1").is_none() && set.get("x-skip").is_none());
542 }
543
544 #[test]
545 fn rejects_doc_with_no_ed25519_keys() {
546 let doc: JwksDoc = serde_json::from_value(json!({"keys": [
547 {"kty": "RSA", "kid": "rsa-1", "principal": "i", "role": "issuer"}
548 ]}))
549 .unwrap();
550 assert!(matches!(KeySet::from_doc(doc), Err(JwksError::Parse(_))));
551 }
552
553 #[test]
554 fn rejects_wrong_length_pubkey() {
555 let short = URL_SAFE_NO_PAD.encode([1u8; 16]);
556 let doc: JwksDoc = serde_json::from_value(json!({"keys": [
557 {"kty": "OKP", "crv": "Ed25519", "x": short, "kid": "ed-bad", "principal": "i", "role": "issuer"}
558 ]}))
559 .unwrap();
560 match KeySet::from_doc(doc).unwrap_err() {
561 JwksError::BadKey { kid, .. } => assert_eq!(kid, "ed-bad"),
562 other => panic!("expected BadKey, got {other:?}"),
563 }
564 }
565
566 #[test]
567 fn role_is_required_on_the_wire() {
568 let x = URL_SAFE_NO_PAD.encode([3u8; 32]);
569 let res: Result<JwksDoc, _> = serde_json::from_value(json!({
570 "keys": [{"kty": "OKP", "crv": "Ed25519", "x": x, "kid": "k", "principal": "i"}]
571 }));
572 assert!(res.is_err(), "a role-less JWK must not parse");
573 }
574
575 #[test]
576 fn atomic_round_trip_through_disk_leaves_no_staging_file() {
577 let tmp = tempfile::tempdir().unwrap();
578 let path = tmp.path().join("nested/jwks.json");
579 let set = KeySet::from_doc(doc_with("ed-rt", &[42u8; 32])).unwrap();
580 write_atomic(&path, &set, SystemTime::now()).unwrap();
581 let staged = std::fs::read_dir(path.parent().unwrap())
582 .unwrap()
583 .filter_map(|e| e.ok())
584 .filter(|e| e.file_name().to_string_lossy().contains(".tmp."))
585 .count();
586 assert_eq!(staged, 0);
587 let (loaded, _) = load_from_disk(&path).unwrap().unwrap();
588 assert_eq!(loaded.get("ed-rt").unwrap().public_key, [42u8; 32]);
589 }
590
591 #[test]
592 fn old_format_cache_is_discarded_not_parsed() {
593 let tmp = tempfile::tempdir().unwrap();
594 let path = tmp.path().join("jwks.json");
595 let x = URL_SAFE_NO_PAD.encode([5u8; 32]);
596 let old = json!({
597 "doc": {"keys": [{"kty": "OKP", "crv": "Ed25519", "x": x, "kid": "old"}]},
598 "last_refresh_secs": 1_700_000_000u64,
599 });
600 std::fs::write(&path, serde_json::to_vec(&old).unwrap()).unwrap();
601 assert!(load_from_disk(&path).unwrap().is_none());
602 assert!(!path.exists(), "old-format cache file must be removed");
603 }
604
605 #[test]
606 fn load_from_disk_missing_returns_none() {
607 let tmp = tempfile::tempdir().unwrap();
608 assert!(load_from_disk(&tmp.path().join("nope.json")).unwrap().is_none());
609 }
610
611 #[test]
612 fn boot_no_cache_unreachable_source_is_fatal() {
613 let tmp = tempfile::tempdir().unwrap();
614 let cfg = JwksCacheConfig::new(tmp.path().join("jwks.json"));
615 let err = pollster::block_on(JwksCache::boot(source(None), cfg)).unwrap_err();
616 assert!(matches!(err, JwksError::BootFetchFatal { .. }));
617 }
618
619 #[test]
620 fn boot_first_fetch_persists_then_restart_serves_from_disk() {
621 let tmp = tempfile::tempdir().unwrap();
622 let cfg = JwksCacheConfig::new(tmp.path().join("jwks.json"));
623 let c = pollster::block_on(JwksCache::boot(source(Some(doc_with("k1", &[1; 32]))), cfg.clone()))
624 .unwrap();
625 assert!(c.keys().get("k1").is_some());
626 let c2 = pollster::block_on(JwksCache::boot(source(None), cfg)).unwrap();
628 assert!(c2.keys().get("k1").is_some());
629 }
630
631 #[test]
632 fn stale_cache_with_unreachable_source_serves_stale() {
633 let tmp = tempfile::tempdir().unwrap();
634 let mut cfg = JwksCacheConfig::new(tmp.path().join("jwks.json"));
635 let set = KeySet::from_doc(doc_with("old", &[2; 32])).unwrap();
636 write_atomic(&cfg.cache_path, &set, UNIX_EPOCH + Duration::from_secs(1)).unwrap();
637 let c = pollster::block_on(JwksCache::boot(source(None), cfg.clone())).unwrap();
638 assert!(c.keys().get("old").is_some());
639 cfg.serve_stale_on_failure = false;
640 assert!(pollster::block_on(JwksCache::boot(source(None), cfg)).is_err());
641 }
642
643 #[test]
646 fn kid_miss_refresh_is_rate_limited() {
647 let tmp = tempfile::tempdir().unwrap();
648 let cfg = JwksCacheConfig::new(tmp.path().join("jwks.json"));
649 let (src, counter) = counted(Some(doc_with("k1", &[1; 32])));
650 let c = pollster::block_on(JwksCache::boot(src, cfg)).unwrap();
651 let calls = || counter.load(Ordering::SeqCst);
652 assert_eq!(calls(), 1);
654 assert!(c.kid_miss_last().is_none());
655 assert!(pollster::block_on(c.keys_for("unknown")).get("unknown").is_none());
656 assert_eq!(calls(), 2);
657 let first = c.kid_miss_last().unwrap();
658 pollster::block_on(c.keys_for("unknown-2"));
659 assert_eq!(calls(), 2, "second miss inside the window must not refetch");
660 assert_eq!(c.kid_miss_last(), Some(first));
661 pollster::block_on(c.keys_for("k1"));
663 assert_eq!(calls(), 2);
664 }
665}