tangled-cli 0.1.0

//! atproto OAuth login for the CLI (RFC 8252 loopback flow).
//!
//! The CLI acts as a "loopback client": `client_id` is the special
//! `http://localhost?...` form that authorization servers must support
//! without fetching hosted client metadata. The flow is:
//!
//! 1. resolve handle -> DID -> PDS host (DID document)
//! 2. discover the PDS's authorization server and its endpoints
//! 3. pushed authorization request (PAR) with PKCE, signed with a fresh
//!    DPoP key
//! 4. open the browser and catch the redirect on a 127.0.0.1 listener
//! 5. exchange the code for DPoP-bound tokens
//!
//! All later PDS requests use the same DPoP key via [`super::auth::PdsAuth`].

use anyhow::{anyhow, bail, Context, Result};
use base64::Engine;
use rand::RngCore;
use sha2::{Digest, Sha256};
use tangled_config::session::{OAuthTokens, Session};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::auth::DpopKey;
use super::http;

const SCOPE: &str = "atproto transition:generic";

fn b64url(data: &[u8]) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

fn urlencode(s: &str) -> String {
    url::form_urlencoded::byte_serialize(s.as_bytes()).collect()
}

fn random_token() -> String {
    let mut bytes = [0u8; 32];
    rand::thread_rng().fill_bytes(&mut bytes);
    b64url(&bytes)
}

/// Resolves a handle or DID to (did, pds_url).
async fn resolve_identity(input: &str) -> Result<(String, String)> {
    let did = if input.starts_with("did:") {
        input.to_string()
    } else {
        #[derive(serde::Deserialize)]
        struct Res {
            did: String,
        }
        crate::progress::with_spinner("Resolving handle...", async {
            let res = http()
                .get(format!(
                    "https://public.api.bsky.app/xrpc/com.atproto.identity.resolveHandle?handle={}",
                    urlencode(input)
                ))
                .send()
                .await?;
            if !res.status().is_success() {
                bail!("could not resolve handle {input}: {}", res.status());
            }
            Ok::<_, anyhow::Error>(res.json::<Res>().await?.did)
        })
        .await?
    };

    let doc_url = if let Some(plc) = did.strip_prefix("did:plc:") {
        format!("https://plc.directory/did:plc:{plc}")
    } else if let Some(host) = did.strip_prefix("did:web:") {
        format!("https://{}/.well-known/did.json", host.replace("%3A", ":"))
    } else {
        bail!("unsupported DID method: {did}");
    };
    let doc: serde_json::Value =
        crate::progress::with_spinner("Fetching DID document...", async {
            Ok::<_, anyhow::Error>(
                http().get(&doc_url).send().await?.json().await?,
            )
        })
        .await?;
    let pds = doc["service"]
        .as_array()
        .and_then(|services| {
            services.iter().find(|s| {
                s["id"].as_str() == Some("#atproto_pds")
                    || s["type"].as_str() == Some("AtprotoPersonalDataServer")
            })
        })
        .and_then(|s| s["serviceEndpoint"].as_str())
        .ok_or_else(|| {
            anyhow!("DID document for {did} has no PDS service endpoint")
        })?
        .to_string();
    Ok((did, pds))
}

#[derive(Debug, serde::Deserialize)]
struct AuthServerMetadata {
    issuer: String,
    authorization_endpoint: String,
    token_endpoint: String,
    pushed_authorization_request_endpoint: String,
}

async fn discover_auth_server(pds: &str) -> Result<AuthServerMetadata> {
    #[derive(serde::Deserialize)]
    struct ProtectedResource {
        authorization_servers: Vec<String>,
    }
    let protected_resource: ProtectedResource =
        crate::progress::with_spinner("Fetching OAuth metadata...", async {
            http()
                .get(format!(
                    "{}/.well-known/oauth-protected-resource",
                    pds.trim_end_matches('/')
                ))
                .send()
                .await?
                .json()
                .await
                .context("fetching PDS protected-resource metadata")
        })
        .await?;
    let issuer = protected_resource
        .authorization_servers
        .first()
        .ok_or_else(|| anyhow!("PDS lists no authorization servers"))?;
    let meta: AuthServerMetadata = crate::progress::with_spinner(
        "Fetching authorization server metadata...",
        async {
            http()
                .get(format!(
                    "{}/.well-known/oauth-authorization-server",
                    issuer.trim_end_matches('/')
                ))
                .send()
                .await?
                .json()
                .await
                .context("fetching authorization server metadata")
        },
    )
    .await?;
    Ok(meta)
}

