1use std::collections::HashMap;
62use std::future::Future;
63use std::pin::Pin;
64use std::sync::Arc;
65use std::task::{Context, Poll};
66use std::time::{Duration, Instant};
67
68use base64::engine::general_purpose::STANDARD as BASE64;
69use base64::Engine as _;
70use http_body_util::{BodyExt as _, Full, Limited};
71use hyper::{Method, Request, Uri};
72use hyper_util::client::legacy::connect::{Connected, Connection as HyperConnection};
73use hyper_util::client::legacy::Client;
74use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
75use rustls::pki_types::{pem::PemObject, CertificateDer, PrivateKeyDer, ServerName};
76use rustls::ClientConfig;
77use tower_service::Service;
78
79use crate::config::{OAuth2ClientCredentials, OAuth2Mtls};
80
81const REFRESH_SKEW: Duration = Duration::from_secs(60);
86
87const TOKEN_BODY_CAP: usize = 1024 * 1024;
91
92const TOKEN_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
96
97#[derive(Clone)]
99struct CachedToken {
100 token: String,
102 expires_at: Instant,
105}
106
107impl CachedToken {
108 fn is_valid(&self) -> bool {
111 Instant::now() < self.expires_at
112 }
113}
114
115#[derive(Debug)]
119pub enum OAuth2Error {
120 InvalidEndpoint(String),
122 EndpointStatus(u16),
124 MalformedResponse(String),
127 Fetch(String),
129 MtlsConfig(String),
131}
132
133impl std::fmt::Display for OAuth2Error {
134 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135 match self {
136 OAuth2Error::InvalidEndpoint(m) => {
137 write!(f, "oauth2 token endpoint is invalid: {m}")
138 }
139 OAuth2Error::EndpointStatus(s) => {
140 write!(f, "oauth2 token endpoint returned status {s}")
141 }
142 OAuth2Error::MalformedResponse(m) => {
143 write!(f, "oauth2 token response is malformed: {m}")
144 }
145 OAuth2Error::Fetch(m) => write!(f, "oauth2 token fetch failed: {m}"),
146 OAuth2Error::MtlsConfig(m) => {
147 write!(f, "oauth2 mtls client certificate error: {m}")
148 }
149 }
150 }
151}
152
153impl std::error::Error for OAuth2Error {}
154
155enum MaybeTls {
162 Plain(hyper_util::rt::tokio::WithHyperIo<tokio::net::TcpStream>),
163 Tls(
164 Box<
165 tokio_rustls::client::TlsStream<
166 hyper_util::rt::tokio::WithHyperIo<tokio::net::TcpStream>,
167 >,
168 >,
169 ),
170}
171
172impl tokio::io::AsyncRead for MaybeTls {
173 fn poll_read(
174 mut self: Pin<&mut Self>,
175 cx: &mut Context<'_>,
176 buf: &mut tokio::io::ReadBuf<'_>,
177 ) -> Poll<std::io::Result<()>> {
178 match &mut *self {
179 MaybeTls::Plain(s) => Pin::new(s).poll_read(cx, buf),
180 MaybeTls::Tls(s) => Pin::new(s.as_mut()).poll_read(cx, buf),
181 }
182 }
183}
184
185impl tokio::io::AsyncWrite for MaybeTls {
186 fn poll_write(
187 mut self: Pin<&mut Self>,
188 cx: &mut Context<'_>,
189 buf: &[u8],
190 ) -> Poll<std::io::Result<usize>> {
191 match &mut *self {
192 MaybeTls::Plain(s) => Pin::new(s).poll_write(cx, buf),
193 MaybeTls::Tls(s) => Pin::new(s.as_mut()).poll_write(cx, buf),
194 }
195 }
196
197 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
198 match &mut *self {
199 MaybeTls::Plain(s) => Pin::new(s).poll_flush(cx),
200 MaybeTls::Tls(s) => Pin::new(s.as_mut()).poll_flush(cx),
201 }
202 }
203
204 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
205 match &mut *self {
206 MaybeTls::Plain(s) => Pin::new(s).poll_shutdown(cx),
207 MaybeTls::Tls(s) => Pin::new(s.as_mut()).poll_shutdown(cx),
208 }
209 }
210}
211
212impl HyperConnection for MaybeTls {
213 fn connected(&self) -> Connected {
214 Connected::new()
215 }
216}
217
218#[derive(Clone)]
224struct OAuth2Connector {
225 http: hyper_util::client::legacy::connect::HttpConnector,
226 tls: tokio_rustls::TlsConnector,
227}
228
229impl OAuth2Connector {
230 fn new(mtls: Option<&OAuth2Mtls>) -> Result<Self, OAuth2Error> {
237 let mut http = hyper_util::client::legacy::connect::HttpConnector::new();
238 http.enforce_http(false); let roots = crate::security::tls::webpki_root_store();
240 let builder = ClientConfig::builder().with_root_certificates(roots);
241 let mut cfg = match mtls {
242 Some(m) => {
243 let certs = load_cert_chain(&m.client_cert)?;
244 let key = load_private_key(&m.client_key)?;
245 builder.with_client_auth_cert(certs, key).map_err(|e| {
246 OAuth2Error::MtlsConfig(format!("building client auth config: {e}"))
247 })?
248 }
249 None => builder.with_no_client_auth(),
250 };
251 cfg.alpn_protocols = vec![b"http/1.1".to_vec()];
252 Ok(OAuth2Connector {
253 http,
254 tls: tokio_rustls::TlsConnector::from(Arc::new(cfg)),
255 })
256 }
257}
258
259impl Service<Uri> for OAuth2Connector {
260 type Response = TokioIo<MaybeTls>;
261 type Error = std::io::Error;
262 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
263
264 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
265 self.http.poll_ready(cx).map_err(std::io::Error::other)
266 }
267
268 fn call(&mut self, uri: Uri) -> Self::Future {
269 let mut http = self.http.clone();
270 let tls = self.tls.clone();
271 Box::pin(async move {
272 let host = uri
273 .host()
274 .ok_or_else(|| std::io::Error::other("oauth2 token url has no host"))?
275 .to_string();
276 let stream = tokio::time::timeout(TOKEN_CONNECT_TIMEOUT, http.call(uri.clone()))
277 .await
278 .map_err(|_| {
279 std::io::Error::new(std::io::ErrorKind::TimedOut, "oauth2 connect timed out")
280 })?
281 .map_err(std::io::Error::other)?
282 .into_inner();
283 match uri.scheme_str() {
284 Some("https") => {
285 let name = ServerName::try_from(host.clone()).map_err(|e| {
286 std::io::Error::other(format!("oauth2 host is not a valid tls name: {e}"))
287 })?;
288 let tls_stream = tokio::time::timeout(
289 TOKEN_CONNECT_TIMEOUT,
290 tls.connect(name, hyper_util::rt::tokio::WithHyperIo::new(stream)),
291 )
292 .await
293 .map_err(|_| {
294 std::io::Error::new(
295 std::io::ErrorKind::TimedOut,
296 "oauth2 tls handshake timed out",
297 )
298 })??;
299 Ok(TokioIo::new(MaybeTls::Tls(Box::new(tls_stream))))
300 }
301 _ => Ok(TokioIo::new(MaybeTls::Plain(
302 hyper_util::rt::tokio::WithHyperIo::new(stream),
303 ))),
304 }
305 })
306 }
307}
308
309fn load_cert_chain(path: &str) -> Result<Vec<CertificateDer<'static>>, OAuth2Error> {
311 let ppath = std::path::PathBuf::from(path);
312 let certs = CertificateDer::pem_file_iter(&ppath)
313 .map_err(|e| OAuth2Error::MtlsConfig(format!("reading client cert '{path}': {e}")))?
314 .collect::<Result<Vec<_>, _>>()
315 .map_err(|e| OAuth2Error::MtlsConfig(format!("parsing client cert '{path}': {e}")))?;
316 if certs.is_empty() {
317 return Err(OAuth2Error::MtlsConfig(format!(
318 "client cert '{path}' contains no certificates"
319 )));
320 }
321 Ok(certs)
322}
323
324fn load_private_key(path: &str) -> Result<PrivateKeyDer<'static>, OAuth2Error> {
328 let ppath = std::path::PathBuf::from(path);
329 let pem = std::fs::read(&ppath)
330 .map_err(|e| OAuth2Error::MtlsConfig(format!("reading client key '{path}': {e}")))?;
331 let key = PrivateKeyDer::pem_slice_iter(&pem)
332 .next()
333 .ok_or_else(|| OAuth2Error::MtlsConfig(format!("client key '{path}' has no private key")))?
334 .map_err(|e| OAuth2Error::MtlsConfig(format!("parsing client key '{path}': {e}")))?;
335 Ok(key)
336}
337
338pub struct OAuth2Client {
343 cfg: OAuth2ClientCredentials,
344 http: Client<OAuth2Connector, Full<bytes::Bytes>>,
345 fetch_lock: tokio::sync::Mutex<()>,
349}
350
351impl OAuth2Client {
352 pub fn build(cfg: OAuth2ClientCredentials) -> Result<Arc<Self>, OAuth2Error> {
356 let connector = OAuth2Connector::new(cfg.mtls.as_ref())?;
357 let mut builder = Client::builder(TokioExecutor::new());
358 builder.pool_timer(TokioTimer::new());
359 Ok(Arc::new(OAuth2Client {
360 cfg,
361 http: builder.build(connector),
362 fetch_lock: tokio::sync::Mutex::new(()),
363 }))
364 }
365
366 pub fn config(&self) -> &OAuth2ClientCredentials {
368 &self.cfg
369 }
370
371 async fn fetch_token(&self) -> Result<CachedToken, OAuth2Error> {
375 let uri: Uri = self.cfg.token_endpoint.parse().map_err(|e| {
376 OAuth2Error::InvalidEndpoint(format!(
377 "'{}' is not a valid URL: {e}",
378 self.cfg.token_endpoint
379 ))
380 })?;
381 let secret = crate::config::credentials::resolve_configured_secret(&self.cfg.client_secret)
386 .map_err(|e| OAuth2Error::Fetch(format!("client_secret unresolvable: {e}")))?;
387 let basic = BASE64.encode(format!("{}:{}", self.cfg.client_id, secret));
388 let mut body = "grant_type=client_credentials".to_string();
390 if !self.cfg.scopes.is_empty() {
391 body.push_str("&scope=");
392 body.push_str(&self.cfg.scopes.join(" "));
393 }
394 let req = Request::builder()
395 .method(Method::POST)
396 .uri(uri)
397 .header(hyper::header::AUTHORIZATION, format!("Basic {basic}"))
398 .header(
399 hyper::header::CONTENT_TYPE,
400 "application/x-www-form-urlencoded",
401 )
402 .body(Full::new(bytes::Bytes::from(body)))
403 .map_err(|e| OAuth2Error::Fetch(format!("building token request: {e}")))?;
404 let resp = self
405 .http
406 .request(req)
407 .await
408 .map_err(|e| OAuth2Error::Fetch(format!("token endpoint request failed: {e:#}")))?;
409 if !resp.status().is_success() {
410 return Err(OAuth2Error::EndpointStatus(resp.status().as_u16()));
411 }
412 let body = Limited::new(resp.into_body(), TOKEN_BODY_CAP)
413 .collect()
414 .await
415 .map_err(|e| OAuth2Error::Fetch(format!("token response body read failed: {e}")))?
416 .to_bytes();
417 let json: serde_json::Value = serde_json::from_slice(&body)
418 .map_err(|e| OAuth2Error::MalformedResponse(format!("not JSON: {e}")))?;
419 let token = json
420 .get("access_token")
421 .and_then(|v| v.as_str())
422 .ok_or_else(|| OAuth2Error::MalformedResponse("missing string access_token".into()))?
423 .to_string();
424 if token.is_empty() {
425 return Err(OAuth2Error::MalformedResponse(
426 "access_token is empty".into(),
427 ));
428 }
429 if let Some(tt) = json.get("token_type").and_then(|v| v.as_str()) {
433 if !tt.eq_ignore_ascii_case("bearer") {
434 return Err(OAuth2Error::MalformedResponse(format!(
435 "token_type '{tt}' is not 'bearer'"
436 )));
437 }
438 }
439 let expires_in = json
440 .get("expires_in")
441 .and_then(|v| v.as_u64())
442 .unwrap_or(3600);
443 let ttl = self.effective_ttl(expires_in);
444 Ok(CachedToken {
445 token,
446 expires_at: Instant::now() + ttl,
447 })
448 }
449
450 fn effective_ttl(&self, expires_in: u64) -> Duration {
455 let skew = REFRESH_SKEW.as_secs();
456 let base = expires_in.saturating_sub(skew).max(1);
457 match self.cfg.token_cache_ttl_s {
458 Some(override_s) => Duration::from_secs(base.min(override_s)),
459 None => Duration::from_secs(base),
460 }
461 }
462
463 #[doc(hidden)]
468 pub fn effective_ttl_for_test(&self, expires_in: u64) -> Duration {
469 self.effective_ttl(expires_in)
470 }
471
472 pub async fn token(&self, cache: &OAuth2TokenCache) -> Result<String, OAuth2Error> {
482 if let Some(token) = cache.get(self).filter(|t| t.is_valid()) {
484 return Ok(token.token.clone());
485 }
486 let _guard = self.fetch_lock.lock().await;
490 if let Some(token) = cache.get(self).filter(|t| t.is_valid()) {
491 return Ok(token.token.clone());
492 }
493 let token = self.fetch_token().await?;
494 let token_string = token.token.clone();
495 cache.put(self, token);
496 Ok(token_string)
497 }
498}
499
500impl std::fmt::Debug for OAuth2Client {
501 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
502 f.debug_struct("OAuth2Client")
503 .field("token_endpoint", &self.cfg.token_endpoint)
504 .field("client_id", &self.cfg.client_id)
505 .field("client_secret", &"[redacted]")
506 .field("scopes", &self.cfg.scopes)
507 .field("mtls", &self.cfg.mtls.is_some())
508 .field("token_cache_ttl_s", &self.cfg.token_cache_ttl_s)
509 .finish_non_exhaustive()
510 }
511}
512
513pub struct OAuth2TokenCache {
519 inner: std::sync::Mutex<HashMap<String, CachedToken>>,
520}
521
522impl OAuth2TokenCache {
523 pub fn new() -> Self {
525 OAuth2TokenCache {
526 inner: std::sync::Mutex::new(HashMap::new()),
527 }
528 }
529
530 fn get(&self, client: &OAuth2Client) -> Option<CachedToken> {
533 self.inner
534 .lock()
535 .expect("oauth2 token cache poisoned")
536 .get(&client.cfg.token_endpoint)
537 .cloned()
538 }
539
540 fn put(&self, client: &OAuth2Client, token: CachedToken) {
543 self.inner
544 .lock()
545 .expect("oauth2 token cache poisoned")
546 .insert(client.cfg.token_endpoint.clone(), token);
547 }
548}
549
550impl Default for OAuth2TokenCache {
551 fn default() -> Self {
552 Self::new()
553 }
554}
555
556impl std::fmt::Debug for OAuth2TokenCache {
557 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
558 f.debug_struct("OAuth2TokenCache").finish_non_exhaustive()
559 }
560}