semora-cli 0.8.0

//! `semora login` (browser approval, PKCE) and `semora logout`.

use std::io::{BufRead, Write};
use std::time::Duration;

use anyhow::{Context, bail};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use rand::RngCore;
use semora_proto::{CliTokenRequest, CliTokenResponse};
use sha2::{Digest, Sha256};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinSet;

use crate::client::{Client, check, http_client, no_server_error};
use crate::config::{Settings, origin};
use crate::credentials::Stored;
use crate::{Ctx, exit, output};

const LOGIN_TIMEOUT: Duration = Duration::from_secs(600);

fn random_b64(n: usize) -> String {
    let mut buf = vec![0u8; n];
    rand::rngs::OsRng.fill_bytes(&mut buf);
    URL_SAFE_NO_PAD.encode(buf)
}

pub(crate) fn hostname() -> String {
    for var in ["COMPUTERNAME", "HOSTNAME"] {
        if let Ok(h) = std::env::var(var)
            && !h.trim().is_empty()
        {
            return h.trim().to_lowercase();
        }
    }
    std::process::Command::new("hostname")
        .output()
        .ok()
        .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_lowercase())
        .filter(|h| !h.is_empty())
        .unwrap_or_else(|| "this-machine".into())
}

const DONE_PAGE: &str = "<!doctype html><meta charset=utf-8><title>Semora</title>\
<body style=\"font:16px system-ui;margin:4rem auto;max-width:32rem\"><h1>Signed in</h1>\
<p>The Semora CLI has its key. You can close this tab.</p>";
const DENIED_PAGE: &str = "<!doctype html><meta charset=utf-8><title>Semora</title>\
<body style=\"font:16px system-ui;margin:4rem auto;max-width:32rem\"><h1>Not signed in</h1>\
<p>The request was denied. You can close this tab.</p>";

/// Answers with a status and no body, and closes the connection.
async fn turn_away(mut sock: TcpStream, status: &str) {
    let resp = format!("HTTP/1.1 {status}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n");
    let _ = sock.write_all(resp.as_bytes()).await;
    let _ = sock.shutdown().await;
}

type Query = Vec<(String, String)>;

/// How many callback connections are open at once. A connection holds a task and at most
/// `MAX_HEAD` bytes, so this bounds what a local process can make the CLI hold. Past it the
/// listener stops accepting until a connection ends, and `HEAD_TIMEOUT` ends a stalled one.
const MAX_OPEN_CONNECTIONS: usize = 32;
/// How long a connection has to send its request head before it is dropped. A browser sends
/// it at once, so this only ever ends a preconnect or a stall.
const HEAD_TIMEOUT: Duration = Duration::from_secs(5);
/// The most of a request that is read while looking for the end of its head.
const MAX_HEAD: usize = 16 * 1024;

/// Reads a request up to the end of its head (or `MAX_HEAD` bytes, or the peer closing).
async fn read_head(sock: &mut TcpStream) -> std::io::Result<Vec<u8>> {
    let mut buf = Vec::with_capacity(2048);
    let mut chunk = [0u8; 1024];
    while !buf.windows(4).any(|w| w == b"\r\n\r\n") && buf.len() < MAX_HEAD {
        let n = sock.read(&mut chunk).await?;
        if n == 0 {
            break;
        }
        buf.extend_from_slice(&chunk[..n]);
    }
    Ok(buf)
}

/// Answers one connection. Returns the query pairs if it was the real callback (a `/callback`
/// request carrying `state`), and `None` for anything else: the connection is turned away, or
/// dropped if it sends no head within `head_timeout` or fails.
async fn serve_connection(
    mut sock: TcpStream,
    state: String,
    head_timeout: Duration,
) -> Option<Query> {
    let buf = match tokio::time::timeout(head_timeout, read_head(&mut sock)).await {
        Ok(Ok(buf)) => buf,
        Ok(Err(_)) | Err(_) => return None,
    };
    let head = String::from_utf8_lossy(&buf);
    let target = head
        .lines()
        .next()
        .and_then(|l| l.split_whitespace().nth(1))
        .unwrap_or("")
        .to_string();
    let Some(query) = target.strip_prefix("/callback?") else {
        turn_away(sock, "404 Not Found").await;
        return None;
    };
    let pairs: Query = url::form_urlencoded::parse(query.as_bytes())
        .into_owned()
        .collect();
    let sent_state = pairs.iter().find(|(k, _)| k == "state").map(|(_, v)| v);
    if sent_state.is_none_or(|s| *s != state) {
        turn_away(sock, "400 Bad Request").await;
        return None;
    }
    let body = if pairs.iter().any(|(k, _)| k == "error") {
        DENIED_PAGE
    } else {
        DONE_PAGE
    };
    let resp = format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    );
    let _ = sock.write_all(resp.as_bytes()).await;
    let _ = sock.shutdown().await;
    Some(pairs)
}

