zilliz 1.4.3

//! Auth0 device-code helpers usable from non-blocking contexts (TUI, tests).
//!
//! These mirror the inner state machine of `src/cli/auth.rs::login_with_browser`,
//! but expose it as pure async functions without `println!` side effects.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use anyhow::{bail, Context, Result};
use serde::Deserialize;

use crate::config::manager::ConfigManager;
use crate::model::types::AuthConfig;

#[derive(Debug, Clone, Deserialize)]
pub struct DeviceCodeResponse {
    pub device_code: String,
    pub user_code: String,
    pub verification_uri_complete: String,
    pub expires_in: u64,
    #[serde(default = "default_interval")]
    pub interval: u64,
}

fn default_interval() -> u64 {
    5
}

/// Auth0's device-flow default poll interval. Used as a floor so a bogus
/// `interval: 0` from the token endpoint can't spin the poll loop.
const MIN_POLL_INTERVAL_SECS: u64 = 5;

/// Request/connect timeouts for the Auth0 + login HTTP calls. Mirrors
/// `ApiClient` so a stalled network/TLS connection can't hang the background
/// sign-in task indefinitely (cancellation is only observed once `send()`
/// returns).
const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);

/// Build a reqwest client with bounded timeouts shared by all device-code
/// HTTP calls.
fn http_client() -> Result<reqwest::Client> {
    reqwest::Client::builder()
        .user_agent(format!("zilliz-cli/{}", env!("CARGO_PKG_VERSION")))
        .timeout(HTTP_REQUEST_TIMEOUT)
        .connect_timeout(HTTP_CONNECT_TIMEOUT)
        .build()
        .context("Failed to build HTTP client")
}

#[derive(Debug, Clone, Deserialize)]
pub struct TokenResponse {
    pub access_token: String,
}

#[derive(Debug, Deserialize)]
struct TokenErrorResponse {
    error: String,
    #[serde(default)]
    #[allow(dead_code)]
    error_description: String,
}

/// Cooperative cancellation token shared between the TUI handler and the
/// background polling task. Hand-rolled (rather than pulling in `tokio_util`)
/// because we only need a single atomic flag.
#[derive(Debug, Clone, Default)]
pub struct CancellationToken {
    flag: Arc<AtomicBool>,
}

impl CancellationToken {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn cancel(&self) {
        self.flag.store(true, Ordering::SeqCst);
    }

    pub fn is_cancelled(&self) -> bool {
        self.flag.load(Ordering::SeqCst)
    }
}

/// Request a device code from Auth0. Pure async, no side effects.
pub async fn request_device_code(auth_config: &AuthConfig) -> Result<DeviceCodeResponse> {
    let client = http_client()?;

    let resp: DeviceCodeResponse = client
        .post(format!("{}/oauth/device/code", auth_config.auth0_domain))
        .form(&[
            ("client_id", auth_config.client_id.as_str()),
            ("scope", "openid email profile"),
        ])
        .send()
        .await
        .context("Failed to request device code")?
        .json()
        .await
        .context("Invalid device code response")?;
    Ok(resp)
}

