aither-codex 0.6.0

//! OAuth authentication for `ChatGPT` subscription accounts.
//!
//! Two login paths are provided:
//!
//! - Browser flow: [`Pkce::authorize_url`] + [`listen_for_callback`] +
//!   [`exchange_code`]. The redirect URI is fixed by the registered client ID
//!   (`http://localhost:1455/auth/callback`).
//! - Device flow (headless/SSH): [`request_device_code`] +
//!   [`poll_device_code`] + [`exchange_code`]. Unlike standard OAuth device
//!   flow, the PKCE verifier is issued by the server during polling, not
//!   generated by the client.
//!
//! Access tokens are short-lived; use [`refresh_tokens`] with the stored
//! refresh token. Note that `OpenAI` rotates refresh tokens: every successful
//! refresh returns a *new* refresh token that must replace the old one.

use std::time::{Duration, SystemTime, UNIX_EPOCH};

use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use zenwave::{Client as _, client};

use crate::constant::{
    AUTHORIZE_URL, CLIENT_ID, DEVICE_CODE_URL, DEVICE_TOKEN_URL, OAUTH_PORT, OAUTH_SCOPES,
    REDIRECT_URI, TOKEN_URL,
};
use crate::error::CodexError;

/// Buffer for reading the OAuth callback HTTP request.
const CALLBACK_BUF: usize = 8192;

/// PKCE verifier/challenge pair plus the OAuth `state` nonce.
#[derive(Debug)]
pub struct Pkce {
    verifier: String,
    challenge: String,
    state: String,
}

impl Pkce {
    /// Generate a fresh PKCE pair and state nonce.
    ///
    /// # Panics
    ///
    /// Panics if the OS secure random source is unavailable.
    #[must_use]
    pub fn generate() -> Self {
        let mut verifier_bytes = [0u8; 32];
        let mut state_bytes = [0u8; 16];
        getrandom::fill(&mut verifier_bytes).expect("secure random source unavailable");
        getrandom::fill(&mut state_bytes).expect("secure random source unavailable");

        let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
        let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
        let state = URL_SAFE_NO_PAD.encode(state_bytes);

        Self {
            verifier,
            challenge,
            state,
        }
    }

    /// The `code_verifier` sent during token exchange.
    #[must_use]
    pub fn verifier(&self) -> &str {
        &self.verifier
    }

    /// The `state` nonce to validate in the callback.
    #[must_use]
    pub fn state(&self) -> &str {
        &self.state
    }

    /// Build the authorization URL the user must open.
    ///
    /// # Panics
    ///
    /// Panics if the constant [`AUTHORIZE_URL`] fails to parse (a build-time
    /// invariant).
    #[must_use]
    pub fn authorize_url(&self) -> String {
        let mut url = url::Url::parse(AUTHORIZE_URL).expect("constant URL parses");
        url.query_pairs_mut()
            .append_pair("response_type", "code")
            .append_pair("client_id", CLIENT_ID)
            .append_pair("redirect_uri", REDIRECT_URI)
            .append_pair("scope", OAUTH_SCOPES)
            .append_pair("code_challenge", &self.challenge)
            .append_pair("code_challenge_method", "S256")
            .append_pair("state", &self.state)
            .append_pair("codex_cli_simplified_flow", "true");
        url.into()
    }
}

