Skip to main content

aauth_protocol/
discovery.rs

1//! Key discovery: `{iss}/.well-known/{dwk}` names a `jwks_uri` (or carries
2//! `jwks`), and the key a token or signature names is found there.
3//!
4//! Discovery runs over a [`Fetch`], so the caller decides which HTTP client
5//! (and which SSRF guard) reaches the network, and tests need none.
6
7use 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
19/// The most a fetched document may weigh. Metadata and key sets are small;
20/// anything larger is a peer misbehaving.
21pub const MAX_DOCUMENT_BYTES: usize = 1 << 20;
22
23/// Finds an issuer's signing keys.
24#[async_trait]
25pub trait KeyResolver: Send + Sync {
26    /// The keys `issuer` publishes through its `dwk` metadata document.
27    async fn jwks(&self, issuer: &str, dwk: &str) -> Result<Jwks>;
28}
29
30/// Resolves nothing: for verifying signatures whose scheme carries its key.
31#[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/// A response to a [`Fetch`].
42#[derive(Clone, Debug)]
43pub struct Fetched {
44    /// The response status.
45    pub status: StatusCode,
46    /// The response headers.
47    pub headers: HeaderMap,
48    /// The response body, at most [`MAX_DOCUMENT_BYTES`].
49    pub body: Bytes,
50}
51
52/// An HTTP GET.
53#[async_trait]
54pub trait Fetch: Send + Sync {
55    /// GET `url` with `headers`. An error is a transport failure; any HTTP
56    /// status is a [`Fetched`].
57    async fn get(&self, url: &str, headers: HeaderMap) -> Result<Fetched>;
58}
59
60const MAX_CACHED: usize = 256;
61
62/// Discovery over a [`Fetch`], caching each issuer's keys for ten minutes
63/// (see [`JwksResolver::with_ttl`]).
64///
65/// Only a server identifier is looked up. Its document must name it as
66/// `issuer`, exactly (RFC 8414 ยง3.3): otherwise a document served at
67/// someone's well-known path by misconfigured hosting could vouch for keys
68/// that are not theirs. A `jwks_uri` must be on the issuer's own origin.
69pub 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    /// Discovery over `fetch`.
85    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    /// Keep each issuer's keys for `ttl` instead.
94    pub fn with_ttl(mut self, ttl: Duration) -> Self {
95        self.ttl = ttl;
96        self
97    }
98
99    /// Forget every cached key set, as after an out-of-band rotation.
100    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/// Read a response body, refusing one larger than `limit`.
181#[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/// [`Fetch`] over a reqwest client.
208#[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    /// Fetch URLs someone else chose: `https` only, never an IP literal or
218    /// `localhost`, no redirects or proxies, and a resolver that refuses any
219    /// name answering with a non-public address or one in `denied`.
220    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    /// Fetch through `client`, unchecked: for URLs you chose yourself, or local development.
228    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/// Serves fixed JSON documents by URL, and `404` for anything else: for
265/// tests, and for parties that live in the same process.
266#[derive(Clone, Debug, Default)]
267pub struct StaticFetch {
268    documents: Arc<Mutex<HashMap<String, serde_json::Value>>>,
269}
270
271impl StaticFetch {
272    /// No documents yet.
273    pub fn new() -> Self {
274        Self::default()
275    }
276
277    /// Serve `document` at `url`.
278    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}