/// POSTs a form to an authorization-server endpoint with a DPoP proof,
/// handling the `use_dpop_nonce` retry. Returns the parsed JSON body.
async fn as_form_post(
    url: &str,
    form: &[(&str, &str)],
    key: &DpopKey,
    nonce: &mut Option<String>,
) -> Result<serde_json::Value> {
    for _ in 0..2 {
        let proof = key.proof("POST", url, nonce.as_deref(), None);
        let (status, body, new_nonce) = crate::progress::with_spinner(
            "Contacting authorization server...",
            async {
                let res = http()
                    .post(url)
                    .header("DPoP", proof)
                    .form(form)
                    .send()
                    .await?;
                let new_nonce = res
                    .headers()
                    .get("dpop-nonce")
                    .and_then(|v| v.to_str().ok())
                    .map(str::to_string);
                let status = res.status();
                let body: serde_json::Value =
                    res.json().await.unwrap_or_default();
                Ok::<_, reqwest::Error>((status, body, new_nonce))
            },
        )
        .await?;
        if let Some(n) = new_nonce {
            *nonce = Some(n);
        }
        if status.is_success() {
            return Ok(body);
        }
        if body["error"].as_str() == Some("use_dpop_nonce") {
            continue; // retry once with the nonce we just stored
        }
        bail!(
            "{url} failed: {status}: {}",
            serde_json::to_string(&body).unwrap_or_default()
        );
    }
    bail!("{url} kept demanding new DPoP nonces");
}

/// Runs the full browser-based OAuth flow and returns a ready Session.
pub async fn login(handle_or_did: &str) -> Result<Session> {
    let (did, pds) = resolve_identity(handle_or_did).await?;
    let meta = discover_auth_server(&pds).await?;
    println!("Authorization server: {}", meta.issuer);

    // Loopback redirect listener on a random port.
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
        .await
        .context("binding loopback listener")?;
    let port = listener.local_addr()?.port();
    let redirect_uri = format!("http://127.0.0.1:{port}/callback");
    let client_id = format!(
        "http://localhost?redirect_uri={}&scope={}",
        urlencode(&redirect_uri),
        urlencode(SCOPE)
    );

    // PKCE + state + DPoP key for this session.
    let verifier = random_token();
    let challenge = b64url(&Sha256::digest(verifier.as_bytes()));
    let state = random_token();
    let key = DpopKey::generate();
    let mut as_nonce: Option<String> = None;

    // Pushed authorization request.
    let par = as_form_post(
        &meta.pushed_authorization_request_endpoint,
        &[
            ("response_type", "code"),
            ("client_id", &client_id),
            ("redirect_uri", &redirect_uri),
            ("scope", SCOPE),
            ("state", &state),
            ("code_challenge", &challenge),
            ("code_challenge_method", "S256"),
            ("login_hint", handle_or_did),
        ],
        &key,
        &mut as_nonce,
    )
    .await
    .context("pushed authorization request")?;
    let request_uri = par["request_uri"]
        .as_str()
        .ok_or_else(|| anyhow!("PAR response missing request_uri"))?;

    let authorize_url = format!(
        "{}?client_id={}&request_uri={}",
        meta.authorization_endpoint,
        urlencode(&client_id),
        urlencode(request_uri)
    );
    println!("Opening browser to authorize; if it does not open, visit:\n  {authorize_url}");
    let _ = open::that(&authorize_url);

    // Wait for the redirect.
    let (code, returned_state, returned_iss) =
        wait_for_callback(listener).await?;
    if returned_state != state {
        bail!("OAuth state mismatch; aborting");
    }
    if let Some(iss) = returned_iss {
        if iss != meta.issuer {
            bail!("OAuth issuer mismatch: expected {}, got {iss}", meta.issuer);
        }
    }

    // Exchange the code for tokens.
    let tokens = as_form_post(
        &meta.token_endpoint,
        &[
            ("grant_type", "authorization_code"),
            ("code", &code),
            ("redirect_uri", &redirect_uri),
            ("client_id", &client_id),
            ("code_verifier", &verifier),
        ],
        &key,
        &mut as_nonce,
    )
    .await
    .context("token exchange")?;

    session_from_token_response(&tokens, &did, &pds, &meta, &client_id, &key)
}

