Skip to main content

aither_codex/
auth.rs

1//! OAuth authentication for `ChatGPT` subscription accounts.
2//!
3//! Two login paths are provided:
4//!
5//! - Browser flow: [`Pkce::authorize_url`] + [`listen_for_callback`] +
6//!   [`exchange_code`]. The redirect URI is fixed by the registered client ID
7//!   (`http://localhost:1455/auth/callback`).
8//! - Device flow (headless/SSH): [`request_device_code`] +
9//!   [`poll_device_code`] + [`exchange_code`]. Unlike standard OAuth device
10//!   flow, the PKCE verifier is issued by the server during polling, not
11//!   generated by the client.
12//!
13//! Access tokens are short-lived; use [`refresh_tokens`] with the stored
14//! refresh token. Note that `OpenAI` rotates refresh tokens: every successful
15//! refresh returns a *new* refresh token that must replace the old one.
16
17use std::time::{Duration, SystemTime, UNIX_EPOCH};
18
19use base64::Engine as _;
20use base64::engine::general_purpose::URL_SAFE_NO_PAD;
21use serde::{Deserialize, Serialize};
22use sha2::{Digest, Sha256};
23use zenwave::{Client as _, client};
24
25use crate::constant::{
26    AUTHORIZE_URL, CLIENT_ID, DEVICE_CODE_URL, DEVICE_TOKEN_URL, OAUTH_PORT, OAUTH_SCOPES,
27    REDIRECT_URI, TOKEN_URL,
28};
29use crate::error::CodexError;
30
31/// Buffer for reading the OAuth callback HTTP request.
32const CALLBACK_BUF: usize = 8192;
33
34/// PKCE verifier/challenge pair plus the OAuth `state` nonce.
35#[derive(Debug)]
36pub struct Pkce {
37    verifier: String,
38    challenge: String,
39    state: String,
40}
41
42impl Pkce {
43    /// Generate a fresh PKCE pair and state nonce.
44    ///
45    /// # Panics
46    ///
47    /// Panics if the OS secure random source is unavailable.
48    #[must_use]
49    pub fn generate() -> Self {
50        let mut verifier_bytes = [0u8; 32];
51        let mut state_bytes = [0u8; 16];
52        getrandom::fill(&mut verifier_bytes).expect("secure random source unavailable");
53        getrandom::fill(&mut state_bytes).expect("secure random source unavailable");
54
55        let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
56        let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
57        let state = URL_SAFE_NO_PAD.encode(state_bytes);
58
59        Self {
60            verifier,
61            challenge,
62            state,
63        }
64    }
65
66    /// The `code_verifier` sent during token exchange.
67    #[must_use]
68    pub fn verifier(&self) -> &str {
69        &self.verifier
70    }
71
72    /// The `state` nonce to validate in the callback.
73    #[must_use]
74    pub fn state(&self) -> &str {
75        &self.state
76    }
77
78    /// Build the authorization URL the user must open.
79    ///
80    /// # Panics
81    ///
82    /// Panics if the constant [`AUTHORIZE_URL`] fails to parse (a build-time
83    /// invariant).
84    #[must_use]
85    pub fn authorize_url(&self) -> String {
86        let mut url = url::Url::parse(AUTHORIZE_URL).expect("constant URL parses");
87        url.query_pairs_mut()
88            .append_pair("response_type", "code")
89            .append_pair("client_id", CLIENT_ID)
90            .append_pair("redirect_uri", REDIRECT_URI)
91            .append_pair("scope", OAUTH_SCOPES)
92            .append_pair("code_challenge", &self.challenge)
93            .append_pair("code_challenge_method", "S256")
94            .append_pair("state", &self.state)
95            .append_pair("codex_cli_simplified_flow", "true");
96        url.into()
97    }
98}
99
100/// Listen on `localhost:1455` for a single OAuth callback and return the
101/// authorization code.
102///
103/// Only the first request on the port is handled; `state` is validated against
104/// `expected_state`. A minimal success page is served to the browser.
105///
106/// # Errors
107///
108/// - [`CodexError::Io`] if the port cannot be bound (e.g. another client is
109///   already listening) or the socket fails.
110/// - [`CodexError::Callback`] on malformed requests or a callback carrying an
111///   `error` parameter.
112/// - [`CodexError::StateMismatch`] if the returned `state` does not match.
113#[cfg(not(target_arch = "wasm32"))]
114pub async fn listen_for_callback(expected_state: &str) -> Result<String, CodexError> {
115    use futures_lite::{AsyncReadExt, AsyncWriteExt};
116
117    let addr = std::net::SocketAddr::from(([127, 0, 0, 1], OAUTH_PORT));
118    let listener = async_io::Async::<std::net::TcpListener>::bind(addr)?;
119    tracing::info!(port = OAUTH_PORT, "waiting for OAuth callback");
120
121    let (mut stream, _) = listener.accept().await?;
122    let mut buf = vec![0u8; CALLBACK_BUF];
123    let n = stream.read(&mut buf).await?;
124    let request = String::from_utf8_lossy(&buf[..n]);
125
126    let (code, error) = match parse_callback_request(&request, expected_state) {
127        Ok(code) => (Some(code), None),
128        Err(err) => (None, Some(err.to_string())),
129    };
130
131    let (status, body) = code.as_ref().map_or(
132        (
133            "400 Bad Request",
134            "<html><body><h1>Login failed</h1><p>Return to the terminal for details.</p></body></html>",
135        ),
136        |_| {
137            (
138                "200 OK",
139                "<html><body><h1>Logged in</h1><p>You may close this tab and return to May.</p></body></html>",
140            )
141        },
142    );
143    let response = format!(
144        "HTTP/1.1 {status}\r\ncontent-type: text/html\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
145        body.len()
146    );
147    let _ = stream.write_all(response.as_bytes()).await;
148    let _ = stream.close().await;
149
150    match (code, error) {
151        (Some(code), _) => Ok(code),
152        (None, Some(err)) => Err(CodexError::Callback(err)),
153        _ => Err(CodexError::Callback("empty callback".into())),
154    }
155}
156
157/// Parse `GET /auth/callback?code=…&state=…` out of the raw HTTP request.
158fn parse_callback_request(request: &str, expected_state: &str) -> Result<String, CodexError> {
159    let path = request
160        .split_whitespace()
161        .nth(1)
162        .ok_or_else(|| CodexError::Callback("malformed HTTP request".into()))?;
163    let url = url::Url::parse(&format!("http://localhost{path}"))
164        .map_err(|e| CodexError::Callback(format!("malformed callback path: {e}")))?;
165
166    let mut code = None;
167    let mut state = None;
168    let mut error = None;
169    for (k, v) in url.query_pairs() {
170        match k.as_ref() {
171            "code" => code = Some(v.into_owned()),
172            "state" => state = Some(v.into_owned()),
173            "error" => error = Some(v.into_owned()),
174            _ => {}
175        }
176    }
177
178    if let Some(error) = error {
179        return Err(CodexError::Callback(format!(
180            "authorization failed: {error}"
181        )));
182    }
183    if state.as_deref() != Some(expected_state) {
184        return Err(CodexError::StateMismatch);
185    }
186    code.ok_or_else(|| CodexError::Callback("callback missing code".into()))
187}
188
189/// Exchange an authorization code + PKCE verifier for credentials.
190///
191/// # Errors
192///
193/// Returns [`CodexError::Http`] on transport errors and [`CodexError::Api`] if
194/// the token endpoint rejects the exchange.
195pub async fn exchange_code(code: &str, verifier: &str) -> Result<CodexCredentials, CodexError> {
196    #[derive(Serialize)]
197    struct ExchangeRequest<'a> {
198        grant_type: &'a str,
199        client_id: &'a str,
200        code: &'a str,
201        code_verifier: &'a str,
202        redirect_uri: &'a str,
203    }
204
205    let mut backend = client();
206    let response: TokenResponse = backend
207        .post(TOKEN_URL)?
208        .json_body(&ExchangeRequest {
209            grant_type: "authorization_code",
210            client_id: CLIENT_ID,
211            code,
212            code_verifier: verifier,
213            redirect_uri: REDIRECT_URI,
214        })?
215        .json()
216        .await?;
217    Ok(response.into_credentials())
218}
219
220/// Exchange a refresh token for a fresh credential set.
221///
222/// The returned credentials contain a **new** refresh token — the caller must
223/// persist it; the old one is invalidated by the rotation.
224///
225/// # Errors
226///
227/// Same as [`exchange_code`].
228pub async fn refresh_tokens(refresh_token: &str) -> Result<CodexCredentials, CodexError> {
229    #[derive(Serialize)]
230    struct RefreshRequest<'a> {
231        grant_type: &'a str,
232        refresh_token: &'a str,
233        client_id: &'a str,
234    }
235
236    let mut backend = client();
237    let response: TokenResponse = backend
238        .post(TOKEN_URL)?
239        .json_body(&RefreshRequest {
240            grant_type: "refresh_token",
241            refresh_token,
242            client_id: CLIENT_ID,
243        })?
244        .json()
245        .await?;
246    Ok(response.into_credentials())
247}
248
249/// Begin the device-authorization flow (for SSH/headless environments where a
250/// browser on `localhost:1455` is unavailable).
251///
252/// # Errors
253///
254/// See [`exchange_code`].
255pub async fn request_device_code() -> Result<DeviceCodeResponse, CodexError> {
256    #[derive(Serialize)]
257    struct DeviceCodeRequest<'a> {
258        client_id: &'a str,
259    }
260
261    let mut backend = client();
262    let response: DeviceCodeResponse = backend
263        .post(DEVICE_CODE_URL)?
264        .json_body(&DeviceCodeRequest {
265            client_id: CLIENT_ID,
266        })?
267        .json()
268        .await?;
269    Ok(response)
270}
271
272/// Poll once for device authorization completion.
273///
274/// On success returns `(authorization_code, server_issued_pkce_verifier)` —
275/// feed both into [`exchange_code`]. While the user has not yet approved,
276/// returns [`CodexError::AuthorizationPending`].
277///
278/// # Errors
279///
280/// - [`CodexError::AuthorizationPending`] while waiting for the user.
281/// - [`CodexError::Api`] on expired or invalid device codes.
282pub async fn poll_device_code(
283    device_auth_id: &str,
284    user_code: &str,
285) -> Result<(String, String), CodexError> {
286    #[derive(Serialize)]
287    struct PollRequest<'a> {
288        device_auth_id: &'a str,
289        user_code: &'a str,
290    }
291    #[derive(Deserialize)]
292    struct PollResponse {
293        authorization_code: String,
294        code_verifier: String,
295    }
296
297    // The endpoint answers 403/404 while authorization is pending.
298    let mut backend = client();
299    let result: Result<PollResponse, zenwave::Error> = backend
300        .post(DEVICE_TOKEN_URL)?
301        .json_body(&PollRequest {
302            device_auth_id,
303            user_code,
304        })?
305        .json()
306        .await;
307
308    match result {
309        Ok(poll) => Ok((poll.authorization_code, poll.code_verifier)),
310        Err(zenwave::Error::Http { status, .. })
311            if status.as_u16() == 403 || status.as_u16() == 404 =>
312        {
313            Err(CodexError::AuthorizationPending)
314        }
315        Err(err) => Err(CodexError::Http(err)),
316    }
317}
318
319/// Persisted `ChatGPT` OAuth credentials.
320///
321/// `account_id` and `residency` are decoded from the JWT claims of the
322/// access/ID tokens at construction time.
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct CodexCredentials {
325    /// Short-lived bearer token sent as `Authorization`.
326    pub access_token: String,
327    /// Long-lived token for [`refresh_tokens`]; rotated on every refresh.
328    pub refresh_token: String,
329    /// `OpenID` ID token; a source of the account ID claim.
330    #[serde(default, skip_serializing_if = "Option::is_none")]
331    pub id_token: Option<String>,
332    /// `ChatGPT-Account-Id` header value (from JWT claims).
333    #[serde(default, skip_serializing_if = "Option::is_none")]
334    pub account_id: Option<String>,
335    /// Optional compute-residency constraint (`no_constraint` means none).
336    #[serde(default, skip_serializing_if = "Option::is_none")]
337    pub residency: Option<String>,
338    /// Access-token expiry as seconds since the Unix epoch.
339    #[serde(default, skip_serializing_if = "Option::is_none")]
340    pub expires_at: Option<u64>,
341    /// Earliest time the server wants a refresh performed (epoch seconds).
342    #[serde(default, skip_serializing_if = "Option::is_none")]
343    pub earliest_refresh_at: Option<u64>,
344}
345
346impl CodexCredentials {
347    /// Build credentials from raw token strings, decoding `account_id` and
348    /// `residency` from the JWT claims.
349    #[must_use]
350    pub fn new(access_token: impl Into<String>, refresh_token: impl Into<String>) -> Self {
351        let mut credentials = Self {
352            access_token: access_token.into(),
353            refresh_token: refresh_token.into(),
354            id_token: None,
355            account_id: None,
356            residency: None,
357            expires_at: None,
358            earliest_refresh_at: None,
359        };
360        credentials.fill_claims();
361        credentials
362    }
363
364    /// Attach an ID token (a source of the account-ID claim) and re-decode.
365    #[must_use]
366    pub fn with_id_token(mut self, id_token: impl Into<String>) -> Self {
367        self.id_token = Some(id_token.into());
368        self.fill_claims();
369        self
370    }
371
372    /// Set the access-token expiry (seconds since the Unix epoch).
373    #[must_use]
374    pub const fn with_expires_at(mut self, expires_at: u64) -> Self {
375        self.expires_at = Some(expires_at);
376        self
377    }
378
379    /// Whether the access token is expired or within `skew` of expiring.
380    #[must_use]
381    pub fn is_expired(&self, skew: Duration) -> bool {
382        self.expires_at.is_some_and(|at| {
383            let now = SystemTime::now()
384                .duration_since(UNIX_EPOCH)
385                .map_or(0, |d| d.as_secs());
386            now.saturating_add(skew.as_secs()) >= at
387        })
388    }
389}
390
391/// Response shape of `POST /api/accounts/deviceauth/usercode`.
392#[derive(Debug, Deserialize)]
393pub struct DeviceCodeResponse {
394    /// Server-side authorization session id used while polling.
395    pub device_auth_id: String,
396    /// The code the user must type at the verification URL.
397    pub user_code: String,
398    /// Suggested poll interval in seconds.
399    #[serde(default)]
400    pub interval: Option<String>,
401}
402
403/// Token endpoint response (both grants share the shape).
404#[derive(Debug, Deserialize)]
405struct TokenResponse {
406    access_token: String,
407    refresh_token: String,
408    #[serde(default)]
409    id_token: Option<String>,
410    #[serde(default)]
411    expires_in: Option<u64>,
412    /// Server hint; numeric epoch seconds or an RFC 3339 string.
413    #[serde(default)]
414    earliest_refresh_at: Option<serde_json::Value>,
415}
416
417impl TokenResponse {
418    fn into_credentials(self) -> CodexCredentials {
419        let now = SystemTime::now()
420            .duration_since(UNIX_EPOCH)
421            .map_or(0, |d| d.as_secs());
422        let expires_at = self.expires_in.map(|secs| now + secs);
423        let earliest_refresh_at = self
424            .earliest_refresh_at
425            .as_ref()
426            .and_then(serde_json::Value::as_u64)
427            .or(expires_at);
428        let mut creds = CodexCredentials {
429            access_token: self.access_token,
430            refresh_token: self.refresh_token,
431            id_token: self.id_token,
432            account_id: None,
433            residency: None,
434            expires_at,
435            earliest_refresh_at,
436        };
437        creds.fill_claims();
438        creds
439    }
440}
441
442impl CodexCredentials {
443    /// Decode `account_id` / `residency` out of the JWT claims.
444    ///
445    /// Checked in order: access token `chatgpt_account_id`, then ID token
446    /// `organizations[0].id` under `https://api.openai.com/auth`.
447    fn fill_claims(&mut self) {
448        if self.account_id.is_none() {
449            self.account_id =
450                decode_claims(&self.access_token).and_then(JwtClaims::chatgpt_account_id);
451        }
452        if self.account_id.is_none() {
453            self.account_id = self
454                .id_token
455                .as_deref()
456                .and_then(decode_claims)
457                .and_then(JwtClaims::chatgpt_account_id);
458        }
459        if self.residency.is_none() {
460            self.residency =
461                decode_claims(&self.access_token).and_then(JwtClaims::compute_residency);
462        }
463    }
464}
465
466/// Partial JWT claims shape we care about.
467#[derive(Debug, Deserialize)]
468struct JwtClaims {
469    /// Flat account-id claim (present on access tokens and many id tokens).
470    #[serde(default)]
471    chatgpt_account_id: Option<String>,
472    /// Residency constraint, e.g. `no_constraint` or a region name.
473    #[serde(default)]
474    chatgpt_compute_residency: Option<String>,
475    /// Namespaced auth claim containing organizations (id tokens).
476    #[serde(default, rename = "https://api.openai.com/auth")]
477    auth: Option<AuthClaims>,
478}
479
480impl JwtClaims {
481    fn chatgpt_account_id(self) -> Option<String> {
482        let flat = self.chatgpt_account_id;
483        match self.auth {
484            Some(mut auth) => flat
485                .or_else(|| auth.chatgpt_account_id.take())
486                .or_else(|| auth.first_organization_id()),
487            None => flat,
488        }
489    }
490
491    fn compute_residency(self) -> Option<String> {
492        let flat = self.chatgpt_compute_residency;
493        match self.auth {
494            Some(auth) => flat.or(auth.chatgpt_compute_residency),
495            None => flat,
496        }
497    }
498}
499
500#[derive(Debug, Deserialize)]
501struct AuthClaims {
502    #[serde(default)]
503    chatgpt_account_id: Option<String>,
504    #[serde(default)]
505    chatgpt_compute_residency: Option<String>,
506    #[serde(default)]
507    organizations: Option<Vec<Organization>>,
508}
509
510impl AuthClaims {
511    fn first_organization_id(self) -> Option<String> {
512        self.organizations?.first().and_then(|org| org.id.clone())
513    }
514}
515
516#[derive(Debug, Deserialize)]
517struct Organization {
518    id: Option<String>,
519}
520
521/// Decode the payload section of a JWT without verifying the signature
522/// (claims are used as routing hints, not as trust decisions).
523fn decode_claims(jwt: &str) -> Option<JwtClaims> {
524    let payload = jwt.split('.').nth(1)?;
525    let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?;
526    serde_json::from_slice(&bytes).ok()
527}
528
529#[cfg(test)]
530mod tests {
531    use super::*;
532
533    #[test]
534    fn pkce_challenge_is_s256_of_verifier() {
535        let pkce = Pkce::generate();
536        let expected = URL_SAFE_NO_PAD.encode(Sha256::digest(pkce.verifier().as_bytes()));
537        assert_eq!(pkce.challenge, expected);
538        assert_eq!(pkce.state().len(), 22); // 16 bytes base64url
539    }
540
541    #[test]
542    fn authorize_url_carries_pkce_params() {
543        let pkce = Pkce::generate();
544        let url = pkce.authorize_url();
545        assert!(url.starts_with(AUTHORIZE_URL));
546        assert!(url.contains("code_challenge_method=S256"));
547        assert!(url.contains(&format!("state={}", pkce.state())));
548        assert!(url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
549    }
550
551    #[test]
552    fn callback_parses_code_and_state() {
553        let request =
554            "GET /auth/callback?code=abc123&state=xyz HTTP/1.1\r\nHost: localhost\r\n\r\n";
555        assert_eq!(parse_callback_request(request, "xyz").unwrap(), "abc123");
556    }
557
558    #[test]
559    fn callback_rejects_wrong_state() {
560        let request = "GET /auth/callback?code=abc&state=xyz HTTP/1.1\r\n\r\n";
561        assert!(matches!(
562            parse_callback_request(request, "other"),
563            Err(CodexError::StateMismatch)
564        ));
565    }
566
567    #[test]
568    fn callback_surfaces_error_param() {
569        let request = "GET /auth/callback?error=access_denied&state=xyz HTTP/1.1\r\n\r\n";
570        assert!(matches!(
571            parse_callback_request(request, "xyz"),
572            Err(CodexError::Callback(_))
573        ));
574    }
575
576    #[test]
577    fn nested_claims_decode() {
578        // header.payload.signature — payload claims live under the
579        // https://api.openai.com/auth namespace on current tokens.
580        let payload = serde_json::json!({
581            "https://api.openai.com/auth": {
582                "chatgpt_account_id": "acct-1",
583                "chatgpt_compute_residency": "no_constraint"
584            }
585        });
586        let jwt = format!(
587            "h.{}.s",
588            URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).unwrap())
589        );
590        let claims = decode_claims(&jwt).expect("claims decode");
591        assert_eq!(claims.chatgpt_account_id().as_deref(), Some("acct-1"));
592        let claims = decode_claims(&jwt).expect("claims decode");
593        assert_eq!(claims.compute_residency().as_deref(), Some("no_constraint"));
594    }
595
596    #[test]
597    fn credentials_roundtrip() {
598        let creds = CodexCredentials::new("at", "rt").with_expires_at(42);
599        let json = serde_json::to_string(&creds).unwrap();
600        let back: CodexCredentials = serde_json::from_str(&json).unwrap();
601        assert_eq!(back.access_token, "at");
602        assert_eq!(back.expires_at, Some(42));
603    }
604}