ai_usagebar/openai/
creds.rs1use 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 #[serde(default)]
31 pub expires_at: Option<String>,
32 #[serde(flatten, default)]
33 pub extra: serde_json::Map<String, serde_json::Value>,
34}
35
36pub 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
52pub 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 pub fn expires_at_secs(&self) -> i64 {
62 parse_jwt_exp(&self.id_token).unwrap_or(0)
63 }
64
65 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
77fn 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 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}