1use std::collections::HashMap;
8use std::sync::{Arc, Mutex};
9use std::time::{Duration, Instant};
10
11use async_trait::async_trait;
12use bytes::Bytes;
13use http::{HeaderMap, StatusCode};
14
15use crate::error::{Error, Result};
16use crate::jwk::Jwks;
17use crate::metadata::{metadata_url, KeyedMetadata};
18
19pub const MAX_DOCUMENT_BYTES: usize = 1 << 20;
22
23#[async_trait]
25pub trait KeyResolver: Send + Sync {
26 async fn jwks(&self, issuer: &str, dwk: &str) -> Result<Jwks>;
28}
29
30#[derive(Clone, Copy, Debug, Default)]
32pub struct NoKeys;
33
34#[async_trait]
35impl KeyResolver for NoKeys {
36 async fn jwks(&self, _issuer: &str, _dwk: &str) -> Result<Jwks> {
37 Err(Error::Discovery("this verifier looks up no keys".into()))
38 }
39}
40
41#[derive(Clone, Debug)]
43pub struct Fetched {
44 pub status: StatusCode,
46 pub headers: HeaderMap,
48 pub body: Bytes,
50}
51
52#[async_trait]
54pub trait Fetch: Send + Sync {
55 async fn get(&self, url: &str, headers: HeaderMap) -> Result<Fetched>;
58}
59
60const MAX_CACHED: usize = 256;
61
62pub struct JwksResolver {
70 fetch: Arc<dyn Fetch>,
71 ttl: Duration,
72 cache: Mutex<HashMap<String, (Instant, Jwks)>>,
73}
74
75impl std::fmt::Debug for JwksResolver {
76 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77 f.debug_struct("JwksResolver")
78 .field("ttl", &self.ttl)
79 .finish_non_exhaustive()
80 }
81}
82
83impl JwksResolver {
84 pub fn new(fetch: Arc<dyn Fetch>) -> Self {
86 Self {
87 fetch,
88 ttl: Duration::from_secs(600),
89 cache: Mutex::new(HashMap::new()),
90 }
91 }
92
93 pub fn with_ttl(mut self, ttl: Duration) -> Self {
95 self.ttl = ttl;
96 self
97 }
98
99 pub fn clear(&self) {
101 if let Ok(mut cache) = self.cache.lock() {
102 cache.clear();
103 }
104 }
105
106 async fn json<T: serde::de::DeserializeOwned>(&self, url: &str) -> Result<T> {
107 let fetched = self.fetch.get(url, HeaderMap::new()).await?;
108 if !fetched.status.is_success() {
109 return Err(Error::Discovery(format!(
110 "{url} answered {}",
111 fetched.status
112 )));
113 }
114 serde_json::from_slice(&fetched.body)
115 .map_err(|e| Error::Discovery(format!("{url} is unreadable: {e}")))
116 }
117}
118
119#[async_trait]
120impl KeyResolver for JwksResolver {
121 async fn jwks(&self, issuer: &str, dwk: &str) -> Result<Jwks> {
122 if !crate::identifier::is_server_identifier(issuer) {
123 return Err(Error::Discovery(format!(
124 "{issuer} is not a server identifier"
125 )));
126 }
127 if !crate::metadata::dwk::is_known(dwk) {
128 return Err(Error::Discovery(format!(
129 "{dwk} is not a metadata document"
130 )));
131 }
132 let url = metadata_url(issuer, dwk);
133 if let Some((at, jwks)) = self.cache.lock().ok().and_then(|c| c.get(&url).cloned()) {
134 if at.elapsed() < self.ttl {
135 return Ok(jwks);
136 }
137 }
138 let metadata: KeyedMetadata = self.json(&url).await?;
139 match &metadata.issuer {
140 Some(named) if named == issuer => {}
141 Some(named) => {
142 return Err(Error::Discovery(format!(
143 "{url} names another issuer, {named}"
144 )))
145 }
146 None => return Err(Error::Discovery(format!("{url} names no issuer"))),
147 }
148 let jwks = match (metadata.jwks, metadata.jwks_uri) {
149 (Some(jwks), _) => jwks,
150 (None, Some(uri)) if crate::identifier::same_origin(issuer, &uri) => {
151 self.json(&uri).await?
152 }
153 (None, Some(uri)) => {
154 return Err(Error::Discovery(format!(
155 "{url} puts its keys on another origin: {uri}"
156 )))
157 }
158 (None, None) => {
159 return Err(Error::Discovery(format!(
160 "{url} has neither jwks nor jwks_uri"
161 )))
162 }
163 };
164 if let Ok(mut cache) = self.cache.lock() {
165 if cache.len() >= MAX_CACHED {
166 let oldest = cache
167 .iter()
168 .min_by_key(|(_, (at, _))| *at)
169 .map(|(k, _)| k.clone());
170 if let Some(oldest) = oldest {
171 cache.remove(&oldest);
172 }
173 }
174 cache.insert(url, (Instant::now(), jwks.clone()));
175 }
176 Ok(jwks)
177 }
178}
179
180#[cfg(feature = "client")]
182pub(crate) async fn read_limited(mut response: reqwest::Response, limit: usize) -> Result<Bytes> {
183 if response
184 .content_length()
185 .is_some_and(|len| len > limit as u64)
186 {
187 return Err(Error::Http(format!(
188 "response is larger than {limit} bytes"
189 )));
190 }
191 let mut body = bytes::BytesMut::new();
192 while let Some(chunk) = response
193 .chunk()
194 .await
195 .map_err(|e| Error::Http(e.to_string()))?
196 {
197 if body.len() + chunk.len() > limit {
198 return Err(Error::Http(format!(
199 "response is larger than {limit} bytes"
200 )));
201 }
202 body.extend_from_slice(&chunk);
203 }
204 Ok(body.freeze())
205}
206
207#[cfg(feature = "client")]
209#[derive(Clone, Debug)]
210pub struct ReqwestFetch {
211 client: reqwest::Client,
212 guarded: bool,
213}
214
215#[cfg(feature = "client")]
216impl ReqwestFetch {
217 pub fn guarded(denied: crate::client::DeniedNetworks) -> Result<Self> {
221 Ok(Self {
222 client: crate::client::guarded_client(denied)?,
223 guarded: true,
224 })
225 }
226
227 pub fn new(client: reqwest::Client) -> Self {
229 Self {
230 client,
231 guarded: false,
232 }
233 }
234}
235
236#[cfg(feature = "client")]
237#[async_trait]
238impl Fetch for ReqwestFetch {
239 async fn get(&self, url: &str, headers: HeaderMap) -> Result<Fetched> {
240 if self.guarded {
241 if !url.starts_with("https://") {
242 return Err(Error::Url(format!("refusing to fetch {url}")));
243 }
244 crate::client::guard::admit_host(url)?;
245 }
246 let response = self
247 .client
248 .get(url)
249 .headers(headers)
250 .send()
251 .await
252 .map_err(|e| Error::Http(format!("{url}: {e}")))?;
253 let status = response.status();
254 let headers = response.headers().clone();
255 let body = read_limited(response, MAX_DOCUMENT_BYTES).await?;
256 Ok(Fetched {
257 status,
258 headers,
259 body,
260 })
261 }
262}
263
264#[derive(Clone, Debug, Default)]
267pub struct StaticFetch {
268 documents: Arc<Mutex<HashMap<String, serde_json::Value>>>,
269}
270
271impl StaticFetch {
272 pub fn new() -> Self {
274 Self::default()
275 }
276
277 pub fn insert(&self, url: impl Into<String>, document: serde_json::Value) {
279 if let Ok(mut documents) = self.documents.lock() {
280 documents.insert(url.into(), document);
281 }
282 }
283}
284
285#[async_trait]
286impl Fetch for StaticFetch {
287 async fn get(&self, url: &str, _headers: HeaderMap) -> Result<Fetched> {
288 let document = self.documents.lock().ok().and_then(|d| d.get(url).cloned());
289 Ok(match document {
290 Some(document) => Fetched {
291 status: StatusCode::OK,
292 headers: HeaderMap::new(),
293 body: Bytes::from(document.to_string()),
294 },
295 None => Fetched {
296 status: StatusCode::NOT_FOUND,
297 headers: HeaderMap::new(),
298 body: Bytes::new(),
299 },
300 })
301 }
302}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307 use serde_json::json;
308
309 #[tokio::test]
310 async fn keys_are_found_through_the_metadata_document() {
311 let fetch = StaticFetch::new();
312 fetch.insert(
313 "https://ap.example/.well-known/aauth-agent.json",
314 json!({ "issuer": "https://ap.example", "jwks_uri": "https://ap.example/jwks.json" }),
315 );
316 fetch.insert(
317 "https://ap.example/jwks.json",
318 json!({ "keys": [{ "kty": "OKP", "kid": "k" }] }),
319 );
320 let resolver = JwksResolver::new(Arc::new(fetch));
321 let jwks = resolver
322 .jwks("https://ap.example", "aauth-agent.json")
323 .await
324 .unwrap();
325 assert_eq!(jwks.keys[0].kid.as_deref(), Some("k"));
326 }
327
328 #[tokio::test]
329 async fn a_document_naming_another_issuer_is_refused() {
330 let fetch = StaticFetch::new();
331 fetch.insert(
332 "https://ap.example/.well-known/aauth-agent.json",
333 json!({ "issuer": "https://evil.example", "jwks": { "keys": [] } }),
334 );
335 let resolver = JwksResolver::new(Arc::new(fetch));
336 let refused = resolver
337 .jwks("https://ap.example", "aauth-agent.json")
338 .await
339 .unwrap_err();
340 assert!(refused.to_string().contains("another issuer"), "{refused}");
341 }
342
343 #[tokio::test]
344 async fn keys_are_cached_until_cleared() {
345 let fetch = StaticFetch::new();
346 let url = "https://ap.example/.well-known/aauth-agent.json";
347 fetch.insert(
348 url,
349 json!({ "issuer": "https://ap.example", "jwks": { "keys": [{ "kty": "OKP", "kid": "a" }] } }),
350 );
351 let resolver = JwksResolver::new(Arc::new(fetch.clone()));
352 resolver
353 .jwks("https://ap.example", "aauth-agent.json")
354 .await
355 .unwrap();
356 fetch.insert(
357 url,
358 json!({ "issuer": "https://ap.example", "jwks": { "keys": [{ "kty": "OKP", "kid": "b" }] } }),
359 );
360 let cached = resolver
361 .jwks("https://ap.example", "aauth-agent.json")
362 .await
363 .unwrap();
364 assert_eq!(cached.keys[0].kid.as_deref(), Some("a"));
365 resolver.clear();
366 let fresh = resolver
367 .jwks("https://ap.example", "aauth-agent.json")
368 .await
369 .unwrap();
370 assert_eq!(fresh.keys[0].kid.as_deref(), Some("b"));
371 }
372}