Skip to main content

ai_usagebar/openai/
creds.rs

1//! Read and write `~/.codex/auth.json` — the OAuth state the OpenAI Codex CLI
2//! maintains. Mirrors codexbar's jq paths.
3
4use std::path::{Path, PathBuf};
5
6use base64::Engine;
7use serde::{Deserialize, Serialize};
8
9use crate::cache::atomic_write;
10use crate::error::{AppError, Result};
11
12#[derive(Debug, Clone, Deserialize, Serialize)]
13pub struct AuthFile {
14    pub tokens: Tokens,
15    #[serde(default, skip_serializing_if = "Option::is_none")]
16    pub last_refresh: Option<String>,
17    #[serde(flatten, default)]
18    pub extra: serde_json::Map<String, serde_json::Value>,
19}
20
21#[derive(Debug, Clone, Deserialize, Serialize)]
22pub struct Tokens {
23    pub access_token: String,
24    pub refresh_token: String,
25    pub id_token: String,
26    #[serde(default)]
27    pub account_id: Option<String>,
28    /// Optional explicit expiry from the OAuth server. When absent, we infer
29    /// from the id_token's `exp` claim.
30    #[serde(default)]
31    pub expires_at: Option<String>,
32    #[serde(flatten, default)]
33    pub extra: serde_json::Map<String, serde_json::Value>,
34}
35
36/// Default location: `~/.codex/auth.json`.
37pub fn default_path() -> Result<PathBuf> {
38    let home = std::env::var_os("HOME").ok_or_else(|| AppError::Other("HOME not set".into()))?;
39    Ok(PathBuf::from(home).join(".codex/auth.json"))
40}
41
42pub fn read_from(path: &Path) -> Result<AuthFile> {
43    let raw = std::fs::read_to_string(path).map_err(|e| AppError::io_at(path, e))?;
44    serde_json::from_str(&raw).map_err(|e| {
45        AppError::Credentials(format!(
46            "could not parse {}: {e}. Run `codex login` to re-authenticate.",
47            path.display()
48        ))
49    })
50}
51
52/// Persist updated tokens, preserving any unknown fields. Atomic.
53pub fn write_back(path: &Path, auth: &AuthFile) -> Result<()> {
54    let bytes = serde_json::to_vec_pretty(auth).map_err(AppError::Json)?;
55    atomic_write(path, &bytes)
56}
57
58impl Tokens {
59    /// Compute the Unix-seconds expiry from the embedded id_token. Returns
60    /// 0 (forcing immediate refresh) when the token isn't parseable.
61    pub fn expires_at_secs(&self) -> i64 {
62        parse_jwt_exp(&self.id_token).unwrap_or(0)
63    }
64
65    /// Plan tier from the id_token's nested claim
66    /// `https://api.openai.com/auth.chatgpt_plan_type`.
67    pub fn plan_type_from_id_token(&self) -> Option<String> {
68        let claims = parse_jwt_claims(&self.id_token)?;
69        claims
70            .get("https://api.openai.com/auth")
71            .and_then(|v| v.get("chatgpt_plan_type"))
72            .and_then(|v| v.as_str())
73            .map(|s| s.to_string())
74    }
75}
76
77/// Parse a JWT's `exp` claim. Returns None for malformed tokens.
78fn parse_jwt_exp(token: &str) -> Option<i64> {
79    let claims = parse_jwt_claims(token)?;
80    claims
81        .get("exp")
82        .and_then(|v| v.as_i64())
83        .or_else(|| claims.get("exp").and_then(|v| v.as_f64()).map(|f| f as i64))
84}
85
86fn parse_jwt_claims(token: &str) -> Option<serde_json::Value> {
87    let mut parts = token.split('.');
88    let _header = parts.next()?;
89    let payload = parts.next()?;
90    let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
91        .decode(payload)
92        .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload))
93        .ok()?;
94    serde_json::from_slice(&decoded).ok()
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use std::io::Write;
101    use tempfile::NamedTempFile;
102
103    fn write_auth(s: &str) -> NamedTempFile {
104        let mut f = NamedTempFile::new().unwrap();
105        f.write_all(s.as_bytes()).unwrap();
106        f.flush().unwrap();
107        f
108    }
109
110    /// Build a fake JWT with the given claims (no signature verification).
111    fn fake_jwt(claims: serde_json::Value) -> String {
112        let header = base64::engine::general_purpose::URL_SAFE_NO_PAD
113            .encode(br#"{"alg":"none","typ":"JWT"}"#);
114        let payload =
115            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string().as_bytes());
116        format!("{header}.{payload}.sig")
117    }
118
119    #[test]
120    fn parses_minimal_auth_file() {
121        let jwt = fake_jwt(serde_json::json!({"exp": 1234567890}));
122        let body = format!(
123            r#"{{"tokens":{{"access_token":"AT","refresh_token":"RT",
124                "id_token":"{jwt}","account_id":"acc"}}}}"#
125        );
126        let f = write_auth(&body);
127        let auth = read_from(f.path()).unwrap();
128        assert_eq!(auth.tokens.access_token, "AT");
129        assert_eq!(auth.tokens.account_id.as_deref(), Some("acc"));
130        assert_eq!(auth.tokens.expires_at_secs(), 1234567890);
131    }
132
133    #[test]
134    fn extracts_plan_type_from_id_token() {
135        let jwt = fake_jwt(serde_json::json!({
136            "exp": 1234567890,
137            "https://api.openai.com/auth": {"chatgpt_plan_type": "plus"}
138        }));
139        let body = format!(
140            r#"{{"tokens":{{"access_token":"AT","refresh_token":"RT","id_token":"{jwt}"}}}}"#
141        );
142        let f = write_auth(&body);
143        let auth = read_from(f.path()).unwrap();
144        assert_eq!(
145            auth.tokens.plan_type_from_id_token().as_deref(),
146            Some("plus")
147        );
148    }
149
150    #[test]
151    fn malformed_jwt_returns_zero_exp() {
152        let body = r#"{"tokens":{"access_token":"x","refresh_token":"y","id_token":"not.a.jwt"}}"#;
153        let f = write_auth(body);
154        let auth = read_from(f.path()).unwrap();
155        assert_eq!(auth.tokens.expires_at_secs(), 0);
156        assert!(auth.tokens.plan_type_from_id_token().is_none());
157    }
158
159    #[test]
160    fn malformed_file_returns_credentials_error() {
161        let f = write_auth("not json");
162        let err = read_from(f.path()).unwrap_err();
163        assert!(matches!(err, AppError::Credentials(_)));
164    }
165
166    #[test]
167    fn write_back_preserves_unknown_fields() {
168        let jwt = fake_jwt(serde_json::json!({"exp": 1234567890}));
169        let body = format!(
170            r#"{{"tokens":{{"access_token":"AT","refresh_token":"RT","id_token":"{jwt}"}},
171                "some_other_field":"keep-me"}}"#
172        );
173        let f = write_auth(&body);
174        let mut auth = read_from(f.path()).unwrap();
175        auth.tokens.access_token = "NEW".into();
176        write_back(f.path(), &auth).unwrap();
177
178        let v: serde_json::Value =
179            serde_json::from_str(&std::fs::read_to_string(f.path()).unwrap()).unwrap();
180        assert_eq!(v["some_other_field"], "keep-me");
181        assert_eq!(v["tokens"]["access_token"], "NEW");
182    }
183}