codex_login/
token_data.rs

1use base64::Engine;
2use chrono::DateTime;
3use chrono::Utc;
4use codex_protocol::auth::PlanType;
5use serde::Deserialize;
6use serde::Serialize;
7use serde::de::DeserializeOwned;
8use thiserror::Error;
9
10#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)]
11pub struct TokenData {
12    /// Flat info parsed from the JWT in auth.json.
13    #[serde(
14        deserialize_with = "deserialize_id_token",
15        serialize_with = "serialize_id_token"
16    )]
17    pub id_token: IdTokenInfo,
18
19    /// This is a JWT.
20    pub access_token: String,
21
22    pub refresh_token: String,
23
24    pub account_id: Option<String>,
25}
26
27/// Flat subset of useful claims in id_token from auth.json.
28#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
29pub struct IdTokenInfo {
30    pub email: Option<String>,
31    /// The ChatGPT subscription plan type
32    /// (e.g., "free", "plus", "pro", "business", "enterprise", "edu").
33    /// (Note: values may vary by backend.)
34    pub chatgpt_plan_type: Option<PlanType>,
35    /// ChatGPT user identifier associated with the token, if present.
36    pub chatgpt_user_id: Option<String>,
37    /// Organization/workspace identifier associated with the token, if present.
38    pub chatgpt_account_id: Option<String>,
39    /// Whether the selected ChatGPT workspace must route through the FedRAMP edge.
40    pub chatgpt_account_is_fedramp: bool,
41    pub raw_jwt: String,
42}
43
44impl IdTokenInfo {
45    pub fn get_chatgpt_plan_type(&self) -> Option<String> {
46        self.chatgpt_plan_type.as_ref().map(|t| match t {
47            PlanType::Known(plan) => plan.display_name().to_string(),
48            PlanType::Unknown(s) => s.clone(),
49        })
50    }
51
52    pub fn get_chatgpt_plan_type_raw(&self) -> Option<String> {
53        self.chatgpt_plan_type.as_ref().map(|t| match t {
54            PlanType::Known(plan) => plan.raw_value().to_string(),
55            PlanType::Unknown(s) => s.clone(),
56        })
57    }
58
59    pub fn is_workspace_account(&self) -> bool {
60        matches!(
61            self.chatgpt_plan_type,
62            Some(PlanType::Known(plan)) if plan.is_workspace_account()
63        )
64    }
65
66    pub fn is_fedramp_account(&self) -> bool {
67        self.chatgpt_account_is_fedramp
68    }
69}
70
71#[derive(Deserialize)]
72struct IdClaims {
73    #[serde(default)]
74    email: Option<String>,
75    #[serde(rename = "https://api.openai.com/profile", default)]
76    profile: Option<ProfileClaims>,
77    #[serde(rename = "https://api.openai.com/auth", default)]
78    auth: Option<AuthClaims>,
79}
80
81#[derive(Deserialize)]
82struct ProfileClaims {
83    #[serde(default)]
84    email: Option<String>,
85}
86
87#[derive(Deserialize)]
88struct AuthClaims {
89    #[serde(default)]
90    chatgpt_plan_type: Option<PlanType>,
91    #[serde(default)]
92    chatgpt_user_id: Option<String>,
93    #[serde(default)]
94    user_id: Option<String>,
95    #[serde(default)]
96    chatgpt_account_id: Option<String>,
97    #[serde(default)]
98    chatgpt_account_is_fedramp: bool,
99}
100
101#[derive(Deserialize)]
102struct StandardJwtClaims {
103    #[serde(default)]
104    exp: Option<i64>,
105}
106
107#[derive(Debug, Error)]
108pub enum IdTokenInfoError {
109    #[error("invalid ID token format")]
110    InvalidFormat,
111    #[error(transparent)]
112    Base64(#[from] base64::DecodeError),
113    #[error(transparent)]
114    Json(#[from] serde_json::Error),
115}
116
117fn decode_jwt_payload<T: DeserializeOwned>(jwt: &str) -> Result<T, IdTokenInfoError> {
118    // JWT format: header.payload.signature
119    let mut parts = jwt.split('.');
120    let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) {
121        (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s),
122        _ => return Err(IdTokenInfoError::InvalidFormat),
123    };
124
125    let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?;
126    let claims = serde_json::from_slice(&payload_bytes)?;
127    Ok(claims)
128}
129
130pub fn parse_jwt_expiration(jwt: &str) -> Result<Option<DateTime<Utc>>, IdTokenInfoError> {
131    let claims: StandardJwtClaims = decode_jwt_payload(jwt)?;
132    Ok(claims
133        .exp
134        .and_then(|exp| DateTime::<Utc>::from_timestamp(exp, 0)))
135}
136
137pub fn parse_chatgpt_jwt_claims(jwt: &str) -> Result<IdTokenInfo, IdTokenInfoError> {
138    let claims: IdClaims = decode_jwt_payload(jwt)?;
139    let email = claims
140        .email
141        .or_else(|| claims.profile.and_then(|profile| profile.email));
142
143    match claims.auth {
144        Some(auth) => Ok(IdTokenInfo {
145            email,
146            raw_jwt: jwt.to_string(),
147            chatgpt_plan_type: auth.chatgpt_plan_type,
148            chatgpt_user_id: auth.chatgpt_user_id.or(auth.user_id),
149            chatgpt_account_id: auth.chatgpt_account_id,
150            chatgpt_account_is_fedramp: auth.chatgpt_account_is_fedramp,
151        }),
152        None => Ok(IdTokenInfo {
153            email,
154            raw_jwt: jwt.to_string(),
155            chatgpt_plan_type: None,
156            chatgpt_user_id: None,
157            chatgpt_account_id: None,
158            chatgpt_account_is_fedramp: false,
159        }),
160    }
161}
162
163fn deserialize_id_token<'de, D>(deserializer: D) -> Result<IdTokenInfo, D::Error>
164where
165    D: serde::Deserializer<'de>,
166{
167    let s = String::deserialize(deserializer)?;
168    parse_chatgpt_jwt_claims(&s).map_err(serde::de::Error::custom)
169}
170
171fn serialize_id_token<S>(id_token: &IdTokenInfo, serializer: S) -> Result<S::Ok, S::Error>
172where
173    S: serde::Serializer,
174{
175    serializer.serialize_str(&id_token.raw_jwt)
176}
177
178#[cfg(test)]
179#[path = "token_data_tests.rs"]
180mod tests;
codex_login/token_data.rs

codex_login/
token_data.rs
