pidge-client 1.5.1

//! OAuth 2.0 authorization code + PKCE flow with a one-shot local HTTP server.
//!
//! Why this and not device-code:
//!
//! Device-code flow works perfectly for work/school M365 accounts, but personal
//! Microsoft accounts (live.com / outlook.com / hotmail.com) consistently hit
//! `invalid_request: response_type missing` errors deep inside Microsoft's MSA
//! pipeline. After significant debugging, and trying multiple redirect-URI
//! shapes (nativeclient, http://localhost, urn:ietf:wg:oauth:2.0:oob), the
//! conclusion was that personal MSA's device-code support is fragile in ways
//! that can't be papered over via app-registration tweaks.
//!
//! The auth-code + PKCE + local-server flow is what modern MSAL libraries use
//! and what Microsoft itself recommends for desktop / CLI apps. Both account
//! types route through the same `/oauth2/v2.0/authorize` endpoint and the
//! same redirect-URI plumbing, so M365 and MSA behave identically.
//!
//! Flow:
//!
//! 1. Bind a `TcpListener` on `127.0.0.1:0`; the OS picks a free port.
//! 2. Generate a 64-char random `code_verifier`, derive `code_challenge =
//!    base64url(SHA256(code_verifier))`, and a random `state` for CSRF.
//! 3. Open the user's browser to
//!    `https://login.microsoftonline.com/common/oauth2/v2.0/authorize?
//!     client_id=…&response_type=code&redirect_uri=http://localhost:{port}&
//!     scope=…&code_challenge=…&code_challenge_method=S256&state=…`
//! 4. The user signs in; Microsoft redirects the browser back to
//!    `http://localhost:{port}/?code=…&state=…`.
//! 5. Our local listener accepts exactly one connection, reads the first
//!    request line, extracts the query params, writes a friendly HTML
//!    response, and closes.
//! 6. We POST `/oauth2/v2.0/token` with the auth code, code_verifier, and
//!    redirect_uri to exchange for an access + refresh token.
//!
//! The local server is single-shot: one connection, one response, then close.
//! No port collisions because we let the OS pick; no listening process left
//! behind; no firewall surprises because nothing outside the loopback
//! interface can reach it.

use std::time::Duration;

use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use rand::distr::Alphanumeric;
use rand::{Rng, RngExt, rng};
use serde::Deserialize;
use sha2::{Digest, Sha256};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

use crate::auth::tokens::TokenSet;
use crate::error::ClientError;

/// The result of a successful sign-in.
pub struct AuthSuccess {
    pub tokens: TokenSet,
    /// The raw `id_token`, if Microsoft returned one. Caller uses this to
    /// extract the tenant ID via the existing `jwt::extract_tenant_id`.
    pub id_token: Option<String>,
}

/// Run the full browser flow to completion. Returns the access+refresh
/// tokens and (optionally) the id_token for tenant extraction.
///
/// `on_open` is called once we know the authorize URL; the caller is
/// expected to print it to the user and best-effort spawn their browser.
pub async fn run<F: FnOnce(&str)>(
    http: &reqwest::Client,
    authority_base: &str,
    client_id: &str,
    scope: &str,
    on_open: F,
) -> Result<AuthSuccess, ClientError> {
    // Bind first so we know which port to embed in the redirect URI.
    let listener = TcpListener::bind("127.0.0.1:0")
        .await
        .map_err(ClientError::Io)?;
    let port = listener.local_addr().map_err(ClientError::Io)?.port();
    let redirect_uri = format!("http://localhost:{port}");

    let verifier = make_code_verifier();
    let challenge = make_code_challenge(&verifier);
    let state = make_random(32);

    let authorize_url = build_authorize_url(
        authority_base,
        client_id,
        &redirect_uri,
        scope,
        &challenge,
        &state,
    );
    on_open(&authorize_url);

    let CallbackParams {
        code,
        state: returned_state,
    } = wait_for_callback(listener, Duration::from_secs(300), "Microsoft sign-in").await?;
    if returned_state != state {
        return Err(ClientError::Graph {
            status: 400,
            message: "OAuth state mismatch: possible CSRF or stale request".to_string(),
        });
    }

    exchange_code_to_success(
        http,
        authority_base,
        client_id,
        &code,
        &verifier,
        &redirect_uri,
    )
    .await
}

/// Redeem an authorization code and shape the response as an [`AuthSuccess`].
/// Shared by the local browser flow and hosted callbacks.
pub(crate) async fn exchange_code_to_success(
    http: &reqwest::Client,
    authority_base: &str,
    client_id: &str,
    code: &str,
    code_verifier: &str,
    redirect_uri: &str,
) -> Result<AuthSuccess, ClientError> {
    let tokens_response = exchange_code(
        http,
        authority_base,
        client_id,
        code,
        code_verifier,
        redirect_uri,
    )
    .await?;

    Ok(AuthSuccess {
        tokens: TokenSet {
            access_token: tokens_response.access_token,
            refresh_token: tokens_response.refresh_token.unwrap_or_default(),
            expires_at: chrono::Utc::now()
                + chrono::Duration::seconds(tokens_response.expires_in.unwrap_or(3600)),
        },
        id_token: tokens_response.id_token,
    })
}

