agcodex_login/
token_data.rs

1use base64::Engine;
2use serde::Deserialize;
3use serde::Serialize;
4use thiserror::Error;
5
6use crate::AuthMode;
7
8#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Default)]
9pub struct TokenData {
10    /// Flat info parsed from the JWT in auth.json.
11    #[serde(
12        deserialize_with = "deserialize_id_token",
13        serialize_with = "serialize_id_token"
14    )]
15    pub id_token: IdTokenInfo,
16
17    /// This is a JWT.
18    pub access_token: String,
19
20    pub refresh_token: String,
21
22    pub account_id: Option<String>,
23}
24
25impl TokenData {
26    /// Returns true if this is a plan that should use the traditional
27    /// "metered" billing via an API key.
28    pub(crate) fn should_use_api_key(&self, preferred_auth_method: AuthMode) -> bool {
29        if preferred_auth_method == AuthMode::ApiKey {
30            return true;
31        }
32
33        self.id_token
34            .chatgpt_plan_type
35            .as_ref()
36            .is_none_or(|plan| plan.is_plan_that_should_use_api_key())
37    }
38}
39
40/// Flat subset of useful claims in id_token from auth.json.
41#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
42pub struct IdTokenInfo {
43    pub email: Option<String>,
44    /// The ChatGPT subscription plan type
45    /// (e.g., "free", "plus", "pro", "business", "enterprise", "edu").
46    /// (Note: ae has not verified that those are the exact values.)
47    pub(crate) chatgpt_plan_type: Option<PlanType>,
48    pub raw_jwt: String,
49}
50
51impl IdTokenInfo {
52    pub fn get_chatgpt_plan_type(&self) -> Option<String> {
53        self.chatgpt_plan_type.as_ref().map(|t| match t {
54            PlanType::Known(plan) => format!("{plan:?}"),
55            PlanType::Unknown(s) => s.clone(),
56        })
57    }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(untagged)]
62pub(crate) enum PlanType {
63    Known(KnownPlan),
64    Unknown(String),
65}
66
67impl PlanType {
68    const fn is_plan_that_should_use_api_key(&self) -> bool {
69        match self {
70            Self::Known(known) => {
71                use KnownPlan::*;
72                !matches!(known, Free | Plus | Pro | Team)
73            }
74            Self::Unknown(_) => {
75                // Unknown plans should use the API key.
76                true
77            }
78        }
79    }
80
81    pub fn as_string(&self) -> String {
82        match self {
83            Self::Known(known) => format!("{known:?}").to_lowercase(),
84            Self::Unknown(s) => s.clone(),
85        }
86    }
87}
88
89#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
90#[serde(rename_all = "lowercase")]
91pub(crate) enum KnownPlan {
92    Free,
93    Plus,
94    Pro,
95    Team,
96    Business,
97    Enterprise,
98    Edu,
99}
100
101#[derive(Deserialize)]
102struct IdClaims {
103    #[serde(default)]
104    email: Option<String>,
105    #[serde(rename = "https://api.openai.com/auth", default)]
106    auth: Option<AuthClaims>,
107}
108
109#[derive(Deserialize)]
110struct AuthClaims {
111    #[serde(default)]
112    chatgpt_plan_type: Option<PlanType>,
113}
114
115#[derive(Debug, Error)]
116pub enum IdTokenInfoError {
117    #[error("invalid ID token format")]
118    InvalidFormat,
119    #[error(transparent)]
120    Base64(#[from] base64::DecodeError),
121    #[error(transparent)]
122    Json(#[from] serde_json::Error),
123}
124
125pub(crate) fn parse_id_token(id_token: &str) -> Result<IdTokenInfo, IdTokenInfoError> {
126    // JWT format: header.payload.signature
127    let mut parts = id_token.split('.');
128    let (_header_b64, payload_b64, _sig_b64) = match (parts.next(), parts.next(), parts.next()) {
129        (Some(h), Some(p), Some(s)) if !h.is_empty() && !p.is_empty() && !s.is_empty() => (h, p, s),
130        _ => return Err(IdTokenInfoError::InvalidFormat),
131    };
132
133    let payload_bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(payload_b64)?;
134    let claims: IdClaims = serde_json::from_slice(&payload_bytes)?;
135
136    Ok(IdTokenInfo {
137        email: claims.email,
138        chatgpt_plan_type: claims.auth.and_then(|a| a.chatgpt_plan_type),
139        raw_jwt: id_token.to_string(),
140    })
141}
142
143fn deserialize_id_token<'de, D>(deserializer: D) -> Result<IdTokenInfo, D::Error>
144where
145    D: serde::Deserializer<'de>,
146{
147    let s = String::deserialize(deserializer)?;
148    parse_id_token(&s).map_err(serde::de::Error::custom)
149}
150
151fn serialize_id_token<S>(id_token: &IdTokenInfo, serializer: S) -> Result<S::Ok, S::Error>
152where
153    S: serde::Serializer,
154{
155    serializer.serialize_str(&id_token.raw_jwt)
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use serde::Serialize;
162
163    #[test]
164    fn id_token_info_parses_email_and_plan() {
165        #[derive(Serialize)]
166        struct Header {
167            alg: &'static str,
168            typ: &'static str,
169        }
170        let header = Header {
171            alg: "none",
172            typ: "JWT",
173        };
174        let payload = serde_json::json!({
175            "email": "user@example.com",
176            "https://api.openai.com/auth": {
177                "chatgpt_plan_type": "pro"
178            }
179        });
180
181        fn b64url_no_pad(bytes: &[u8]) -> String {
182            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)
183        }
184
185        let header_b64 = b64url_no_pad(&serde_json::to_vec(&header).unwrap());
186        let payload_b64 = b64url_no_pad(&serde_json::to_vec(&payload).unwrap());
187        let signature_b64 = b64url_no_pad(b"sig");
188        let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}");
189
190        let info = parse_id_token(&fake_jwt).expect("should parse");
191        assert_eq!(info.email.as_deref(), Some("user@example.com"));
192        assert_eq!(
193            info.chatgpt_plan_type,
194            Some(PlanType::Known(KnownPlan::Pro))
195        );
196    }
197}
agcodex_login/token_data.rs

agcodex_login/
token_data.rs
