korax-public/mailcore/src/oauth.rs

use crate::types::MailError;
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};

struct Credentials {
    client_id: String,
    client_secret: String,
}

static CREDENTIALS: OnceLock<Credentials> = OnceLock::new();
static DESKTOP_CREDENTIALS: OnceLock<Credentials> = OnceLock::new();

/// Must be called once at startup (from mailmore) before any OAuth operation.
pub fn init_credentials(client_id: &str, client_secret: &str) {
    let _ = CREDENTIALS.set(Credentials {
        client_id: client_id.to_string(),
        client_secret: client_secret.to_string(),
    });
}

/// Desktop/mobile PKCE client credentials (loopback redirect, any port).
pub fn init_desktop_credentials(client_id: &str, client_secret: &str) {
    let _ = DESKTOP_CREDENTIALS.set(Credentials {
        client_id: client_id.to_string(),
        client_secret: client_secret.to_string(),
    });
}

fn get_client_id() -> &'static str {
    CREDENTIALS
        .get()
        .map(|c| c.client_id.as_str())
        .unwrap_or("")
}

fn get_client_secret() -> &'static str {
    CREDENTIALS
        .get()
        .map(|c| c.client_secret.as_str())
        .unwrap_or("")
}

fn get_desktop_client_id() -> &'static str {
    DESKTOP_CREDENTIALS
        .get()
        .map(|c| c.client_id.as_str())
        .unwrap_or_else(|| get_client_id())
}

fn get_desktop_client_secret() -> &'static str {
    DESKTOP_CREDENTIALS
        .get()
        .map(|c| c.client_secret.as_str())
        .unwrap_or_else(|| get_client_secret())
}

struct TokenEntry {
    access_token: String,
    refresh_token: String,
    expires_at: i64,
}

static TOKEN_STORE: OnceLock<Mutex<HashMap<String, TokenEntry>>> = OnceLock::new();

fn get_store() -> &'static Mutex<HashMap<String, TokenEntry>> {
    TOKEN_STORE.get_or_init(|| Mutex::new(HashMap::new()))
}

/// Populate the in-memory token store from Flutter secure storage.
/// Call once at startup for every OAuth2 account, passing the persisted
/// refresh token (access token and expiry start zeroed; first IMAP connect
/// will refresh automatically).
pub fn init_session(email: &str, refresh_token: &str) {
    let mut store = get_store().lock().unwrap();
    let entry = store.entry(email.to_string()).or_insert(TokenEntry {
        access_token: String::new(),
        refresh_token: String::new(),
        expires_at: 0,
    });
    // Only update if the refresh token is non-empty (don't clobber a live entry).
    if !refresh_token.is_empty() {
        entry.refresh_token = refresh_token.to_string();
    }
}

/// Write all three token fields to the in-memory store.
pub fn store_tokens(
    email: &str,
    access_token: &str,
    refresh_token: &str,
    expires_at: i64,
) {
    let mut store = get_store().lock().unwrap();
    store.insert(
        email.to_string(),
        TokenEntry {
            access_token: access_token.to_string(),
            refresh_token: refresh_token.to_string(),
            expires_at,
        },
    );
}

/// Remove the token entry for an account.
pub fn delete_tokens(email: &str) {
    get_store().lock().unwrap().remove(email);
}

/// Return a valid access token, refreshing if within 60 seconds of expiry.
pub async fn get_valid_token(email: &str) -> Result<String, MailError> {
    let (access, expires_at) = {
        let store = get_store().lock().unwrap();
        let entry = store.get(email).ok_or_else(|| {
            MailError::Auth(format!("No OAuth2 session for {}", email))
        })?;
        (entry.access_token.clone(), entry.expires_at)
    };

    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;

    if !access.is_empty() && expires_at - now > 60 {
        return Ok(access);
    }

    tracing::debug!("OAuth2 token for {} is expired or missing, refreshing", email);
    refresh_access_token(email).await
}

/// Refresh the access token using the stored refresh token.
/// Updates the in-memory store and returns the new access token.
pub async fn refresh_access_token(email: &str) -> Result<String, MailError> {
    let refresh = {
        let store = get_store().lock().unwrap();
        store
            .get(email)
            .map(|e| e.refresh_token.clone())
            .filter(|r| !r.is_empty())
            .ok_or_else(|| {
                MailError::Auth(format!("No OAuth2 refresh token for {}", email))
            })?
    };

    let params = [
        ("grant_type", "refresh_token"),
        ("refresh_token", &refresh),
        ("client_id", get_desktop_client_id()),
        ("client_secret", get_desktop_client_secret()),
    ];

    let resp = reqwest::Client::new()
        .post("https://oauth2.googleapis.com/token")
        .form(&params)
        .send()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let body: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let new_access = body["access_token"]
        .as_str()
        .ok_or_else(|| {
            MailError::TokenRefreshFailed(format!(
                "No access_token in refresh response: {}",
                body
            ))
        })?
        .to_string();

    let expires_in = body["expires_in"].as_i64().unwrap_or(3600);
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;

    store_tokens(email, &new_access, &refresh, now + expires_in);

    Ok(new_access)
}