/// Poll the Auth0 token endpoint until success, failure, or the supplied
/// cancellation token fires.
pub async fn poll_for_token(
    auth_config: &AuthConfig,
    device_code: &str,
    interval_secs: u64,
    timeout_secs: u64,
    cancel: CancellationToken,
) -> Result<TokenResponse> {
    let client = http_client()?;
    let url = format!("{}/oauth/token", auth_config.auth0_domain);
    let start = std::time::Instant::now();
    let timeout = Duration::from_secs(timeout_secs);
    // Clamp to a sane floor: Auth0 occasionally returns `interval: 0`, which
    // would otherwise make `target` zero and spin the poll loop.
    let mut interval = interval_secs.max(MIN_POLL_INTERVAL_SECS);

    loop {
        if cancel.is_cancelled() {
            bail!("Cancelled by user");
        }
        if start.elapsed() > timeout {
            bail!("Authentication timed out. Please try again.");
        }

        // Sleep in short slices so cancellation responds quickly.
        let slice = Duration::from_millis(250);
        let mut slept = Duration::ZERO;
        let target = Duration::from_secs(interval);
        while slept < target {
            if cancel.is_cancelled() {
                bail!("Cancelled by user");
            }
            tokio::time::sleep(slice).await;
            slept += slice;
        }

        let resp = client
            .post(&url)
            .form(&[
                ("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
                ("device_code", device_code),
                ("client_id", auth_config.client_id.as_str()),
            ])
            .send()
            .await
            .context("Token poll request failed")?;

        if resp.status().is_success() {
            return resp
                .json::<TokenResponse>()
                .await
                .context("Invalid token response");
        }

        let error_resp: TokenErrorResponse = resp.json().await.unwrap_or(TokenErrorResponse {
            error: "unknown".to_string(),
            error_description: "Unknown error".to_string(),
        });

        match error_resp.error.as_str() {
            "authorization_pending" => continue,
            "slow_down" => {
                interval = (interval + 5).min(60);
                continue;
            }
            "expired_token" => bail!("Authentication timed out. Please try again."),
            "access_denied" => bail!("Authentication was denied."),
            other => bail!("Authentication failed: {}", other),
        }
    }
}

/// Persist an API key into `~/.zilliz/credentials`. Thin wrapper that exists
/// only so callers don't import `ConfigManager` just to call one method.
pub fn save_api_key(config_mgr: &ConfigManager, api_key: &str) -> Result<()> {
    config_mgr.save_api_key_only(api_key)
}

/// Result of exchanging an Auth0 access token for CLI credentials via the
/// `/account/v1/cli/login` endpoint. Mirrors the JSON returned by the API.
#[derive(Clone)]
pub struct LoginPayload {
    pub user_id: String,
    pub email: String,
    pub name: String,
    /// Raw org objects from the API. Passed through to
    /// `ConfigManager::save_login_data` which knows how to extract api keys.
    pub orgs: Vec<serde_json::Value>,
}

// Manual, redacted `Debug` so `{:?}` on a `LoginPayload` (or the `WizardMsg`
// that wraps it) can never leak the API keys carried inside `orgs`.
impl std::fmt::Debug for LoginPayload {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LoginPayload")
            .field("user_id", &self.user_id)
            .field("email", &self.email)
            .field("name", &self.name)
            .field(
                "orgs",
                &format_args!("<{} org(s) redacted>", self.orgs.len()),
            )
            .finish()
    }
}

/// Exchange an Auth0 access token for CLI credentials. Pure async, no
/// `println!` / no `ConfigManager` access — callers persist the result via
/// `ConfigManager::save_login_data`.
pub async fn exchange_token(auth_config: &AuthConfig, access_token: &str) -> Result<LoginPayload> {
    let client = http_client()?;

    let url = format!("{}/account/v1/cli/login", auth_config.login_api);
    let resp = client
        .post(&url)
        .header("Authorization", format!("Bearer {}", access_token))
        .send()
        .await
        .context("Failed to exchange token")?;

    let status = resp.status();
    let body_text = resp.text().await.context("Failed to read login response")?;
    // Do not embed the raw response body in the error: a login endpoint could
    // include credentials/API keys in an error-shaped payload, and TUI/CLI
    // error paths surface this to stderr.
    let body: serde_json::Value = serde_json::from_str(&body_text)
        .with_context(|| format!("Invalid login response (HTTP {})", status.as_u16()))?;

    let code = body
        .get("code")
        .or_else(|| body.get("Code"))
        .and_then(|v| v.as_i64())
        .unwrap_or(-1);

    if !status.is_success() || (code != 0 && code != 200) {
        let msg = body
            .get("msg")
            .or_else(|| body.get("Message"))
            .and_then(|v| v.as_str())
            .unwrap_or("Login failed");
        bail!("Login failed ({}): {}", status.as_u16(), msg);
    }

    let data = body
        .get("data")
        .or_else(|| body.get("Data"))
        .cloned()
        .unwrap_or_default();
    let user = data.get("user").cloned().unwrap_or_default();
    let orgs = data
        .get("orgs")
        .and_then(|v| v.as_array())
        .cloned()
        .unwrap_or_default();

    Ok(LoginPayload {
        user_id: user
            .get("userId")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string(),
        email: user
            .get("email")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string(),
        name: user
            .get("name")
            .and_then(|v| v.as_str())
            .unwrap_or("")
            .to_string(),
        orgs,
    })
}

/// Best-effort browser opener. Errors are intentionally swallowed by callers;
/// the TUI surfaces a hint to copy the URL manually instead.
pub fn open_browser(url: &str) -> std::io::Result<()> {
    #[cfg(target_os = "macos")]
    {
        std::process::Command::new("open").arg(url).spawn()?;
    }
    #[cfg(target_os = "linux")]
    {
        std::process::Command::new("xdg-open").arg(url).spawn()?;
    }
    #[cfg(target_os = "windows")]
    {
        // Use the OS URL handler directly rather than `cmd /c start <url>`,
        // which would route a device-code URL through the Windows shell where
        // metacharacters could be interpreted.
        std::process::Command::new("explorer.exe")
            .arg(url)
            .spawn()?;
    }
    Ok(())
}