From 3ba4ac8ee6f1a50c74ef40b66bb4330622519cb0 Mon Sep 17 00:00:00 2001 From: Eason Goodale Date: Fri, 8 Aug 2025 14:08:24 -0700 Subject: [PATCH] split --- codex-rs/login/src/auth_file.rs | 47 +++++ codex-rs/login/src/jwt_utils.rs | 23 +++ codex-rs/login/src/lib.rs | 5 + codex-rs/login/src/pkce.rs | 23 +++ codex-rs/login/src/redeem.rs | 157 ++++++++++++++++ codex-rs/login/src/server.rs | 297 ++++-------------------------- codex-rs/login/src/success_url.rs | 32 ++++ codex-rs/login/tests/headless.rs | 1 + 8 files changed, 325 insertions(+), 260 deletions(-) create mode 100644 codex-rs/login/src/auth_file.rs create mode 100644 codex-rs/login/src/jwt_utils.rs create mode 100644 codex-rs/login/src/pkce.rs create mode 100644 codex-rs/login/src/redeem.rs create mode 100644 codex-rs/login/src/success_url.rs diff --git a/codex-rs/login/src/auth_file.rs b/codex-rs/login/src/auth_file.rs new file mode 100644 index 0000000000..7cd7c0bfed --- /dev/null +++ b/codex-rs/login/src/auth_file.rs @@ -0,0 +1,47 @@ +use chrono::Utc; +use serde_json::json; +use std::fs::OpenOptions; +#[cfg(unix)] +use std::os::unix::fs::OpenOptionsExt; +use std::path::Path; + +pub(crate) fn now_rfc3339_z() -> String { + Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true) +} + +pub(crate) fn write_auth_file( + codex_home: &Path, + api_key: Option, + id_token: &str, + access_token: &str, + refresh_token: &str, + account_id: Option, +) -> std::io::Result<()> { + std::fs::create_dir_all(codex_home)?; + let auth_path = codex_home.join("auth.json"); + + let contents = serde_json::to_string_pretty(&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": now_rfc3339_z(), + })) + .unwrap_or_else(|_| "{}".to_string()); + + let mut opts = OpenOptions::new(); + opts.create(true).truncate(true).write(true); + #[cfg(unix)] + { + opts.mode(0o600); + } + let mut f = opts.open(auth_path)?; + use std::io::Write; + f.write_all(contents.as_bytes())?; + f.flush() +} + + diff --git a/codex-rs/login/src/jwt_utils.rs b/codex-rs/login/src/jwt_utils.rs new file mode 100644 index 0000000000..d14964d2ec --- /dev/null +++ b/codex-rs/login/src/jwt_utils.rs @@ -0,0 +1,23 @@ +use base64::Engine as _; + +pub(crate) fn parse_jwt_claims(token: &str) -> serde_json::Value { + let mut parts = token.split('.'); + let _header = parts.next(); + let payload = parts.next(); + let _sig = parts.next(); + match payload { + Some(p) if !p.is_empty() => decode_jwt_payload_segment(p), + _ => serde_json::Value::Object(Default::default()), + } +} + +fn decode_jwt_payload_segment(segment_b64: &str) -> serde_json::Value { + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(segment_b64) + .ok(); + decoded + .and_then(|bytes| serde_json::from_slice::(&bytes).ok()) + .unwrap_or(serde_json::Value::Object(Default::default())) +} + + diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 1151ded007..7f1f02bdc6 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -25,6 +25,11 @@ use crate::token_data::parse_id_token; mod token_data; mod server; +mod pkce; +mod jwt_utils; +mod auth_file; +mod success_url; +mod redeem; pub use server::LoginServerOptions; pub use server::run_local_login_server_with_options; pub use server::process_callback_headless; diff --git a/codex-rs/login/src/pkce.rs b/codex-rs/login/src/pkce.rs new file mode 100644 index 0000000000..a5bb52e2f1 --- /dev/null +++ b/codex-rs/login/src/pkce.rs @@ -0,0 +1,23 @@ +use base64::Engine as _; +use rand::RngCore; +use sha2::{Digest, Sha256}; + +#[derive(Debug, Clone)] +pub(crate) struct PkceCodes { + pub(crate) code_verifier: String, + pub(crate) code_challenge: String, +} + +pub(crate) fn generate_pkce() -> PkceCodes { + let mut bytes = [0u8; 64]; + rand::thread_rng().fill_bytes(&mut bytes); + let code_verifier = hex::encode(bytes); + let digest = Sha256::digest(code_verifier.as_bytes()); + let code_challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest); + PkceCodes { + code_verifier, + code_challenge, + } +} + + diff --git a/codex-rs/login/src/redeem.rs b/codex-rs/login/src/redeem.rs new file mode 100644 index 0000000000..77b443d5ff --- /dev/null +++ b/codex-rs/login/src/redeem.rs @@ -0,0 +1,157 @@ +use crate::auth_file::now_rfc3339_z; +use crate::jwt_utils::parse_jwt_claims; +use chrono::{DateTime, Duration, Utc}; +use reqwest::blocking::Client; +use serde_json::json; +use std::time::Duration as StdDuration; + +const DEFAULT_ISSUER: &str = "https://auth.openai.com"; + +pub(crate) fn maybe_redeem_credits( + issuer: &str, + client_id: &str, + id_token_opt: Option<&str>, + refresh_token: &str, + codex_home: &std::path::Path, +) { + let client = Client::builder() + .timeout(StdDuration::from_secs(30)) + .build(); + let Ok(client) = client else { return }; + + // Parse initial ID token claims and check expiration. + let mut id_token = id_token_opt.unwrap_or("").to_string(); + let mut claims = parse_jwt_claims(&id_token); + + let mut token_expired = true; + if let Some(exp) = claims.get("exp").and_then(|v| v.as_i64()) { + let now_ms = (Utc::now().timestamp_millis()) as i64; + token_expired = now_ms >= exp * 1000; + } + + if token_expired { + eprintln!("Refreshing credentials..."); + #[derive(serde::Serialize)] + struct RefreshReq<'a> { + client_id: &'a str, + grant_type: &'a str, + refresh_token: &'a str, + scope: &'a str, + } + let body = RefreshReq { + client_id, + grant_type: "refresh_token", + refresh_token, + scope: "openid profile email", + }; + let resp = client + .post("https://auth.openai.com/oauth/token") + .json(&body) + .send(); + let Ok(resp) = resp else { return }; + let Ok(val) = resp.json::() else { return }; + let new_id_token = val.get("id_token").and_then(|v| v.as_str()).map(|s| s.to_string()); + let new_refresh_token = val + .get("refresh_token") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + if let (Some(new_id), Some(new_refresh)) = (new_id_token, new_refresh_token) { + // Update file on disk with new tokens. + // Read, modify, write. + let path = codex_home.join("auth.json"); + if let Ok(mut existing) = std::fs::read_to_string(&path) { + if let Ok(mut obj) = serde_json::from_str::(&existing) { + obj["tokens"]["id_token"] = serde_json::Value::String(new_id.clone()); + obj["tokens"]["refresh_token"] = serde_json::Value::String(new_refresh.clone()); + // last_refresh is top-level + obj["last_refresh"] = serde_json::Value::String(now_rfc3339_z()); + existing = serde_json::to_string_pretty(&obj).unwrap_or(existing); + let _ = std::fs::write(&path, existing); + id_token = new_id; + claims = parse_jwt_claims(&id_token); + } + } + } else { + return; + } + } + + // Eligibility checks + let auth_claims = claims + .get("https://api.openai.com/auth") + .cloned() + .unwrap_or(serde_json::Value::Object(Default::default())); + + // Subscription active > 7 days check (parity with Python script) + if let Some(sub_start_str) = auth_claims + .get("chatgpt_subscription_active_start") + .and_then(|v| v.as_str()) + { + if let Ok(sub_start) = DateTime::parse_from_rfc3339(sub_start_str) + .map(|dt| dt.with_timezone(&Utc)) + { + if Utc::now() - sub_start < Duration::days(7) { + eprintln!( + "Sorry, your subscription must be active for more than 7 days to redeem credits." + ); + return; + } + } + } + + let needs_setup = { + let completed = auth_claims + .get("completed_platform_onboarding") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let is_owner = auth_claims + .get("is_org_owner") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + !completed && is_owner + }; + if needs_setup { + eprintln!("Only users with Plus or Pro subscriptions can redeem free API credits."); + return; + } + let plan_type = auth_claims + .get("chatgpt_plan_type") + .and_then(|v| v.as_str()) + .unwrap_or(""); + if plan_type != "plus" && plan_type != "pro" { + eprintln!("Only users with Plus or Pro subscriptions can redeem free API credits."); + return; + } + + let api_host = if issuer == DEFAULT_ISSUER { + "https://api.openai.com" + } else { + "https://api.openai.org" + }; + + let payload = json!({"id_token": id_token}); + let resp = client + .post(format!("{api_host}/v1/billing/redeem_credits")) + .json(&payload) + .send(); + if let Ok(r) = resp { + if let Ok(val) = r.json::() { + let granted = val + .get("granted_chatgpt_subscriber_api_credits") + .and_then(|v| v.as_i64()) + .unwrap_or(0); + if granted > 0 { + let amount = if plan_type == "plus" { "$5" } else { "$50" }; + eprintln!( + "Thanks for being a ChatGPT {} subscriber! If you haven't already redeemed, you should receive {} in API credits.", + if plan_type == "plus" { "Plus" } else { "Pro" }, + amount + ); + } else { + eprintln!("It looks like no credits were granted: {}", val); + } + } + } +} + + diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index f125a7a064..ca1a0a817a 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -1,44 +1,29 @@ -use base64::Engine as _; use chrono::Utc; use rand::RngCore; use reqwest::blocking::Client; use serde::Deserialize; use serde_json::json; -use sha2::{Digest, Sha256}; use std::collections::HashMap; -use std::fs::{OpenOptions}; -#[cfg(unix)] -use std::os::unix::fs::OpenOptionsExt; use std::path::Path; use std::time::Duration; use tiny_http::{Header, Method, Response, Server}; use url::Url; +use url::form_urlencoded; use std::path::PathBuf; +use crate::auth_file::write_auth_file; +use crate::jwt_utils::parse_jwt_claims; +use crate::pkce::generate_pkce; +use crate::redeem::maybe_redeem_credits; +use crate::success_url::build_success_url; + const DEFAULT_PORT: u16 = 1455; const DEFAULT_ISSUER: &str = "https://auth.openai.com"; // Copied from the Python HTML to keep UX consistent. pub const LOGIN_SUCCESS_HTML: &str = include_str!("./success_page.html"); -#[derive(Debug, Clone)] -pub struct PkceCodes { - code_verifier: String, - code_challenge: String, -} - -fn generate_pkce() -> PkceCodes { - // Equivalent to Python's secrets.token_hex(64) - let mut bytes = [0u8; 64]; - rand::thread_rng().fill_bytes(&mut bytes); - let code_verifier = hex::encode(bytes); - let digest = Sha256::digest(code_verifier.as_bytes()); - let code_challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest); - PkceCodes { - code_verifier, - code_challenge, - } -} +// PKCE helpers are in crate::pkce #[derive(Debug, Deserialize)] struct CodeExchangeResponse { @@ -52,198 +37,11 @@ struct TokenExchangeResponse { access_token: String, } -fn now_rfc3339_z() -> String { - Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true) -} +// JWT helpers are in crate::jwt_utils -fn decode_jwt_payload_segment(segment_b64: &str) -> serde_json::Value { - let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(segment_b64) - .ok(); - decoded - .and_then(|bytes| serde_json::from_slice::(&bytes).ok()) - .unwrap_or(serde_json::Value::Object(Default::default())) -} +// Auth file writer is in crate::auth_file -fn parse_jwt_claims(token: &str) -> serde_json::Value { - // Expect three segments: header.payload.signature - let mut parts = token.split('.'); - let _header = parts.next(); - let payload = parts.next(); - let _sig = parts.next(); - match payload { - Some(p) if !p.is_empty() => decode_jwt_payload_segment(p), - _ => serde_json::Value::Object(Default::default()), - } -} - -fn write_auth_file( - codex_home: &Path, - api_key: Option, - id_token: &str, - access_token: &str, - refresh_token: &str, - account_id: Option, -) -> std::io::Result<()> { - std::fs::create_dir_all(codex_home)?; - let auth_path = codex_home.join("auth.json"); - - let contents = 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": now_rfc3339_z(), - }) - .to_string(); - - let mut opts = OpenOptions::new(); - opts.create(true).truncate(true).write(true); - #[cfg(unix)] - { - opts.mode(0o600); - } - let mut f = opts.open(auth_path)?; - use std::io::Write; - f.write_all(contents.as_bytes())?; - f.flush() -} - -fn maybe_redeem_credits( - issuer: &str, - client_id: &str, - id_token_opt: Option<&str>, - refresh_token: &str, - codex_home: &Path, -) { - // Best-effort: any error should not abort the login flow. - let client = Client::builder() - .timeout(Duration::from_secs(30)) - .build(); - let Ok(client) = client else { return }; - - // Parse initial ID token claims and check expiration. - let mut id_token = id_token_opt.unwrap_or("").to_string(); - let mut claims = parse_jwt_claims(&id_token); - - let mut token_expired = true; - if let Some(exp) = claims.get("exp").and_then(|v| v.as_i64()) { - let now_ms = (Utc::now().timestamp_millis()) as i64; - token_expired = now_ms >= exp * 1000; - } - - if token_expired { - eprintln!("Refreshing credentials..."); - #[derive(serde::Serialize)] - struct RefreshReq<'a> { - client_id: &'a str, - grant_type: &'a str, - refresh_token: &'a str, - scope: &'a str, - } - let body = RefreshReq { - client_id, - grant_type: "refresh_token", - refresh_token, - scope: "openid profile email", - }; - let resp = client - .post("https://auth.openai.com/oauth/token") - .json(&body) - .send(); - let Ok(resp) = resp else { return }; - let Ok(val) = resp.json::() else { return }; - let new_id_token = val.get("id_token").and_then(|v| v.as_str()).map(|s| s.to_string()); - let new_refresh_token = val - .get("refresh_token") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - if let (Some(new_id), Some(new_refresh)) = (new_id_token, new_refresh_token) { - // Update file on disk with new tokens. - // Read, modify, write. - let path = codex_home.join("auth.json"); - if let Ok(mut existing) = std::fs::read_to_string(&path) { - if let Ok(mut obj) = serde_json::from_str::(&existing) { - obj["tokens"]["id_token"] = serde_json::Value::String(new_id.clone()); - obj["tokens"]["refresh_token"] = serde_json::Value::String(new_refresh.clone()); - obj["tokens"]["last_refresh"] = serde_json::Value::String(now_rfc3339_z()); - existing = serde_json::to_string_pretty(&obj).unwrap_or(existing); - let _ = std::fs::write(&path, existing); - id_token = new_id; - claims = parse_jwt_claims(&id_token); - } - } - } else { - return; - } - } - - // Eligibility checks - let auth_claims = claims - .get("https://api.openai.com/auth") - .cloned() - .unwrap_or(serde_json::Value::Object(Default::default())); - let needs_setup = { - let completed = auth_claims - .get("completed_platform_onboarding") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - let is_owner = auth_claims - .get("is_org_owner") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - !completed && is_owner - }; - if needs_setup { - eprintln!("Only users with Plus or Pro subscriptions can redeem free API credits."); - return; - } - let plan_type = auth_claims - .get("chatgpt_plan_type") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if plan_type != "plus" && plan_type != "pro" { - eprintln!("Only users with Plus or Pro subscriptions can redeem free API credits."); - return; - } - - let api_host = if issuer == DEFAULT_ISSUER { - "https://api.openai.com" - } else { - "https://api.openai.org" - }; - - let payload = json!({"id_token": id_token}); - let resp = client - .post(format!("{api_host}/v1/billing/redeem_credits")) - .json(&payload) - .send(); - match resp { - Ok(r) => match r.json::() { - Ok(val) => { - let granted = val - .get("granted_chatgpt_subscriber_api_credits") - .and_then(|v| v.as_i64()) - .unwrap_or(0); - if granted > 0 { - let amount = if plan_type == "plus" { "$5" } else { "$50" }; - eprintln!( - "Thanks for being a ChatGPT {} subscriber! If you haven't already redeemed, you should receive {} in API credits.", - if plan_type == "plus" { "Plus" } else { "Pro" }, - amount - ); - } else { - eprintln!("It looks like no credits were granted: {}", val); - } - } - Err(_) => {} - }, - Err(_) => {} - } -} +// Credit redemption logic is in crate::redeem #[derive(Debug, Clone)] pub struct LoginServerOptions { @@ -362,16 +160,8 @@ pub fn run_local_login_server_with_options(opts: LoginServerOptions) -> std::io: } (Method::Get, "/auth/callback") => { // Parse query params - let params: HashMap = query - .as_deref() - .unwrap_or("") - .split('&') - .filter_map(|kv| kv.split_once('=')) - .filter_map(|(k, v)| { - let kk = urlencoding::decode(k).ok()?.into_owned(); - let vv = urlencoding::decode(v).ok()?.into_owned(); - Some((kk, vv)) - }) + let params: HashMap = form_urlencoded::parse(query.as_deref().unwrap_or("").as_bytes()) + .into_owned() .collect(); if params.get("state").map(|s| s.as_str()) != Some(state.as_str()) { @@ -498,17 +288,15 @@ pub fn run_local_login_server_with_options(opts: LoginServerOptions) -> std::io: } else { "https://platform.api.openai.org" }; - let mut success_url = Url::parse(&format!("{}/success", url_base)).unwrap(); - if let Some(id_tok) = Some(&tokens.id_token) { - success_url.query_pairs_mut().append_pair("id_token", id_tok); - } - if let Some(org) = org_id { success_url.query_pairs_mut().append_pair("org_id", org); } - if let Some(proj) = project_id { success_url.query_pairs_mut().append_pair("project_id", proj); } - if let Some(pt) = plan_type { success_url.query_pairs_mut().append_pair("plan_type", pt); } - success_url - .query_pairs_mut() - .append_pair("needs_setup", if needs_setup { "true" } else { "false" }) - .append_pair("platform_url", platform_url); + let success_url = build_success_url( + &url_base, + Some(&tokens.id_token), + org_id, + project_id, + plan_type, + needs_setup, + platform_url, + ); let mut resp = Response::empty(302); resp.add_header(Header::from_bytes(&b"Location"[..], success_url.as_str()).unwrap()); @@ -692,32 +480,20 @@ pub fn process_callback_headless( // Build success URL let base = default_url_base(opts.port); - let mut success_url = Url::parse(&format!("{base}/success")).unwrap(); - success_url.query_pairs_mut().append_pair("id_token", &id_token); - if let Some(org) = org_id { - success_url.query_pairs_mut().append_pair("org_id", org); - } - if let Some(proj) = project_id { - success_url - .query_pairs_mut() - .append_pair("project_id", proj); - } - if !plan_type.is_empty() { - success_url - .query_pairs_mut() - .append_pair("plan_type", plan_type); - } - success_url - .query_pairs_mut() - .append_pair("needs_setup", if needs_setup { "true" } else { "false" }) - .append_pair( - "platform_url", - if opts.issuer == DEFAULT_ISSUER { - "https://platform.openai.com" - } else { - "https://platform.api.openai.org" - }, - ); + let platform_url = if opts.issuer == DEFAULT_ISSUER { + "https://platform.openai.com" + } else { + "https://platform.api.openai.org" + }; + let success_url = build_success_url( + &base, + Some(&id_token), + org_id, + project_id, + if plan_type.is_empty() { None } else { Some(plan_type) }, + needs_setup, + platform_url, + ); Ok(HeadlessOutcome { success_url: success_url.into_string(), @@ -726,3 +502,4 @@ pub fn process_callback_headless( } +// Success URL builder is in crate::success_url diff --git a/codex-rs/login/src/success_url.rs b/codex-rs/login/src/success_url.rs new file mode 100644 index 0000000000..c1dea345dc --- /dev/null +++ b/codex-rs/login/src/success_url.rs @@ -0,0 +1,32 @@ +use url::Url; + +pub(crate) fn build_success_url( + url_base: &str, + id_token: Option<&str>, + org_id: Option<&str>, + project_id: Option<&str>, + plan_type: Option<&str>, + needs_setup: bool, + platform_url: &str, +) -> Url { + let mut success_url = Url::parse(&format!("{}/success", url_base)).expect("valid base url"); + if let Some(id) = id_token { + success_url.query_pairs_mut().append_pair("id_token", id); + } + if let Some(org) = org_id { + success_url.query_pairs_mut().append_pair("org_id", org); + } + if let Some(proj) = project_id { + success_url.query_pairs_mut().append_pair("project_id", proj); + } + if let Some(pt) = plan_type { + success_url.query_pairs_mut().append_pair("plan_type", pt); + } + success_url + .query_pairs_mut() + .append_pair("needs_setup", if needs_setup { "true" } else { "false" }) + .append_pair("platform_url", platform_url); + success_url +} + + diff --git a/codex-rs/login/tests/headless.rs b/codex-rs/login/tests/headless.rs index 56fca3b815..5196718497 100644 --- a/codex-rs/login/tests/headless.rs +++ b/codex-rs/login/tests/headless.rs @@ -66,6 +66,7 @@ fn default_opts(tmp: &TempDir) -> LoginServerOptions { open_browser: false, redeem_credits: true, expose_state_endpoint: false, + testing_timeout_secs: None, } }