1use std::net::{IpAddr, SocketAddr};
31use std::sync::Arc;
32use std::time::{Duration, Instant};
33
34use dashmap::DashMap;
35use serde::Deserialize;
36use tokio::sync::Mutex;
37use tracing::warn;
38
39use crate::cimd::ssrf;
40
41const MAX_DOCUMENT_BYTES: usize = 64 * 1024;
46
47const FETCH_TIMEOUT: Duration = Duration::from_secs(5);
55
56const DNS_TIMEOUT: Duration = Duration::from_secs(3);
62
63const CACHE_TTL: Duration = Duration::from_secs(300);
65
66const NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(60);
70
71const MAX_CACHE_ENTRIES: usize = 10_000;
77
78#[derive(Debug, Clone, Deserialize)]
79pub struct ClientMetadataDocument {
80 pub client_id: String,
81 pub client_name: String,
82 #[serde(default)]
83 pub redirect_uris: Vec<String>,
84 #[serde(default = "default_token_endpoint_auth_method")]
85 pub token_endpoint_auth_method: String,
86 #[serde(default)]
87 pub jwks: Option<serde_json::Value>,
88}
89
90fn default_token_endpoint_auth_method() -> String {
91 "none".to_string()
92}
93
94#[derive(Debug, Clone, thiserror::Error)]
95pub enum CimdError {
96 #[error(transparent)]
97 Ssrf(#[from] ssrf::SsrfError),
98 #[error("dns resolution failed for `{0}`: {1}")]
103 DnsResolutionFailed(String, String),
104 #[error(
110 "`{0}` resolved to at least one private/loopback/link-local address; blocked to prevent SSRF"
111 )]
112 DnsBlocked(String),
113 #[error("fetch failed: {0}")]
114 Fetch(String),
115 #[error(
119 "peer address {actual} did not match the validated address {expected}; possible proxy interception or DNS-rebinding attempt"
120 )]
121 PeerMismatch {
122 expected: SocketAddr,
123 actual: SocketAddr,
124 },
125 #[error("invalid client metadata document: {0}")]
126 InvalidDocument(String),
127 #[error(
128 "client metadata document client_id `{document_client_id}` does not match the requested URL `{requested_url}`"
129 )]
130 ClientIdMismatch {
131 document_client_id: String,
132 requested_url: String,
133 },
134}
135
136impl CimdError {
137 #[must_use]
147 pub fn kind(&self) -> &'static str {
148 match self {
149 Self::Ssrf(e) => e.kind(),
150 Self::DnsResolutionFailed(..) => "dns_resolution_failed",
151 Self::DnsBlocked(_) => "ssrf_blocked",
152 Self::Fetch(_) => "cimd_fetch_failed",
153 Self::PeerMismatch { .. } => "ssrf_blocked",
154 Self::InvalidDocument(_) => "invalid_client_metadata",
155 Self::ClientIdMismatch { .. } => "invalid_client_metadata",
156 }
157 }
158}
159
160#[must_use]
165pub fn is_cimd_client_id(client_id: &str) -> bool {
166 client_id.starts_with("https://")
167}
168
169pub async fn resolve_and_validate_address(host: &str, port: u16) -> Result<SocketAddr, CimdError> {
182 let lookup = tokio::time::timeout(DNS_TIMEOUT, tokio::net::lookup_host((host, port)))
183 .await
184 .map_err(|_| {
185 CimdError::DnsResolutionFailed(
186 host.to_string(),
187 format!("timed out after {DNS_TIMEOUT:?}"),
188 )
189 })?
190 .map_err(|e| CimdError::DnsResolutionFailed(host.to_string(), e.to_string()))?;
191 let addrs: Vec<SocketAddr> = lookup.collect();
192 if addrs.is_empty() {
193 return Err(CimdError::DnsResolutionFailed(
194 host.to_string(),
195 "resolved to no addresses".to_string(),
196 ));
197 }
198 if addrs
199 .iter()
200 .any(|addr| ssrf::check_ip_not_private(addr.ip(), host).is_err())
201 {
202 return Err(CimdError::DnsBlocked(host.to_string()));
203 }
204 Ok(addrs[0])
205}
206
207pub(crate) async fn fetch_via_pinned_address(
224 url: &str,
225 addr: SocketAddr,
226) -> Result<ClientMetadataDocument, CimdError> {
227 let parsed =
228 url::Url::parse(url).map_err(|e| CimdError::Fetch(format!("parse `{url}`: {e}")))?;
229 let host = parsed
230 .host_str()
231 .ok_or_else(|| CimdError::Fetch(format!("no host in `{url}`")))?;
232 let client = reqwest::Client::builder()
233 .resolve(host, addr)
234 .no_proxy()
239 .redirect(reqwest::redirect::Policy::none())
244 .timeout(FETCH_TIMEOUT)
245 .build()
246 .map_err(|e| CimdError::Fetch(format!("build pinned client for `{url}`: {e}")))?;
247 fetch_document_at(&client, url, addr).await
248}
249
250pub(crate) async fn fetch_document_at(
267 client: &reqwest::Client,
268 url: &str,
269 pinned_addr: SocketAddr,
270) -> Result<ClientMetadataDocument, CimdError> {
271 let mut response = client
272 .get(url)
273 .send()
274 .await
275 .map_err(|e| CimdError::Fetch(format!("GET `{url}`: {e}")))?;
276
277 match response.remote_addr() {
291 Some(peer) if peer == pinned_addr => {}
292 Some(peer) => {
293 return Err(CimdError::PeerMismatch {
294 expected: pinned_addr,
295 actual: peer,
296 });
297 }
298 None => {
299 return Err(CimdError::Fetch(format!(
300 "no remote peer address available for `{url}`; refusing to trust an unverified connection"
301 )));
302 }
303 }
304
305 if !response.status().is_success() {
306 return Err(CimdError::Fetch(format!(
307 "GET `{url}` returned HTTP {}",
308 response.status()
309 )));
310 }
311
312 let mut buf: Vec<u8> = Vec::new();
313 while let Some(chunk) = response
314 .chunk()
315 .await
316 .map_err(|e| CimdError::Fetch(format!("read body from `{url}`: {e}")))?
317 {
318 buf.extend_from_slice(&chunk);
319 if buf.len() > MAX_DOCUMENT_BYTES {
320 return Err(CimdError::InvalidDocument(format!(
321 "document at `{url}` exceeds the {MAX_DOCUMENT_BYTES}-byte limit"
322 )));
323 }
324 }
325
326 let document: ClientMetadataDocument = serde_json::from_slice(&buf).map_err(|e| {
327 CimdError::InvalidDocument(format!("document at `{url}` is not valid JSON: {e}"))
328 })?;
329 if document.client_id.is_empty() || document.client_name.is_empty() {
330 return Err(CimdError::InvalidDocument(format!(
331 "document at `{url}` is missing required client_id or client_name"
332 )));
333 }
334 match document.token_endpoint_auth_method.as_str() {
335 "none" if document.jwks.is_none() => {}
336 "private_key_jwt" if document.jwks.is_some() => {}
337 "none" => {
338 return Err(CimdError::InvalidDocument(
339 "public clients must not declare jwks".to_string(),
340 ));
341 }
342 "private_key_jwt" => {
343 return Err(CimdError::InvalidDocument(
344 "private_key_jwt clients require jwks".to_string(),
345 ));
346 }
347 _ => {
348 return Err(CimdError::InvalidDocument(
349 "unsupported token_endpoint_auth_method".to_string(),
350 ));
351 }
352 }
353 if document.redirect_uris.is_empty() {
354 return Err(CimdError::InvalidDocument(format!(
355 "document at `{url}` declares no redirect_uris"
356 )));
357 }
358 if document.client_id != url {
359 return Err(CimdError::ClientIdMismatch {
360 document_client_id: document.client_id,
361 requested_url: url.to_string(),
362 });
363 }
364 Ok(document)
365}
366
367struct CacheEntry {
368 result: Result<ClientMetadataDocument, CimdError>,
374 fetched_at: Instant,
375 ttl: Duration,
376}
377
378pub struct DocumentCache {
398 entries: DashMap<String, CacheEntry>,
399 build_locks: DashMap<String, Arc<Mutex<()>>>,
400}
401
402impl DocumentCache {
403 #[must_use]
404 pub fn new() -> Self {
405 Self {
406 entries: DashMap::new(),
407 build_locks: DashMap::new(),
408 }
409 }
410
411 fn get_fresh(&self, url: &str) -> Option<Result<ClientMetadataDocument, CimdError>> {
412 let entry = self.entries.get(url)?;
413 if entry.fetched_at.elapsed() >= entry.ttl {
414 return None;
415 }
416 Some(entry.result.clone())
417 }
418
419 fn lock_for(&self, url: String) -> Arc<Mutex<()>> {
425 if self.build_locks.len() >= MAX_CACHE_ENTRIES {
426 self.build_locks
427 .retain(|_, lock| Arc::strong_count(lock) > 1);
428 }
429 self.build_locks
430 .entry(url)
431 .or_insert_with(|| Arc::new(Mutex::new(())))
432 .clone()
433 }
434
435 fn insert(
436 &self,
437 url: String,
438 result: &Result<ClientMetadataDocument, CimdError>,
439 ttl: Duration,
440 ) {
441 if self.entries.len() >= MAX_CACHE_ENTRIES {
442 self.entries.retain(|_, e| e.fetched_at.elapsed() < e.ttl);
443 if self.entries.len() >= MAX_CACHE_ENTRIES {
444 warn!(
449 cache_size = self.entries.len(),
450 "CIMD document cache at capacity with no expired entries to evict; not caching this result"
451 );
452 return;
453 }
454 }
455 self.entries.insert(
456 url,
457 CacheEntry {
458 result: result.clone(),
459 fetched_at: Instant::now(),
460 ttl,
461 },
462 );
463 }
464}
465
466impl Default for DocumentCache {
467 fn default() -> Self {
468 Self::new()
469 }
470}
471
472pub async fn fetch_and_validate_client_metadata(
482 cache: &DocumentCache,
483 url: &str,
484) -> Result<ClientMetadataDocument, CimdError> {
485 if let Some(cached) = cache.get_fresh(url) {
486 return cached;
487 }
488
489 let parsed = ssrf::validate_url_shape(url)?;
495
496 let lock = cache.lock_for(url.to_string());
497 let _guard = lock.lock().await;
498
499 if let Some(cached) = cache.get_fresh(url) {
502 return cached;
503 }
504
505 let result: Result<ClientMetadataDocument, CimdError> = async {
506 let host = parsed
507 .host_str()
508 .ok_or_else(|| CimdError::DnsResolutionFailed(url.to_string(), "no host".to_string()))?
509 .to_string();
510 let port = parsed.port_or_known_default().unwrap_or(443);
511
512 let addr = match parsed.host() {
513 Some(url::Host::Domain(_)) => resolve_and_validate_address(&host, port).await?,
514 Some(url::Host::Ipv4(ip)) => SocketAddr::new(IpAddr::V4(ip), port),
517 Some(url::Host::Ipv6(ip)) => SocketAddr::new(IpAddr::V6(ip), port),
518 None => unreachable!("validate_url_shape guarantees a host"),
519 };
520
521 fetch_via_pinned_address(url, addr).await
522 }
523 .await;
524
525 let ttl = if result.is_ok() {
526 CACHE_TTL
527 } else {
528 NEGATIVE_CACHE_TTL
529 };
530 cache.insert(url.to_string(), &result, ttl);
531 result
532}
533
534#[cfg(test)]
535mod tests {
536 use super::*;
537 use wiremock::matchers::{method, path};
538 use wiremock::{Mock, MockServer, ResponseTemplate};
539
540 #[test]
541 fn is_cimd_client_id_detects_https_urls_only() {
542 assert!(is_cimd_client_id(
543 "https://app.example.com/oauth/client-metadata.json"
544 ));
545 assert!(!is_cimd_client_id("abcDEF123opaque-token"));
546 assert!(!is_cimd_client_id("http://app.example.com/client.json"));
547 }
548
549 #[tokio::test]
550 async fn resolve_and_validate_address_rejects_loopback_host() {
551 let err = resolve_and_validate_address("localhost", 443)
552 .await
553 .unwrap_err();
554 assert_eq!(err.kind(), "ssrf_blocked");
555 }
556
557 #[tokio::test]
558 async fn resolve_and_validate_address_reports_dns_failure_distinctly_from_ssrf_block() {
559 let err = resolve_and_validate_address("definitely-does-not-exist.invalid", 443)
564 .await
565 .unwrap_err();
566 assert_eq!(err.kind(), "dns_resolution_failed");
567 }
568
569 #[tokio::test]
570 async fn fetch_via_pinned_address_succeeds_for_matching_client_id() {
571 let server = MockServer::start().await;
572 let addr = server.address();
573 let url = format!("{}/client.json", server.uri());
574 Mock::given(method("GET"))
575 .and(path("/client.json"))
576 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
577 "client_id": url,
578 "client_name": "Example MCP Client",
579 "redirect_uris": ["http://127.0.0.1:3000/callback"],
580 })))
581 .mount(&server)
582 .await;
583
584 let document = fetch_via_pinned_address(&url, *addr)
585 .await
586 .expect("fetch ok");
587 assert_eq!(document.client_id, url);
588 assert_eq!(document.client_name, "Example MCP Client");
589 assert_eq!(
590 document.redirect_uris,
591 vec!["http://127.0.0.1:3000/callback"]
592 );
593 }
594
595 #[tokio::test]
596 async fn fetch_document_at_rejects_peer_mismatch() {
597 let server = MockServer::start().await;
603 let real_addr = *server.address();
604 let url = format!("{}/client.json", server.uri());
605 Mock::given(method("GET"))
606 .and(path("/client.json"))
607 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
608 "client_id": url,
609 "client_name": "Example",
610 "redirect_uris": ["http://127.0.0.1:3000/callback"],
611 })))
612 .mount(&server)
613 .await;
614
615 let wrong_addr = SocketAddr::new(real_addr.ip(), real_addr.port().wrapping_add(1).max(1));
616 let client = reqwest::Client::builder()
617 .no_proxy()
618 .redirect(reqwest::redirect::Policy::none())
619 .build()
620 .unwrap();
621 let err = fetch_document_at(&client, &url, wrong_addr)
622 .await
623 .unwrap_err();
624 assert!(matches!(err, CimdError::PeerMismatch { .. }));
625 assert_eq!(err.kind(), "ssrf_blocked");
626 }
627
628 #[tokio::test]
629 async fn fetch_via_pinned_address_rejects_client_id_mismatch() {
630 let server = MockServer::start().await;
631 let addr = *server.address();
632 let url = format!("{}/client.json", server.uri());
633 Mock::given(method("GET"))
634 .and(path("/client.json"))
635 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
636 "client_id": "https://attacker.example/spoofed.json",
637 "client_name": "Spoofed Client",
638 "redirect_uris": ["http://127.0.0.1:9999/callback"],
639 })))
640 .mount(&server)
641 .await;
642
643 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
644 assert!(matches!(err, CimdError::ClientIdMismatch { .. }));
645 assert_eq!(err.kind(), "invalid_client_metadata");
646 }
647
648 #[tokio::test]
649 async fn fetch_via_pinned_address_rejects_missing_required_fields() {
650 let server = MockServer::start().await;
651 let addr = *server.address();
652 let url = format!("{}/client.json", server.uri());
653 Mock::given(method("GET"))
654 .and(path("/client.json"))
655 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
656 "client_id": url,
657 "redirect_uris": ["http://127.0.0.1:3000/callback"],
658 })))
659 .mount(&server)
660 .await;
661
662 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
663 assert!(matches!(err, CimdError::InvalidDocument(_)));
664 assert_eq!(err.kind(), "invalid_client_metadata");
665 }
666
667 #[tokio::test]
668 async fn fetch_via_pinned_address_rejects_malformed_json() {
669 let server = MockServer::start().await;
670 let addr = *server.address();
671 let url = format!("{}/client.json", server.uri());
672 Mock::given(method("GET"))
673 .and(path("/client.json"))
674 .respond_with(ResponseTemplate::new(200).set_body_string("{not valid json"))
675 .mount(&server)
676 .await;
677
678 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
679 assert!(matches!(err, CimdError::InvalidDocument(_)));
680 assert_eq!(err.kind(), "invalid_client_metadata");
681 }
682
683 #[tokio::test]
684 async fn fetch_via_pinned_address_rejects_empty_redirect_uris() {
685 let server = MockServer::start().await;
686 let addr = *server.address();
687 let url = format!("{}/client.json", server.uri());
688 Mock::given(method("GET"))
689 .and(path("/client.json"))
690 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
691 "client_id": url,
692 "client_name": "Example",
693 "redirect_uris": [],
694 })))
695 .mount(&server)
696 .await;
697
698 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
699 assert!(matches!(err, CimdError::InvalidDocument(_)));
700 assert_eq!(err.kind(), "invalid_client_metadata");
701 }
702
703 #[tokio::test]
717 async fn fetch_via_pinned_address_rejects_an_unsupported_token_endpoint_auth_method() {
718 let server = MockServer::start().await;
719 let addr = *server.address();
720 let url = format!("{}/client.json", server.uri());
721 Mock::given(method("GET"))
722 .and(path("/client.json"))
723 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
724 "client_id": url,
725 "client_name": "Example",
726 "redirect_uris": ["http://127.0.0.1:3000/callback"],
727 "token_endpoint_auth_method": "client_secret_basic",
728 })))
729 .mount(&server)
730 .await;
731
732 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
733 let CimdError::InvalidDocument(message) = &err else {
734 panic!("expected an invalid-document rejection, got {err:?}");
735 };
736 assert_eq!(message, "unsupported token_endpoint_auth_method");
737 assert_eq!(err.kind(), "invalid_client_metadata");
738 }
739
740 #[tokio::test]
744 async fn fetch_via_pinned_address_defaults_a_missing_auth_method_to_public() {
745 let server = MockServer::start().await;
746 let addr = *server.address();
747 let url = format!("{}/client.json", server.uri());
748 Mock::given(method("GET"))
749 .and(path("/client.json"))
750 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
751 "client_id": url,
752 "client_name": "Example",
753 "redirect_uris": ["http://127.0.0.1:3000/callback"],
754 })))
755 .mount(&server)
756 .await;
757
758 let document = fetch_via_pinned_address(&url, addr)
759 .await
760 .expect("fetch ok");
761 assert_eq!(document.token_endpoint_auth_method, "none");
762 assert!(document.jwks.is_none());
763 }
764
765 #[tokio::test]
766 async fn fetch_via_pinned_address_rejects_non_success_status() {
767 let server = MockServer::start().await;
768 let addr = *server.address();
769 let url = format!("{}/missing.json", server.uri());
770 Mock::given(method("GET"))
771 .and(path("/missing.json"))
772 .respond_with(ResponseTemplate::new(404))
773 .mount(&server)
774 .await;
775
776 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
777 assert!(matches!(err, CimdError::Fetch(_)));
778 assert_eq!(err.kind(), "cimd_fetch_failed");
779 }
780
781 #[tokio::test]
782 async fn fetch_via_pinned_address_does_not_follow_redirects() {
783 let server = MockServer::start().await;
784 let addr = *server.address();
785 let url = format!("{}/redirecting.json", server.uri());
786 Mock::given(method("GET"))
787 .and(path("/redirecting.json"))
788 .respond_with(
789 ResponseTemplate::new(302)
790 .insert_header("Location", "https://attacker.example/elsewhere.json"),
791 )
792 .mount(&server)
793 .await;
794
795 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
796 assert!(matches!(err, CimdError::Fetch(_)));
797 assert_eq!(err.kind(), "cimd_fetch_failed");
798 }
799
800 #[tokio::test]
801 async fn fetch_via_pinned_address_rejects_oversized_body_without_hanging() {
802 let server = MockServer::start().await;
803 let addr = *server.address();
804 let url = format!("{}/big.json", server.uri());
805 let oversized = "x".repeat(MAX_DOCUMENT_BYTES + 1024);
806 Mock::given(method("GET"))
807 .and(path("/big.json"))
808 .respond_with(ResponseTemplate::new(200).set_body_string(oversized))
809 .mount(&server)
810 .await;
811
812 let err = fetch_via_pinned_address(&url, addr).await.unwrap_err();
813 assert!(matches!(err, CimdError::InvalidDocument(_)));
814 assert_eq!(err.kind(), "invalid_client_metadata");
815 }
816
817 #[test]
818 fn cache_returns_none_when_expired() {
819 let cache = DocumentCache::new();
820 let doc = ClientMetadataDocument {
821 client_id: "https://app.example.com/client.json".to_string(),
822 client_name: "Example".to_string(),
823 redirect_uris: vec!["http://127.0.0.1:3000/callback".to_string()],
824 token_endpoint_auth_method: "none".to_string(),
825 jwks: None,
826 };
827 cache.insert(
828 "https://app.example.com/client.json".to_string(),
829 &Ok(doc),
830 Duration::from_millis(1),
831 );
832 std::thread::sleep(Duration::from_millis(20));
833 assert!(
834 cache
835 .get_fresh("https://app.example.com/client.json")
836 .is_none()
837 );
838 }
839
840 #[test]
841 fn cache_returns_document_when_fresh() {
842 let cache = DocumentCache::new();
843 let doc = ClientMetadataDocument {
844 client_id: "https://app.example.com/client.json".to_string(),
845 client_name: "Example".to_string(),
846 redirect_uris: vec!["http://127.0.0.1:3000/callback".to_string()],
847 token_endpoint_auth_method: "none".to_string(),
848 jwks: None,
849 };
850 cache.insert(
851 "https://app.example.com/client.json".to_string(),
852 &Ok(doc),
853 CACHE_TTL,
854 );
855 assert!(
856 cache
857 .get_fresh("https://app.example.com/client.json")
858 .is_some()
859 );
860 }
861
862 #[test]
863 fn cache_caches_negative_results_too() {
864 let cache = DocumentCache::new();
865 let err = CimdError::DnsBlocked("app.example.com".to_string());
866 cache.insert(
867 "https://app.example.com/client.json".to_string(),
868 &Err(err),
869 NEGATIVE_CACHE_TTL,
870 );
871 let cached = cache
872 .get_fresh("https://app.example.com/client.json")
873 .expect("negative result should be cached");
874 assert!(matches!(cached, Err(CimdError::DnsBlocked(_))));
879 assert_eq!(cached.unwrap_err().kind(), "ssrf_blocked");
880 }
881
882 #[test]
883 fn entries_insert_skips_caching_once_at_capacity_with_nothing_expired() {
884 let cache = DocumentCache::new();
885 for i in 0..MAX_CACHE_ENTRIES {
886 let url = format!("https://app.example.com/{i}.json");
887 let doc = ClientMetadataDocument {
888 client_id: url.clone(),
889 client_name: "Example".to_string(),
890 redirect_uris: vec!["http://127.0.0.1:3000/callback".to_string()],
891 token_endpoint_auth_method: "none".to_string(),
892 jwks: None,
893 };
894 cache.insert(url, &Ok(doc), CACHE_TTL);
895 }
896 assert_eq!(cache.entries.len(), MAX_CACHE_ENTRIES);
897
898 let overflow_url = "https://app.example.com/overflow.json".to_string();
899 let doc = ClientMetadataDocument {
900 client_id: overflow_url.clone(),
901 client_name: "Example".to_string(),
902 redirect_uris: vec!["http://127.0.0.1:3000/callback".to_string()],
903 token_endpoint_auth_method: "none".to_string(),
904 jwks: None,
905 };
906 cache.insert(overflow_url.clone(), &Ok(doc), CACHE_TTL);
907
908 assert_eq!(cache.entries.len(), MAX_CACHE_ENTRIES);
912 assert!(cache.get_fresh(&overflow_url).is_none());
913 }
914
915 #[test]
916 fn lock_for_sweeps_idle_locks_through_the_real_entry_point_when_at_capacity() {
917 let cache = DocumentCache::new();
918 for i in 0..MAX_CACHE_ENTRIES {
919 cache.build_locks.insert(
920 format!("https://idle.example/{i}.json"),
921 Arc::new(Mutex::new(())),
922 );
923 }
924 assert_eq!(cache.build_locks.len(), MAX_CACHE_ENTRIES);
925
926 let _ = cache.lock_for("https://new.example/client.json".to_string());
927
928 assert_eq!(cache.build_locks.len(), 1);
932 assert!(
933 cache
934 .build_locks
935 .get("https://new.example/client.json")
936 .is_some()
937 );
938 }
939
940 #[tokio::test]
941 async fn build_locks_coalesces_two_concurrent_misses_for_the_same_url_to_one_fetch() {
942 use std::sync::atomic::{AtomicUsize, Ordering};
943
944 async fn simulate_fetch(
955 cache: &DocumentCache,
956 url: &str,
957 fetch_count: &AtomicUsize,
958 ) -> ClientMetadataDocument {
959 if let Some(Ok(doc)) = cache.get_fresh(url) {
960 return doc;
961 }
962 let lock = cache.lock_for(url.to_string());
963 let _guard = lock.lock().await;
964 if let Some(Ok(doc)) = cache.get_fresh(url) {
965 return doc;
966 }
967 fetch_count.fetch_add(1, Ordering::SeqCst);
968 tokio::time::sleep(Duration::from_millis(30)).await;
971 let doc = ClientMetadataDocument {
972 client_id: url.to_string(),
973 client_name: "Example".to_string(),
974 redirect_uris: vec!["http://127.0.0.1:3000/callback".to_string()],
975 token_endpoint_auth_method: "none".to_string(),
976 jwks: None,
977 };
978 cache.insert(url.to_string(), &Ok(doc.clone()), CACHE_TTL);
979 doc
980 }
981
982 let cache = DocumentCache::new();
983 let url = "https://app.example.com/client.json";
984 let fetch_count = AtomicUsize::new(0);
985 let (a, b) = tokio::join!(
986 simulate_fetch(&cache, url, &fetch_count),
987 simulate_fetch(&cache, url, &fetch_count),
988 );
989 assert_eq!(a.client_id, url);
990 assert_eq!(b.client_id, url);
991 assert_eq!(
992 fetch_count.load(Ordering::SeqCst),
993 1,
994 "two concurrent misses for the same URL must coalesce to exactly one fetch"
995 );
996 }
997
998 #[test]
999 fn build_locks_retain_keeps_only_locks_with_an_active_holder() {
1000 let cache = DocumentCache::new();
1005 let idle = Arc::new(Mutex::new(()));
1006 let busy = Arc::new(Mutex::new(()));
1007 let _busy_holder = busy.clone(); cache.build_locks.insert("idle.example".to_string(), idle);
1009 cache.build_locks.insert("busy.example".to_string(), busy);
1010 cache
1011 .build_locks
1012 .retain(|_, lock| Arc::strong_count(lock) > 1);
1013 assert!(cache.build_locks.get("idle.example").is_none());
1014 assert!(cache.build_locks.get("busy.example").is_some());
1015 }
1016}