/// Listen on `localhost:1455` for a single OAuth callback and return the
/// authorization code.
///
/// Only the first request on the port is handled; `state` is validated against
/// `expected_state`. A minimal success page is served to the browser.
///
/// # Errors
///
/// - [`CodexError::Io`] if the port cannot be bound (e.g. another client is
///   already listening) or the socket fails.
/// - [`CodexError::Callback`] on malformed requests or a callback carrying an
///   `error` parameter.
/// - [`CodexError::StateMismatch`] if the returned `state` does not match.
#[cfg(not(target_arch = "wasm32"))]
pub async fn listen_for_callback(expected_state: &str) -> Result<String, CodexError> {
    use futures_lite::{AsyncReadExt, AsyncWriteExt};

    let addr = std::net::SocketAddr::from(([127, 0, 0, 1], OAUTH_PORT));
    let listener = async_io::Async::<std::net::TcpListener>::bind(addr)?;
    tracing::info!(port = OAUTH_PORT, "waiting for OAuth callback");

    let (mut stream, _) = listener.accept().await?;
    let mut buf = vec![0u8; CALLBACK_BUF];
    let n = stream.read(&mut buf).await?;
    let request = String::from_utf8_lossy(&buf[..n]);

    let (code, error) = match parse_callback_request(&request, expected_state) {
        Ok(code) => (Some(code), None),
        Err(err) => (None, Some(err.to_string())),
    };

    let (status, body) = code.as_ref().map_or(
        (
            "400 Bad Request",
            "<html><body><h1>Login failed</h1><p>Return to the terminal for details.</p></body></html>",
        ),
        |_| {
            (
                "200 OK",
                "<html><body><h1>Logged in</h1><p>You may close this tab and return to May.</p></body></html>",
            )
        },
    );
    let response = format!(
        "HTTP/1.1 {status}\r\ncontent-type: text/html\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
        body.len()
    );
    let _ = stream.write_all(response.as_bytes()).await;
    let _ = stream.close().await;

    match (code, error) {
        (Some(code), _) => Ok(code),
        (None, Some(err)) => Err(CodexError::Callback(err)),
        _ => Err(CodexError::Callback("empty callback".into())),
    }
}

/// Parse `GET /auth/callback?code=…&state=…` out of the raw HTTP request.
fn parse_callback_request(request: &str, expected_state: &str) -> Result<String, CodexError> {
    let path = request
        .split_whitespace()
        .nth(1)
        .ok_or_else(|| CodexError::Callback("malformed HTTP request".into()))?;
    let url = url::Url::parse(&format!("http://localhost{path}"))
        .map_err(|e| CodexError::Callback(format!("malformed callback path: {e}")))?;

    let mut code = None;
    let mut state = None;
    let mut error = None;
    for (k, v) in url.query_pairs() {
        match k.as_ref() {
            "code" => code = Some(v.into_owned()),
            "state" => state = Some(v.into_owned()),
            "error" => error = Some(v.into_owned()),
            _ => {}
        }
    }

    if let Some(error) = error {
        return Err(CodexError::Callback(format!(
            "authorization failed: {error}"
        )));
    }
    if state.as_deref() != Some(expected_state) {
        return Err(CodexError::StateMismatch);
    }
    code.ok_or_else(|| CodexError::Callback("callback missing code".into()))
}

/// Exchange an authorization code + PKCE verifier for credentials.
///
/// # Errors
///
/// Returns [`CodexError::Http`] on transport errors and [`CodexError::Api`] if
/// the token endpoint rejects the exchange.
pub async fn exchange_code(code: &str, verifier: &str) -> Result<CodexCredentials, CodexError> {
    #[derive(Serialize)]
    struct ExchangeRequest<'a> {
        grant_type: &'a str,
        client_id: &'a str,
        code: &'a str,
        code_verifier: &'a str,
        redirect_uri: &'a str,
    }

    let mut backend = client();
    let response: TokenResponse = backend
        .post(TOKEN_URL)?
        .json_body(&ExchangeRequest {
            grant_type: "authorization_code",
            client_id: CLIENT_ID,
            code,
            code_verifier: verifier,
            redirect_uri: REDIRECT_URI,
        })?
        .json()
        .await?;
    Ok(response.into_credentials())
}

