fix: bound Codex OAuth refresh bodies and deadline

This commit is contained in:
2026-10-12 02:30:16 +09:00
parent 98023b9872
commit 9e1de93a04
5 changed files with 571 additions and 62 deletions
+1
View File
@@ -37,6 +37,7 @@ chrono = { version = "0.4", default-features = false, features = ["serde", "cloc
agen-macros = { workspace = true } agen-macros = { workspace = true }
[dev-dependencies] [dev-dependencies]
tokio = { workspace = true, features = ["test-util"] }
clap = { version = "4.5", features = ["derive", "env"] } clap = { version = "4.5", features = ["derive", "env"] }
tempfile = { workspace = true } tempfile = { workspace = true }
dotenv = "0.15" dotenv = "0.15"
+20 -5
View File
@@ -108,13 +108,21 @@ pub async fn load(path: &Path) -> Result<AuthSnapshot, CodexAuthError> {
/// ///
/// 並行する Codex CLI / 別プロセスが先に refresh していた場合の /// 並行する Codex CLI / 別プロセスが先に refresh していた場合の
/// fields を保護するため、書込前に再 load して merge する。 /// fields を保護するため、書込前に再 load して merge する。
/// Account / refresh token must still match the request's input: unrelated
/// fields can be merged, but another login or token rotation must not be overwritten.
pub async fn persist_refreshed( pub async fn persist_refreshed(
path: &Path, path: &Path,
refresh_input: &AuthSnapshot,
new_id_token: Option<String>, new_id_token: Option<String>,
new_access_token: Option<String>, new_access_token: Option<String>,
new_refresh_token: Option<String>, new_refresh_token: Option<String>,
) -> Result<AuthSnapshot, CodexAuthError> { ) -> Result<AuthSnapshot, CodexAuthError> {
let mut current = load(path).await?; let mut current = load(path).await?;
if current.account_id != refresh_input.account_id
|| current.refresh_token != refresh_input.refresh_token
{
return Ok(current);
}
let raw = &mut current.raw; let raw = &mut current.raw;
let tokens = raw let tokens = raw
.get_mut("tokens") .get_mut("tokens")
@@ -243,10 +251,16 @@ mod tests {
"agent_identity":{"workspace_id":"w","agent_runtime_id":"r","agent_private_key":"k","registered_at":"x"} "agent_identity":{"workspace_id":"w","agent_runtime_id":"r","agent_private_key":"k","registered_at":"x"}
}"#, }"#,
); );
let updated = let input = load(&path).await.unwrap();
persist_refreshed(&path, None, Some("new-acc".into()), Some("new-ref".into())) let updated = persist_refreshed(
.await &path,
.unwrap(); &input,
None,
Some("new-acc".into()),
Some("new-ref".into()),
)
.await
.unwrap();
assert_eq!(updated.access_token, "new-acc"); assert_eq!(updated.access_token, "new-acc");
assert_eq!(updated.refresh_token, "new-ref"); assert_eq!(updated.refresh_token, "new-ref");
// 未知フィールド agent_identity が保たれる // 未知フィールド agent_identity が保たれる
@@ -266,7 +280,8 @@ mod tests {
); );
// 既存ファイルを 644 に変えてから persist → 600 に直るか // 既存ファイルを 644 に変えてから persist → 600 に直るか
std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap(); std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)).unwrap();
persist_refreshed(&path, None, Some("a2".into()), None) let input = load(&path).await.unwrap();
persist_refreshed(&path, &input, None, Some("a2".into()), None)
.await .await
.unwrap(); .unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777; let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
+22 -2
View File
@@ -17,6 +17,8 @@ mod auth_json;
mod error; mod error;
mod jwt; mod jwt;
mod refresh; mod refresh;
#[cfg(test)]
mod refresh_tests;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::Arc;
@@ -40,6 +42,9 @@ pub struct CodexAuthProvider {
auth_path: PathBuf, auth_path: PathBuf,
refresh_endpoint: String, refresh_endpoint: String,
http: reqwest::Client, http: reqwest::Client,
refresh_limits: refresh::RefreshLimits,
#[cfg(test)]
before_guarded_reload: Option<Arc<tokio::sync::Barrier>>,
state: Arc<Mutex<State>>, state: Arc<Mutex<State>>,
} }
@@ -76,6 +81,9 @@ impl CodexAuthProvider {
auth_path: codex_home.join("auth.json"), auth_path: codex_home.join("auth.json"),
refresh_endpoint: refresh::REFRESH_URL.to_string(), refresh_endpoint: refresh::REFRESH_URL.to_string(),
http: reqwest::Client::new(), http: reqwest::Client::new(),
refresh_limits: refresh::RefreshLimits::default(),
#[cfg(test)]
before_guarded_reload: None,
state: Arc::new(Mutex::new(State::default())), state: Arc::new(Mutex::new(State::default())),
} }
} }
@@ -108,6 +116,11 @@ impl CodexAuthProvider {
// 3. Refresh 直前に再 load し account_id 一致を確認(guarded reload)。 // 3. Refresh 直前に再 load し account_id 一致を確認(guarded reload)。
// 一致しなければ他プロセスが先に更新済 → 自分は refresh しない // 一致しなければ他プロセスが先に更新済 → 自分は refresh しない
#[cfg(test)]
if let Some(barrier) = &self.before_guarded_reload {
barrier.wait().await;
barrier.wait().await;
}
let pre_refresh = auth_json::load(&self.auth_path).await?; let pre_refresh = auth_json::load(&self.auth_path).await?;
if !is_stale(&pre_refresh) { if !is_stale(&pre_refresh) {
state.cached = Some(pre_refresh.clone()); state.cached = Some(pre_refresh.clone());
@@ -124,12 +137,14 @@ impl CodexAuthProvider {
&self.http, &self.http,
&self.refresh_endpoint, &self.refresh_endpoint,
&pre_refresh.refresh_token, &pre_refresh.refresh_token,
&self.refresh_limits,
) )
.await?; .await?;
// 5. 書き戻し(書込前に再 load + merge は persist_refreshed 内で実施) // 5. 書き戻し(書込前に再 load + merge は persist_refreshed 内で実施)
let new_snap = auth_json::persist_refreshed( let new_snap = auth_json::persist_refreshed(
&self.auth_path, &self.auth_path,
&pre_refresh,
refreshed.id_token, refreshed.id_token,
refreshed.access_token, refreshed.access_token,
refreshed.refresh_token, refreshed.refresh_token,
@@ -210,13 +225,18 @@ mod tests {
use wiremock::matchers::{method, path as url_path}; use wiremock::matchers::{method, path as url_path};
use wiremock::{Mock, MockServer, ResponseTemplate}; use wiremock::{Mock, MockServer, ResponseTemplate};
fn make_jwt(payload: serde_json::Value) -> String { pub(super) fn make_jwt(payload: serde_json::Value) -> String {
let p = base64::engine::general_purpose::URL_SAFE_NO_PAD let p = base64::engine::general_purpose::URL_SAFE_NO_PAD
.encode(serde_json::to_vec(&payload).unwrap()); .encode(serde_json::to_vec(&payload).unwrap());
format!("h.{p}.s") format!("h.{p}.s")
} }
fn write_auth(dir: &std::path::Path, exp: i64, fedramp: bool, refresh: &str) -> PathBuf { pub(super) fn write_auth(
dir: &std::path::Path,
exp: i64,
fedramp: bool,
refresh: &str,
) -> PathBuf {
let id_token = make_jwt(json!({ let id_token = make_jwt(json!({
"https://api.openai.com/auth": { "https://api.openai.com/auth": {
"chatgpt_account_id": "acc-xyz", "chatgpt_account_id": "acc-xyz",
+97 -55
View File
@@ -10,7 +10,26 @@ use super::error::{CodexAuthError, PermanentReason};
pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann";
pub const REFRESH_URL: &str = "https://auth.openai.com/oauth/token"; pub const REFRESH_URL: &str = "https://auth.openai.com/oauth/token";
pub const DEFAULT_REFRESH_TIMEOUT: Duration = Duration::from_secs(30); /// Limits cover send, both response bodies, and decode, not auth file I/O.
/// OAuth responses contain only a few tokens; 64 KiB leaves ample JWT headroom.
#[derive(Clone)]
pub(super) struct RefreshLimits {
pub timeout: Duration,
pub max_body_bytes: usize,
#[cfg(test)]
pub body_started: Option<std::sync::Arc<tokio::sync::Notify>>,
}
impl Default for RefreshLimits {
fn default() -> Self {
Self {
timeout: Duration::from_secs(30),
max_body_bytes: 64 * 1024,
#[cfg(test)]
body_started: None,
}
}
}
#[derive(Serialize)] #[derive(Serialize)]
struct RefreshRequest<'a> { struct RefreshRequest<'a> {
@@ -37,54 +56,94 @@ pub async fn request_refresh(
client: &reqwest::Client, client: &reqwest::Client,
endpoint: &str, endpoint: &str,
refresh_token: &str, refresh_token: &str,
limits: &RefreshLimits,
) -> Result<RefreshResponse, CodexAuthError> {
let deadline = tokio::time::Instant::now() + limits.timeout;
let result = tokio::time::timeout_at(
deadline,
request_refresh_bounded(client, endpoint, refresh_token, limits),
)
.await;
// JSON decode is synchronous. The byte bound limits its work, and this
// post-check prevents a decode that crosses the deadline from being persisted.
if tokio::time::Instant::now() >= deadline {
return Err(CodexAuthError::RefreshTransient(
"oauth_token_refresh timed out".into(),
));
}
result.map_err(|_| CodexAuthError::RefreshTransient("oauth_token_refresh timed out".into()))?
}
async fn request_refresh_bounded(
client: &reqwest::Client,
endpoint: &str,
refresh_token: &str,
limits: &RefreshLimits,
) -> Result<RefreshResponse, CodexAuthError> { ) -> Result<RefreshResponse, CodexAuthError> {
let body = RefreshRequest { let body = RefreshRequest {
client_id: CLIENT_ID, client_id: CLIENT_ID,
grant_type: "refresh_token", grant_type: "refresh_token",
refresh_token, refresh_token,
}; };
let response = response_with_timeout( let mut response = client
client .post(endpoint)
.post(endpoint) .header("Content-Type", "application/json")
.header("Content-Type", "application/json") .json(&body)
.json(&body) .send()
.send(), .await
DEFAULT_REFRESH_TIMEOUT, .map_err(|_| CodexAuthError::RefreshTransient("send failed".into()))?;
)
.await?;
let status = response.status(); let status = response.status();
if status.is_success() { if response
response .content_length()
.json::<RefreshResponse>() .is_some_and(|len| len > limits.max_body_bytes as u64)
.await {
.map_err(|e| CodexAuthError::RefreshTransient(format!("parse response: {e}"))) return Err(body_too_large());
} else { }
let body = response.text().await.unwrap_or_default(); #[cfg(test)]
if status == reqwest::StatusCode::UNAUTHORIZED { if let Some(started) = &limits.body_started {
let (reason, message) = classify_permanent(&body); started.notify_one();
Err(CodexAuthError::RefreshPermanent { reason, message }) }
} else { let mut bytes = Vec::new();
Err(CodexAuthError::RefreshTransient(format!( while let Some(chunk) = response
"{status}: {body}" .chunk()
))) .await
.map_err(|_| CodexAuthError::RefreshTransient("read response body failed".into()))?
{
// Content-Length can be absent or untrustworthy. Bound actual bytes too,
// before copying a chunk into the accumulated body.
if chunk.len() > limits.max_body_bytes.saturating_sub(bytes.len()) {
return Err(body_too_large());
} }
bytes.extend_from_slice(&chunk);
}
if status.is_success() {
let tokens: RefreshResponse = serde_json::from_slice(&bytes).map_err(|_| {
// Do not expose untrusted body values (which can contain secrets).
CodexAuthError::RefreshTransient("decode response failed".into())
})?;
if tokens.access_token.as_deref().is_none_or(str::is_empty)
|| tokens.id_token.as_deref().is_some_and(str::is_empty)
|| tokens.refresh_token.as_deref().is_some_and(str::is_empty)
{
return Err(CodexAuthError::RefreshTransient(
"incomplete token response".into(),
));
}
Ok(tokens)
} else if status == reqwest::StatusCode::UNAUTHORIZED {
let (reason, message) = classify_permanent(&String::from_utf8_lossy(&bytes));
Err(CodexAuthError::RefreshPermanent { reason, message })
} else {
// Error bodies are consumed within the same limits but never echoed.
Err(CodexAuthError::RefreshTransient(format!(
"refresh endpoint returned {status}"
)))
} }
} }
async fn response_with_timeout( fn body_too_large() -> CodexAuthError {
future: impl std::future::Future<Output = Result<reqwest::Response, reqwest::Error>>, CodexAuthError::RefreshTransient("response body exceeds size limit".into())
timeout: Duration,
) -> Result<reqwest::Response, CodexAuthError> {
tokio::time::timeout(timeout, future)
.await
.map_err(|_| {
CodexAuthError::RefreshTransient(format!(
"oauth_token_refresh timed out after {}s",
timeout.as_secs()
))
})?
.map_err(|e| CodexAuthError::RefreshTransient(format!("send: {e}")))
} }
fn classify_permanent(body: &str) -> (PermanentReason, String) { fn classify_permanent(body: &str) -> (PermanentReason, String) {
@@ -99,7 +158,7 @@ fn classify_permanent(body: &str) -> (PermanentReason, String) {
PermanentReason::Expired => "Your refresh token has expired".to_string(), PermanentReason::Expired => "Your refresh token has expired".to_string(),
PermanentReason::Reused => "Your refresh token was already used".to_string(), PermanentReason::Reused => "Your refresh token was already used".to_string(),
PermanentReason::Revoked => "Your refresh token was revoked".to_string(), PermanentReason::Revoked => "Your refresh token was revoked".to_string(),
PermanentReason::Other => format!("Unknown 401 from refresh endpoint: {body}"), PermanentReason::Other => "Unknown 401 from refresh endpoint".to_string(),
}; };
(reason, message) (reason, message)
} }
@@ -126,23 +185,6 @@ fn extract_error_code(body: &str) -> Option<String> {
mod tests { mod tests {
use super::*; use super::*;
#[tokio::test]
async fn refresh_response_timeout_is_transient() {
let err = match response_with_timeout(
std::future::pending::<Result<reqwest::Response, reqwest::Error>>(),
Duration::from_millis(5),
)
.await
{
Ok(_) => panic!("expected refresh timeout"),
Err(err) => err,
};
assert!(
matches!(err, CodexAuthError::RefreshTransient(message) if message.contains("timed out"))
);
}
#[test] #[test]
fn classify_expired() { fn classify_expired() {
let body = r#"{"error":{"code":"refresh_token_expired"}}"#; let body = r#"{"error":{"code":"refresh_token_expired"}}"#;
@@ -0,0 +1,431 @@
//! Real AuthProvider calls against owned auth files and a loopback HTTP peer.
//! Gates place callers at response-body / reload boundaries; paused Tokio time
//! advances the injected HTTP deadline only after headers have been received.
use super::tests::{make_jwt, write_auth};
use super::*;
use serde_json::{Value, json};
use std::time::Duration as StdDuration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::{Notify, mpsc};
use tokio::task::{JoinHandle, JoinSet};
struct Reply {
status: u16,
body: Vec<u8>,
chunked: bool,
gate: Option<Arc<Notify>>,
truncate: bool,
}
impl Reply {
fn body(status: u16, body: impl Into<Vec<u8>>) -> Self {
Self {
status,
body: body.into(),
chunked: false,
gate: None,
truncate: false,
}
}
fn gated(status: u16, body: impl Into<Vec<u8>>, gate: &Arc<Notify>) -> Self {
Self {
gate: Some(gate.clone()),
..Self::body(status, body)
}
}
}
struct Stub {
endpoint: String,
requests: mpsc::UnboundedReceiver<Value>,
task: JoinHandle<()>,
}
impl Drop for Stub {
fn drop(&mut self) {
// Dropping the server's JoinSet also aborts stalled connection handlers.
self.task.abort();
}
}
impl Stub {
async fn start(replies: Vec<Reply>) -> Self {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let endpoint = format!("http://{}/oauth/token", listener.local_addr().unwrap());
let (tx, requests) = mpsc::unbounded_channel();
let task = tokio::spawn(async move {
let mut handlers = JoinSet::new();
for reply in replies {
let (mut stream, _) = listener.accept().await.unwrap();
let tx = tx.clone();
handlers.spawn(async move {
let mut request = Vec::new();
let header_end = loop {
let byte = stream.read_u8().await.unwrap();
request.push(byte);
if request.ends_with(b"\r\n\r\n") { break request.len(); }
assert!(request.len() < 16 * 1024, "fixture request headers too large");
};
let headers = std::str::from_utf8(&request).unwrap();
assert!(headers.starts_with("POST /oauth/token HTTP/1.1\r\n"));
let len: usize = headers.lines().find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length").then(|| value.trim().parse().unwrap())
}).unwrap();
request.resize(header_end + len, 0);
stream.read_exact(&mut request[header_end..]).await.unwrap();
tx.send(serde_json::from_slice(&request[header_end..]).unwrap()).unwrap();
let framing = if reply.chunked {
"Transfer-Encoding: chunked".to_string()
} else {
format!("Content-Length: {}", reply.body.len())
};
stream.write_all(format!(
"HTTP/1.1 {} Fixture\r\nContent-Type: application/json\r\n{framing}\r\nConnection: close\r\n\r\n", reply.status
).as_bytes()).await.unwrap();
if let Some(gate) = reply.gate { gate.notified().await; }
let body = if reply.truncate { &reply.body[..reply.body.len() / 2] } else { &reply.body[..] };
if reply.chunked {
for chunk in body.chunks(512) {
let _ = stream.write_all(format!("{:x}\r\n", chunk.len()).as_bytes()).await;
let _ = stream.write_all(chunk).await;
let _ = stream.write_all(b"\r\n").await;
}
let _ = stream.write_all(b"0\r\n\r\n").await;
} else {
let _ = stream.write_all(body).await;
}
});
}
while let Some(result) = handlers.join_next().await {
result.unwrap();
}
});
Self {
endpoint,
requests,
task,
}
}
}
fn rotated_body() -> Vec<u8> {
serde_json::to_vec(&json!({
"access_token": make_jwt(json!({"exp": Utc::now().timestamp() + 3600})),
"refresh_token": "fixture-rotated-secret",
}))
.unwrap()
}
fn provider(dir: &std::path::Path, stub: &Stub, started: &Arc<Notify>) -> Arc<CodexAuthProvider> {
let mut provider = CodexAuthProvider::new(dir.to_path_buf())
.with_refresh_endpoint(&stub.endpoint)
.with_http_client(reqwest::Client::builder().no_proxy().build().unwrap());
provider.refresh_limits = refresh::RefreshLimits {
timeout: StdDuration::from_secs(10),
max_body_bytes: 1024,
body_started: Some(started.clone()),
};
Arc::new(provider)
}
async fn bounded<F: std::future::Future>(future: F) -> F::Output {
tokio::time::timeout(StdDuration::from_secs(5), future)
.await
.expect("fixture/caller did not progress")
}
fn assert_rotated(path: &std::path::Path, headers: &[(HeaderName, HeaderValue)]) {
let value: Value = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
assert_eq!(value["tokens"]["refresh_token"], "fixture-rotated-secret");
assert_eq!(value["fixture_extra"], "preserve me");
let auth = headers
.iter()
.find(|(name, _)| name == "authorization")
.unwrap();
assert_eq!(
auth.1.to_str().unwrap(),
format!(
"Bearer {}",
value["tokens"]["access_token"].as_str().unwrap()
)
);
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
assert_eq!(
std::fs::metadata(path).unwrap().permissions().mode() & 0o777,
0o600
);
}
}
fn stale_auth(dir: &std::path::Path) -> PathBuf {
let path = write_auth(
dir,
Utc::now().timestamp() - 60,
false,
"fixture-old-secret",
);
let mut value: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
value["fixture_extra"] = json!("preserve me");
std::fs::write(&path, serde_json::to_vec(&value).unwrap()).unwrap();
path
}
#[tokio::test]
async fn stalled_success_and_error_bodies_timeout_without_persist_and_release_waiting_caller() {
for status in [200, 401, 503] {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let original = std::fs::read(&path).unwrap();
let gate = Arc::new(Notify::new());
let started = Arc::new(Notify::new());
let mut stub = Stub::start(vec![
Reply::gated(status, b"unfinished".to_vec(), &gate),
Reply::body(200, rotated_body()),
])
.await;
let provider = provider(dir.path(), &stub, &started);
let first = tokio::spawn({
let provider = provider.clone();
async move { provider.headers().await }
});
bounded(started.notified()).await; // provider has received headers and is consuming the body
let request = bounded(stub.requests.recv()).await.unwrap();
assert_eq!(request["refresh_token"], "fixture-old-secret");
let second = provider.headers();
tokio::pin!(second);
assert!(futures::poll!(&mut second).is_pending());
assert!(provider.state.try_lock().is_err());
assert!(
stub.requests.try_recv().is_err(),
"waiter must not refresh concurrently"
);
tokio::time::pause();
tokio::time::advance(StdDuration::from_secs(11)).await;
let error = first.await.unwrap().unwrap_err();
tokio::time::resume();
assert!(error.to_string().contains("timed out"), "{error}");
assert_eq!(std::fs::read(&path).unwrap(), original);
let headers = bounded(second).await.unwrap();
assert_rotated(&path, &headers);
let retry = bounded(stub.requests.recv()).await.unwrap();
assert_eq!(retry["refresh_token"], "fixture-old-secret");
}
}
#[tokio::test]
async fn oversized_malformed_and_incomplete_bodies_do_not_persist_and_allow_later_refresh() {
let mut cases = vec![
(Reply::body(200, vec![b'x'; 1025]), "size limit"),
(
Reply {
chunked: true,
..Reply::body(200, vec![b'x'; 1025])
},
"size limit",
),
(
Reply {
chunked: true,
..Reply::body(401, vec![b'x'; 1025])
},
"size limit",
),
(Reply::body(503, vec![b'x'; 1025]), "size limit"),
(
Reply::body(
200,
br#"{"access_token": {"secret":"fixture-body-secret"}}"#.to_vec(),
),
"decode",
),
(Reply::body(200, b"{".to_vec()), "decode"),
(Reply::body(200, b"{}".to_vec()), "incomplete"),
(
Reply::body(
200,
br#"{"access_token":"","refresh_token":"fixture-body-secret"}"#.to_vec(),
),
"incomplete",
),
(
Reply {
truncate: true,
..Reply::body(200, rotated_body())
},
"read response body",
),
(
Reply {
truncate: true,
..Reply::body(401, b"fixture-body-secret".to_vec())
},
"read response body",
),
(
Reply::body(401, b"fixture-body-secret".to_vec()),
"codex login",
),
(Reply::body(503, b"fixture-body-secret".to_vec()), "503"),
];
for (reply, expected) in cases.drain(..) {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let original = std::fs::read(&path).unwrap();
let started = Arc::new(Notify::new());
let mut stub = Stub::start(vec![reply, Reply::body(200, rotated_body())]).await;
let provider = provider(dir.path(), &stub, &started);
let error = bounded(provider.headers()).await.unwrap_err();
let message = error.to_string();
assert!(message.contains(expected), "expected {expected}: {message}");
assert!(!message.contains("fixture-body-secret"));
assert!(!message.contains("fixture-old-secret"));
assert_eq!(std::fs::read(&path).unwrap(), original);
let headers = bounded(provider.headers()).await.unwrap();
assert_rotated(&path, &headers);
for _ in 0..2 {
assert_eq!(
bounded(stub.requests.recv()).await.unwrap()["refresh_token"],
"fixture-old-secret"
);
}
}
}
#[tokio::test]
async fn cancelled_body_wait_leaves_auth_unchanged_and_releases_waiting_caller() {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let original = std::fs::read(&path).unwrap();
let started = Arc::new(Notify::new());
let gate = Arc::new(Notify::new());
let mut stub = Stub::start(vec![
Reply::gated(200, rotated_body(), &gate),
Reply::body(200, rotated_body()),
])
.await;
let provider = provider(dir.path(), &stub, &started);
let first = tokio::spawn({
let provider = provider.clone();
async move { provider.headers().await }
});
bounded(started.notified()).await;
bounded(stub.requests.recv()).await.unwrap();
let second = provider.headers();
tokio::pin!(second);
assert!(futures::poll!(&mut second).is_pending());
assert!(stub.requests.try_recv().is_err());
first.abort();
assert!(first.await.unwrap_err().is_cancelled());
assert_eq!(std::fs::read(&path).unwrap(), original);
assert_rotated(&path, &bounded(second).await.unwrap());
assert_eq!(
bounded(stub.requests.recv()).await.unwrap()["refresh_token"],
"fixture-old-secret"
);
}
#[tokio::test]
async fn shared_callers_refresh_once_and_use_persisted_rotated_tokens() {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let started = Arc::new(Notify::new());
let gate = Arc::new(Notify::new());
let mut stub = Stub::start(vec![Reply::gated(200, rotated_body(), &gate)]).await;
let provider = provider(dir.path(), &stub, &started);
let first = tokio::spawn({
let provider = provider.clone();
async move { provider.headers().await }
});
bounded(started.notified()).await;
let second = provider.headers();
tokio::pin!(second);
assert!(futures::poll!(&mut second).is_pending());
assert_eq!(
bounded(stub.requests.recv()).await.unwrap()["refresh_token"],
"fixture-old-secret"
);
let mut external: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
external["concurrent_extra"] = json!("merge me");
std::fs::write(&path, serde_json::to_vec(&external).unwrap()).unwrap();
gate.notify_one();
let first_headers = bounded(first).await.unwrap().unwrap();
let second_headers = bounded(second).await.unwrap();
assert_eq!(first_headers, second_headers);
assert_rotated(&path, &second_headers);
let persisted: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
assert_eq!(persisted["concurrent_extra"], "merge me");
assert!(
stub.requests.try_recv().is_err(),
"single-flight must make only one request"
);
}
#[tokio::test]
async fn account_switch_before_guarded_reload_skips_old_account_refresh() {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let barrier = Arc::new(tokio::sync::Barrier::new(2));
let mut provider = CodexAuthProvider::new(dir.path().to_path_buf());
let mut stub = Stub::start(vec![]).await;
provider.refresh_endpoint = stub.endpoint.clone();
provider.before_guarded_reload = Some(barrier.clone());
let task = tokio::spawn(async move { provider.headers().await });
bounded(barrier.wait()).await;
let mut switched: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
switched["tokens"]["account_id"] = json!("switched-account");
switched["tokens"]["refresh_token"] = json!("fixture-switched-secret");
std::fs::write(&path, serde_json::to_vec(&switched).unwrap()).unwrap();
let original = std::fs::read(&path).unwrap();
barrier.wait().await;
let headers = bounded(task).await.unwrap().unwrap();
assert!(
headers
.iter()
.any(|(name, value)| name == "chatgpt-account-id" && value == "switched-account")
);
assert_eq!(std::fs::read(&path).unwrap(), original);
assert!(stub.requests.try_recv().is_err());
}
#[tokio::test]
async fn account_or_token_switch_during_refresh_is_not_overwritten_by_old_response() {
for account_changed in [false, true] {
let dir = tempfile::tempdir().unwrap();
let path = stale_auth(dir.path());
let started = Arc::new(Notify::new());
let gate = Arc::new(Notify::new());
let mut stub = Stub::start(vec![Reply::gated(200, rotated_body(), &gate)]).await;
let provider = provider(dir.path(), &stub, &started);
let task = tokio::spawn({
let provider = provider.clone();
async move { provider.headers().await }
});
bounded(started.notified()).await;
bounded(stub.requests.recv()).await.unwrap();
let mut switched: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
if account_changed {
switched["tokens"]["account_id"] = json!("switched-account");
} else {
switched["tokens"]["refresh_token"] = json!("fixture-switched-secret");
}
switched["tokens"]["access_token"] =
json!(make_jwt(json!({"exp": Utc::now().timestamp() + 7200})));
std::fs::write(&path, serde_json::to_vec(&switched).unwrap()).unwrap();
let original = std::fs::read(&path).unwrap();
gate.notify_one();
let headers = bounded(task).await.unwrap().unwrap();
assert_eq!(std::fs::read(&path).unwrap(), original);
let expected = format!(
"Bearer {}",
switched["tokens"]["access_token"].as_str().unwrap()
);
assert!(headers.iter().any(|(name, value)| name == "authorization" && value.to_str().unwrap() == expected));
assert_eq!(bounded(provider.headers()).await.unwrap(), headers);
assert!(stub.requests.try_recv().is_err());
}
}