/// Exchange an authorization code from the desktop PKCE flow for tokens.
/// Also resolves the user email from the id_token or userinfo endpoint.
/// Returns (email, access_token, refresh_token).
pub async fn exchange_auth_code_desktop(
    code: &str,
    redirect_uri: &str,
    code_verifier: &str,
) -> Result<(String, String, String), MailError> {
    let client = reqwest::Client::new();

    let params = [
        ("grant_type", "authorization_code"),
        ("code", code),
        ("client_id", get_desktop_client_id()),
        ("client_secret", get_desktop_client_secret()),
        ("redirect_uri", redirect_uri),
        ("code_verifier", code_verifier),
    ];

    let resp = client
        .post("https://oauth2.googleapis.com/token")
        .form(&params)
        .send()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let body: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let access_token = body["access_token"]
        .as_str()
        .ok_or_else(|| {
            MailError::TokenRefreshFailed(format!("No access_token in desktop exchange: {}", body))
        })?
        .to_string();

    let refresh_token = body["refresh_token"]
        .as_str()
        .map(|s| s.to_string())
        .unwrap_or_default();

    // Try to extract email from the id_token JWT payload (base64url-decode the middle part)
    let email = body["id_token"]
        .as_str()
        .and_then(|jwt| {
            let parts: Vec<&str> = jwt.split('.').collect();
            parts.get(1).and_then(|payload| {
                use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
                URL_SAFE_NO_PAD
                    .decode(payload)
                    .ok()
                    .and_then(|bytes| serde_json::from_slice::<serde_json::Value>(&bytes).ok())
                    .and_then(|json| json["email"].as_str().map(|s| s.to_string()))
            })
        });

    let email = match email {
        Some(e) => e,
        None => {
            // Fallback: call userinfo endpoint
            let userinfo: serde_json::Value = client
                .get("https://www.googleapis.com/oauth2/v1/userinfo")
                .bearer_auth(&access_token)
                .send()
                .await
                .map_err(|e| MailError::TokenRefreshFailed(format!("userinfo fetch failed: {}", e)))?
                .json()
                .await
                .map_err(|e| MailError::TokenRefreshFailed(format!("userinfo parse failed: {}", e)))?;
            userinfo["email"]
                .as_str()
                .ok_or_else(|| {
                    MailError::TokenRefreshFailed("No email in userinfo response".to_string())
                })?
                .to_string()
        }
    };

    let expires_in = body["expires_in"].as_i64().unwrap_or(3600);
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;

    store_tokens(&email, &access_token, &refresh_token, now + expires_in);

    Ok((email, access_token, refresh_token))
}

/// Exchange a server-side authorization code for access + refresh tokens.
/// Returns (access_token, refresh_token) so the caller can persist the
/// refresh token via Flutter secure storage.
pub async fn exchange_auth_code(
    email: &str,
    server_auth_code: &str,
) -> Result<(String, String), MailError> {
    let params = [
        ("grant_type", "authorization_code"),
        ("code", server_auth_code),
        ("client_id", get_client_id()),
        ("client_secret", get_client_secret()),
        ("redirect_uri", ""),
    ];

    let resp = reqwest::Client::new()
        .post("https://oauth2.googleapis.com/token")
        .form(&params)
        .send()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let body: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| MailError::TokenRefreshFailed(e.to_string()))?;

    let access_token = body["access_token"]
        .as_str()
        .ok_or_else(|| {
            MailError::TokenRefreshFailed(format!(
                "No access_token in auth code exchange response: {}",
                body
            ))
        })?
        .to_string();

    // Google only returns refresh_token on first consent. Fall back to any
    // previously stored refresh token if the exchange omits it.
    let refresh_token = body["refresh_token"].as_str().map(|s| s.to_string())
        .or_else(|| {
            get_store().lock().unwrap()
                .get(email)
                .map(|e| e.refresh_token.clone())
                .filter(|r| !r.is_empty())
        })
        .unwrap_or_default();

    let expires_in = body["expires_in"].as_i64().unwrap_or(3600);
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_secs() as i64;

    store_tokens(email, &access_token, &refresh_token, now + expires_in);

    Ok((access_token, refresh_token))
}