fix: bound Codex OAuth refresh bodies and deadline
This commit is contained in:
@@ -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"
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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());
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user