/// Exchange a refresh token for a fresh credential set.
///
/// The returned credentials contain a **new** refresh token — the caller must
/// persist it; the old one is invalidated by the rotation.
///
/// # Errors
///
/// Same as [`exchange_code`].
pub async fn refresh_tokens(refresh_token: &str) -> Result<CodexCredentials, CodexError> {
    #[derive(Serialize)]
    struct RefreshRequest<'a> {
        grant_type: &'a str,
        refresh_token: &'a str,
        client_id: &'a str,
    }

    let mut backend = client();
    let response: TokenResponse = backend
        .post(TOKEN_URL)?
        .json_body(&RefreshRequest {
            grant_type: "refresh_token",
            refresh_token,
            client_id: CLIENT_ID,
        })?
        .json()
        .await?;
    Ok(response.into_credentials())
}

/// Begin the device-authorization flow (for SSH/headless environments where a
/// browser on `localhost:1455` is unavailable).
///
/// # Errors
///
/// See [`exchange_code`].
pub async fn request_device_code() -> Result<DeviceCodeResponse, CodexError> {
    #[derive(Serialize)]
    struct DeviceCodeRequest<'a> {
        client_id: &'a str,
    }

    let mut backend = client();
    let response: DeviceCodeResponse = backend
        .post(DEVICE_CODE_URL)?
        .json_body(&DeviceCodeRequest {
            client_id: CLIENT_ID,
        })?
        .json()
        .await?;
    Ok(response)
}

/// Poll once for device authorization completion.
///
/// On success returns `(authorization_code, server_issued_pkce_verifier)` —
/// feed both into [`exchange_code`]. While the user has not yet approved,
/// returns [`CodexError::AuthorizationPending`].
///
/// # Errors
///
/// - [`CodexError::AuthorizationPending`] while waiting for the user.
/// - [`CodexError::Api`] on expired or invalid device codes.
pub async fn poll_device_code(
    device_auth_id: &str,
    user_code: &str,
) -> Result<(String, String), CodexError> {
    #[derive(Serialize)]
    struct PollRequest<'a> {
        device_auth_id: &'a str,
        user_code: &'a str,
    }
    #[derive(Deserialize)]
    struct PollResponse {
        authorization_code: String,
        code_verifier: String,
    }

    // The endpoint answers 403/404 while authorization is pending.
    let mut backend = client();
    let result: Result<PollResponse, zenwave::Error> = backend
        .post(DEVICE_TOKEN_URL)?
        .json_body(&PollRequest {
            device_auth_id,
            user_code,
        })?
        .json()
        .await;

    match result {
        Ok(poll) => Ok((poll.authorization_code, poll.code_verifier)),
        Err(zenwave::Error::Http { status, .. })
            if status.as_u16() == 403 || status.as_u16() == 404 =>
        {
            Err(CodexError::AuthorizationPending)
        }
        Err(err) => Err(CodexError::Http(err)),
    }
}

/// Persisted `ChatGPT` OAuth credentials.
///
/// `account_id` and `residency` are decoded from the JWT claims of the
/// access/ID tokens at construction time.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodexCredentials {
    /// Short-lived bearer token sent as `Authorization`.
    pub access_token: String,
    /// Long-lived token for [`refresh_tokens`]; rotated on every refresh.
    pub refresh_token: String,
    /// `OpenID` ID token; a source of the account ID claim.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub id_token: Option<String>,
    /// `ChatGPT-Account-Id` header value (from JWT claims).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub account_id: Option<String>,
    /// Optional compute-residency constraint (`no_constraint` means none).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub residency: Option<String>,
    /// Access-token expiry as seconds since the Unix epoch.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<u64>,
    /// Earliest time the server wants a refresh performed (epoch seconds).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub earliest_refresh_at: Option<u64>,
}

impl CodexCredentials {
    /// Build credentials from raw token strings, decoding `account_id` and
    /// `residency` from the JWT claims.
    #[must_use]
    pub fn new(access_token: impl Into<String>, refresh_token: impl Into<String>) -> Self {
        let mut credentials = Self {
            access_token: access_token.into(),
            refresh_token: refresh_token.into(),
            id_token: None,
            account_id: None,
            residency: None,
            expires_at: None,
            earliest_refresh_at: None,
        };
        credentials.fill_claims();
        credentials
    }

