mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
feat: verify agent identity JWTs with JWKS (#19764)
This commit is contained in:
committed by
GitHub
Unverified
parent
6138063656
commit
f6797c3ac6
Generated
+1
@@ -2850,6 +2850,7 @@ dependencies = [
|
||||
"codex-terminal-detection",
|
||||
"codex-utils-template",
|
||||
"core_test_support",
|
||||
"jsonwebtoken",
|
||||
"keyring",
|
||||
"once_cell",
|
||||
"os_info",
|
||||
|
||||
@@ -19,6 +19,9 @@ use ed25519_dalek::pkcs8::EncodePrivateKey;
|
||||
use jsonwebtoken::Algorithm;
|
||||
use jsonwebtoken::DecodingKey;
|
||||
use jsonwebtoken::Validation;
|
||||
use jsonwebtoken::decode;
|
||||
use jsonwebtoken::decode_header;
|
||||
use jsonwebtoken::jwk::JwkSet;
|
||||
use rand::TryRngCore;
|
||||
use rand::rngs::OsRng;
|
||||
use serde::Deserialize;
|
||||
@@ -28,6 +31,9 @@ use sha2::Digest as _;
|
||||
use sha2::Sha512;
|
||||
|
||||
const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const AGENT_IDENTITY_JWKS_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const AGENT_IDENTITY_JWT_AUDIENCE: &str = "codex-app-server";
|
||||
const AGENT_IDENTITY_JWT_ISSUER: &str = "https://chatgpt.com/codex-backend/agent-identity";
|
||||
|
||||
/// Stored key material for a registered agent identity.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
@@ -58,6 +64,10 @@ pub struct GeneratedAgentKeyMaterial {
|
||||
/// Claims carried by an Agent Identity JWT.
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
|
||||
pub struct AgentIdentityJwtClaims {
|
||||
pub iss: String,
|
||||
pub aud: String,
|
||||
pub iat: usize,
|
||||
pub exp: usize,
|
||||
pub agent_runtime_id: String,
|
||||
pub agent_private_key: String,
|
||||
pub account_id: String,
|
||||
@@ -115,27 +125,49 @@ pub fn authorization_header_for_agent_task(
|
||||
Ok(format!("AgentAssertion {serialized_assertion}"))
|
||||
}
|
||||
|
||||
pub async fn fetch_agent_identity_jwks(
|
||||
client: &reqwest::Client,
|
||||
chatgpt_base_url: &str,
|
||||
) -> Result<JwkSet> {
|
||||
let response = client
|
||||
.get(agent_identity_jwks_url(chatgpt_base_url))
|
||||
.timeout(AGENT_IDENTITY_JWKS_TIMEOUT)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to request agent identity JWKS")?
|
||||
.error_for_status()
|
||||
.context("agent identity JWKS endpoint returned an error")?;
|
||||
|
||||
response
|
||||
.json()
|
||||
.await
|
||||
.context("failed to decode agent identity JWKS")
|
||||
}
|
||||
|
||||
pub fn decode_agent_identity_jwt(
|
||||
jwt: &str,
|
||||
public_key_base64: Option<&str>,
|
||||
jwks: Option<&JwkSet>,
|
||||
) -> Result<AgentIdentityJwtClaims> {
|
||||
let Some(public_key_base64) = public_key_base64 else {
|
||||
let Some(jwks) = jwks else {
|
||||
return decode_agent_identity_jwt_payload(jwt);
|
||||
};
|
||||
|
||||
let mut validation = Validation::new(Algorithm::EdDSA);
|
||||
validation.required_spec_claims.clear();
|
||||
validation.validate_exp = false;
|
||||
validation.validate_aud = false;
|
||||
|
||||
let public_key = BASE64_STANDARD
|
||||
.decode(public_key_base64)
|
||||
.context("agent identity JWT public key is not valid base64")?;
|
||||
let decoding_key = DecodingKey::from_ed_der(&public_key);
|
||||
|
||||
jsonwebtoken::decode::<AgentIdentityJwtClaims>(jwt, &decoding_key, &validation)
|
||||
let header = decode_header(jwt).context("failed to decode agent identity JWT header")?;
|
||||
let kid = header
|
||||
.kid
|
||||
.context("agent identity JWT header does not include a kid")?;
|
||||
let jwk = jwks
|
||||
.find(&kid)
|
||||
.with_context(|| format!("agent identity JWT kid {kid} is not trusted"))?;
|
||||
let decoding_key = DecodingKey::from_jwk(jwk).context("failed to build JWT decoding key")?;
|
||||
let mut validation = Validation::new(Algorithm::RS256);
|
||||
validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]);
|
||||
validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]);
|
||||
validation.required_spec_claims.insert("iss".to_string());
|
||||
validation.required_spec_claims.insert("aud".to_string());
|
||||
decode::<AgentIdentityJwtClaims>(jwt, &decoding_key, &validation)
|
||||
.map(|data| data.claims)
|
||||
.context("failed to decode agent identity JWT")
|
||||
.context("failed to verify agent identity JWT")
|
||||
}
|
||||
|
||||
fn decode_agent_identity_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T> {
|
||||
@@ -279,6 +311,15 @@ pub fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String {
|
||||
format!("{trimmed}/authenticate_app_v2")
|
||||
}
|
||||
|
||||
pub fn agent_identity_jwks_url(chatgpt_base_url: &str) -> String {
|
||||
let trimmed = chatgpt_base_url.trim_end_matches('/');
|
||||
if trimmed.contains("/backend-api") {
|
||||
format!("{trimmed}/wham/agent-identities/jwks")
|
||||
} else {
|
||||
format!("{trimmed}/agent-identities/jwks")
|
||||
}
|
||||
}
|
||||
|
||||
pub fn agent_identity_request_id() -> Result<String> {
|
||||
let mut request_id_bytes = [0u8; 16];
|
||||
OsRng
|
||||
@@ -290,29 +331,6 @@ pub fn agent_identity_request_id() -> Result<String> {
|
||||
))
|
||||
}
|
||||
|
||||
pub fn normalize_chatgpt_base_url(chatgpt_base_url: &str) -> String {
|
||||
let mut base_url = chatgpt_base_url.trim_end_matches('/').to_string();
|
||||
for suffix in [
|
||||
"/wham/remote/control/server/enroll",
|
||||
"/wham/remote/control/server",
|
||||
] {
|
||||
if let Some(stripped) = base_url.strip_suffix(suffix) {
|
||||
base_url = stripped.to_string();
|
||||
break;
|
||||
}
|
||||
}
|
||||
if let Some(stripped) = base_url.strip_suffix("/codex") {
|
||||
base_url = stripped.to_string();
|
||||
}
|
||||
if (base_url.starts_with("https://chatgpt.com")
|
||||
|| base_url.starts_with("https://chat.openai.com"))
|
||||
&& !base_url.contains("/backend-api")
|
||||
{
|
||||
base_url = format!("{base_url}/backend-api");
|
||||
}
|
||||
base_url
|
||||
}
|
||||
|
||||
pub fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials {
|
||||
AgentBillOfMaterials {
|
||||
agent_version: env!("CARGO_PKG_VERSION").to_string(),
|
||||
@@ -472,6 +490,10 @@ mod tests {
|
||||
#[test]
|
||||
fn decode_agent_identity_jwt_reads_claims() {
|
||||
let jwt = jwt_with_payload(serde_json::json!({
|
||||
"iss": AGENT_IDENTITY_JWT_ISSUER,
|
||||
"aud": AGENT_IDENTITY_JWT_AUDIENCE,
|
||||
"iat": 1_700_000_000usize,
|
||||
"exp": 4_000_000_000usize,
|
||||
"agent_runtime_id": "agent-runtime-id",
|
||||
"agent_private_key": "private-key",
|
||||
"account_id": "account-id",
|
||||
@@ -481,12 +503,15 @@ mod tests {
|
||||
"chatgpt_account_is_fedramp": false,
|
||||
}));
|
||||
|
||||
let claims =
|
||||
decode_agent_identity_jwt(&jwt, /*public_key_base64*/ None).expect("JWT should decode");
|
||||
let claims = decode_agent_identity_jwt(&jwt, /*jwks*/ None).expect("JWT should decode");
|
||||
|
||||
assert_eq!(
|
||||
claims,
|
||||
AgentIdentityJwtClaims {
|
||||
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
|
||||
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
|
||||
iat: 1_700_000_000,
|
||||
exp: 4_000_000_000,
|
||||
agent_runtime_id: "agent-runtime-id".to_string(),
|
||||
agent_private_key: "private-key".to_string(),
|
||||
account_id: "account-id".to_string(),
|
||||
@@ -499,15 +524,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_agent_identity_jwt_verifies_when_public_key_is_present() {
|
||||
let mut secret_key_bytes = [0u8; 32];
|
||||
secret_key_bytes[0] = 1;
|
||||
let signing_key = SigningKey::from_bytes(&secret_key_bytes);
|
||||
let private_key_pkcs8 = signing_key
|
||||
.to_pkcs8_der()
|
||||
.expect("private key should encode");
|
||||
let public_key_base64 = BASE64_STANDARD.encode(signing_key.verifying_key().as_bytes());
|
||||
fn decode_agent_identity_jwt_verifies_when_jwks_is_present() {
|
||||
let jwks = test_jwks("test-key");
|
||||
let claims = AgentIdentityJwtClaims {
|
||||
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
|
||||
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
|
||||
iat: 1_700_000_000,
|
||||
exp: 4_000_000_000,
|
||||
agent_runtime_id: "agent-runtime-id".to_string(),
|
||||
agent_private_key: "private-key".to_string(),
|
||||
account_id: "account-id".to_string(),
|
||||
@@ -517,8 +540,12 @@ mod tests {
|
||||
chatgpt_account_is_fedramp: false,
|
||||
};
|
||||
let jwt = jsonwebtoken::encode(
|
||||
&Header::new(Algorithm::EdDSA),
|
||||
&test_jwt_header("test-key"),
|
||||
&serde_json::json!({
|
||||
"iss": claims.iss,
|
||||
"aud": claims.aud,
|
||||
"iat": claims.iat,
|
||||
"exp": claims.exp,
|
||||
"agent_runtime_id": claims.agent_runtime_id,
|
||||
"agent_private_key": claims.agent_private_key,
|
||||
"account_id": claims.account_id,
|
||||
@@ -527,11 +554,15 @@ mod tests {
|
||||
"plan_type": "pro",
|
||||
"chatgpt_account_is_fedramp": claims.chatgpt_account_is_fedramp,
|
||||
}),
|
||||
&EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()),
|
||||
&test_rsa_encoding_key(),
|
||||
)
|
||||
.expect("JWT should encode");
|
||||
|
||||
let expected_claims = AgentIdentityJwtClaims {
|
||||
iss: AGENT_IDENTITY_JWT_ISSUER.to_string(),
|
||||
aud: AGENT_IDENTITY_JWT_AUDIENCE.to_string(),
|
||||
iat: 1_700_000_000,
|
||||
exp: 4_000_000_000,
|
||||
agent_runtime_id: "agent-runtime-id".to_string(),
|
||||
agent_private_key: "private-key".to_string(),
|
||||
account_id: "account-id".to_string(),
|
||||
@@ -541,31 +572,22 @@ mod tests {
|
||||
chatgpt_account_is_fedramp: false,
|
||||
};
|
||||
assert_eq!(
|
||||
decode_agent_identity_jwt(&jwt, Some(&public_key_base64)).expect("JWT should verify"),
|
||||
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect("JWT should verify"),
|
||||
expected_claims
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decode_agent_identity_jwt_rejects_wrong_public_key() {
|
||||
let mut signing_secret_key_bytes = [0u8; 32];
|
||||
signing_secret_key_bytes[0] = 1;
|
||||
let signing_key = SigningKey::from_bytes(&signing_secret_key_bytes);
|
||||
let private_key_pkcs8 = signing_key
|
||||
.to_pkcs8_der()
|
||||
.expect("private key should encode");
|
||||
|
||||
let mut other_secret_key_bytes = [0u8; 32];
|
||||
other_secret_key_bytes[0] = 2;
|
||||
let other_public_key_base64 = BASE64_STANDARD.encode(
|
||||
SigningKey::from_bytes(&other_secret_key_bytes)
|
||||
.verifying_key()
|
||||
.as_bytes(),
|
||||
);
|
||||
fn decode_agent_identity_jwt_rejects_untrusted_kid() {
|
||||
let jwks = test_jwks("other-key");
|
||||
|
||||
let jwt = jsonwebtoken::encode(
|
||||
&Header::new(Algorithm::EdDSA),
|
||||
&test_jwt_header("test-key"),
|
||||
&serde_json::json!({
|
||||
"iss": AGENT_IDENTITY_JWT_ISSUER,
|
||||
"aud": AGENT_IDENTITY_JWT_AUDIENCE,
|
||||
"iat": 1_700_000_000,
|
||||
"exp": 4_000_000_000usize,
|
||||
"agent_runtime_id": "agent-runtime-id",
|
||||
"agent_private_key": "private-key",
|
||||
"account_id": "account-id",
|
||||
@@ -574,19 +596,111 @@ mod tests {
|
||||
"plan_type": "pro",
|
||||
"chatgpt_account_is_fedramp": false,
|
||||
}),
|
||||
&EncodingKey::from_ed_der(private_key_pkcs8.as_bytes()),
|
||||
&test_rsa_encoding_key(),
|
||||
)
|
||||
.expect("JWT should encode");
|
||||
|
||||
decode_agent_identity_jwt(&jwt, Some(&other_public_key_base64))
|
||||
.expect_err("JWT should not verify");
|
||||
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_chatgpt_base_url_strips_codex_before_backend_api() {
|
||||
fn decode_agent_identity_jwt_requires_issuer_and_audience() {
|
||||
let jwks = test_jwks("test-key");
|
||||
let jwt = jsonwebtoken::encode(
|
||||
&test_jwt_header("test-key"),
|
||||
&serde_json::json!({
|
||||
"iat": 1_700_000_000,
|
||||
"exp": 4_000_000_000usize,
|
||||
"agent_runtime_id": "agent-runtime-id",
|
||||
"agent_private_key": "private-key",
|
||||
"account_id": "account-id",
|
||||
"chatgpt_user_id": "user-id",
|
||||
"email": "user@example.com",
|
||||
"plan_type": "pro",
|
||||
"chatgpt_account_is_fedramp": false,
|
||||
}),
|
||||
&test_rsa_encoding_key(),
|
||||
)
|
||||
.expect("JWT should encode");
|
||||
|
||||
decode_agent_identity_jwt(&jwt, Some(&jwks)).expect_err("JWT should not verify");
|
||||
}
|
||||
|
||||
fn test_jwt_header(kid: &str) -> Header {
|
||||
let mut header = Header::new(Algorithm::RS256);
|
||||
header.kid = Some(kid.to_string());
|
||||
header
|
||||
}
|
||||
|
||||
fn test_rsa_encoding_key() -> EncodingKey {
|
||||
EncodingKey::from_rsa_pem(
|
||||
br#"-----BEGIN PRIVATE KEY-----
|
||||
MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
|
||||
bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
|
||||
FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
|
||||
mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
|
||||
0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
|
||||
zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
|
||||
uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
|
||||
1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
|
||||
RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
|
||||
NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
|
||||
rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
|
||||
Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
|
||||
aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
|
||||
iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
|
||||
YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
|
||||
3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
|
||||
sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
|
||||
gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
|
||||
sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
|
||||
BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
|
||||
KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
|
||||
r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
|
||||
29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
|
||||
J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
|
||||
5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
|
||||
8U4M2TSWCKUY/A6sT4W8+mT9
|
||||
-----END PRIVATE KEY-----"#,
|
||||
)
|
||||
.expect("test RSA key should parse")
|
||||
}
|
||||
|
||||
fn test_jwks(kid: &str) -> jsonwebtoken::jwk::JwkSet {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"keys": [{
|
||||
"kty": "RSA",
|
||||
"kid": kid,
|
||||
"use": "sig",
|
||||
"alg": "RS256",
|
||||
"n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
|
||||
"e": "AQAB",
|
||||
}]
|
||||
}))
|
||||
.expect("test JWKS should parse")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_jwks_url_uses_backend_api_base_url() {
|
||||
assert_eq!(
|
||||
normalize_chatgpt_base_url("https://chatgpt.com/codex"),
|
||||
"https://chatgpt.com/backend-api"
|
||||
agent_identity_jwks_url("https://chatgpt.com/backend-api"),
|
||||
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
|
||||
);
|
||||
assert_eq!(
|
||||
agent_identity_jwks_url("https://chatgpt.com/backend-api/"),
|
||||
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_identity_jwks_url_uses_codex_api_base_url() {
|
||||
assert_eq!(
|
||||
agent_identity_jwks_url("http://localhost:8080/api/codex"),
|
||||
"http://localhost:8080/api/codex/agent-identities/jwks"
|
||||
);
|
||||
assert_eq!(
|
||||
agent_identity_jwks_url("http://localhost:8080/api/codex/"),
|
||||
"http://localhost:8080/api/codex/agent-identities/jwks"
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -207,7 +207,10 @@ pub async fn run_login_with_agent_identity(
|
||||
&config.codex_home,
|
||||
&agent_identity,
|
||||
config.cli_auth_credentials_store_mode,
|
||||
) {
|
||||
Some(&config.chatgpt_base_url),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
eprintln!("{LOGIN_SUCCESS_MESSAGE}");
|
||||
std::process::exit(0);
|
||||
@@ -362,8 +365,12 @@ pub async fn run_login_with_device_code_fallback_to_browser(
|
||||
pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! {
|
||||
let config = load_config_or_exit(cli_config_overrides).await;
|
||||
|
||||
match CodexAuth::from_auth_storage(&config.codex_home, config.cli_auth_credentials_store_mode)
|
||||
.await
|
||||
match CodexAuth::from_auth_storage(
|
||||
&config.codex_home,
|
||||
config.cli_auth_credentials_store_mode,
|
||||
Some(&config.chatgpt_base_url),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(auth)) => match auth.auth_mode() {
|
||||
AuthMode::ApiKey => match auth.get_token() {
|
||||
|
||||
@@ -6,8 +6,6 @@ use pretty_assertions::assert_eq;
|
||||
use serde_json::Value;
|
||||
use tempfile::TempDir;
|
||||
|
||||
const FAKE_AGENT_IDENTITY_JWT: &str = "eyJhbGciOiJFZERTQSIsInR5cCI6IkpXVCJ9.eyJhZ2VudF9ydW50aW1lX2lkIjoiYWdlbnQtcnVudGltZS1pZCIsImFnZW50X3ByaXZhdGVfa2V5IjoicHJpdmF0ZS1rZXkiLCJhY2NvdW50X2lkIjoiYWNjb3VudC0xMjMiLCJjaGF0Z3B0X3VzZXJfaWQiOiJ1c2VyLWlkIiwiZW1haWwiOiJ1c2VyQGV4YW1wbGUuY29tIiwicGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lzX2ZlZHJhbXAiOmZhbHNlfQ.c2ln";
|
||||
|
||||
fn codex_command(codex_home: &Path) -> Result<assert_cmd::Command> {
|
||||
let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?);
|
||||
cmd.env("CODEX_HOME", codex_home);
|
||||
@@ -53,22 +51,16 @@ fn login_with_api_key_reads_stdin_and_writes_auth_json() -> Result<()> {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn login_with_agent_identity_reads_stdin_and_writes_auth_json() -> Result<()> {
|
||||
fn login_with_agent_identity_rejects_invalid_jwt() -> Result<()> {
|
||||
let codex_home = TempDir::new()?;
|
||||
write_file_auth_config(codex_home.path())?;
|
||||
|
||||
let mut cmd = codex_command(codex_home.path())?;
|
||||
cmd.args(["login", "--with-agent-identity"])
|
||||
.write_stdin(format!("{FAKE_AGENT_IDENTITY_JWT}\n"))
|
||||
.write_stdin("not-a-jwt\n")
|
||||
.assert()
|
||||
.success()
|
||||
.stderr(contains("Successfully logged in"));
|
||||
|
||||
let auth = read_auth_json(codex_home.path())?;
|
||||
assert_eq!(auth["auth_mode"], "agentIdentity");
|
||||
assert_eq!(auth["agent_identity"], FAKE_AGENT_IDENTITY_JWT);
|
||||
assert!(auth["OPENAI_API_KEY"].is_null());
|
||||
assert!(auth.get("tokens").is_none());
|
||||
.failure()
|
||||
.stderr(contains("Error logging in with Agent Identity"));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1090,13 +1090,17 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() {
|
||||
let mut config = load_default_config_for_test(&codex_home).await;
|
||||
config.model_provider = model_provider;
|
||||
|
||||
let auth_manager =
|
||||
match CodexAuth::from_auth_storage(codex_home.path(), AuthCredentialsStoreMode::File).await
|
||||
{
|
||||
Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth),
|
||||
Ok(None) => panic!("No CodexAuth found in codex_home"),
|
||||
Err(e) => panic!("Failed to load CodexAuth: {e}"),
|
||||
};
|
||||
let auth_manager = match CodexAuth::from_auth_storage(
|
||||
codex_home.path(),
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(auth)) => codex_core::test_support::auth_manager_from_auth(auth),
|
||||
Ok(None) => panic!("No CodexAuth found in codex_home"),
|
||||
Err(e) => panic!("Failed to load CodexAuth: {e}"),
|
||||
};
|
||||
let thread_manager = ThreadManager::new(
|
||||
&config,
|
||||
auth_manager,
|
||||
|
||||
@@ -439,6 +439,7 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result
|
||||
auth_credentials_store_mode: config.cli_auth_credentials_store_mode,
|
||||
forced_login_method: config.forced_login_method,
|
||||
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
|
||||
chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
|
||||
})
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -45,6 +45,7 @@ webbrowser = { workspace = true }
|
||||
[dev-dependencies]
|
||||
anyhow = { workspace = true }
|
||||
core_test_support = { workspace = true }
|
||||
jsonwebtoken = { workspace = true }
|
||||
keyring = { workspace = true }
|
||||
pretty_assertions = { workspace = true }
|
||||
regex-lite = { workspace = true }
|
||||
|
||||
@@ -82,7 +82,7 @@ mod tests {
|
||||
use serial_test::serial;
|
||||
|
||||
#[test]
|
||||
#[serial(agent_identity_authapi_base_url_env)]
|
||||
#[serial(codex_auth_env)]
|
||||
fn agent_identity_authapi_base_url_prefers_env_value() {
|
||||
let _guard = EnvVarGuard::set(
|
||||
CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR,
|
||||
@@ -95,7 +95,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[serial(agent_identity_authapi_base_url_env)]
|
||||
#[serial(codex_auth_env)]
|
||||
fn agent_identity_authapi_base_url_uses_prod_authapi_by_default() {
|
||||
let _guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL_ENV_VAR);
|
||||
assert_eq!(
|
||||
|
||||
@@ -16,6 +16,11 @@ use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tempfile::TempDir;
|
||||
use tempfile::tempdir;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_without_id_token() {
|
||||
@@ -78,15 +83,29 @@ fn login_with_api_key_overwrites_existing_auth_json() {
|
||||
assert!(auth.tokens.is_none(), "tokens should be cleared");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn login_with_agent_identity_writes_only_token() {
|
||||
#[tokio::test]
|
||||
async fn login_with_agent_identity_writes_only_token() {
|
||||
let dir = tempdir().unwrap();
|
||||
let auth_path = dir.path().join("auth.json");
|
||||
let record = agent_identity_record("account-123");
|
||||
let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
|
||||
let agent_identity = signed_agent_identity_jwt(&record).expect("signed agent identity");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/wham/agent-identities/jwks"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let chatgpt_base_url = format!("{}/backend-api", server.uri());
|
||||
|
||||
super::login_with_agent_identity(dir.path(), &agent_identity, AuthCredentialsStoreMode::File)
|
||||
.expect("login_with_agent_identity should succeed");
|
||||
super::login_with_agent_identity(
|
||||
dir.path(),
|
||||
&agent_identity,
|
||||
AuthCredentialsStoreMode::File,
|
||||
Some(&chatgpt_base_url),
|
||||
)
|
||||
.await
|
||||
.expect("login_with_agent_identity should succeed");
|
||||
|
||||
let storage = FileAuthStorage::new(dir.path().to_path_buf());
|
||||
let auth = storage
|
||||
@@ -99,15 +118,21 @@ fn login_with_agent_identity_writes_only_token() {
|
||||
);
|
||||
assert!(auth.tokens.is_none(), "tokens should be cleared");
|
||||
assert!(auth.openai_api_key.is_none(), "API key should be cleared");
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn login_with_agent_identity_rejects_invalid_jwt() {
|
||||
#[tokio::test]
|
||||
async fn login_with_agent_identity_rejects_invalid_jwt() {
|
||||
let dir = tempdir().unwrap();
|
||||
|
||||
let err =
|
||||
super::login_with_agent_identity(dir.path(), "not-a-jwt", AuthCredentialsStoreMode::File)
|
||||
.expect_err("invalid Agent Identity token should fail");
|
||||
let err = super::login_with_agent_identity(
|
||||
dir.path(),
|
||||
"not-a-jwt",
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect_err("invalid Agent Identity token should fail");
|
||||
|
||||
assert_eq!(err.kind(), std::io::ErrorKind::Other);
|
||||
assert!(
|
||||
@@ -117,11 +142,47 @@ fn login_with_agent_identity_rejects_invalid_jwt() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn login_with_agent_identity_rejects_unsigned_jwt() {
|
||||
let dir = tempdir().unwrap();
|
||||
let record = agent_identity_record("account-123");
|
||||
let agent_identity = fake_agent_identity_jwt(&record).expect("fake agent identity");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/wham/agent-identities/jwks"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let chatgpt_base_url = format!("{}/backend-api", server.uri());
|
||||
|
||||
super::login_with_agent_identity(
|
||||
dir.path(),
|
||||
&agent_identity,
|
||||
AuthCredentialsStoreMode::File,
|
||||
Some(&chatgpt_base_url),
|
||||
)
|
||||
.await
|
||||
.expect_err("unsigned Agent Identity token should fail");
|
||||
|
||||
assert!(
|
||||
!get_auth_file(dir.path()).exists(),
|
||||
"unsigned Agent Identity token should not write auth.json"
|
||||
);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn missing_auth_json_returns_none() {
|
||||
let dir = tempdir().unwrap();
|
||||
let auth = CodexAuth::from_auth_storage(dir.path(), AuthCredentialsStoreMode::File)
|
||||
.await
|
||||
.expect("call should succeed");
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let auth = CodexAuth::from_auth_storage(
|
||||
dir.path(),
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("call should succeed");
|
||||
assert_eq!(auth, None);
|
||||
}
|
||||
|
||||
@@ -129,6 +190,7 @@ async fn missing_auth_json_returns_none() {
|
||||
#[serial(codex_auth_env)]
|
||||
async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let fake_jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -143,6 +205,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
@@ -186,6 +249,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() {
|
||||
#[serial(codex_auth_env)]
|
||||
async fn loads_api_key_from_auth_json() {
|
||||
let dir = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let auth_file = dir.path().join("auth.json");
|
||||
std::fs::write(
|
||||
auth_file,
|
||||
@@ -197,6 +261,7 @@ async fn loads_api_key_from_auth_json() {
|
||||
dir.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
@@ -255,8 +320,10 @@ async fn unauthorized_recovery_reports_mode_and_step_names() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -271,6 +338,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
@@ -288,6 +356,7 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() {
|
||||
codex_home.path(),
|
||||
updated_auth_dot_json,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("updated auth should parse");
|
||||
@@ -591,6 +660,7 @@ async fn build_config(
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode::File,
|
||||
forced_login_method,
|
||||
forced_chatgpt_workspace_id,
|
||||
chatgpt_base_url: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -611,6 +681,14 @@ impl EnvVarGuard {
|
||||
}
|
||||
Self { key, original }
|
||||
}
|
||||
|
||||
fn remove(key: &'static str) -> Self {
|
||||
let original = env::var_os(key);
|
||||
unsafe {
|
||||
env::remove_var(key);
|
||||
}
|
||||
Self { key, original }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -625,6 +703,55 @@ impl Drop for EnvVarGuard {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn load_auth_reads_agent_identity_from_env() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let expected_record = agent_identity_record("account-123");
|
||||
let agent_identity =
|
||||
signed_agent_identity_jwt(&expected_record).expect("signed agent identity");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/wham/agent-identities/jwks"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(test_jwks_body()))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/backend-api/v1/agent/agent-runtime-id/task/register"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"task_id": "task-123",
|
||||
})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let _agent_guard = EnvVarGuard::set(CODEX_AGENT_IDENTITY_ENV_VAR, &agent_identity);
|
||||
|
||||
let chatgpt_base_url = format!("{}/backend-api", server.uri());
|
||||
let _authapi_guard =
|
||||
EnvVarGuard::set("CODEX_AGENT_IDENTITY_AUTHAPI_BASE_URL", &chatgpt_base_url);
|
||||
let auth = super::load_auth(
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
Some(&chatgpt_base_url),
|
||||
)
|
||||
.await
|
||||
.expect("env auth should load")
|
||||
.expect("env auth should be present");
|
||||
|
||||
let CodexAuth::AgentIdentity(agent_identity) = auth else {
|
||||
panic!("env auth should load as agent identity");
|
||||
};
|
||||
assert_eq!(agent_identity.record(), &expected_record);
|
||||
assert_eq!(agent_identity.process_task_id(), "task-123");
|
||||
assert!(
|
||||
!get_auth_file(codex_home.path()).exists(),
|
||||
"env auth should not write auth.json"
|
||||
);
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn load_auth_keeps_codex_api_key_env_precedence() {
|
||||
@@ -638,6 +765,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ true,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("env auth should load")
|
||||
@@ -650,6 +778,7 @@ async fn load_auth_keeps_codex_api_key_env_precedence() {
|
||||
#[serial(codex_auth_env)]
|
||||
async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
|
||||
.expect("seed api key");
|
||||
|
||||
@@ -674,6 +803,7 @@ async fn enforce_login_restrictions_logs_out_for_method_mismatch() {
|
||||
#[serial(codex_auth_env)]
|
||||
async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -705,6 +835,7 @@ async fn enforce_login_restrictions_logs_out_for_workspace_mismatch() {
|
||||
#[serial(codex_auth_env)]
|
||||
async fn enforce_login_restrictions_allows_matching_workspace() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -736,6 +867,7 @@ async fn enforce_login_restrictions_allows_matching_workspace() {
|
||||
async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_forced_chatgpt_workspace_id_is_set()
|
||||
{
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
login_with_api_key(codex_home.path(), "sk-test", AuthCredentialsStoreMode::File)
|
||||
.expect("seed api key");
|
||||
|
||||
@@ -759,6 +891,7 @@ async fn enforce_login_restrictions_allows_api_key_if_login_method_not_set_but_f
|
||||
#[serial(codex_auth_env)]
|
||||
async fn enforce_login_restrictions_blocks_env_api_key_when_chatgpt_required() {
|
||||
let _guard = EnvVarGuard::set(CODEX_API_KEY_ENV_VAR, "sk-env");
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let codex_home = tempdir().unwrap();
|
||||
|
||||
let config = build_config(
|
||||
@@ -795,6 +928,10 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
|
||||
let encode = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
|
||||
let header_b64 = encode(br#"{"alg":"EdDSA","typ":"JWT"}"#);
|
||||
let payload = json!({
|
||||
"iss": "https://chatgpt.com/codex-backend/agent-identity",
|
||||
"aud": "codex-app-server",
|
||||
"iat": 1_700_000_000usize,
|
||||
"exp": 4_000_000_000usize,
|
||||
"agent_runtime_id": record.agent_runtime_id,
|
||||
"agent_private_key": record.agent_private_key,
|
||||
"account_id": record.account_id,
|
||||
@@ -808,9 +945,77 @@ fn fake_agent_identity_jwt(record: &AgentIdentityAuthRecord) -> std::io::Result<
|
||||
Ok(format!("{header_b64}.{payload_b64}.{signature_b64}"))
|
||||
}
|
||||
|
||||
fn signed_agent_identity_jwt(
|
||||
record: &AgentIdentityAuthRecord,
|
||||
) -> jsonwebtoken::errors::Result<String> {
|
||||
let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256);
|
||||
header.kid = Some("test-key".to_string());
|
||||
jsonwebtoken::encode(
|
||||
&header,
|
||||
&json!({
|
||||
"iss": "https://chatgpt.com/codex-backend/agent-identity",
|
||||
"aud": "codex-app-server",
|
||||
"iat": 1_700_000_000usize,
|
||||
"exp": 4_000_000_000usize,
|
||||
"agent_runtime_id": record.agent_runtime_id,
|
||||
"agent_private_key": record.agent_private_key,
|
||||
"account_id": record.account_id,
|
||||
"chatgpt_user_id": record.chatgpt_user_id,
|
||||
"email": record.email,
|
||||
"plan_type": record.plan_type,
|
||||
"chatgpt_account_is_fedramp": record.chatgpt_account_is_fedramp,
|
||||
}),
|
||||
&jsonwebtoken::EncodingKey::from_rsa_pem(TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM)?,
|
||||
)
|
||||
}
|
||||
|
||||
fn test_jwks_body() -> serde_json::Value {
|
||||
json!({
|
||||
"keys": [{
|
||||
"kty": "RSA",
|
||||
"kid": "test-key",
|
||||
"use": "sig",
|
||||
"alg": "RS256",
|
||||
"n": "1qQF2MqTrGAMDm7wXbjJP5sWqGA83tAGUs2ksy7iJXLJdhCg4AtwGm4SFl4f6kxhCSzlN1QdXuZjvRT2wZZiGUi9xUE28rf4WLrTxSnwqLuTy5knMP08yC0t_0YU_FGPZMcWb14hG05IvZr8UbmRaVagxSR8H4rSIymRoVwwmFSrqz068XrWGSYNIfLEASyo5GdAaqmk1JALINHgYGQJVxMxtwcvDxoVKmC7eltUNymMNBZhsv4E8sx9YNLpBoEibznfEpDU_DGzrM5eZCsQzaqbhBOlGd427ifud_Nnd9cPqzgCUc23-0FXSPfpbgksCXAwAmD0OFjQWrgqVdKL6Q",
|
||||
"e": "AQAB",
|
||||
}]
|
||||
})
|
||||
}
|
||||
|
||||
const TEST_AGENT_IDENTITY_RSA_PRIVATE_KEY_PEM: &[u8] = br#"-----BEGIN PRIVATE KEY-----
|
||||
MIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggSkAgEAAoIBAQDWpAXYypOsYAwO
|
||||
bvBduMk/mxaoYDze0AZSzaSzLuIlcsl2EKDgC3AabhIWXh/qTGEJLOU3VB1e5mO9
|
||||
FPbBlmIZSL3FQTbyt/hYutPFKfCou5PLmScw/TzILS3/RhT8UY9kxxZvXiEbTki9
|
||||
mvxRuZFpVqDFJHwfitIjKZGhXDCYVKurPTrxetYZJg0h8sQBLKjkZ0BqqaTUkAsg
|
||||
0eBgZAlXEzG3By8PGhUqYLt6W1Q3KYw0FmGy/gTyzH1g0ukGgSJvOd8SkNT8MbOs
|
||||
zl5kKxDNqpuEE6UZ3jbuJ+5382d31w+rOAJRzbf7QVdI9+luCSwJcDACYPQ4WNBa
|
||||
uCpV0ovpAgMBAAECggEAVu84LwZdqYN9XpswX8VoPYrjMm9IODapWQBRpQFoNyK2
|
||||
1ksF3bjEPvA2Azk8U/l7k+vLKw22l6lY3EyRZPcz5GnB8xLm3ogE3mtNOp4yCyVu
|
||||
RxhQ91aaN7mU17/a4BdorLi2LYVCg3zBmYociD1Q2AluNGsCmwPu+K7tfR2J0Sg8
|
||||
NjqiTbDG1XDpR/icwgC9t6vh8lZpCHDhF4tbQfLLVLeA/OdcuzXDyMCXbmdVIdBQ
|
||||
rm4aIFmr2e1/2ctTbCg85S6AGFTH+pSLjrwTzyvf+F6NW5uNjLQAQLFj+EznBDxj
|
||||
Xdx90cySrjsKK6PVWQF4RiTvkSW8eWL7R6B2FZbGwQKBgQDuVQRj72hWloR7mbEL
|
||||
aUEEv3pIXTMXWEsoMBNczos/1L1RnAN1AI44TurznasPZAWvQj+kVbLDR+TAeZrL
|
||||
iA8HIWswQUI18hFmgKzSkwIXGtubcKVrgsKeS4lMDKCM/Ef6WAYdeq6ronoY5lCN
|
||||
YrJFmGp81W5zcV7lyiycgbSiGwKBgQDmjWYf6pZjrK7Z+OJ3X1AZfi2vss15SCvL
|
||||
3fPgzIDbViztpGyQhc3DQZIsBNIu0xZp/veGce9TEeTds2ro9NfdJFeou8+fC7Pq
|
||||
sOsM3amGFFi+ZW/9BWyjZEM88bgWWAjqLHbpfHDxjAf5CSxddqxgHlbP0Ytyb1Vg
|
||||
gmPDn9YKSwKBgQDbTi3hC35WFuDHn0/zcSHcDZmnFuOZeqyFyV83yfMGhGrEuqvP
|
||||
sPgtRikajJ3IZsB4WZyYSidZXEFY/0z6NjOl2xF38MTNQPbT/FmK1q1Yt2UWrlv5
|
||||
BvSwlk87RG9D7C0LZo4R+D7cPoDdgqjiwMvMEIkEX5zn641oI1ZTmWKuuwKBgQCD
|
||||
KF+3unnRvHRAVoFnTZbA2fJdqMeRvogD04GhGlYX8V9f1hFY6nXTJaNlXVzA/J8c
|
||||
r8ra9kgjJuPfZ+ljG58OFFW2DRohLcQtuHYPfK6rMzoFHqnl9EcIcMp7ijuionR3
|
||||
29HOJFgQYgxLFXfit9d6WugiE+BTupiEbckZif13HwKBgE/lAlkVHP6YahOO2Ljc
|
||||
J1bwkqKZTB5dHolX9A58e/xXnfZ5P8f3Z83+Izap3FwqQulk7b1WO1MQcHuVg2NN
|
||||
5da0D4h2rYOXnbYIg0BVu4spQbaM6ewsp66b8+MzLOBvj8SzWdt1Oyw0q/MRyQAR
|
||||
8U4M2TSWCKUY/A6sT4W8+mT9
|
||||
-----END PRIVATE KEY-----"#;
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn plan_type_maps_known_plan() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -825,6 +1030,7 @@ async fn plan_type_maps_known_plan() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
@@ -834,8 +1040,10 @@ async fn plan_type_maps_known_plan() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn plan_type_maps_self_serve_business_usage_based_plan() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -850,6 +1058,7 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
@@ -862,8 +1071,10 @@ async fn plan_type_maps_self_serve_business_usage_based_plan() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -878,6 +1089,7 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
@@ -890,8 +1102,10 @@ async fn plan_type_maps_enterprise_cbp_usage_based_plan() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn plan_type_maps_unknown_to_unknown() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -906,6 +1120,7 @@ async fn plan_type_maps_unknown_to_unknown() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
@@ -915,8 +1130,10 @@ async fn plan_type_maps_unknown_to_unknown() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[serial(codex_auth_env)]
|
||||
async fn missing_plan_type_maps_to_unknown() {
|
||||
let codex_home = tempdir().unwrap();
|
||||
let _agent_guard = EnvVarGuard::remove(CODEX_AGENT_IDENTITY_ENV_VAR);
|
||||
let _jwt = write_auth_file(
|
||||
AuthFileParams {
|
||||
openai_api_key: None,
|
||||
@@ -931,6 +1148,7 @@ async fn missing_plan_type_maps_to_unknown() {
|
||||
codex_home.path(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("load auth")
|
||||
|
||||
@@ -16,6 +16,8 @@ use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
use codex_agent_identity::decode_agent_identity_jwt;
|
||||
use codex_agent_identity::fetch_agent_identity_jwks;
|
||||
use codex_app_server_protocol::AuthMode;
|
||||
use codex_app_server_protocol::AuthMode as ApiAuthMode;
|
||||
use codex_protocol::config_types::ForcedLoginMethod;
|
||||
@@ -29,6 +31,7 @@ pub use crate::auth::storage::AuthDotJson;
|
||||
use crate::auth::storage::AuthStorageBackend;
|
||||
use crate::auth::storage::create_auth_storage;
|
||||
use crate::auth::util::try_parse_error_message;
|
||||
use crate::default_client::build_reqwest_client;
|
||||
use crate::default_client::create_client;
|
||||
use crate::token_data::TokenData;
|
||||
use crate::token_data::parse_chatgpt_jwt_claims;
|
||||
@@ -88,6 +91,7 @@ const REFRESH_TOKEN_INVALIDATED_MESSAGE: &str = "Your access token could not be
|
||||
const REFRESH_TOKEN_UNKNOWN_MESSAGE: &str =
|
||||
"Your access token could not be refreshed. Please log out and sign in again.";
|
||||
const REFRESH_TOKEN_ACCOUNT_MISMATCH_MESSAGE: &str = "Your access token could not be refreshed because you have since logged out or signed in to another account. Please sign in again.";
|
||||
const DEFAULT_CHATGPT_BACKEND_BASE_URL: &str = "https://chatgpt.com/backend-api";
|
||||
const REFRESH_TOKEN_URL: &str = "https://auth.openai.com/oauth/token";
|
||||
pub(super) const REVOKE_TOKEN_URL: &str = "https://auth.openai.com/oauth/revoke";
|
||||
pub const REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR: &str = "CODEX_REFRESH_TOKEN_URL_OVERRIDE";
|
||||
@@ -197,6 +201,7 @@ impl CodexAuth {
|
||||
codex_home: &Path,
|
||||
auth_dot_json: AuthDotJson,
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode,
|
||||
chatgpt_base_url: Option<&str>,
|
||||
) -> std::io::Result<Self> {
|
||||
let auth_mode = auth_dot_json.resolved_mode();
|
||||
let client = create_client();
|
||||
@@ -212,7 +217,7 @@ impl CodexAuth {
|
||||
"agent identity auth is missing an agent identity token.",
|
||||
));
|
||||
};
|
||||
return Self::from_agent_identity_jwt(&agent_identity).await;
|
||||
return Self::from_agent_identity_jwt(&agent_identity, chatgpt_base_url).await;
|
||||
}
|
||||
|
||||
let storage_mode = auth_dot_json.storage_mode(auth_credentials_store_mode);
|
||||
@@ -237,17 +242,26 @@ impl CodexAuth {
|
||||
pub async fn from_auth_storage(
|
||||
codex_home: &Path,
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode,
|
||||
chatgpt_base_url: Option<&str>,
|
||||
) -> std::io::Result<Option<Self>> {
|
||||
load_auth(
|
||||
codex_home,
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
auth_credentials_store_mode,
|
||||
chatgpt_base_url,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
|
||||
let record = AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
|
||||
pub async fn from_agent_identity_jwt(
|
||||
jwt: &str,
|
||||
chatgpt_base_url: Option<&str>,
|
||||
) -> std::io::Result<Self> {
|
||||
let base_url = chatgpt_base_url
|
||||
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
|
||||
.trim_end_matches('/')
|
||||
.to_string();
|
||||
let record = verified_agent_identity_record(jwt, &base_url).await?;
|
||||
Ok(Self::AgentIdentity(AgentIdentityAuth::load(record).await?))
|
||||
}
|
||||
|
||||
@@ -493,6 +507,18 @@ pub fn read_codex_agent_identity_from_env() -> Option<String> {
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
async fn verified_agent_identity_record(
|
||||
jwt: &str,
|
||||
chatgpt_base_url: &str,
|
||||
) -> std::io::Result<AgentIdentityAuthRecord> {
|
||||
AgentIdentityAuthRecord::from_agent_identity_jwt(jwt)?;
|
||||
let jwks = fetch_agent_identity_jwks(&build_reqwest_client(), chatgpt_base_url)
|
||||
.await
|
||||
.map_err(std::io::Error::other)?;
|
||||
let claims = decode_agent_identity_jwt(jwt, Some(&jwks)).map_err(std::io::Error::other)?;
|
||||
Ok(claims.into())
|
||||
}
|
||||
|
||||
/// Delete the auth.json file inside `codex_home` if it exists. Returns `Ok(true)`
|
||||
/// if a file was removed, `Ok(false)` if no auth file was present.
|
||||
pub fn logout(
|
||||
@@ -535,12 +561,17 @@ pub fn login_with_api_key(
|
||||
}
|
||||
|
||||
/// Writes an `auth.json` that contains only the Agent Identity token.
|
||||
pub fn login_with_agent_identity(
|
||||
pub async fn login_with_agent_identity(
|
||||
codex_home: &Path,
|
||||
agent_identity: &str,
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode,
|
||||
chatgpt_base_url: Option<&str>,
|
||||
) -> std::io::Result<()> {
|
||||
AgentIdentityAuthRecord::from_agent_identity_jwt(agent_identity)?;
|
||||
let base_url = chatgpt_base_url
|
||||
.unwrap_or(DEFAULT_CHATGPT_BACKEND_BASE_URL)
|
||||
.trim_end_matches('/')
|
||||
.to_string();
|
||||
verified_agent_identity_record(agent_identity, &base_url).await?;
|
||||
let auth_dot_json = AuthDotJson {
|
||||
auth_mode: Some(ApiAuthMode::AgentIdentity),
|
||||
openai_api_key: None,
|
||||
@@ -599,6 +630,7 @@ pub struct AuthConfig {
|
||||
pub auth_credentials_store_mode: AuthCredentialsStoreMode,
|
||||
pub forced_login_method: Option<ForcedLoginMethod>,
|
||||
pub forced_chatgpt_workspace_id: Option<String>,
|
||||
pub chatgpt_base_url: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<()> {
|
||||
@@ -606,6 +638,7 @@ pub async fn enforce_login_restrictions(config: &AuthConfig) -> std::io::Result<
|
||||
&config.codex_home,
|
||||
/*enable_codex_api_key_env*/ true,
|
||||
config.auth_credentials_store_mode,
|
||||
config.chatgpt_base_url.as_deref(),
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
@@ -711,6 +744,7 @@ async fn load_auth(
|
||||
codex_home: &Path,
|
||||
enable_codex_api_key_env: bool,
|
||||
auth_credentials_store_mode: AuthCredentialsStoreMode,
|
||||
chatgpt_base_url: Option<&str>,
|
||||
) -> std::io::Result<Option<CodexAuth>> {
|
||||
// API key via env var takes precedence over any other auth method.
|
||||
if enable_codex_api_key_env && let Some(api_key) = read_codex_api_key_from_env() {
|
||||
@@ -728,6 +762,7 @@ async fn load_auth(
|
||||
codex_home,
|
||||
auth_dot_json,
|
||||
AuthCredentialsStoreMode::Ephemeral,
|
||||
chatgpt_base_url,
|
||||
)
|
||||
.await?;
|
||||
return Ok(Some(auth));
|
||||
@@ -739,7 +774,7 @@ async fn load_auth(
|
||||
}
|
||||
|
||||
if let Some(agent_identity) = read_codex_agent_identity_from_env() {
|
||||
return CodexAuth::from_agent_identity_jwt(&agent_identity)
|
||||
return CodexAuth::from_agent_identity_jwt(&agent_identity, chatgpt_base_url)
|
||||
.await
|
||||
.map(Some);
|
||||
}
|
||||
@@ -751,9 +786,13 @@ async fn load_auth(
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let auth =
|
||||
CodexAuth::from_auth_dot_json(codex_home, auth_dot_json, auth_credentials_store_mode)
|
||||
.await?;
|
||||
let auth = CodexAuth::from_auth_dot_json(
|
||||
codex_home,
|
||||
auth_dot_json,
|
||||
auth_credentials_store_mode,
|
||||
chatgpt_base_url,
|
||||
)
|
||||
.await?;
|
||||
Ok(Some(auth))
|
||||
}
|
||||
|
||||
@@ -1288,6 +1327,7 @@ impl AuthManager {
|
||||
&codex_home,
|
||||
enable_codex_api_key_env,
|
||||
auth_credentials_store_mode,
|
||||
chatgpt_base_url.as_deref(),
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
@@ -1491,6 +1531,7 @@ impl AuthManager {
|
||||
&self.codex_home,
|
||||
self.enable_codex_api_key_env,
|
||||
self.auth_credentials_store_mode,
|
||||
self.chatgpt_base_url.as_deref(),
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
|
||||
@@ -19,6 +19,7 @@ use std::sync::Mutex;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::token_data::TokenData;
|
||||
use codex_agent_identity::AgentIdentityJwtClaims;
|
||||
use codex_agent_identity::decode_agent_identity_jwt;
|
||||
use codex_app_server_protocol::AuthMode;
|
||||
use codex_config::types::AuthCredentialsStoreMode;
|
||||
@@ -59,10 +60,16 @@ pub struct AgentIdentityAuthRecord {
|
||||
|
||||
impl AgentIdentityAuthRecord {
|
||||
pub(crate) fn from_agent_identity_jwt(jwt: &str) -> std::io::Result<Self> {
|
||||
let claims = decode_agent_identity_jwt(jwt, /*public_key_base64*/ None)
|
||||
.map_err(std::io::Error::other)?;
|
||||
let claims =
|
||||
decode_agent_identity_jwt(jwt, /*jwks*/ None).map_err(std::io::Error::other)?;
|
||||
|
||||
Ok(Self {
|
||||
Ok(claims.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AgentIdentityJwtClaims> for AgentIdentityAuthRecord {
|
||||
fn from(claims: AgentIdentityJwtClaims) -> Self {
|
||||
Self {
|
||||
agent_runtime_id: claims.agent_runtime_id,
|
||||
agent_private_key: claims.agent_private_key,
|
||||
account_id: claims.account_id,
|
||||
@@ -70,7 +77,7 @@ impl AgentIdentityAuthRecord {
|
||||
email: claims.email,
|
||||
plan_type: claims.plan_type,
|
||||
chatgpt_account_is_fedramp: claims.chatgpt_account_is_fedramp,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -228,10 +228,14 @@ c2ln",
|
||||
)
|
||||
.expect("auth.json should be written");
|
||||
|
||||
CodexAuth::from_auth_storage(codex_home, AuthCredentialsStoreMode::File)
|
||||
.await
|
||||
.expect("auth should load")
|
||||
.expect("auth should be present")
|
||||
CodexAuth::from_auth_storage(
|
||||
codex_home,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*chatgpt_base_url*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("auth should load")
|
||||
.expect("auth should be present")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -900,6 +900,7 @@ pub async fn run_main(
|
||||
auth_credentials_store_mode: config.cli_auth_credentials_store_mode,
|
||||
forced_login_method: config.forced_login_method,
|
||||
forced_chatgpt_workspace_id: config.forced_chatgpt_workspace_id.clone(),
|
||||
chatgpt_base_url: Some(config.chatgpt_base_url.clone()),
|
||||
})
|
||||
.await
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user