AkurAI-Drive / client/src/auth.rs

Latest change feb7ac86b5e60d55b429eacb7146cf0dc32c6628 - Add native Linux sync client (akurai-drive-sync) by olibuijr

//! Device-token login: spin up a loopback HTTP server on a random free port,
//! open the browser to the backend's `/auth/device` endpoint, and capture the
//! `?token=` the backend redirects back with.

use anyhow::{bail, Result};
use std::time::{Duration, Instant};

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

/// Run the full login handshake and return the captured device token.
pub fn login(api_base: &str) -> Result<String> {
    // Bind to port 0 → the OS picks a free loopback port.
    let server = tiny_http::Server::http("127.0.0.1:0")
        .map_err(|e| anyhow::anyhow!("starting loopback server: {e}"))?;
    let port = match server.server_addr() {
        tiny_http::ListenAddr::IP(addr) => addr.port(),
        other => bail!("unexpected listen address: {other:?}"),
    };
    let redirect = format!("http://127.0.0.1:{port}/cb");
    let auth_url = format!("{api_base}/auth/device?redirect={}", urlencode(&redirect));

    println!("Opening your browser to log in...");
    println!("If it doesn't open, visit:\n  {auth_url}");
    open_browser(&auth_url);

    // Wait for the backend's redirect to hit our /cb endpoint.
    let deadline = Instant::now() + LOGIN_TIMEOUT;
    loop {
        let remaining = deadline.saturating_duration_since(Instant::now());
        if remaining.is_zero() {
            bail!("login timed out after {}s", LOGIN_TIMEOUT.as_secs());
        }
        match server.recv_timeout(remaining.min(Duration::from_secs(1))) {
            Ok(Some(request)) => {
                if let Some(token) = token_from_path(request.url()) {
                    respond_ok(request);
                    return Ok(token);
                }
                // Some other path (favicon, etc.) — answer and keep waiting.
                respond_wait(request);
            }
            Ok(None) => continue, // timeout tick; re-check deadline
            Err(e) => bail!("loopback server error: {e}"),
        }
    }
}

/// Extract the `token` query parameter from a request path like
/// `/cb?token=abc123`.
fn token_from_path(path: &str) -> Option<String> {
    let query = path.split_once('?').map(|(_, q)| q)?;
    for pair in query.split('&') {
        if let Some((k, v)) = pair.split_once('=') {
            if k == "token" {
                let decoded = urldecode(v);
                if !decoded.is_empty() {
                    return Some(decoded);
                }
            }
        }
    }
    None
}

fn respond_ok(request: tiny_http::Request) {
    let html = "<!doctype html><html><head><meta charset=utf-8>\
        <title>AkurAI Drive</title></head><body style=\"font-family:system-ui;\
        text-align:center;padding-top:4rem;\"><h2>Logged in to AkurAI Drive</h2>\
        <p>You can close this tab and return to the terminal.</p></body></html>";
    let resp = tiny_http::Response::from_string(html).with_header(
        tiny_http::Header::from_bytes(&b"Content-Type"[..], &b"text/html; charset=utf-8"[..])
            .unwrap(),
    );
    let _ = request.respond(resp);
}

fn respond_wait(request: tiny_http::Request) {
    let resp = tiny_http::Response::from_string("waiting for token...");
    let _ = request.respond(resp);
}

/// Try the `open` crate; fall back to `xdg-open`.
fn open_browser(url: &str) {
    if open::that(url).is_ok() {
        return;
    }
    let _ = std::process::Command::new("xdg-open").arg(url).spawn();
}

/// Minimal percent-encoding for the redirect URL we embed in the query string.
fn urlencode(s: &str) -> String {
    let mut out = String::with_capacity(s.len() * 3);
    for b in s.bytes() {
        match b {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                out.push(b as char)
            }
            _ => out.push_str(&format!("%{b:02X}")),
        }
    }
    out
}

/// Minimal percent-decoding for the captured token value.
fn urldecode(s: &str) -> String {
    let bytes = s.as_bytes();
    let mut out = Vec::with_capacity(bytes.len());
    let mut i = 0;
    while i < bytes.len() {
        match bytes[i] {
            b'%' if i + 2 < bytes.len() => {
                let hi = hexval(bytes[i + 1]);
                let lo = hexval(bytes[i + 2]);
                if let (Some(hi), Some(lo)) = (hi, lo) {
                    out.push(hi * 16 + lo);
                    i += 3;
                    continue;
                }
                out.push(bytes[i]);
                i += 1;
            }
            b'+' => {
                out.push(b' ');
                i += 1;
            }
            b => {
                out.push(b);
                i += 1;
            }
        }
    }
    String::from_utf8_lossy(&out).into_owned()
}

fn hexval(b: u8) -> Option<u8> {
    match b {
        b'0'..=b'9' => Some(b - b'0'),
        b'a'..=b'f' => Some(b - b'a' + 10),
        b'A'..=b'F' => Some(b - b'A' + 10),
        _ => None,
    }
}

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

    #[test]
    fn extracts_token() {
        assert_eq!(token_from_path("/cb?token=abc123"), Some("abc123".into()));
        assert_eq!(
            token_from_path("/cb?foo=1&token=xy%2Dz&bar=2"),
            Some("xy-z".into())
        );
    }

    #[test]
    fn no_token_is_none() {
        assert_eq!(token_from_path("/cb"), None);
        assert_eq!(token_from_path("/favicon.ico"), None);
        assert_eq!(token_from_path("/cb?token="), None);
    }

    #[test]
    fn roundtrip_encoding() {
        let original = "http://127.0.0.1:54321/cb";
        assert_eq!(urldecode(&urlencode(original)), original);
    }
}