write_auth_json

This commit is contained in:
Eason Goodale
2025-08-12 00:20:54 -07:00
parent 7cad669e09
commit 701c2f0802
5 changed files with 240 additions and 98 deletions

View File

@@ -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()),

View File

@@ -79,26 +79,69 @@ pub(crate) fn update_tokens(
access_token: Option<String>,
refresh_token: Option<String>,
) -> std::io::Result<AuthDotJson> {
// 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<String> = None;
let mut prior_refresh: Option<String> = 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::<serde_json::Value>(&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)
}

View File

@@ -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"));
}

View File

@@ -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<String>,
}
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<String>,
) -> Result<Self, IdTokenInfoError> {
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<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
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<String>,
}
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<D>(deserializer: D) -> Result<Self, D::Error>
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<String>,
}
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<IdTokenInfo, IdTokenInfoE
})
}
fn deserialize_id_token<'de, D>(deserializer: D) -> Result<IdTokenInfo, D::Error>
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)]