Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions codex-rs/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

240 changes: 169 additions & 71 deletions codex-rs/agent-identity/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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)]
Expand Down Expand Up @@ -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,
Comment thread
efrazer-oai marked this conversation as resolved.
Comment thread
efrazer-oai marked this conversation as resolved.
pub iat: usize,
pub exp: usize,
pub agent_runtime_id: String,
pub agent_private_key: String,
pub account_id: String,
Expand Down Expand Up @@ -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))
Comment thread
shijie-oai marked this conversation as resolved.
.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);
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

validation.set_audience(&[AGENT_IDENTITY_JWT_AUDIENCE]);
validation.set_issuer(&[AGENT_IDENTITY_JWT_ISSUER]);
Comment thread
efrazer-oai marked this conversation as resolved.
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Validate issuer against the configured backend, not a fixed host

decode_agent_identity_jwt verifies signatures using JWKS fetched from the caller-provided chatgpt_base_url, but still hard-codes issuer to https://chatgpt.com/.... In non-prod/custom deployments (e.g. localhost or staging), correctly signed tokens from that backend can be rejected solely due to issuer mismatch, breaking AgentIdentity auth load/login flows.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

codex-backend currently always sets issuer to chatgpt.com regardless of the deployment, this should be fine.

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> {
Expand Down Expand Up @@ -279,6 +311,11 @@ 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('/');
format!("{trimmed}/wham/agent-identities/jwks")
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Build JWKS URL for /api/codex bases correctly

agent_identity_jwks_url always appends /wham/agent-identities/jwks to the configured base. With a Codex-API base like https://.../api/codex, this produces .../api/codex/wham/agent-identities/jwks, but JWKS is served at .../api/codex/agent-identities/jwks. AgentIdentity JWT verification will fail in Codex-API style deployments, breaking auth load/login.

Useful? React with 👍 / 👎.

}

pub fn agent_identity_request_id() -> Result<String> {
let mut request_id_bytes = [0u8; 16];
OsRng
Expand All @@ -290,29 +327,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(),
Expand Down Expand Up @@ -471,6 +485,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",
Expand All @@ -480,12 +498,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(),
Expand All @@ -498,15 +519,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(),
Expand All @@ -516,8 +535,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,
Expand All @@ -526,11 +549,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(),
Expand All @@ -540,31 +567,45 @@ 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");
fn decode_agent_identity_jwt_rejects_untrusted_kid() {
let jwks = test_jwks("other-key");

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(),
);
let jwt = jsonwebtoken::encode(
&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",
"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");
}

#[test]
fn decode_agent_identity_jwt_requires_issuer_and_audience() {
let jwks = test_jwks("test-key");
let jwt = jsonwebtoken::encode(
&Header::new(Algorithm::EdDSA),
&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",
Expand All @@ -573,19 +614,76 @@ 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");
}

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 normalize_chatgpt_base_url_strips_codex_before_backend_api() {
fn agent_identity_jwks_url_uses_backend_api_base_url() {
assert_eq!(
agent_identity_jwks_url("https://chatgpt.com/backend-api"),
"https://chatgpt.com/backend-api/wham/agent-identities/jwks"
);
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"
);
}

Expand Down
6 changes: 5 additions & 1 deletion codex-rs/cli/src/login.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -366,6 +369,7 @@ pub async fn run_login_status(cli_config_overrides: CliConfigOverrides) -> ! {
&config.codex_home,
config.cli_auth_credentials_store_mode,
config.agent_identity_authapi_base_url.as_deref(),
Some(&config.chatgpt_base_url),
)
.await
{
Expand Down
Loading
Loading