Skip to main content

alibi_plugins/oauth_provider/
client_jwks.rs

1//! Bounded remote client metadata transport and provider-scoped public-key caching.
2use alibi_core::store::OAuthClient;
3use alibi_core::{AuthContext, AuthError, AuthResult, AuthSchema, Headers};
4use async_trait::async_trait;
5use indexmap::IndexMap;
6use serde_json::Value;
7use std::sync::Arc;
8use std::sync::Mutex;
9use tokio::sync::watch;
10
11/// Cancellation exposed to an application-owned client metadata transport.
12#[derive(Clone, Debug)]
13pub struct OAuthClientMetadataAbortSignal {
14    receiver: watch::Receiver<bool>,
15    _keepalive: watch::Sender<bool>,
16}
17impl OAuthClientMetadataAbortSignal {
18    #[must_use]
19    pub fn aborted(&self) -> bool {
20        *self.receiver.borrow()
21    }
22    /// Wait until the resource request's abort signal fires.
23    pub async fn cancelled(&self) {
24        let mut receiver = self.receiver.clone();
25        loop {
26            let aborted = *receiver.borrow_and_update();
27            if aborted || receiver.changed().await.is_err() {
28                return;
29            }
30        }
31    }
32}
33/// Actual options supplied to a configured metadata resource fetch callback.
34#[derive(Clone, Debug)]
35pub struct OAuthClientMetadataFetchOptions {
36    pub headers: Headers,
37    pub redirect: &'static str,
38    pub signal: OAuthClientMetadataAbortSignal,
39}
40/// A streaming HTTP response and its redirect history, matching a fetch Response.
41#[derive(Debug)]
42pub struct OAuthClientMetadataFetchResponse {
43    pub response: reqwest::Response,
44    pub redirected: bool,
45}
46/// Application-owned network boundary for a discovered client's referenced resources.
47#[async_trait]
48pub trait OAuthClientMetadataResourceFetch: Send + Sync + std::fmt::Debug {
49    async fn fetch(
50        &self,
51        url: &str,
52        options: OAuthClientMetadataFetchOptions,
53    ) -> AuthResult<OAuthClientMetadataFetchResponse>;
54}
55
56const FETCH_ERROR: &str = "failed to fetch client JWKS";
57const MAX_BYTES: usize = 64 * 1024;
58#[derive(Clone, Debug)]
59struct Entry {
60    jwks: Value,
61    fetched_at: i64,
62}
63#[derive(Debug, Default)]
64pub(super) struct Cache {
65    entries: Mutex<IndexMap<String, Entry>>,
66}
67impl Cache {
68    fn get(&self, key: &str) -> AuthResult<Option<Entry>> {
69        Ok(self
70            .entries
71            .lock()
72            .map_err(|_| AuthError::internal("JWKS cache unavailable"))?
73            .get(key)
74            .cloned())
75    }
76    fn set(&self, key: String, jwks: Value, fetched_at: i64) -> AuthResult<()> {
77        let mut entries = self
78            .entries
79            .lock()
80            .map_err(|_| AuthError::internal("JWKS cache unavailable"))?;
81        _ = entries.insert(key, Entry { jwks, fetched_at });
82        if entries.len() > 500 {
83            _ = entries.shift_remove_index(0);
84        }
85        Ok(())
86    }
87}
88fn owner(client: &OAuthClient) -> Option<&str> {
89    client
90        .metadata
91        .get("client_discovery_id")
92        .and_then(Value::as_str)
93        .filter(|id| !id.is_empty())
94}
95fn uri(client: &OAuthClient) -> Option<&str> {
96    client
97        .metadata
98        .get("jwks_uri")
99        .and_then(Value::as_str)
100        .filter(|uri| !uri.is_empty())
101}
102fn cache_key(client: &OAuthClient, uri: &str) -> String {
103    format!(
104        "{}:{uri}",
105        client
106            .metadata
107            .get("client_discovery_id")
108            .and_then(Value::as_str)
109            .unwrap_or("managed")
110    )
111}
112fn transport(
113    provider: &super::OAuthProviderPlugin,
114    client: &OAuthClient,
115) -> Option<Arc<dyn OAuthClientMetadataResourceFetch>> {
116    let id = owner(client)?;
117    provider
118        .config
119        .client_discoveries
120        .iter()
121        .find(|source| source.id() == id)?
122        .metadata_resource_fetch()
123}
124fn validate_uri<S: AuthSchema>(
125    ctx: &AuthContext<S>,
126    client: &OAuthClient,
127    uri: &str,
128) -> AuthResult<Result<(), &'static str>> {
129    let parsed =
130        url::Url::parse(uri).map_err(|_| AuthError::internal("Invalid client JWKS URL"))?;
131    if parsed.scheme() != "https" {
132        return Ok(Err("jwks_uri must use HTTPS"));
133    }
134    if !parsed.username().is_empty()
135        || parsed
136            .password()
137            .is_some_and(|password| !password.is_empty())
138    {
139        return Ok(Err("jwks_uri must not contain credentials"));
140    }
141    if parsed.fragment().is_some() {
142        return Ok(Err("jwks_uri must not include a fragment component"));
143    }
144    if !crate::sso::endpoint_host::public(parsed.host_str().unwrap_or_default()) {
145        return Ok(Err(
146            "jwks_uri must not point to a private or reserved address",
147        ));
148    }
149    if owner(client).is_some()
150        && url::Url::parse(&client.client_id).is_ok_and(|id| {
151            matches!(id.scheme(), "https" | "http") && id.origin() == parsed.origin()
152        })
153    {
154        return Ok(Ok(()));
155    }
156    if !ctx.config.is_redirect_target_trusted(uri) {
157        return Ok(Err("client jwks_uri is not trusted"));
158    }
159    Ok(Ok(()))
160}
161pub(super) async fn get<S: AuthSchema>(
162    ctx: &AuthContext<S>,
163    provider: &super::OAuthProviderPlugin,
164    client: &OAuthClient,
165) -> AuthResult<Result<Value, &'static str>> {
166    if let Some(inline) = client
167        .metadata
168        .get("jwks")
169        .filter(|value| !value.is_null() && *value != "")
170    {
171        return Ok(Ok(if let Some(raw) = inline.as_str() {
172            alibi_core::utils::json::parse_value(raw)
173                .map_err(|_| AuthError::internal("Invalid stored client JWKS"))?
174                .to_json_value()?
175        } else {
176            inline.clone()
177        }));
178    }
179    let Some(uri) = uri(client) else {
180        return Ok(Err("client has no JWKS configured"));
181    };
182    let fetcher = transport(provider, client);
183    if owner(client).is_some() && fetcher.is_none() {
184        return Ok(Err(
185            "client discovery does not provide a metadata resource transport",
186        ));
187    }
188    if let Err(description) = validate_uri(ctx, client, uri)? {
189        return Ok(Err(description));
190    }
191    let now = chrono::Utc::now().timestamp_millis();
192    let key = cache_key(client, uri);
193    let cache = provider.client_jwks.clone();
194    let cached = cache.get(&key)?;
195    if let Some(entry) = &cached
196        && now.saturating_sub(entry.fetched_at) < 300_000
197    {
198        return Ok(Ok(entry.jwks.clone()));
199    }
200    match fetch(uri, fetcher).await {
201        Ok(Some(jwks)) => {
202            cache.set(key, jwks.clone(), now)?;
203            Ok(Ok(jwks))
204        }
205        Ok(None) => Ok(Err(FETCH_ERROR)),
206        Err(()) => Ok(
207            if let Some(entry) =
208                cached.filter(|entry| now.saturating_sub(entry.fetched_at) < 600_000)
209            {
210                Ok(entry.jwks)
211            } else {
212                Err(FETCH_ERROR)
213            },
214        ),
215    }
216}
217pub(super) async fn refetch(
218    provider: &super::OAuthProviderPlugin,
219    client: &OAuthClient,
220) -> AuthResult<Option<Value>> {
221    let Some(uri) = uri(client) else {
222        return Ok(None);
223    };
224    let fetcher = transport(provider, client);
225    if owner(client).is_some() && fetcher.is_none() {
226        return Ok(None);
227    }
228    let Ok(Some(jwks)) = fetch(uri, fetcher).await else {
229        return Ok(None);
230    };
231    provider.client_jwks.set(
232        cache_key(client, uri),
233        jwks.clone(),
234        chrono::Utc::now().timestamp_millis(),
235    )?;
236    Ok(Some(jwks))
237}
238struct AbortTimer(tokio::task::JoinHandle<()>);
239impl Drop for AbortTimer {
240    fn drop(&mut self) {
241        self.0.abort();
242    }
243}
244async fn fetch(
245    uri: &str,
246    transport: Option<Arc<dyn OAuthClientMetadataResourceFetch>>,
247) -> Result<Option<Value>, ()> {
248    let (sender, receiver) = watch::channel(false);
249    let signal = OAuthClientMetadataAbortSignal {
250        receiver,
251        _keepalive: sender.clone(),
252    };
253    let _timer = AbortTimer(tokio::spawn(async move {
254        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
255        _ = sender.send(true);
256    }));
257    let mut response = if let Some(transport) = transport {
258        let mut headers = Headers::new();
259        _ = headers.insert("accept", "application/json");
260        transport
261            .fetch(
262                uri,
263                OAuthClientMetadataFetchOptions {
264                    headers,
265                    redirect: "error",
266                    signal,
267                },
268            )
269            .await
270            .map_err(|_| ())?
271    } else {
272        let client = reqwest::Client::builder()
273            .redirect(reqwest::redirect::Policy::none())
274            .build()
275            .map_err(|_| ())?;
276        OAuthClientMetadataFetchResponse {
277            response: client
278                .get(uri)
279                .header("accept", "application/json")
280                .timeout(std::time::Duration::from_secs(5))
281                .send()
282                .await
283                .map_err(|_| ())?,
284            redirected: false,
285        }
286    };
287    if response.redirected || response.response.status() != reqwest::StatusCode::OK {
288        return Err(());
289    }
290    let Some(content_type) = header(&response.response, "content-type") else {
291        return Err(());
292    };
293    if !json_content_type(&content_type) {
294        return Err(());
295    }
296    if header(&response.response, "content-length")
297        .as_deref()
298        .and_then(alibi_core::utils::javascript::string_to_number)
299        .is_some_and(|length| length.is_finite() && length > 65_536.0)
300    {
301        return Err(());
302    }
303    let mut bytes = Vec::new();
304    while let Some(chunk) = response.response.chunk().await.map_err(|_| ())? {
305        if bytes.len().saturating_add(chunk.len()) > MAX_BYTES {
306            return Err(());
307        }
308        bytes.extend_from_slice(&chunk);
309    }
310    let text = String::from_utf8_lossy(&bytes);
311    let text = text.strip_prefix('\u{feff}').unwrap_or(&text);
312    let Ok(parsed) =
313        alibi_core::utils::json::parse_value(text).and_then(|value| value.to_json_value())
314    else {
315        return Ok(None);
316    };
317    let valid = parsed
318        .get("keys")
319        .and_then(Value::as_array)
320        .is_some_and(|keys| !keys.is_empty() && keys.iter().all(Value::is_object))
321        && super::clients::public_jwks_error(&parsed).is_none();
322    Ok(valid.then_some(parsed))
323}
324fn header(response: &reqwest::Response, name: &str) -> Option<String> {
325    let values = response
326        .headers()
327        .get_all(name)
328        .iter()
329        .map(|value| {
330            value.to_str().map_or_else(
331                |_| {
332                    value
333                        .as_bytes()
334                        .iter()
335                        .map(|byte| char::from(*byte))
336                        .collect()
337                },
338                str::to_owned,
339            )
340        })
341        .collect::<Vec<_>>();
342    (!values.is_empty()).then(|| {
343        values
344            .join(", ")
345            .trim_matches([' ', '\t', '\r', '\n'])
346            .to_owned()
347    })
348}
349fn json_content_type(value: &str) -> bool {
350    let value = value.to_ascii_lowercase();
351    let Some(subtype) = value.strip_prefix("application/") else {
352        return false;
353    };
354    let subtype = subtype
355        .split(';')
356        .next()
357        .unwrap_or_default()
358        .trim_end_matches(alibi_core::utils::javascript::is_whitespace);
359    subtype == "json"
360        || subtype.strip_suffix("+json").is_some_and(|prefix| {
361            !prefix.is_empty()
362                && prefix
363                    .bytes()
364                    .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.'))
365        })
366}