pub(crate) fn build_authorize_url(
    authority_base: &str,
    client_id: &str,
    redirect_uri: &str,
    scope: &str,
    challenge: &str,
    state: &str,
) -> String {
    build_authorize_url_with_hint(
        authority_base,
        client_id,
        redirect_uri,
        scope,
        challenge,
        state,
        None,
    )
}

/// [`build_authorize_url`] with an optional `login_hint`: the address
/// Microsoft should preselect in its account picker, for flows where the
/// caller already knows which mailbox the user is expected to sign in with.
pub(crate) fn build_authorize_url_with_hint(
    authority_base: &str,
    client_id: &str,
    redirect_uri: &str,
    scope: &str,
    challenge: &str,
    state: &str,
    login_hint: Option<&str>,
) -> String {
    let mut url = url::Url::parse(&format!("{authority_base}/oauth2/v2.0/authorize"))
        .expect("authority_base is a valid URL");
    {
        let mut q = url.query_pairs_mut();
        q.append_pair("client_id", client_id)
            .append_pair("response_type", "code")
            .append_pair("redirect_uri", redirect_uri)
            .append_pair("response_mode", "query")
            .append_pair("scope", scope)
            .append_pair("state", state)
            .append_pair("code_challenge", challenge)
            .append_pair("code_challenge_method", "S256")
            // `prompt=select_account` forces Microsoft to show the account picker
            // even if the user is already signed in to *some* account. This is
            // what stops "browser is already signed in to my M365 account so
            // pidge auto-grabs that one when I wanted my live.com account".
            .append_pair("prompt", "select_account");
        if let Some(hint) = login_hint.map(str::trim).filter(|h| !h.is_empty()) {
            q.append_pair("login_hint", hint);
        }
    }
    url.into()
}

pub(crate) struct CallbackParams {
    pub(crate) code: String,
    pub(crate) state: String,
}

/// Accept a single connection on the listener, parse the request line for
/// query parameters, write a success/error response, close. Single-shot.
///
/// `timeout` bounds how long we'll wait for the browser to hit the callback,
/// covering both the initial connection *and* the request read that follows,
/// so a client that opens the socket and never sends a request line still
/// hits the deadline rather than hanging forever. The Microsoft flow uses 5
/// minutes (matching Microsoft's own OAuth code TTL); the MCP sign-in flow
/// (see `crate::mcp::oauth`) uses a longer window since it's a separate,
/// often manually-triggered, sign-in step.
///
/// `label` identifies who this callback is for (e.g. `"Microsoft sign-in"`
/// or `"MCP sign-in"`); it prefixes the error message when the remote party
/// reports `error`/`error_description`, so an MCP-server error isn't
/// misattributed to Microsoft.
/// A timeout for the user to read: whole minutes as "5 min", else seconds.
fn timeout_label(timeout: Duration) -> String {
    let secs = timeout.as_secs();
    if secs >= 60 && secs.is_multiple_of(60) {
        format!("{} min", secs / 60)
    } else {
        format!("{secs}s")
    }
}