/// Serves `/callback` requests on the loopback listener until one carries `state`, the value
/// this login sent; returns that request's query pairs. Any other request, to another path or
/// with another state (or none), is turned away and the wait goes on: a stray or forged local
/// request must not end the login. Each connection is served in its own task, so one that
/// sends nothing (a browser preconnect) or sends slowly cannot hold up the real callback. The
/// caller's timeout ends the wait.
async fn wait_for_callback(listener: TcpListener, state: &str) -> anyhow::Result<Query> {
    serve_callbacks(listener, state, MAX_OPEN_CONNECTIONS, HEAD_TIMEOUT).await
}

async fn serve_callbacks(
    listener: TcpListener,
    state: &str,
    max_open: usize,
    head_timeout: Duration,
) -> anyhow::Result<Query> {
    // Dropping the set, on return or when the caller's timeout cancels this, aborts the
    // connections still open and so closes them.
    let mut open = JoinSet::new();
    loop {
        tokio::select! {
            accepted = listener.accept(), if open.len() < max_open => {
                let (sock, _) = accepted?;
                open.spawn(serve_connection(sock, state.to_string(), head_timeout));
            }
            Some(done) = open.join_next() => {
                if let Some(pairs) = done? {
                    return Ok(pairs);
                }
            }
        }
    }
}

pub async fn login(ctx: &Ctx, no_browser: bool, name: Option<String>) -> anyhow::Result<i32> {
    let server = ctx.settings.server.clone().ok_or_else(no_server_error)?;
    let base = origin(&server);
    let name = name.unwrap_or_else(|| format!("cli@{}", hostname()));
    let verifier = random_b64(32);
    let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
    let state = random_b64(16);

    let listener = if no_browser {
        None
    } else {
        Some(
            TcpListener::bind("127.0.0.1:0")
                .await
                .context("opening a local port for the sign-in callback")?,
        )
    };
    let port = match &listener {
        Some(l) => l.local_addr()?.port(),
        None => 0,
    };
    let query = url::form_urlencoded::Serializer::new(String::new())
        .append_pair("port", &port.to_string())
        .append_pair("state", &state)
        .append_pair("code_challenge", &challenge)
        .append_pair("name", &name)
        .finish();
    let url = format!("{base}/cli/authorize?{query}");

    eprintln!("Sign in to {base} in your browser and allow the CLI:");
    println!("{url}");
    std::io::stdout().flush()?;

    let code = match listener {
        Some(listener) => {
            if ctx.settings.open_browser && open::that(&url).is_err() {
                output::warn("could not open a browser; open the URL above yourself");
            }
            // `wait_for_callback` has already checked the state.
            let pairs = tokio::time::timeout(LOGIN_TIMEOUT, wait_for_callback(listener, &state))
                .await
                .map_err(|_| exit::timeout("no answer from the browser within 10 minutes"))??;
            let get = |k: &str| {
                pairs
                    .iter()
                    .find(|(key, _)| key == k)
                    .map(|(_, v)| v.clone())
            };
            if let Some(err) = get("error") {
                return Err(exit::auth(format!("sign-in was not allowed ({err})")));
            }
            get("code").context("the sign-in callback had no code")?
        }
        None => {
            eprint!("After you allow it, paste the code shown in the browser: ");
            std::io::stderr().flush()?;
            let line = tokio::task::spawn_blocking(|| {
                let mut s = String::new();
                std::io::stdin().lock().read_line(&mut s).map(|_| s)
            })
            .await??;
            let code = line.trim().to_string();
            if code.is_empty() {
                bail!("no code entered");
            }
            code
        }
    };

    let resp = http_client()?
        .post(format!("{base}/api/v1/cli/token"))
        .json(&CliTokenRequest {
            code,
            code_verifier: verifier,
        })
        .send()
        .await
        .with_context(|| format!("could not reach {base}"))?;
    let token: CliTokenResponse = check(resp, "sign-in").await?.json().await?;
    // The key this CLI had stored for this server, read before the new one replaces it.
    // The environment's key is not in the store and is never revoked here.
    let previous = replaced_key(ctx.store.get(&base), &token.api_key);
    let stored = ctx.store.set(&base, &token.api_key, &token.name)?;
    if let Some(old) = previous
        && let Err(e) = revoke_stored_key(&ctx.settings, old).await
    {
        output::warn(format!(
            "could not revoke the previous key on the server: {e:#}"
        ));
    }
    let place = match stored {
        Stored::Keyring => "the system keyring".to_string(),
        Stored::File => ctx.store.path().display().to_string(),
    };
    eprintln!(
        "Signed in to {base} as {}. Key \"{}\" saved in {place}.",
        token.user.email, token.name
    );
    if ctx.settings.key_source == crate::config::KeySource::Env {
        output::warn("SEMORA_API_KEY is set, so it wins over this login until you unset it");
    }
    Ok(exit::OK)
}

