From 8cac68328a9d74679c19923f2aa784915ef0e20b Mon Sep 17 00:00:00 2001 From: ee Date: Mon, 19 Oct 2026 14:55:49 +0000 Subject: [PATCH] Add renewable federated auth primitives to the Responses API proxy (#55298) ## What changed Add a federated `Auth` variant that exchanges JWT subject tokens for OpenAI access tokens using identity provider and service account IDs. Serialize concurrent exchanges, cache credentials, and renew proactively with a fresh subject token. Validate token types, bearer credentials, payload sizes, and lifetimes. Bound caching by both relative and absolute expiry, discard cached credentials before renewal, and disable redirects for federated requests. Keep credential buffers zeroizing and diagnostics free of secrets. ## Testing Add HTTP fixture tests for concurrent cache reuse, renewal with fresh subjects, expiry across host suspension, failed renewal, oversized payloads, invalid subjects, and unsupported or missing issued token types. GitOrigin-RevId: e7e6dacd4a834fb73166caf146d4cdc50a895f56 --- codex-rs/Cargo.lock | 1 + codex-rs/responses-api-proxy/Cargo.toml | 1 + codex-rs/responses-api-proxy/src/auth.rs | 238 ++++++++++- .../responses-api-proxy/src/auth_tests.rs | 383 +++++++++++++++++- 4 files changed, 604 insertions(+), 19 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 20fcf90a18..33f0811a10 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -4476,6 +4476,7 @@ name = "codex-responses-api-proxy" version = "0.0.0" dependencies = [ "anyhow", + "bytes", "clap", "codex-http-client", "codex-process-hardening", diff --git a/codex-rs/responses-api-proxy/Cargo.toml b/codex-rs/responses-api-proxy/Cargo.toml index 15107789b6..82a9086783 100644 --- a/codex-rs/responses-api-proxy/Cargo.toml +++ b/codex-rs/responses-api-proxy/Cargo.toml @@ -18,6 +18,7 @@ workspace = true [dependencies] anyhow = { workspace = true } +bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } codex-http-client = { workspace = true } codex-process-hardening = { workspace = true } diff --git a/codex-rs/responses-api-proxy/src/auth.rs b/codex-rs/responses-api-proxy/src/auth.rs index 5e95ce63ce..1b61b5c34f 100644 --- a/codex-rs/responses-api-proxy/src/auth.rs +++ b/codex-rs/responses-api-proxy/src/auth.rs @@ -1,12 +1,19 @@ -//! Supplies sensitive request headers without copying the locked API key. +//! Authenticates proxy requests using a locked API key or renewable OpenAI credentials. +//! Serializes exchanges and never reuses expired credentials on renewal errors. //! Token clients never follow redirects or emit credential-bearing URL diagnostics. use std::fmt; use std::future::Future; use std::pin::Pin; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; use anyhow::Result; use anyhow::anyhow; +use bytes::Bytes; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClient; use codex_http_client::HttpClientBuilder; @@ -15,22 +22,26 @@ use codex_http_client::HttpResponse; use codex_http_client::OutboundProxyPolicy; use codex_http_client::RouteAwareClientPool; use http::HeaderValue; +use http::header::CONTENT_TYPE; +use serde::Deserialize; +use serde::Serialize; use serde::de::DeserializeOwned; -use std::time::Duration; +use tokio::runtime::Runtime; +use url::Url; use zeroize::Zeroizing; -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] +const OPENAI_TOKEN_URL: &str = "https://auth.openai.com/oauth/token"; pub(crate) const MAX_TOKEN_BYTES: usize = 16 * 1024; -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] const MAX_RESPONSE_BYTES: usize = 64 * 1024; -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] +// Bound the complete serialized token and provider metadata, including JSON escaping. +const MAX_EXCHANGE_REQUEST_BYTES: usize = 64 * 1024; pub(crate) const TOKEN_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const MAX_FEDERATED_TOKEN_LIFETIME: Duration = Duration::from_secs(3600); /// Supplies a fresh JWT whenever an OpenAI credential needs to be issued. /// Acquisition stays asynchronous so callers can compose it with the exchange. /// Implementations must not include tokens, credential-bearing URLs, or upstream /// response bodies in errors. The exchange deliberately has no issuer-specific config. -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] pub(crate) trait SubjectTokenSource: Send { fn subject_token<'a>(&'a mut self, client: &'a TokenClient) -> SubjectTokenFuture<'a>; } @@ -41,15 +52,9 @@ pub(crate) type SubjectTokenFuture<'a> = /// Credential requests must not follow redirects or log URLs or response headers. /// The private field prevents other modules from wrapping arbitrary HTTP clients. -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] pub(crate) struct TokenClient(HttpClient); -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] impl TokenClient { - #[allow( - dead_code, - reason = "federation creates the routed client in the next stack layer" - )] pub(crate) fn new() -> Self { // The standalone proxy does not load Codex's session config. For its // opt-in auth flow, resolve system and environment proxy settings per URL. @@ -74,31 +79,67 @@ impl TokenClient { pub(crate) enum Auth { ApiKey(&'static str), + Federated { state: Box> }, } impl fmt::Debug for Auth { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::ApiKey(_) => f.debug_tuple("ApiKey").field(&"[REDACTED]").finish(), + Self::Federated { .. } => f + .debug_struct("Federated") + .field("state", &"[REDACTED]") + .finish(), } } } +/// The sensitive header used for an individual request. #[derive(Clone, Debug)] pub(crate) struct AuthHeader { pub(crate) value: HeaderValue, } impl Auth { + #[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] + pub(crate) fn federated( + source: impl SubjectTokenSource + 'static, + provider: String, + service_account: String, + ) -> Result { + let client = TokenClient::new(); + // Only credential acquisition needs an async runtime. The proxy's + // blocking request threads drive it while holding the credential mutex. + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|_| anyhow!("could not start workload identity HTTP runtime"))?; + Ok(Self::Federated { + state: Box::new(Mutex::new(FederatedAuth { + client, + runtime, + token_url: Url::parse(OPENAI_TOKEN_URL)?, + source: Box::new(source), + provider, + service_account, + cached: None, + })), + }) + } + pub(crate) fn redirect_policy(&self) -> reqwest::redirect::Policy { match self { Self::ApiKey(_) => reqwest::redirect::Policy::default(), + // A federated credential must never be sent to a redirect target. + Self::Federated { .. } => reqwest::redirect::Policy::none(), } } - /// Returns a sensitive authorization header. API keys are validated when read, - /// so this currently always succeeds. The fallible return type allows later - /// credential providers to report acquisition failures. + /// Returns a sensitive authorization header. + /// + /// API keys are validated when read. Federated authentication can fail if the + /// credential cache cannot be locked, or if token acquisition, exchange, or + /// validation fails. pub(crate) fn header(&self) -> Result { match self { Self::ApiKey(value) => { @@ -107,11 +148,155 @@ impl Auth { header.set_sensitive(true); Ok(AuthHeader { value: header }) } + Self::Federated { state } => state + .lock() + .map_err(|_| anyhow!("workload identity credential cache unavailable"))? + .header(), } } } -#[cfg_attr(not(test), allow(dead_code, reason = "used by a later stack layer"))] +pub(crate) struct FederatedAuth { + client: TokenClient, + runtime: Runtime, + token_url: Url, + source: Box, + provider: String, + service_account: String, + cached: Option, +} + +struct CachedToken { + header: AuthHeader, + renew_at: Instant, + expires_at: SystemTime, +} + +impl FederatedAuth { + fn header(&mut self) -> Result { + if let Some(cached) = &self.cached + && Instant::now() < cached.renew_at + // Monotonic time can stop while a host is suspended. + && SystemTime::now() < cached.expires_at + { + return Ok(cached.header.clone()); + } + // Discard the old credential before a fallible exchange. Never reuse it + // when acquisition fails. + self.cached = None; + let cached = self.exchange()?; + let header = cached.header.clone(); + self.cached = Some(cached); + Ok(header) + } + + fn exchange(&mut self) -> Result { + let (started, exchange): (Instant, ExchangeResponse) = self.runtime.block_on(async { + let subject = self.source.subject_token(&self.client).await?; + if subject.is_empty() || subject.len() > MAX_TOKEN_BYTES { + return Err(anyhow!( + "workload identity source returned an invalid subject token" + )); + } + let request = ExchangeRequest { + grant_type: "urn:ietf:params:oauth:grant-type:token-exchange", + subject_token_type: "urn:ietf:params:oauth:token-type:jwt", + subject_token: subject.as_str(), + identity_provider_id: &self.provider, + service_account_id: &self.service_account, + }; + // Serialize through a fixed slice so secret-bearing storage cannot grow. + // The HTTP body owns this buffer, including any request clones used for retries. + let mut body = Zeroizing::new(vec![0u8; MAX_EXCHANGE_REQUEST_BYTES]); + let written = { + let mut writer = std::io::Cursor::new(body.as_mut_slice()); + serde_json::to_writer(&mut writer, &request).map_err(|_| { + anyhow!( + "OpenAI workload identity exchange request exceeds {MAX_EXCHANGE_REQUEST_BYTES} bytes" + ) + })?; + writer.position() as usize + }; + body.truncate(written); + // Start the lifetime clock before the OpenAI exchange, after obtaining + // the subject. expires_in is measured from issuance, not receipt. + let started = Instant::now(); + let response = self + .client + .http() + .post(self.token_url.clone()) + .timeout(TOKEN_REQUEST_TIMEOUT) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) + .body(Bytes::from_owner(body)) + .send() + .await + .map_err(|_| anyhow!("OpenAI workload identity token exchange failed"))?; + let exchange = read_response(response, "OpenAI workload identity").await?; + Ok((started, exchange)) + })?; + if !exchange.token_type.eq_ignore_ascii_case("bearer") + || exchange.issued_token_type != "urn:ietf:params:oauth:token-type:access_token" + { + return Err(anyhow!( + "OpenAI returned an unsupported workload identity token type" + )); + } + if exchange.expires_in > MAX_FEDERATED_TOKEN_LIFETIME.as_secs() { + return Err(anyhow!("OpenAI returned an invalid token lifetime")); + } + let header = bearer_header(exchange.access_token.as_bytes())?; + let expires_at = UNIX_EPOCH + .checked_add(Duration::from_secs(exchange.expires_at)) + .ok_or_else(|| anyhow!("OpenAI returned an expired workload identity token"))?; + let now = SystemTime::now(); + let remaining_absolute = expires_at + .duration_since(now) + .map_err(|_| anyhow!("OpenAI returned an expired workload identity token"))?; + let remaining_relative = Duration::from_secs(exchange.expires_in) + .checked_sub(started.elapsed()) + .ok_or_else(|| anyhow!("OpenAI returned an expired workload identity token"))?; + let remaining = remaining_absolute.min(remaining_relative); + if remaining.is_zero() { + return Err(anyhow!( + "OpenAI returned an expired workload identity token" + )); + } + // Bound the wall-clock deadline too, even when the issuer's clock is + // ahead. Otherwise a host suspend could outlast the shorter relative TTL. + let expires_at = now + .checked_add(remaining) + .ok_or_else(|| anyhow!("OpenAI returned an invalid token lifetime"))?; + // Proactive renewal without a fixed skew that rejects short-lived subject tokens. + let skew = (remaining / 10).min(Duration::from_secs(60)); + let renew_at = Instant::now() + .checked_add(remaining - skew) + .ok_or_else(|| anyhow!("OpenAI returned an invalid token lifetime"))?; + Ok(CachedToken { + header: AuthHeader { value: header }, + renew_at, + expires_at, + }) + } +} + +#[derive(Serialize)] +struct ExchangeRequest<'a> { + grant_type: &'static str, + subject_token_type: &'static str, + subject_token: &'a str, + identity_provider_id: &'a str, + service_account_id: &'a str, +} + +#[derive(Deserialize)] +struct ExchangeResponse { + access_token: Zeroizing, + token_type: String, + issued_token_type: String, + expires_in: u64, + expires_at: u64, +} + pub(crate) async fn read_response( mut response: HttpResponse, source: &str, @@ -140,6 +325,27 @@ pub(crate) async fn read_response( .map_err(|_| anyhow!("{source} returned an invalid token response")) } +// RFC 6750 bearer token characters, rather than the API key's narrower alphabet. +// Accept JWT periods, but never control bytes, whitespace, or header injection. +fn bearer_header(token: &[u8]) -> Result { + if token.is_empty() + || token.len() > MAX_TOKEN_BYTES + || !token.iter().all(|byte| { + byte.is_ascii_alphanumeric() + || matches!(byte, b'-' | b'.' | b'_' | b'~' | b'+' | b'/' | b'=') + }) + { + return Err(anyhow!("invalid workload identity bearer credential")); + } + let mut bytes = Zeroizing::new(Vec::with_capacity(7 + token.len())); + bytes.extend_from_slice(b"Bearer "); + bytes.extend_from_slice(token); + let mut header = HeaderValue::from_bytes(&bytes) + .map_err(|_| anyhow!("invalid workload identity bearer credential"))?; + header.set_sensitive(true); + Ok(header) +} + #[cfg(test)] #[path = "auth_tests.rs"] pub(crate) mod tests; diff --git a/codex-rs/responses-api-proxy/src/auth_tests.rs b/codex-rs/responses-api-proxy/src/auth_tests.rs index ecc31985f1..b9df94a72b 100644 --- a/codex-rs/responses-api-proxy/src/auth_tests.rs +++ b/codex-rs/responses-api-proxy/src/auth_tests.rs @@ -1,7 +1,13 @@ -//! Credential-safe diagnostics and transport fixtures for the production token-client policy. +//! Tests credential-safe diagnostics, federated exchanges, expiry, and transport. -use super::Auth; -use super::TokenClient; +use super::*; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::sync::Arc; +use std::sync::mpsc; +use std::thread; +use tiny_http::Header; +use tiny_http::Server; pub(crate) fn direct_token_client() -> TokenClient { TokenClient(TokenClient::builder().build_direct().unwrap()) @@ -17,3 +23,374 @@ fn debug_redacts_api_keys_in_both_formats() { assert!(output.contains("[REDACTED]")); } } + +// A non-GitHub issuer: each renewal must ask the source for a fresh subject. +struct RotatingSource { + calls: usize, +} + +impl SubjectTokenSource for RotatingSource { + fn subject_token<'a>(&'a mut self, _client: &'a TokenClient) -> SubjectTokenFuture<'a> { + Box::pin(async move { + self.calls += 1; + Ok(Zeroizing::new(format!("issuer.subject.{}", self.calls))) + }) + } +} + +fn server() -> (Server, Url) { + let server = Server::http("127.0.0.1:0").unwrap(); + let url = Url::parse(&format!("http://{}", server.server_addr())).unwrap(); + (server, url) +} + +fn auth_at(url: &Url) -> Auth { + let auth = Auth::federated( + RotatingSource { calls: 0 }, + "ip-test".to_owned(), + "sa-test".to_owned(), + ) + .unwrap(); + let Auth::Federated { state, .. } = &auth else { + unreachable!() + }; + let mut state = state.lock().unwrap(); + state.token_url = url.join("/oauth/token").unwrap(); + // Fixtures must not use a developer's system or environment proxy. + state.client = direct_token_client(); + drop(state); + auth +} + +fn recv(server: &Server) -> tiny_http::Request { + server + .recv_timeout(Duration::from_secs(5)) + .unwrap() + .expect("expected HTTP request") +} + +fn recv_exchange(server: &Server, subject: &str) -> tiny_http::Request { + let mut req = recv(server); + assert_eq!(req.url(), "/oauth/token"); + assert_eq!(req.method(), &tiny_http::Method::Post); + assert!(!req.headers().iter().any(|h| h.field.equiv("Authorization"))); + assert_eq!( + req.headers() + .iter() + .find(|h| h.field.equiv("Content-Type")) + .map(|h| h.value.as_str()), + Some("application/json") + ); + let mut body = String::new(); + req.as_reader().read_to_string(&mut body).unwrap(); + assert_eq!( + serde_json::from_str::(&body).unwrap(), + json!({ + "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "subject_token": subject, + "identity_provider_id": "ip-test", + "service_account_id": "sa-test", + }) + ); + req +} + +fn serve_exchange(server: &Server, subject: &str, token: &str, expires_in: u64, expires_at: u64) { + respond_exchange( + recv_exchange(server, subject), + token, + expires_in, + expires_at, + ); +} + +fn respond_exchange(req: tiny_http::Request, token: &str, expires_in: u64, expires_at: u64) { + req.respond( + tiny_http::Response::from_string( + json!({ + "access_token": token, "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "expires_in": expires_in, "expires_at": expires_at, + }) + .to_string(), + ) + .with_header(Header::from_bytes("Content-Type", "application/json").unwrap()), + ) + .unwrap(); +} + +fn epoch_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() +} + +#[test] +fn oversized_serialized_exchange_request_is_rejected_before_sending() { + let (server, url) = server(); + let auth = auth_at(&url); + let Auth::Federated { state, .. } = &auth else { + unreachable!() + }; + // Fits unescaped, but JSON escaping makes the serialized request exceed the limit. + state.lock().unwrap().provider = "\0".repeat(MAX_EXCHANGE_REQUEST_BYTES / 6 + 1); + + assert_eq!( + auth.header().unwrap_err().to_string(), + format!( + "OpenAI workload identity exchange request exceeds {MAX_EXCHANGE_REQUEST_BYTES} bytes" + ) + ); + assert!(state.lock().unwrap().cached.is_none()); + assert!(server.try_recv().unwrap().is_none()); +} + +#[test] +fn oversized_exchange_never_becomes_a_cached_credential() { + let (server, url) = server(); + let worker = thread::spawn(move || { + let body = json!({ + "access_token": "secret".repeat(MAX_RESPONSE_BYTES), + "token_type": "Bearer", + "issued_token_type": "urn:ietf:params:oauth:token-type:access_token", + "expires_in": 300, + "expires_at": epoch_seconds() + 300, + }); + let _ = recv_exchange(&server, "issuer.subject.1").respond( + tiny_http::Response::from_string(body.to_string()) + .with_header(Header::from_bytes("Content-Type", "application/json").unwrap()), + ); + }); + let auth = auth_at(&url); + let error = format!("{:#}", auth.header().unwrap_err()); + assert_eq!( + error, + "OpenAI workload identity returned an oversized token response" + ); + let Auth::Federated { state, .. } = &auth else { + unreachable!() + }; + assert!(state.lock().unwrap().cached.is_none()); + worker.join().unwrap(); +} + +#[test] +fn exchanges_once_for_concurrent_requests_and_renews_with_a_new_subject() { + let (server, url) = server(); + let auth = Arc::new(auth_at(&url)); + let (attempt_tx, attempt_rx) = mpsc::channel(); + let (in_flight_tx, in_flight_rx) = mpsc::channel(); + let worker = thread::spawn(move || { + let pending = recv_exchange(&server, "issuer.subject.1"); + in_flight_tx.send(()).unwrap(); + // Each other caller must observe the lock held during the first exchange. + for _ in 0..7 { + attempt_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + } + assert!(server.try_recv().unwrap().is_none()); + respond_exchange( + pending, + "openai.first.jwt", + /*expires_in*/ 300, + epoch_seconds() + 300, + ); + serve_exchange( + &server, + "issuer.subject.2", + "openai.second.jwt", + /*expires_in*/ 300, + epoch_seconds() + 300, + ); + }); + let first = { + let auth = Arc::clone(&auth); + thread::spawn(move || auth.header().unwrap()) + }; + in_flight_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + let callers: Vec<_> = (0..7) + .map(|_| { + let auth = Arc::clone(&auth); + let attempted = attempt_tx.clone(); + thread::spawn(move || { + let Auth::Federated { state } = auth.as_ref() else { + unreachable!() + }; + assert!(matches!( + state.try_lock(), + Err(std::sync::TryLockError::WouldBlock) + )); + attempted.send(()).unwrap(); + auth.header().unwrap() + }) + }) + .collect(); + for caller in std::iter::once(first).chain(callers) { + let header = caller.join().unwrap(); + assert!(header.value.is_sensitive()); + assert_eq!(header.value.to_str().unwrap(), "Bearer openai.first.jwt"); + } + let Auth::Federated { state, .. } = auth.as_ref() else { + unreachable!() + }; + state.lock().unwrap().cached.as_mut().unwrap().renew_at = Instant::now(); + assert_eq!( + auth.header().unwrap().value.to_str().unwrap(), + "Bearer openai.second.jwt" + ); + worker.join().unwrap(); +} + +#[test] +fn exchange_is_never_cached_if_expired_or_lifetime_exceeds_the_contract() { + for (expires_in, expires_at, expected_error) in [ + ( + 3600, + epoch_seconds() - 1, + "OpenAI returned an expired workload identity token", + ), + ( + 0, + epoch_seconds() + 3600, + "OpenAI returned an expired workload identity token", + ), + ( + 3601, + epoch_seconds() + 7200, + "OpenAI returned an invalid token lifetime", + ), + ] { + let (server, url) = server(); + let worker = thread::spawn(move || { + serve_exchange( + &server, + "issuer.subject.1", + "expired.secret.token", + expires_in, + expires_at, + ); + }); + let auth = auth_at(&url); + assert_eq!(auth.header().unwrap_err().to_string(), expected_error); + let Auth::Federated { state, .. } = auth else { + unreachable!() + }; + assert!(state.lock().unwrap().cached.is_none()); + worker.join().unwrap(); + } +} + +#[test] +fn unsupported_or_missing_issued_token_type_is_never_cached_or_echoed() { + for (issued_type, expected_error) in [ + ( + Some("urn:ietf:params:oauth:token-type:id_token"), + "OpenAI returned an unsupported workload identity token type", + ), + ( + Some("echoed.secret.in.type"), + "OpenAI returned an unsupported workload identity token type", + ), + ( + None, + "OpenAI workload identity returned an invalid token response", + ), + ] { + let (server, url) = server(); + let worker = thread::spawn(move || { + let mut body = json!({ + "access_token": "openai.secret.jwt", "token_type": "Bearer", + "expires_in": 300, "expires_at": epoch_seconds() + 300, + }); + if let Some(issued_type) = issued_type { + body["issued_token_type"] = json!(issued_type); + } + recv_exchange(&server, "issuer.subject.1") + .respond(tiny_http::Response::from_string(body.to_string())) + .unwrap(); + }); + let auth = auth_at(&url); + assert_eq!(format!("{:#}", auth.header().unwrap_err()), expected_error); + let Auth::Federated { state, .. } = auth else { + unreachable!() + }; + assert!(state.lock().unwrap().cached.is_none()); + worker.join().unwrap(); + } +} + +#[test] +fn absolute_expiry_forces_renewal_and_discards_cached_credential_on_failure() { + let (server, url) = server(); + let worker = thread::spawn(move || { + serve_exchange( + &server, + "issuer.subject.1", + "openai.first.jwt", + /*expires_in*/ 300, + // The issuer's clock is ahead. Accept the token, but don't cache it + // past the shorter relative TTL, including across host suspension. + epoch_seconds() + 7200, + ); + serve_exchange( + &server, + "issuer.subject.2", + "openai.second.jwt", + /*expires_in*/ 300, + epoch_seconds() + 300, + ); + recv_exchange(&server, "issuer.subject.3") + .respond(tiny_http::Response::empty(tiny_http::StatusCode(503))) + .unwrap(); + }); + let auth = auth_at(&url); + assert_eq!( + auth.header().unwrap().value.to_str().unwrap(), + "Bearer openai.first.jwt" + ); + let Auth::Federated { state, .. } = &auth else { + unreachable!() + }; + { + let mut state = state.lock().unwrap(); + let cached = state.cached.as_mut().unwrap(); + // Model a suspend: wall time has crossed expiry but Instant has not advanced. + assert!(Instant::now() < cached.renew_at); + assert!(cached.expires_at <= SystemTime::now() + Duration::from_secs(300)); + cached.expires_at = UNIX_EPOCH; + } + assert_eq!( + auth.header().unwrap().value.to_str().unwrap(), + "Bearer openai.second.jwt" + ); + state.lock().unwrap().cached.as_mut().unwrap().expires_at = UNIX_EPOCH; + assert_eq!( + auth.header().unwrap_err().to_string(), + "OpenAI workload identity token request rejected (HTTP 503)" + ); + assert!(state.lock().unwrap().cached.is_none()); + worker.join().unwrap(); +} + +#[test] +fn invalid_subjects_are_rejected_before_exchange() { + struct InvalidSource(String); + impl SubjectTokenSource for InvalidSource { + fn subject_token<'a>(&'a mut self, _client: &'a TokenClient) -> SubjectTokenFuture<'a> { + Box::pin(async move { Ok(Zeroizing::new(self.0.clone())) }) + } + } + for token in [String::new(), "a".repeat(MAX_TOKEN_BYTES + 1)] { + let auth = Auth::federated( + InvalidSource(token), + "ip-test".to_owned(), + "sa-test".to_owned(), + ) + .unwrap(); + assert_eq!( + auth.header().unwrap_err().to_string(), + "workload identity source returned an invalid subject token" + ); + } +}