pub(crate) async fn wait_for_callback(
    listener: TcpListener,
    timeout: Duration,
    label: &str,
) -> Result<CallbackParams, ClientError> {
    let deadline = tokio::time::Instant::now() + timeout;
    let timed_out = || ClientError::Graph {
        status: 408,
        message: format!(
            "timed out waiting for browser sign-in ({})",
            timeout_label(timeout)
        ),
    };

    let (mut stream, _) = tokio::time::timeout_at(deadline, listener.accept())
        .await
        .map_err(|_| timed_out())?
        .map_err(ClientError::Io)?;

    // We only need the first ~1KB to parse the request line "GET /?...".
    let mut buf = [0u8; 2048];
    let n = tokio::time::timeout_at(deadline, stream.read(&mut buf))
        .await
        .map_err(|_| timed_out())?
        .map_err(ClientError::Io)?;
    let request = std::str::from_utf8(&buf[..n]).unwrap_or("");
    let first_line = request.lines().next().unwrap_or("");
    let path_and_query =
        first_line
            .split_whitespace()
            .nth(1)
            .ok_or_else(|| ClientError::Graph {
                status: 400,
                message: "malformed browser callback request".to_string(),
            })?;

    // Parse "/?code=...&state=..." (or "/?error=...&error_description=...").
    let query_start = path_and_query.find('?').unwrap_or(path_and_query.len());
    let query = &path_and_query[query_start.saturating_add(1)..];
    let pairs: Vec<(String, String)> = url::form_urlencoded::parse(query.as_bytes())
        .map(|(k, v)| (k.into_owned(), v.into_owned()))
        .collect();

    let mut code: Option<String> = None;
    let mut state: Option<String> = None;
    let mut err: Option<String> = None;
    let mut err_description: Option<String> = None;
    for (k, v) in pairs {
        match k.as_str() {
            "code" => code = Some(v),
            "state" => state = Some(v),
            "error" => err = Some(v),
            "error_description" => err_description = Some(v),
            _ => {}
        }
    }

    if let Some(e) = err {
        let detail = err_description.unwrap_or_default();
        write_html(&mut stream, &error_page_html(&e, &detail), "Sign-in failed")
            .await
            .ok();
        return Err(ClientError::Graph {
            status: 400,
            message: format!("{label}: {e} ({detail})"),
        });
    }

    let code = code.ok_or_else(|| ClientError::Graph {
        status: 400,
        message: "browser callback missing `code` parameter".to_string(),
    })?;
    let state = state.unwrap_or_default();

    write_html(&mut stream, SUCCESS_HTML, "Signed in")
        .await
        .ok();

    Ok(CallbackParams { code, state })
}

async fn write_html(
    stream: &mut tokio::net::TcpStream,
    body: &str,
    title: &str,
) -> std::io::Result<()> {
    let body_bytes = body.as_bytes();
    let response = format!(
        "HTTP/1.1 200 OK\r\n\
         Content-Type: text/html; charset=utf-8\r\n\
         Content-Length: {}\r\n\
         Connection: close\r\n\
         X-Title: {}\r\n\
         \r\n",
        body_bytes.len(),
        title,
    );
    stream.write_all(response.as_bytes()).await?;
    stream.write_all(body_bytes).await?;
    stream.shutdown().await
}

const SUCCESS_HTML: &str = r#"<!doctype html><html><head><meta charset="utf-8"><title>Signed in</title>
<style>
body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", system-ui, sans-serif;
       max-width: 480px; margin: 80px auto; text-align: center; color: #1d1d1f; }
