github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refresh rejected proxy credentials and back off failed renewals (#55299)

Browse Source 
## Why

The Responses API proxy could reuse a federated credential after an upstream `401`, and concurrent requests could each retry a failed token exchange.

## What changed

- Invalidate the rejected credential so the next request obtains a fresh token. Forward the `401` without replaying the request or waiting for an ongoing renewal; leave credentials intact on `403`.
- Track rejection per token issuance so a late `401` cannot invalidate a newer credential.
- Back off failed exchanges exponentially from 1 to 30 seconds, starting when the exchange fails. Reject requests during backoff without reusing the old token, and reset the delay after a successful exchange.

## Testing

Add tests for response forwarding without replay, late and concurrent rejections, shared renewal failures, backoff timing and reset, and credential-safe renewal errors without following redirects.

GitOrigin-RevId: eb1ed6848d74fdba1534919cace5126f4a0595cf
This commit is contained in:
ee
2026-10-19 15:08:50 +00:00
committed by copyberry
parent 8cac68328a
commit 54f99017f2
3 changed files with 537 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/responses-api-proxy/src/auth.rs
+75 -13
View File
@@ -1,11 +1,14 @@
//! Authenticates proxy requests using a locked API key or renewable OpenAI credentials.
//! Serializes exchanges and never reuses expired credentials on renewal errors.
//! Federated exchanges are serialized, fail closed on renewal errors, and never replay requests.
//! 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::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;
@@ -37,6 +40,8 @@ const MAX_RESPONSE_BYTES: usize = 64 * 1024;
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);
const INITIAL_RENEWAL_BACKOFF: Duration = Duration::from_secs(1);
const MAX_RENEWAL_BACKOFF: Duration = Duration::from_secs(30);
/// Supplies a fresh JWT whenever an OpenAI credential needs to be issued.
/// Acquisition stays asynchronous so callers can compose it with the exchange.
@@ -94,10 +99,21 @@ impl fmt::Debug for Auth {
}
}
/// The sensitive header used for an individual request.
/// The sensitive header and, for federation, the rejection state of this issuance.
/// In-flight requests share only their own issuance's flag, never a replacement's.
#[derive(Clone, Debug)]
pub(crate) struct AuthHeader {
pub(crate) value: HeaderValue,
rejected: Option<Arc<AtomicBool>>,
}
impl AuthHeader {
pub(crate) fn invalidate(&self) {
if let Some(rejected) = &self.rejected {
// This flag carries no other state; credentials remain mutex-protected.
rejected.store(true, Ordering::Relaxed);
}
}
}
impl Auth {
@@ -123,6 +139,7 @@ impl Auth {
provider,
service_account,
cached: None,
renewal_failure: None,
})),
})
}
@@ -146,12 +163,15 @@ impl Auth {
// Keep the existing locked, static API key allocation.
let mut header = HeaderValue::from_static(value);
header.set_sensitive(true);
Ok(AuthHeader { value: header })
Ok(AuthHeader {
value: header,
rejected: None,
})
}
Self::Federated { state } => state
.lock()
.map_err(|_| anyhow!("workload identity credential cache unavailable"))?
.header(),
.header(Instant::now),
}
}
}
@@ -164,6 +184,7 @@ pub(crate) struct FederatedAuth {
provider: String,
service_account: String,
cached: Option<CachedToken>,
renewal_failure: Option<RenewalFailure>,
}
struct CachedToken {
@@ -172,22 +193,60 @@ struct CachedToken {
expires_at: SystemTime,
}
struct RenewalFailure {
retry_at: Instant,
delay: Duration,
}
impl FederatedAuth {
fn header(&mut self) -> Result<AuthHeader> {
fn header(&mut self, clock: impl Fn() -> Instant) -> Result<AuthHeader> {
let now = clock();
if let Some(cached) = &self.cached
&& Instant::now() < cached.renew_at
&& now < cached.renew_at
// Monotonic time can stop while a host is suspended.
&& SystemTime::now() < cached.expires_at
&& cached
.header
.rejected
.as_ref()
.is_none_or(|rejected| !rejected.load(Ordering::Relaxed))
{
return Ok(cached.header.clone());
}
// Discard the old credential before a fallible exchange. Never reuse it
// when acquisition fails.
if let Some(failure) = &self.renewal_failure
&& now < failure.retry_at
{
// Keep neither the failed token nor an error that could hold source
// credentials. The initiating request reports the original safe error.
return Err(anyhow!(
"workload identity renewal temporarily unavailable; retry later"
));
}
// One exchange at a time (under Auth's mutex), including during startup.
// On failure, fail closed instead of reusing an expired/revoked token.
self.cached = None;
let cached = self.exchange()?;
let header = cached.header.clone();
self.cached = Some(cached);
Ok(header)
match self.exchange() {
Ok(cached) => {
let header = cached.header.clone();
self.cached = Some(cached);
self.renewal_failure = None;
Ok(header)
}
Err(err) => {
// Waiters share this failure instead of each retrying the token
// endpoints. Later requests may retry after a bounded delay.
let delay = self
.renewal_failure
.as_ref()
.map(|failure| (failure.delay * 2).min(MAX_RENEWAL_BACKOFF))
.unwrap_or(INITIAL_RENEWAL_BACKOFF);
self.renewal_failure = Some(RenewalFailure {
retry_at: clock() + delay,
delay,
});
Err(err)
}
}
}
fn exchange(&mut self) -> Result<CachedToken> {
@@ -272,7 +331,10 @@ impl FederatedAuth {
.checked_add(remaining - skew)
.ok_or_else(|| anyhow!("OpenAI returned an invalid token lifetime"))?;
Ok(CachedToken {
header: AuthHeader { value: header },
header: AuthHeader {
value: header,
rejected: Some(Arc::new(AtomicBool::new(/*v*/ false))),
},
renew_at,
expires_at,
})
codex-rs/responses-api-proxy/src/auth_tests.rs
+456
View File
@@ -2,12 +2,15 @@
use super::*;
use pretty_assertions::assert_eq;
use reqwest::blocking::Client;
use serde_json::json;
use std::sync::Arc;
use std::sync::Barrier;
use std::sync::mpsc;
use std::thread;
use tiny_http::Header;
use tiny_http::Server;
use tiny_http::StatusCode;
pub(crate) fn direct_token_client() -> TokenClient {
TokenClient(TokenClient::builder().build_direct().unwrap())
@@ -96,6 +99,22 @@ fn recv_exchange(server: &Server, subject: &str) -> tiny_http::Request {
req
}
fn recv_response_request(
server: &Server,
expected_token: &str,
expected_body: &str,
) -> tiny_http::Request {
let mut req = recv(server);
assert_eq!(req.url(), "/v1/responses");
assert_eq!(req.method(), &tiny_http::Method::Post);
assert!(req.headers().iter().any(|h| h.field.equiv("Authorization")
&& h.value.as_str() == format!("Bearer {expected_token}")));
let mut body = String::new();
req.as_reader().read_to_string(&mut body).unwrap();
assert_eq!(body, expected_body);
req
}
fn serve_exchange(server: &Server, subject: &str, token: &str, expires_in: u64, expires_at: u64) {
respond_exchange(
recv_exchange(server, subject),
@@ -373,6 +392,50 @@ fn absolute_expiry_forces_renewal_and_discards_cached_credential_on_failure() {
worker.join().unwrap();
}
#[test]
fn rejected_renewal_discards_old_credential_and_does_not_log_or_follow_redirect() {
let (server, url) = server();
let target = url.join("/credential-leak").unwrap();
let worker = thread::spawn(move || {
serve_exchange(
&server,
"issuer.subject.1",
"original.token",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
recv_exchange(&server, "issuer.subject.2")
.respond(
tiny_http::Response::from_string("issuer.subject.2 original.token")
.with_status_code(StatusCode(302))
.with_header(Header::from_bytes("Location", target.as_str()).unwrap()),
)
.unwrap();
assert!(
server
.recv_timeout(Duration::from_millis(200))
.unwrap()
.is_none()
);
});
let auth = auth_at(&url);
assert_eq!(
auth.header().unwrap().value.to_str().unwrap(),
"Bearer original.token"
);
let Auth::Federated { ref state, .. } = auth else {
unreachable!()
};
state.lock().unwrap().cached.as_mut().unwrap().renew_at = Instant::now();
let error = format!("{:#}", auth.header().unwrap_err());
assert!(error.contains("302"), "{error}");
for secret in ["issuer.subject.2", "original.token", url.as_str()] {
assert!(!error.contains(secret), "{error}");
}
assert!(state.lock().unwrap().cached.is_none());
worker.join().unwrap();
}
#[test]
fn invalid_subjects_are_rejected_before_exchange() {
struct InvalidSource(String);
@@ -394,3 +457,396 @@ fn invalid_subjects_are_rejected_before_exchange() {
);
}
}
#[test]
fn forwards_401_without_replay_then_reacquires_but_keeps_token_on_403() {
let (upstream, upstream_url) = server();
let auth = auth_at(&upstream_url);
let (proxy, proxy_url) = server();
let upstream_worker = thread::spawn(move || {
serve_exchange(
&upstream,
"issuer.subject.1",
"openai.first.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
recv_response_request(&upstream, "openai.first.jwt", "first")
.respond(
tiny_http::Response::from_string("rejected")
.with_status_code(StatusCode(401))
.with_header(Header::from_bytes("WWW-Authenticate", "Bearer").unwrap()),
)
.unwrap();
// No Responses replay or speculative exchange: the caller submits a
// distinct POST. A permission failure on that POST must not evict auth.
serve_exchange(
&upstream,
"issuer.subject.2",
"openai.second.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
recv_response_request(&upstream, "openai.second.jwt", "second")
.respond(
tiny_http::Response::from_string("forbidden").with_status_code(StatusCode(403)),
)
.unwrap();
recv_response_request(&upstream, "openai.second.jwt", "third")
.respond(tiny_http::Response::from_string("success"))
.unwrap();
});
let proxy_worker = thread::spawn(move || {
let client = Client::builder()
.no_proxy()
.timeout(Duration::from_secs(5))
.redirect(auth.redirect_policy())
.build()
.unwrap();
let config = crate::ForwardConfig {
host_header: HeaderValue::from_str(upstream_url.authority()).unwrap(),
upstream_url: upstream_url.join("/v1/responses").unwrap(),
};
for _ in 0..3 {
crate::forward_request(
&client,
&auth,
&config,
/*dump_dir*/ None,
recv(&proxy),
)
.unwrap();
}
});
let client = Client::builder()
.no_proxy()
.timeout(Duration::from_secs(5))
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap();
let response_url = proxy_url.join("/v1/responses").unwrap();
let response = client
.post(response_url.clone())
.body("first")
.send()
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::UNAUTHORIZED);
assert_eq!(
response.headers().get("www-authenticate").unwrap(),
"Bearer"
);
assert_eq!(response.text().unwrap(), "rejected");
let response = client
.post(response_url.clone())
.body("second")
.send()
.unwrap();
assert_eq!(response.status(), reqwest::StatusCode::FORBIDDEN);
assert_eq!(response.text().unwrap(), "forbidden");
let response = client.post(response_url).body("third").send().unwrap();
assert_eq!(response.status(), reqwest::StatusCode::OK);
assert_eq!(response.text().unwrap(), "success");
proxy_worker.join().unwrap();
upstream_worker.join().unwrap();
}
#[test]
fn late_rejection_does_not_invalidate_a_newer_credential() {
let (server, url) = server();
let worker = thread::spawn(move || {
serve_exchange(
&server,
"issuer.subject.1",
"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 auth = auth_at(&url);
let original = auth.header().unwrap();
original.invalidate();
let renewed = auth.header().unwrap();
assert_eq!(renewed.value.to_str().unwrap(), "Bearer openai.second.jwt");
worker.join().unwrap();
// Model an in-flight request made with the old token completing after
// the refresh. The server is closed, so an accidental third exchange fails.
original.invalidate();
assert_eq!(auth.header().unwrap().value, renewed.value);
}
#[test]
fn forwards_401_before_an_inflight_renewal_finishes() {
let (upstream, upstream_url) = server();
let auth = Arc::new(auth_at(&upstream_url));
let acquiring = Arc::clone(&auth);
let first = thread::spawn(move || acquiring.header().unwrap());
serve_exchange(
&upstream,
"issuer.subject.1",
"openai.first.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
let original = first.join().unwrap();
let (proxy, proxy_url) = server();
let forwarding = Arc::clone(&auth);
let config = crate::ForwardConfig {
host_header: HeaderValue::from_str(upstream_url.authority()).unwrap(),
upstream_url: upstream_url.join("/v1/responses").unwrap(),
};
let proxy_worker = thread::spawn(move || {
let client = Client::builder()
.no_proxy()
.redirect(forwarding.redirect_policy())
.build()
.unwrap();
crate::forward_request(
&client,
&forwarding,
&config,
/*dump_dir*/ None,
recv(&proxy),
)
.unwrap();
});
let (sent, received) = mpsc::sync_channel(1);
let caller = thread::spawn(move || {
let response = Client::builder()
.no_proxy()
.build()
.unwrap()
.post(proxy_url.join("/v1/responses").unwrap())
.body("first")
.send()
.unwrap();
sent.send((response.status(), response.text().unwrap()))
.unwrap();
});
// The request is already at upstream with the first credential. Hold its
// response until a concurrent rejection has started a new exchange.
let response_request = recv_response_request(&upstream, "openai.first.jwt", "first");
original.invalidate();
let acquiring = Arc::clone(&auth);
let renewal = thread::spawn(move || acquiring.header().unwrap());
let exchange = recv_exchange(&upstream, "issuer.subject.2");
response_request
.respond(tiny_http::Response::from_string("rejected").with_status_code(StatusCode(401)))
.unwrap();
// This must arrive before we let the exchange finish and release the mutex.
// Always release it before asserting so a regression cannot strand workers.
let response = received.recv_timeout(Duration::from_secs(5));
respond_exchange(
exchange,
"openai.second.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
let renewed = renewal.join().unwrap();
caller.join().unwrap();
proxy_worker.join().unwrap();
assert_eq!(
response.unwrap(),
(reqwest::StatusCode::UNAUTHORIZED, "rejected".to_string())
);
assert_eq!(auth.header().unwrap().value, renewed.value);
}
fn header_at(auth: &Auth, now: Instant) -> Result<AuthHeader> {
let Auth::Federated { state } = auth else {
unreachable!()
};
state.lock().unwrap().header(|| now)
}
// Advance to the recorded retry deadline without waiting on real time.
fn elapse_backoff(auth: &Auth, now: &mut Instant) -> Duration {
let Auth::Federated { state } = auth else {
unreachable!()
};
let state = state.lock().unwrap();
assert!(state.cached.is_none());
let failure = state.renewal_failure.as_ref().unwrap();
*now = failure.retry_at;
failure.delay
}
#[test]
fn renewal_backoff_starts_after_a_slow_exchange_fails() {
let (server, url) = server();
let started = Instant::now();
let failed = started + Duration::from_secs(10);
let clock = Arc::new(Mutex::new(started));
let upstream_clock = Arc::clone(&clock);
let worker = thread::spawn(move || {
serve_exchange(
&server,
"issuer.subject.1",
"openai.first.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
let exchange = recv_exchange(&server, "issuer.subject.2");
// Simulate time spent waiting on the failed renewal, longer than its cooldown.
*upstream_clock.lock().unwrap() = failed;
exchange
.respond(
tiny_http::Response::from_string("temporarily unavailable")
.with_status_code(StatusCode(503)),
)
.unwrap();
});
let auth = auth_at(&url);
auth.header().unwrap().invalidate();
let Auth::Federated { state } = &auth else {
unreachable!()
};
let error = state
.lock()
.unwrap()
.header(|| *clock.lock().unwrap())
.unwrap_err();
worker.join().unwrap();
assert_eq!(
error.to_string(),
"OpenAI workload identity token request rejected (HTTP 503)"
);
let retry_at = failed + Duration::from_secs(1);
assert_eq!(
state
.lock()
.unwrap()
.renewal_failure
.as_ref()
.unwrap()
.retry_at,
retry_at
);
assert_eq!(
header_at(&auth, retry_at - Duration::from_nanos(1))
.unwrap_err()
.to_string(),
"workload identity renewal temporarily unavailable; retry later"
);
}
#[test]
fn concurrent_renewal_failure_backs_off_caps_and_resets_after_success() {
let (server, url) = server();
let worker = thread::spawn(move || {
serve_exchange(
&server,
"issuer.subject.1",
"openai.first.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
for subject in 2..=8 {
recv_exchange(&server, &format!("issuer.subject.{subject}"))
.respond(
tiny_http::Response::from_string("private-issuer-token openai.first.jwt")
.with_status_code(StatusCode(503)),
)
.unwrap();
}
serve_exchange(
&server,
"issuer.subject.9",
"openai.recovered.jwt",
/*expires_in*/ 300,
epoch_seconds() + 300,
);
recv_exchange(&server, "issuer.subject.10")
.respond(
tiny_http::Response::from_string("private-issuer-token")
.with_status_code(StatusCode(503)),
)
.unwrap();
});
let auth = Arc::new(auth_at(&url));
let original = auth.header().unwrap();
let Auth::Federated { state, .. } = auth.as_ref() else {
unreachable!()
};
let mut now = Instant::now();
state.lock().unwrap().cached.as_mut().unwrap().renew_at = now;
let barrier = Arc::new(Barrier::new(8));
// Start every participant before joining any of them.
let mut callers = Vec::with_capacity(8);
for _ in 0..8 {
let auth = Arc::clone(&auth);
let barrier = Arc::clone(&barrier);
callers.push(thread::spawn(move || {
barrier.wait();
format!("{:#}", header_at(&auth, now).unwrap_err())
}));
}
let errors: Vec<_> = callers
.into_iter()
.map(|caller| caller.join().unwrap())
.collect();
assert_eq!(errors.iter().filter(|e| e.contains("503")).count(), 1);
assert_eq!(
errors.iter().filter(|e| e.contains("retry later")).count(),
7
);
for error in &errors {
for secret in ["private-issuer-token", "openai.first.jwt", url.as_str()] {
assert!(!error.contains(secret), "{error}");
}
}
// An older 401 arriving during backoff cannot trigger another renewal.
original.invalidate();
assert!(
header_at(&auth, now)
.unwrap_err()
.to_string()
.contains("retry later")
);
for seconds in [1, 2, 4, 8, 16, 30] {
assert_eq!(
elapse_backoff(&auth, &mut now),
Duration::from_secs(seconds)
);
assert!(
header_at(&auth, now)
.unwrap_err()
.to_string()
.contains("503")
);
assert!(
header_at(&auth, now)
.unwrap_err()
.to_string()
.contains("retry later")
);
}
assert_eq!(elapse_backoff(&auth, &mut now), Duration::from_secs(30));
let renewed = header_at(&auth, now).unwrap();
assert_eq!(
renewed.value.to_str().unwrap(),
"Bearer openai.recovered.jwt"
);
assert!(state.lock().unwrap().renewal_failure.is_none());
renewed.invalidate();
assert!(
header_at(&auth, now)
.unwrap_err()
.to_string()
.contains("503")
);
assert!(
header_at(&auth, now)
.unwrap_err()
.to_string()
.contains("retry later")
);
assert_eq!(elapse_backoff(&auth, &mut now), Duration::from_secs(1));
worker.join().unwrap();
}
codex-rs/responses-api-proxy/src/lib.rs
+6 -1
View File
@@ -226,7 +226,7 @@ fn forward_request(
return Err(err);
}
};
headers.insert(AUTHORIZATION, auth_header.value);
headers.insert(AUTHORIZATION, auth_header.value.clone());
headers.insert(HOST, config.host_header.clone());
@@ -242,6 +242,11 @@ fn forward_request(
// implements `Read`, so we can use it directly as the body of the
// `tiny_http::Response`.
let status = upstream_resp.status();
if status == reqwest::StatusCode::UNAUTHORIZED {
// Do not replay a Responses POST. A subsequent request can obtain a
// fresh credential, but a late 401 must not evict one renewed in-flight.
auth_header.invalidate();
}
let mut response_headers = Vec::new();
for (name, value) in upstream_resp.headers().iter() {
// Skip headers that tiny_http manages itself.