    /// Attach an ID token (a source of the account-ID claim) and re-decode.
    #[must_use]
    pub fn with_id_token(mut self, id_token: impl Into<String>) -> Self {
        self.id_token = Some(id_token.into());
        self.fill_claims();
        self
    }

    /// Set the access-token expiry (seconds since the Unix epoch).
    #[must_use]
    pub const fn with_expires_at(mut self, expires_at: u64) -> Self {
        self.expires_at = Some(expires_at);
        self
    }

    /// Whether the access token is expired or within `skew` of expiring.
    #[must_use]
    pub fn is_expired(&self, skew: Duration) -> bool {
        self.expires_at.is_some_and(|at| {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .map_or(0, |d| d.as_secs());
            now.saturating_add(skew.as_secs()) >= at
        })
    }
}

/// Response shape of `POST /api/accounts/deviceauth/usercode`.
#[derive(Debug, Deserialize)]
pub struct DeviceCodeResponse {
    /// Server-side authorization session id used while polling.
    pub device_auth_id: String,
    /// The code the user must type at the verification URL.
    pub user_code: String,
    /// Suggested poll interval in seconds.
    #[serde(default)]
    pub interval: Option<String>,
}

/// Token endpoint response (both grants share the shape).
#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
    refresh_token: String,
    #[serde(default)]
    id_token: Option<String>,
    #[serde(default)]
    expires_in: Option<u64>,
    /// Server hint; numeric epoch seconds or an RFC 3339 string.
    #[serde(default)]
    earliest_refresh_at: Option<serde_json::Value>,
}

impl TokenResponse {
    fn into_credentials(self) -> CodexCredentials {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());
        let expires_at = self.expires_in.map(|secs| now + secs);
        let earliest_refresh_at = self
            .earliest_refresh_at
            .as_ref()
            .and_then(serde_json::Value::as_u64)
            .or(expires_at);
        let mut creds = CodexCredentials {
            access_token: self.access_token,
            refresh_token: self.refresh_token,
            id_token: self.id_token,
            account_id: None,
            residency: None,
            expires_at,
            earliest_refresh_at,
        };
        creds.fill_claims();
        creds
    }
}

impl CodexCredentials {
    /// Decode `account_id` / `residency` out of the JWT claims.
    ///
    /// Checked in order: access token `chatgpt_account_id`, then ID token
    /// `organizations[0].id` under `https://api.openai.com/auth`.
    fn fill_claims(&mut self) {
        if self.account_id.is_none() {
            self.account_id =
                decode_claims(&self.access_token).and_then(JwtClaims::chatgpt_account_id);
        }
        if self.account_id.is_none() {
            self.account_id = self
                .id_token
                .as_deref()
                .and_then(decode_claims)
                .and_then(JwtClaims::chatgpt_account_id);
        }
        if self.residency.is_none() {
            self.residency =
                decode_claims(&self.access_token).and_then(JwtClaims::compute_residency);
        }
    }
}

/// Partial JWT claims shape we care about.
#[derive(Debug, Deserialize)]
struct JwtClaims {
    /// Flat account-id claim (present on access tokens and many id tokens).
    #[serde(default)]
    chatgpt_account_id: Option<String>,
    /// Residency constraint, e.g. `no_constraint` or a region name.
    #[serde(default)]
    chatgpt_compute_residency: Option<String>,
    /// Namespaced auth claim containing organizations (id tokens).
    #[serde(default, rename = "https://api.openai.com/auth")]
    auth: Option<AuthClaims>,
}

impl JwtClaims {
    fn chatgpt_account_id(self) -> Option<String> {
        let flat = self.chatgpt_account_id;
        match self.auth {
            Some(mut auth) => flat
                .or_else(|| auth.chatgpt_account_id.take())
                .or_else(|| auth.first_organization_id()),
            None => flat,
        }
    }

