From 701c2f0802a62e8552b326bb89f02edbd63f92ce Mon Sep 17 00:00:00 2001 From: Eason Goodale Date: Tue, 12 Aug 2025 00:20:54 -0700 Subject: [PATCH] write_auth_json --- codex-rs/login/src/auth.rs | 1 + codex-rs/login/src/auth_store.rs | 118 ++++++++++------- codex-rs/login/src/lib_tests.rs | 119 +++++++++++------- codex-rs/login/src/old_login_python_script.py | 0 codex-rs/login/src/token_data.rs | 100 ++++++++++++--- 5 files changed, 240 insertions(+), 98 deletions(-) create mode 100644 codex-rs/login/src/old_login_python_script.py diff --git a/codex-rs/login/src/auth.rs b/codex-rs/login/src/auth.rs index 7baa704847..4cd4f16894 100644 --- a/codex-rs/login/src/auth.rs +++ b/codex-rs/login/src/auth.rs @@ -128,6 +128,7 @@ impl CodexAuth { openai_api_key: None, tokens: Some(TokenData { id_token: Default::default(), + id_token_raw: String::new(), access_token: "Access Token".to_string(), refresh_token: "test".to_string(), account_id: Some("account_id".to_string()), diff --git a/codex-rs/login/src/auth_store.rs b/codex-rs/login/src/auth_store.rs index 99051144c0..ea5e0f1633 100644 --- a/codex-rs/login/src/auth_store.rs +++ b/codex-rs/login/src/auth_store.rs @@ -79,26 +79,69 @@ pub(crate) fn update_tokens( access_token: Option, refresh_token: Option, ) -> std::io::Result { - // Read, modify, write raw JSON to preserve id_token as a string on disk. - let mut contents = String::new(); - { - let mut f = File::open(auth_file)?; - use std::io::Read as _; - f.read_to_string(&mut contents)?; - } - let mut obj: serde_json::Value = serde_json::from_str(&contents)?; - obj["tokens"]["id_token"] = serde_json::Value::String(id_token); - if let Some(a) = access_token { - obj["tokens"]["access_token"] = serde_json::Value::String(a); - } - if let Some(r) = refresh_token { - obj["tokens"]["refresh_token"] = serde_json::Value::String(r); - } - obj["last_refresh"] = - serde_json::Value::String(Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true)); - let updated = serde_json::to_string_pretty(&obj)?; - std::fs::write(auth_file, updated)?; - // Return parsed structure + let mut prior_access: Option = None; + let mut prior_refresh: Option = None; + let mut auth = match try_read_auth_json(auth_file) { + Ok(a) => a, + Err(_) => { + // Try to salvage existing access/refresh from raw JSON on disk + if let Ok(mut f) = File::open(auth_file) { + let mut contents = String::new(); + use std::io::Read as _; + if f.read_to_string(&mut contents).is_ok() { + if let Ok(val) = serde_json::from_str::(&contents) { + prior_access = val + .get("tokens") + .and_then(|t| t.get("access_token")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + prior_refresh = val + .get("tokens") + .and_then(|t| t.get("refresh_token")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + } + } + } + AuthDotJson { + openai_api_key: None, + tokens: None, + last_refresh: None, + } + } + }; + let now = Utc::now(); + auth.last_refresh = Some(now); + + let new_tokens = match auth.tokens.take() { + Some(mut tokens) => { + tokens.id_token_raw = id_token; + if let Some(a) = access_token.clone() { + tokens.access_token = a; + } + if let Some(r) = refresh_token.clone() { + tokens.refresh_token = r; + } + // Re-parse id_token_raw into parsed fields + tokens.id_token = crate::token_data::parse_id_token(&tokens.id_token_raw) + .map_err(std::io::Error::other)?; + tokens + } + None => { + // Construct fresh TokenData from provided values + let a = access_token + .or_else(|| prior_access.clone()) + .ok_or_else(|| std::io::Error::other("missing access_token"))?; + let r = refresh_token + .or_else(|| prior_refresh.clone()) + .ok_or_else(|| std::io::Error::other("missing refresh_token"))?; + crate::token_data::TokenData::from_raw(id_token, a, r, None) + .map_err(std::io::Error::other)? + } + }; + + auth.tokens = Some(new_tokens); + write_auth_json(auth_file, &auth)?; try_read_auth_json(auth_file) } @@ -114,26 +157,17 @@ pub(crate) fn write_new_auth_json( ) -> std::io::Result<()> { std::fs::create_dir_all(codex_home)?; let auth_file = get_auth_file(codex_home); - // Write explicit JSON preserving raw JWT in id_token - let json_data = serde_json::json!({ - "OPENAI_API_KEY": api_key, - "tokens": { - "id_token": id_token, - "access_token": access_token, - "refresh_token": refresh_token, - "account_id": account_id, - }, - "last_refresh": Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true), - }); - let contents = serde_json::to_string_pretty(&json_data)?; - let mut options = OpenOptions::new(); - options.truncate(true).write(true).create(true); - #[cfg(unix)] - { - options.mode(0o600); - } - let mut file = options.open(&auth_file)?; - use std::io::Write as _; - file.write_all(contents.as_bytes())?; - file.flush() + let tokens = crate::token_data::TokenData::from_raw( + id_token.to_string(), + access_token.to_string(), + refresh_token.to_string(), + account_id, + ) + .map_err(std::io::Error::other)?; + let auth_dot_json = AuthDotJson { + openai_api_key: api_key, + tokens: Some(tokens), + last_refresh: Some(Utc::now()), + }; + write_auth_json(&auth_file, &auth_dot_json) } diff --git a/codex-rs/login/src/lib_tests.rs b/codex-rs/login/src/lib_tests.rs index 989c50ce9d..ef7bb17e12 100644 --- a/codex-rs/login/src/lib_tests.rs +++ b/codex-rs/login/src/lib_tests.rs @@ -1,9 +1,7 @@ #![expect(clippy::expect_used, clippy::unwrap_used)] -use super::*; use crate::auth::AuthMode; use crate::auth::CodexAuth; use crate::auth::load_auth; -use crate::auth_store::AuthDotJson; use crate::auth_store::get_auth_file; use crate::auth_store::logout; use crate::token_data::IdTokenInfo; @@ -54,27 +52,27 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { assert_eq!(AuthMode::ChatGPT, mode); let guard = auth_dot_json.lock().unwrap(); - let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); + let actual = guard.as_ref().expect("AuthDotJson should exist"); + assert_eq!(actual.openai_api_key, None); + let tokens = actual.tokens.as_ref().expect("tokens should exist"); assert_eq!( - &AuthDotJson { - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: Some("user@example.com".to_string()), - chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), - }, - access_token: "test-access-token".to_string(), - refresh_token: "test-refresh-token".to_string(), - account_id: None, - }), - last_refresh: Some( - chrono::DateTime::parse_from_rfc3339(LAST_REFRESH) - .unwrap() - .with_timezone(&chrono::Utc), - ), - }, - auth_dot_json - ) + tokens.id_token, + IdTokenInfo { + email: Some("user@example.com".to_string()), + chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), + } + ); + assert_eq!(tokens.access_token, "test-access-token".to_string()); + assert_eq!(tokens.refresh_token, "test-refresh-token".to_string()); + assert_eq!(tokens.account_id, None); + assert_eq!( + actual.last_refresh, + Some( + chrono::DateTime::parse_from_rfc3339(LAST_REFRESH) + .unwrap() + .with_timezone(&chrono::Utc) + ) + ); } /// Even if the OPENAI_API_KEY is set in auth.json, if the plan is not in @@ -102,27 +100,27 @@ async fn pro_account_with_api_key_still_uses_chatgpt_auth() { assert_eq!(AuthMode::ChatGPT, mode); let guard = auth_dot_json.lock().unwrap(); - let auth_dot_json = guard.as_ref().expect("AuthDotJson should exist"); + let actual = guard.as_ref().expect("AuthDotJson should exist"); + assert_eq!(actual.openai_api_key, None); + let tokens = actual.tokens.as_ref().expect("tokens should exist"); assert_eq!( - &AuthDotJson { - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: Some("user@example.com".to_string()), - chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), - }, - access_token: "test-access-token".to_string(), - refresh_token: "test-refresh-token".to_string(), - account_id: None, - }), - last_refresh: Some( - chrono::DateTime::parse_from_rfc3339(LAST_REFRESH) - .unwrap() - .with_timezone(&chrono::Utc), - ), - }, - auth_dot_json - ) + tokens.id_token, + IdTokenInfo { + email: Some("user@example.com".to_string()), + chatgpt_plan_type: Some(PlanType::Known(KnownPlan::Pro)), + } + ); + assert_eq!(tokens.access_token, "test-access-token".to_string()); + assert_eq!(tokens.refresh_token, "test-refresh-token".to_string()); + assert_eq!(tokens.account_id, None); + assert_eq!( + actual.last_refresh, + Some( + chrono::DateTime::parse_from_rfc3339(LAST_REFRESH) + .unwrap() + .with_timezone(&chrono::Utc) + ) + ); } /// If the OPENAI_API_KEY is set in auth.json and it is an enterprise @@ -289,3 +287,40 @@ fn update_tokens_preserves_id_token_as_string() { let val: serde_json::Value = serde_json::from_str(&raw).unwrap(); assert_eq!(val["tokens"]["id_token"].as_str(), Some(new_id.as_str())); } + +#[test] +fn write_new_auth_json_is_python_compatible_shape() { + let dir = tempdir().unwrap(); + let id_token = { + #[derive(Serialize)] + struct Header { alg: &'static str, typ: &'static str } + let header = Header { alg: "none", typ: "JWT" }; + let payload = serde_json::json!({"sub": "123"}); + let b64 = |b: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(b); + let header_b64 = b64(&serde_json::to_vec(&header).unwrap()); + let payload_b64 = b64(&serde_json::to_vec(&payload).unwrap()); + let signature_b64 = b64(b"sig"); + format!("{header_b64}.{payload_b64}.{signature_b64}") + }; + + crate::auth_store::write_new_auth_json( + dir.path(), + Some("sk-test".to_string()), + &id_token, + "a1", + "r1", + Some("acc".to_string()), + ) + .unwrap(); + + let raw = std::fs::read_to_string(dir.path().join("auth.json")).unwrap(); + let val: serde_json::Value = serde_json::from_str(&raw).unwrap(); + + assert_eq!(val["OPENAI_API_KEY"].as_str(), Some("sk-test")); + assert!(val["last_refresh"].as_str().is_some()); + assert!(val["tokens"].is_object()); + assert_eq!(val["tokens"]["id_token"].as_str(), Some(id_token.as_str())); + assert_eq!(val["tokens"]["access_token"].as_str(), Some("a1")); + assert_eq!(val["tokens"]["refresh_token"].as_str(), Some("r1")); + assert_eq!(val["tokens"]["account_id"].as_str(), Some("acc")); +} diff --git a/codex-rs/login/src/old_login_python_script.py b/codex-rs/login/src/old_login_python_script.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/codex-rs/login/src/token_data.rs b/codex-rs/login/src/token_data.rs index e66dcb1a46..13a9eaea48 100644 --- a/codex-rs/login/src/token_data.rs +++ b/codex-rs/login/src/token_data.rs @@ -3,22 +3,47 @@ use serde::Deserialize; use serde::Serialize; use thiserror::Error; -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)] +#[derive(Clone, Debug, Default)] pub struct TokenData { - /// Flat info parsed from the JWT in auth.json. - #[serde(deserialize_with = "deserialize_id_token")] + /// Flat info parsed from the JWT in auth.json (not serialized). pub id_token: IdTokenInfo, - + /// Raw JWT string used for serialization as `tokens.id_token` on disk. + pub id_token_raw: String, /// This is a JWT. pub access_token: String, - pub refresh_token: String, - pub account_id: Option, } + +impl PartialEq for TokenData { + fn eq(&self, other: &Self) -> bool { + self.id_token == other.id_token + && self.access_token == other.access_token + && self.refresh_token == other.refresh_token + && self.account_id == other.account_id + } +} + +impl Eq for TokenData {} /// Returns true if this is a plan that should use the traditional /// "metered" billing via an API key. impl TokenData { + pub fn from_raw( + id_token_raw: String, + access_token: String, + refresh_token: String, + account_id: Option, + ) -> Result { + let id_token = parse_id_token(&id_token_raw)?; + Ok(Self { + id_token, + id_token_raw, + access_token, + refresh_token, + account_id, + }) + } + pub(crate) fn is_plan_that_should_use_api_key(&self) -> bool { self.id_token .chatgpt_plan_type @@ -27,6 +52,61 @@ impl TokenData { } } +impl Serialize for TokenData { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + #[derive(Serialize)] + struct Ser<'a> { + #[serde(rename = "id_token")] + id_token_raw: &'a str, + access_token: &'a str, + refresh_token: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + account_id: &'a Option, + } + let helper = Ser { + id_token_raw: &self.id_token_raw, + access_token: &self.access_token, + refresh_token: &self.refresh_token, + account_id: &self.account_id, + }; + helper.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for TokenData { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + #[derive(Deserialize)] + struct De { + #[serde(rename = "id_token")] + id_token_raw: String, + access_token: String, + refresh_token: String, + #[serde(default)] + account_id: Option, + } + let De { + id_token_raw, + access_token, + refresh_token, + account_id, + } = De::deserialize(deserializer)?; + let id_token = parse_id_token(&id_token_raw).map_err(serde::de::Error::custom)?; + Ok(TokenData { + id_token, + id_token_raw, + access_token, + refresh_token, + account_id, + }) + } +} + /// Flat subset of useful claims in id_token from auth.json. #[derive(Debug, Clone, PartialEq, Eq, Default, Serialize)] pub struct IdTokenInfo { @@ -107,14 +187,6 @@ pub(crate) fn parse_id_token(id_token: &str) -> Result(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - let s = String::deserialize(deserializer)?; - parse_id_token(&s).map_err(serde::de::Error::custom) -} - // -------- Helpers for parsing OpenAI auth claims from arbitrary JWTs -------- #[derive(Default, Deserialize)]