pub async fn logout(ctx: &Ctx) -> anyhow::Result<i32> {
    let server = ctx.settings.server.clone().ok_or_else(no_server_error)?;
    let base = origin(&server);
    let Some(key) = ctx.store.get(&base) else {
        eprintln!("No stored login for {base}.");
        if ctx.settings.key_source == crate::config::KeySource::Env {
            eprintln!("SEMORA_API_KEY comes from the environment or a .env file; remove it there.");
        }
        return Ok(exit::OK);
    };
    if let Err(e) = revoke_stored_key(&ctx.settings, key).await {
        output::warn(format!("could not revoke the key on the server: {e:#}"));
    }
    ctx.store.remove(&base)?;
    eprintln!("Signed out of {base}.");
    Ok(exit::OK)
}

/// The key a login replaces: the one stored for the server, unless it is the new key itself.
fn replaced_key(stored: Option<String>, new: &str) -> Option<String> {
    stored.filter(|old| old != new)
}

/// Revokes on the server a key this CLI had stored: `whoami` names the key, then it is
/// revoked by id. `key` must come from the store. A key from `SEMORA_API_KEY` or `--api-key`
/// is not this CLI's to revoke. `Ok` means the key no longer works, whether it was revoked
/// now or already was (a 401 means it was revoked or has expired, which is what we want).
/// `Err` is any other failure; callers warn and carry on.
async fn revoke_stored_key(settings: &Settings, key: String) -> anyhow::Result<()> {
    let mut settings = settings.clone();
    settings.api_key = Some(key);
    settings.key_source = crate::config::KeySource::Stored;
    let client = Client::new(&settings)?;
    match client.whoami().await {
        Ok(w) => match w.key {
            Some(k) => client.revoke_key(k.id).await,
            None => Ok(()),
        },
        Err(e) if exit::code_of(&e) == exit::AUTH => Ok(()),
        Err(e) => Err(e),
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::{TcpListener, TcpStream};

    use super::{replaced_key, serve_callbacks};

    /// A listener on a free loopback port, served with the given limits; the handle ends
    /// with the query pairs of the real callback.
    async fn listen(
        max_open: usize,
        head_timeout: Duration,
    ) -> (u16, tokio::task::JoinHandle<anyhow::Result<super::Query>>) {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let served = tokio::spawn(serve_callbacks(listener, "right", max_open, head_timeout));
        (port, served)
    }

    async fn connect(port: u16) -> TcpStream {
        TcpStream::connect(("127.0.0.1", port)).await.unwrap()
    }

    async fn send_callback(sock: &mut TcpStream) {
        sock.write_all(b"GET /callback?code=C&state=right HTTP/1.1\r\nHost: x\r\n\r\n")
            .await
            .unwrap();
    }

    /// Reads until the peer closes; `None` if it neither answers nor closes within `wait`.
    async fn read_all(sock: &mut TcpStream, wait: Duration) -> Option<String> {
        let mut out = Vec::new();
        let read = tokio::time::timeout(wait, sock.read_to_end(&mut out)).await;
        // A reset after the answer is as good as a close.
        read.ok()
            .map(|_| String::from_utf8_lossy(&out).into_owned())
    }

    #[tokio::test]
    async fn a_connection_that_sends_no_head_is_dropped_after_the_timeout() {
        let (port, served) = listen(8, Duration::from_millis(300)).await;
        let mut silent = connect(port).await;
        // Nothing was sent, so it ends with no answer once the timeout has passed ...
        assert_eq!(
            read_all(&mut silent, Duration::from_secs(5))
                .await
                .as_deref(),
            Some(""),
            "the silent connection is closed by the CLI"
        );
        // ... and the listener is still serving.
        let mut real = connect(port).await;
        send_callback(&mut real).await;
        let answer = read_all(&mut real, Duration::from_secs(5)).await.unwrap();
        assert!(answer.starts_with("HTTP/1.1 200"), "{answer}");
        let pairs = served.await.unwrap().unwrap();
        assert!(pairs.contains(&("code".to_string(), "C".to_string())));
    }

    #[tokio::test]
    async fn at_most_max_open_connections_are_served_at_once() {
        let (port, served) = listen(2, Duration::from_millis(1500)).await;
        let _stalled = [connect(port).await, connect(port).await];
        let mut real = connect(port).await;
        send_callback(&mut real).await;
        // Both slots are held by the stalled connections, so the real one is not read yet.
        assert_eq!(
            read_all(&mut real, Duration::from_millis(200)).await,
            None,
            "a third connection is not served while two are open"
        );
        // When the stalled ones time out, it is.
        let answer = read_all(&mut real, Duration::from_secs(10)).await.unwrap();
        assert!(answer.starts_with("HTTP/1.1 200"), "{answer}");
        served.await.unwrap().unwrap();
    }

    #[test]
    fn a_login_replaces_a_different_stored_key_only() {
        assert_eq!(replaced_key(None, "sm_new"), None);
        assert_eq!(replaced_key(Some("sm_new".into()), "sm_new"), None);
        assert_eq!(
            replaced_key(Some("sm_old".into()), "sm_new").as_deref(),
            Some("sm_old")
        );
    }
}