    fn compute_residency(self) -> Option<String> {
        let flat = self.chatgpt_compute_residency;
        match self.auth {
            Some(auth) => flat.or(auth.chatgpt_compute_residency),
            None => flat,
        }
    }
}

#[derive(Debug, Deserialize)]
struct AuthClaims {
    #[serde(default)]
    chatgpt_account_id: Option<String>,
    #[serde(default)]
    chatgpt_compute_residency: Option<String>,
    #[serde(default)]
    organizations: Option<Vec<Organization>>,
}

impl AuthClaims {
    fn first_organization_id(self) -> Option<String> {
        self.organizations?.first().and_then(|org| org.id.clone())
    }
}

#[derive(Debug, Deserialize)]
struct Organization {
    id: Option<String>,
}

/// Decode the payload section of a JWT without verifying the signature
/// (claims are used as routing hints, not as trust decisions).
fn decode_claims(jwt: &str) -> Option<JwtClaims> {
    let payload = jwt.split('.').nth(1)?;
    let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?;
    serde_json::from_slice(&bytes).ok()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pkce_challenge_is_s256_of_verifier() {
        let pkce = Pkce::generate();
        let expected = URL_SAFE_NO_PAD.encode(Sha256::digest(pkce.verifier().as_bytes()));
        assert_eq!(pkce.challenge, expected);
        assert_eq!(pkce.state().len(), 22); // 16 bytes base64url
    }

    #[test]
    fn authorize_url_carries_pkce_params() {
        let pkce = Pkce::generate();
        let url = pkce.authorize_url();
        assert!(url.starts_with(AUTHORIZE_URL));
        assert!(url.contains("code_challenge_method=S256"));
        assert!(url.contains(&format!("state={}", pkce.state())));
        assert!(url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
    }

    #[test]
    fn callback_parses_code_and_state() {
        let request =
            "GET /auth/callback?code=abc123&state=xyz HTTP/1.1\r\nHost: localhost\r\n\r\n";
        assert_eq!(parse_callback_request(request, "xyz").unwrap(), "abc123");
    }

    #[test]
    fn callback_rejects_wrong_state() {
        let request = "GET /auth/callback?code=abc&state=xyz HTTP/1.1\r\n\r\n";
        assert!(matches!(
            parse_callback_request(request, "other"),
            Err(CodexError::StateMismatch)
        ));
    }

    #[test]
    fn callback_surfaces_error_param() {
        let request = "GET /auth/callback?error=access_denied&state=xyz HTTP/1.1\r\n\r\n";
        assert!(matches!(
            parse_callback_request(request, "xyz"),
            Err(CodexError::Callback(_))
        ));
    }

    #[test]
    fn nested_claims_decode() {
        // header.payload.signature — payload claims live under the
        // https://api.openai.com/auth namespace on current tokens.
        let payload = serde_json::json!({
            "https://api.openai.com/auth": {
                "chatgpt_account_id": "acct-1",
                "chatgpt_compute_residency": "no_constraint"
            }
        });
        let jwt = format!(
            "h.{}.s",
            URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).unwrap())
        );
        let claims = decode_claims(&jwt).expect("claims decode");
        assert_eq!(claims.chatgpt_account_id().as_deref(), Some("acct-1"));
        let claims = decode_claims(&jwt).expect("claims decode");
        assert_eq!(claims.compute_residency().as_deref(), Some("no_constraint"));
    }

    #[test]
    fn credentials_roundtrip() {
        let creds = CodexCredentials::new("at", "rt").with_expires_at(42);
        let json = serde_json::to_string(&creds).unwrap();
        let back: CodexCredentials = serde_json::from_str(&json).unwrap();
        assert_eq!(back.access_token, "at");
        assert_eq!(back.expires_at, Some(42));
    }
}