diff --git a/crates/agen/Cargo.toml b/crates/agen/Cargo.toml index e01b2ca40..0d667331c 100644 --- a/crates/agen/Cargo.toml +++ b/crates/agen/Cargo.toml @@ -37,6 +37,7 @@ chrono = { version = "0.4", default-features = false, features = ["serde", "cloc agen-macros = { workspace = true } [dev-dependencies] +tokio = { workspace = true, features = ["test-util"] } clap = { version = "4.5", features = ["derive", "env"] } tempfile = { workspace = true } dotenv = "0.15" diff --git a/crates/agen/src/providers/codex/auth_json.rs b/crates/agen/src/providers/codex/auth_json.rs index 374224730..212c75d34 100644 --- a/crates/agen/src/providers/codex/auth_json.rs +++ b/crates/agen/src/providers/codex/auth_json.rs @@ -108,13 +108,21 @@ pub async fn load(path: &Path) -> Result { /// /// 並行する Codex CLI / 別プロセスが先に refresh していた場合の /// 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( path: &Path, + refresh_input: &AuthSnapshot, new_id_token: Option, new_access_token: Option, new_refresh_token: Option, ) -> Result { 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 tokens = raw .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"} }"#, ); - let updated = - persist_refreshed(&path, None, Some("new-acc".into()), Some("new-ref".into())) - .await - .unwrap(); + let input = load(&path).await.unwrap(); + let updated = persist_refreshed( + &path, + &input, + None, + Some("new-acc".into()), + Some("new-ref".into()), + ) + .await + .unwrap(); assert_eq!(updated.access_token, "new-acc"); assert_eq!(updated.refresh_token, "new-ref"); // 未知フィールド agent_identity が保たれる @@ -266,7 +280,8 @@ mod tests { ); // 既存ファイルを 644 に変えてから persist → 600 に直るか 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 .unwrap(); let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777; diff --git a/crates/agen/src/providers/codex/mod.rs b/crates/agen/src/providers/codex/mod.rs index a3a5ad2f0..4123f4925 100644 --- a/crates/agen/src/providers/codex/mod.rs +++ b/crates/agen/src/providers/codex/mod.rs @@ -17,6 +17,8 @@ mod auth_json; mod error; mod jwt; mod refresh; +#[cfg(test)] +mod refresh_tests; use std::path::PathBuf; use std::sync::Arc; @@ -40,6 +42,9 @@ pub struct CodexAuthProvider { auth_path: PathBuf, refresh_endpoint: String, http: reqwest::Client, + refresh_limits: refresh::RefreshLimits, + #[cfg(test)] + before_guarded_reload: Option>, state: Arc>, } @@ -76,6 +81,9 @@ impl CodexAuthProvider { auth_path: codex_home.join("auth.json"), refresh_endpoint: refresh::REFRESH_URL.to_string(), http: reqwest::Client::new(), + refresh_limits: refresh::RefreshLimits::default(), + #[cfg(test)] + before_guarded_reload: None, state: Arc::new(Mutex::new(State::default())), } } @@ -108,6 +116,11 @@ impl CodexAuthProvider { // 3. Refresh 直前に再 load し account_id 一致を確認(guarded reload)。 // 一致しなければ他プロセスが先に更新済 → 自分は 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?; if !is_stale(&pre_refresh) { state.cached = Some(pre_refresh.clone()); @@ -124,12 +137,14 @@ impl CodexAuthProvider { &self.http, &self.refresh_endpoint, &pre_refresh.refresh_token, + &self.refresh_limits, ) .await?; // 5. 書き戻し(書込前に再 load + merge は persist_refreshed 内で実施) let new_snap = auth_json::persist_refreshed( &self.auth_path, + &pre_refresh, refreshed.id_token, refreshed.access_token, refreshed.refresh_token, @@ -210,13 +225,18 @@ mod tests { use wiremock::matchers::{method, path as url_path}; 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 .encode(serde_json::to_vec(&payload).unwrap()); 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!({ "https://api.openai.com/auth": { "chatgpt_account_id": "acc-xyz", diff --git a/crates/agen/src/providers/codex/refresh.rs b/crates/agen/src/providers/codex/refresh.rs index 1b526eeb9..a9844d66a 100644 --- a/crates/agen/src/providers/codex/refresh.rs +++ b/crates/agen/src/providers/codex/refresh.rs @@ -10,7 +10,26 @@ use super::error::{CodexAuthError, PermanentReason}; pub const CLIENT_ID: &str = "app_EMoamEEZ73f0CkXaXp7hrann"; 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>, +} + +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)] struct RefreshRequest<'a> { @@ -37,54 +56,94 @@ pub async fn request_refresh( client: &reqwest::Client, endpoint: &str, refresh_token: &str, + limits: &RefreshLimits, +) -> Result { + 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 { let body = RefreshRequest { client_id: CLIENT_ID, grant_type: "refresh_token", refresh_token, }; - let response = response_with_timeout( - client - .post(endpoint) - .header("Content-Type", "application/json") - .json(&body) - .send(), - DEFAULT_REFRESH_TIMEOUT, - ) - .await?; + let mut response = client + .post(endpoint) + .header("Content-Type", "application/json") + .json(&body) + .send() + .await + .map_err(|_| CodexAuthError::RefreshTransient("send failed".into()))?; let status = response.status(); - if status.is_success() { - response - .json::() - .await - .map_err(|e| CodexAuthError::RefreshTransient(format!("parse response: {e}"))) - } else { - let body = response.text().await.unwrap_or_default(); - if status == reqwest::StatusCode::UNAUTHORIZED { - let (reason, message) = classify_permanent(&body); - Err(CodexAuthError::RefreshPermanent { reason, message }) - } else { - Err(CodexAuthError::RefreshTransient(format!( - "{status}: {body}" - ))) + if response + .content_length() + .is_some_and(|len| len > limits.max_body_bytes as u64) + { + return Err(body_too_large()); + } + #[cfg(test)] + if let Some(started) = &limits.body_started { + started.notify_one(); + } + let mut bytes = Vec::new(); + while let Some(chunk) = response + .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( - future: impl std::future::Future>, - timeout: Duration, -) -> Result { - 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 body_too_large() -> CodexAuthError { + CodexAuthError::RefreshTransient("response body exceeds size limit".into()) } 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::Reused => "Your refresh token was already used".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) } @@ -126,23 +185,6 @@ fn extract_error_code(body: &str) -> Option { mod tests { use super::*; - #[tokio::test] - async fn refresh_response_timeout_is_transient() { - let err = match response_with_timeout( - std::future::pending::>(), - 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] fn classify_expired() { let body = r#"{"error":{"code":"refresh_token_expired"}}"#; diff --git a/crates/agen/src/providers/codex/refresh_tests.rs b/crates/agen/src/providers/codex/refresh_tests.rs new file mode 100644 index 000000000..9030f9722 --- /dev/null +++ b/crates/agen/src/providers/codex/refresh_tests.rs @@ -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, + chunked: bool, + gate: Option>, + truncate: bool, +} + +impl Reply { + fn body(status: u16, body: impl Into>) -> Self { + Self { + status, + body: body.into(), + chunked: false, + gate: None, + truncate: false, + } + } + + fn gated(status: u16, body: impl Into>, gate: &Arc) -> Self { + Self { + gate: Some(gate.clone()), + ..Self::body(status, body) + } + } +} + +struct Stub { + endpoint: String, + requests: mpsc::UnboundedReceiver, + 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) -> 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 { + 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) -> Arc { + 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(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()); + } +}