fn session_from_token_response(
    tokens: &serde_json::Value,
    expected_did: &str,
    pds: &str,
    meta: &AuthServerMetadata,
    client_id: &str,
    key: &DpopKey,
) -> Result<Session> {
    let access_token = tokens["access_token"]
        .as_str()
        .ok_or_else(|| anyhow!("token response missing access_token"))?;
    let refresh_token = tokens["refresh_token"]
        .as_str()
        .ok_or_else(|| anyhow!("token response missing refresh_token"))?;
    let sub = tokens["sub"].as_str().unwrap_or(expected_did);
    if sub != expected_did {
        bail!("logged-in DID {sub} does not match requested identity {expected_did}");
    }
    let expires_in = tokens["expires_in"].as_i64().unwrap_or(300);

    Ok(Session {
        did: sub.to_string(),
        pds: Some(pds.to_string()),
        oauth: Some(OAuthTokens {
            access_token: access_token.to_string(),
            refresh_token: refresh_token.to_string(),
            issuer: meta.issuer.clone(),
            token_endpoint: meta.token_endpoint.clone(),
            client_id: client_id.to_string(),
            dpop_key: key.to_b64(),
            expires_at: chrono::Utc::now().timestamp() + expires_in,
        }),
        ..Default::default()
    })
}

/// Refreshes an OAuth session, rotating the refresh token.
pub async fn refresh(session: &Session) -> Result<Session> {
    let oauth = session
        .oauth
        .as_ref()
        .ok_or_else(|| anyhow!("not an OAuth session"))?;
    let key = DpopKey::from_b64(&oauth.dpop_key)?;
    let mut nonce: Option<String> = None;
    let tokens = as_form_post(
        &oauth.token_endpoint,
        &[
            ("grant_type", "refresh_token"),
            ("refresh_token", &oauth.refresh_token),
            ("client_id", &oauth.client_id),
        ],
        &key,
        &mut nonce,
    )
    .await
    .context("refreshing OAuth session")?;

    let meta = AuthServerMetadata {
        issuer: oauth.issuer.clone(),
        authorization_endpoint: String::new(),
        token_endpoint: oauth.token_endpoint.clone(),
        pushed_authorization_request_endpoint: String::new(),
    };
    let mut refreshed = session_from_token_response(
        &tokens,
        &session.did,
        session.pds.as_deref().unwrap_or_default(),
        &meta,
        &oauth.client_id,
        &key,
    )?;
    refreshed.handle = session.handle.clone();
    refreshed.pds = session.pds.clone();
    Ok(refreshed)
}

/// Accepts one HTTP request on the listener and extracts `code`, `state`,
/// and `iss` from the callback query string.
async fn wait_for_callback(
    listener: tokio::net::TcpListener,
) -> Result<(String, String, Option<String>)> {
    tokio::time::timeout(std::time::Duration::from_secs(300), async {
        loop {
            let (mut stream, _) = listener.accept().await?;
            let mut buf = vec![0u8; 8192];
            let n = stream.read(&mut buf).await?;
            let request = String::from_utf8_lossy(&buf[..n]).into_owned();
            let Some(query) = request
                .lines()
                .next()
                .and_then(|line| line.split_whitespace().nth(1))
                .filter(|path| path.starts_with("/callback"))
                .and_then(|path| path.split_once('?').map(|(_, q)| q.to_string()))
            else {
                // Not the callback (e.g. favicon request); answer and keep waiting.
                let _ = stream
                    .write_all(b"HTTP/1.1 404 Not Found\r\ncontent-length: 0\r\n\r\n")
                    .await;
                continue;
            };

            let mut code = None;
            let mut state = None;
            let mut iss = None;
            let mut error = None;
            for (k, v) in url::form_urlencoded::parse(query.as_bytes()) {
                match k.as_ref() {
                    "code" => code = Some(v.into_owned()),
                    "state" => state = Some(v.into_owned()),
                    "iss" => iss = Some(v.into_owned()),
                    "error" => error = Some(v.into_owned()),
                    _ => {}
                }
            }

            let page = "<html><body><h2>tangled-cli: login complete</h2>\
                        You can close this tab and return to the terminal.</body></html>";
            let response = format!(
                "HTTP/1.1 200 OK\r\ncontent-type: text/html\r\ncontent-length: {}\r\n\r\n{page}",
                page.len()
            );
            let _ = stream.write_all(response.as_bytes()).await;

            if let Some(error) = error {
                bail!("authorization was denied: {error}");
            }
            return Ok((
                code.ok_or_else(|| anyhow!("callback missing code"))?,
                state.ok_or_else(|| anyhow!("callback missing state"))?,
                iss,
            ));
        }
    })
    .await
    .map_err(|_| anyhow!("timed out waiting for browser authorization (5 minutes)"))?
}