.check { font-size: 48px; color: #34c759; }
h1 { font-size: 24px; margin: 16px 0 8px; }
p { color: #6e6e73; }
</style></head>
<body>
  <div class="check">✓</div>
  <h1>Signed in to pidge</h1>
  <p>You can close this window and return to the terminal.</p>
</body></html>"#;

fn error_page_html(err: &str, description: &str) -> String {
    let err = html_escape(err);
    let description = html_escape(description);
    format!(
        r#"<!doctype html><html><head><meta charset="utf-8"><title>Sign-in failed</title>
<style>
body {{ font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", system-ui, sans-serif;
       max-width: 540px; margin: 80px auto; color: #1d1d1f; }}
.x {{ font-size: 48px; color: #ff3b30; text-align: center; }}
h1 {{ font-size: 22px; margin: 16px 0 8px; text-align: center; }}
.detail {{ background: #f5f5f7; padding: 16px; border-radius: 8px; font-family: ui-monospace, monospace;
         font-size: 13px; white-space: pre-wrap; word-break: break-word; }}
</style></head>
<body>
  <div class="x">✕</div>
  <h1>Sign-in failed</h1>
  <p class="detail"><strong>{err}</strong>
{description}</p>
  <p>You can close this window. Return to the terminal for next steps.</p>
</body></html>"#
    )
}

/// Escape the five HTML-significant characters. `err` and `error_description`
/// are fixed strings from pidge's own server, but the MCP flow (see
/// `crate::mcp::oauth`) can point this callback at any server the user
/// configures, so this page must not trust its input to be markup-free.
fn html_escape(s: &str) -> String {
    s.chars()
        .fold(String::with_capacity(s.len()), |mut acc, c| {
            match c {
                '&' => acc.push_str("&amp;"),
                '<' => acc.push_str("&lt;"),
                '>' => acc.push_str("&gt;"),
                '"' => acc.push_str("&quot;"),
                '\'' => acc.push_str("&#39;"),
                _ => acc.push(c),
            }
            acc
        })
}

// --- PKCE helpers ----------------------------------------------------------

/// RFC 7636 says the code_verifier is "a high-entropy cryptographic random
/// STRING, using the unreserved characters … with a minimum length of 43
/// characters and a maximum length of 128 characters." 64 alphanumerics is
/// comfortably inside the spec and gives ~380 bits of entropy.
pub(crate) fn make_code_verifier() -> String {
    let mut rng = rng();
    (0..64).map(|_| rng.sample(Alphanumeric) as char).collect()
}

/// `base64url(SHA256(code_verifier))` per RFC 7636 §4.2. URL_SAFE_NO_PAD is
/// the exact encoding the OAuth spec requires.
pub(crate) fn make_code_challenge(verifier: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(verifier.as_bytes());
    URL_SAFE_NO_PAD.encode(hasher.finalize())
}

/// 32 bytes of OS random → URL-safe base64. Used for the `state` CSRF nonce.
pub(crate) fn make_random(byte_len: usize) -> String {
    let mut buf = vec![0u8; byte_len];
    rng().fill_bytes(&mut buf);
    URL_SAFE_NO_PAD.encode(&buf)
}

// --- token exchange --------------------------------------------------------

#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
    refresh_token: Option<String>,
    expires_in: Option<i64>,
    id_token: Option<String>,
}

async fn exchange_code(
    http: &reqwest::Client,
    authority_base: &str,
    client_id: &str,
    code: &str,
    code_verifier: &str,
    redirect_uri: &str,
) -> Result<TokenResponse, ClientError> {
    let url = format!("{authority_base}/oauth2/v2.0/token");
    let params = [
        ("client_id", client_id),
        ("grant_type", "authorization_code"),
        ("code", code),
        ("code_verifier", code_verifier),
        ("redirect_uri", redirect_uri),
    ];
    let resp = http.post(&url).form(&params).send().await?;
    let status = resp.status();
    if !status.is_success() {
        let text = resp.text().await.unwrap_or_default();
        return Err(ClientError::Graph {
            status: status.as_u16(),
            message: text,
        });
    }
    Ok(resp.json().await?)
}

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

    #[test]
    fn timeout_label_uses_minutes_when_whole() {
        assert_eq!(timeout_label(Duration::from_secs(300)), "5 min");
        assert_eq!(timeout_label(Duration::from_secs(600)), "10 min");
        assert_eq!(timeout_label(Duration::from_secs(45)), "45s");
    }

    #[test]
    fn code_verifier_is_64_alphanumerics() {
        let v = make_code_verifier();
        assert_eq!(v.len(), 64);
        assert!(v.chars().all(|c| c.is_ascii_alphanumeric()));
    }

    #[test]
    fn challenge_matches_rfc_7636_example() {
        // RFC 7636 §4 example:
        //   verifier  = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk"
        //   challenge = "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
        let verifier = "dBjftJeZ4CVP-mB92K27uhbUJU1p1r_wW1gFWFOEjXk";
        assert_eq!(
            make_code_challenge(verifier),
            "E9Melhoa2OwvFrEMTJguCHaoeK1t8URWbuGJSstw-cM"
        );
    }

    #[test]
    fn authorize_url_contains_required_params() {
        let url = build_authorize_url(
            "https://login.microsoftonline.com/common",
            "client-id-here",
            "http://localhost:47821",
            "User.Read offline_access",
            "challenge-here",
            "state-here",
        );
        assert!(url.contains("client_id=client-id-here"));
        assert!(url.contains("response_type=code"));
        assert!(url.contains("code_challenge=challenge-here"));
        assert!(url.contains("code_challenge_method=S256"));
        assert!(url.contains("state=state-here"));
        // redirect_uri is percent-encoded inside the query string.
        assert!(url.contains("redirect_uri=http%3A%2F%2Flocalhost%3A47821"));
        assert!(url.contains("prompt=select_account"));
        assert!(!url.contains("login_hint"));
    }

    #[test]
    fn authorize_url_carries_the_login_hint_when_given() {
        let url = build_authorize_url_with_hint(
            "https://login.microsoftonline.com/common",
            "client-id-here",
            "http://localhost:47821",
            "scope-here",
            "challenge-here",
            "state-here",
            Some("jane.doe@example.com"),
        );
        assert!(url.contains("login_hint=jane.doe%40example.com"), "{url}");
        assert!(url.contains("prompt=select_account"));
        let blank = build_authorize_url_with_hint(
            "https://login.microsoftonline.com/common",
            "c",
            "http://localhost:1",
            "s",
            "ch",
            "st",
            Some("  "),
        );
        assert!(!blank.contains("login_hint"));
    }

    #[test]
    fn random_state_is_unique_per_call() {
        let a = make_random(32);
        let b = make_random(32);
        assert_ne!(a, b);
        assert_eq!(URL_SAFE_NO_PAD.decode(&a).unwrap().len(), 32);
    }

    #[test]
    fn error_page_html_escapes_html_special_characters() {
        let page = error_page_html("<script>alert(1)</script>", "quote\" apostrophe' amp&");
        assert!(!page.contains("<script>"));
        assert!(page.contains("&lt;script&gt;alert(1)&lt;/script&gt;"));
        assert!(page.contains("quote&quot; apostrophe&#39; amp&amp;"));
    }
}