proofborne 0.2.0

A local coding agent that proves its work
use std::{
    io::{Read, Write},
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener, TcpStream},
    process::Command,
    thread,
    time::{Duration, Instant},
};

use anyhow::{Context, Result, anyhow, bail};
use proofborne_adapters::{
    OAuthProviderConfig, PkcePair, build_authorize_url, generate_pkce, generate_state,
};
use url::Url;

/// Runs a browser authorization-code flow against a configured loopback redirect.
///
/// The listener is bound only to a loopback address. The callback path and CSRF
/// state must match exactly; authorization codes are returned to the caller but
/// never printed or included in errors.
pub(super) async fn authorize_code(
    provider: &OAuthProviderConfig,
    no_browser: bool,
    timeout: Duration,
) -> Result<(PkcePair, String, String, OAuthProviderConfig)> {
    let endpoint = LoopbackEndpoint::bind(provider.redirect_uri.as_deref())?;
    let mut effective_provider = provider.clone();
    effective_provider.redirect_uri = Some(endpoint.redirect_uri());
    let pkce = generate_pkce();
    let state = generate_state();
    let authorize_url =
        build_authorize_url(&effective_provider, &pkce, &state).map_err(|error| anyhow!(error))?;

    if no_browser {
        println!("Open this URL to authenticate:\n{authorize_url}");
    } else if let Err(error) = open_browser(&authorize_url) {
        eprintln!("Could not open a browser automatically: {error:#}");
        println!("Open this URL to authenticate:\n{authorize_url}");
    } else {
        eprintln!("Browser opened for OAuth authentication; waiting for callback...");
    }

    let listener = endpoint.listener;
    let path = endpoint.path;
    let callback_state = state.clone();
    let code = tokio::task::spawn_blocking(move || {
        wait_for_callback(listener, &path, &callback_state, timeout)
    })
    .await
    .context("OAuth callback worker failed")??;
    Ok((pkce, code, state, effective_provider))
}

#[derive(Debug)]
struct LoopbackEndpoint {
    listener: TcpListener,
    ip: IpAddr,
    host: String,
    path: String,
}

impl LoopbackEndpoint {
    fn bind(redirect_uri: Option<&str>) -> Result<Self> {
        let redirect_uri = redirect_uri
            .ok_or_else(|| anyhow!("OAuth profile requires redirect_uri for browser login"))?;
        let parsed = Url::parse(redirect_uri).context("OAuth redirect_uri is not a URL")?;
        if parsed.scheme() != "http"
            || parsed.username() != ""
            || parsed.password().is_some()
            || parsed.query().is_some()
            || parsed.fragment().is_some()
        {
            bail!(
                "OAuth redirect_uri must be a plain HTTP loopback URL without credentials, query, or fragment"
            );
        }
        let host = parsed
            .host_str()
            .ok_or_else(|| anyhow!("OAuth redirect_uri has no host"))?;
        let (host, ip) = match host {
            "localhost" => ("localhost".to_owned(), IpAddr::V4(Ipv4Addr::LOCALHOST)),
            "127.0.0.1" => ("127.0.0.1".to_owned(), IpAddr::V4(Ipv4Addr::LOCALHOST)),
            "::1" => ("::1".to_owned(), IpAddr::V6(Ipv6Addr::LOCALHOST)),
            _ => bail!("OAuth redirect_uri host must be localhost, 127.0.0.1, or ::1"),
        };
        let port = parsed
            .port()
            .ok_or_else(|| anyhow!("OAuth redirect_uri must include an explicit loopback port"))?;
        let path = parsed.path().to_owned();
        if path.is_empty() || !path.starts_with('/') {
            bail!("OAuth redirect_uri must include an absolute callback path");
        }
        let address = SocketAddr::new(ip, port);
        let listener = TcpListener::bind(address)
            .with_context(|| format!("could not bind OAuth callback at {address}"))?;
        listener
            .set_nonblocking(true)
            .context("could not configure OAuth callback listener")?;
        Ok(Self {
            listener,
            ip,
            host,
            path,
        })
    }

    fn redirect_uri(&self) -> String {
        let port = self
            .listener
            .local_addr()
            .expect("bound OAuth callback listener has a local address")
            .port();
        if self.ip.is_ipv6() {
            format!("http://[{}]:{port}{}", self.host, self.path)
        } else {
            format!("http://{}:{port}{}", self.host, self.path)
        }
    }
}

fn wait_for_callback(
    listener: TcpListener,
    expected_path: &str,
    expected_state: &str,
    timeout: Duration,
) -> Result<String> {
    let deadline = Instant::now() + timeout;
    loop {
        match listener.accept() {
            Ok((stream, _peer)) => match handle_callback(stream, expected_path, expected_state) {
                Ok(code) => return Ok(code),
                Err(error)
                    if error
                        .downcast_ref::<std::io::Error>()
                        .is_some_and(|source| {
                            matches!(
                                source.kind(),
                                std::io::ErrorKind::WouldBlock
                                    | std::io::ErrorKind::TimedOut
                                    | std::io::ErrorKind::UnexpectedEof
                            )
                        }) => {}
                Err(error) => return Err(error),
            },
            Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                if Instant::now() >= deadline {
                    bail!("timed out waiting for OAuth callback");
                }
                thread::sleep(Duration::from_millis(25));
            }
            Err(error) => return Err(error).context("OAuth callback listener failed"),
        }
    }
}

fn handle_callback(
    mut stream: TcpStream,
    expected_path: &str,
    expected_state: &str,
) -> Result<String> {
    stream
        .set_nonblocking(false)
        .context("could not configure OAuth callback connection")?;
    stream
        .set_read_timeout(Some(Duration::from_secs(2)))
        .context("could not bound OAuth callback read")?;
    let mut buffer = [0u8; 8192];
    let count = stream
        .read(&mut buffer)
        .context("could not read OAuth callback request")?;
    if count == 0 {
        return Err(std::io::Error::new(
            std::io::ErrorKind::UnexpectedEof,
            "OAuth callback request was empty",
        )
        .into());
    }
    let request = std::str::from_utf8(&buffer[..count]).context("OAuth callback was not UTF-8")?;
    let request_line = request
        .lines()
        .next()
        .ok_or_else(|| anyhow!("OAuth callback request had no request line"))?;
    let mut fields = request_line.split_ascii_whitespace();
    let method = fields.next().unwrap_or_default();
    let target = fields.next().unwrap_or_default();
    if method != "GET" {
        let _ = respond(&mut stream, 405, "Authentication callback must use GET.");
        bail!("OAuth callback used an unsupported method");
    }
    let callback_url = Url::parse(&format!("http://127.0.0.1{target}"))
        .context("OAuth callback request target was invalid")?;
    if callback_url.path() != expected_path {
        let _ = respond(
            &mut stream,
            404,
            "Authentication callback path was not recognized.",
        );
        bail!("OAuth callback path did not match configured redirect_uri");
    }
    if callback_url.query_pairs().any(|(key, _)| key == "error") {
        let _ = respond(&mut stream, 400, "The provider rejected authentication.");
        bail!("OAuth provider rejected authentication");
    }
    let state = callback_url
        .query_pairs()
        .find(|(key, _)| key == "state")
        .map(|(_, value)| value.into_owned())
        .ok_or_else(|| anyhow!("OAuth callback did not include state"))?;
    if state != expected_state {
        let _ = respond(&mut stream, 400, "Authentication state did not match.");
        bail!("OAuth callback state did not match");
    }
    let code = callback_url
        .query_pairs()
        .find(|(key, _)| key == "code")
        .map(|(_, value)| value.into_owned())
        .filter(|value| !value.is_empty())
        .ok_or_else(|| anyhow!("OAuth callback did not include an authorization code"))?;
    let _ = respond(
        &mut stream,
        200,
        "Authentication complete. You can close this window.",
    );
    Ok(code)
}

fn respond(stream: &mut TcpStream, status: u16, body: &str) -> Result<()> {
    let reason = match status {
        200 => "OK",
        400 => "Bad Request",
        404 => "Not Found",
        405 => "Method Not Allowed",
        _ => "Error",
    };
    let response = format!(
        "HTTP/1.1 {status} {reason}\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
        body.len()
    );
    stream
        .write_all(response.as_bytes())
        .context("could not respond to OAuth callback")
}

fn open_browser(url: &str) -> Result<()> {
    #[cfg(windows)]
    let mut command = {
        let mut command = Command::new("rundll32");
        command.args(["url.dll,FileProtocolHandler", url]);
        command
    };
    #[cfg(target_os = "macos")]
    let mut command = {
        let mut command = Command::new("open");
        command.arg(url);
        command
    };
    #[cfg(all(unix, not(target_os = "macos")))]
    let mut command = {
        let mut command = Command::new("xdg-open");
        command.arg(url);
        command
    };
    let status = command.status().context("browser launcher failed")?;
    if !status.success() {
        bail!("browser launcher exited unsuccessfully");
    }
    Ok(())
}

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

    #[test]
    fn loopback_endpoint_rejects_non_loopback_hosts() {
        let error = LoopbackEndpoint::bind(Some("http://0.0.0.0:1455/callback")).unwrap_err();
        assert!(error.to_string().contains("localhost, 127.0.0.1, or ::1"));
    }

    #[test]
    fn loopback_endpoint_rejects_missing_explicit_port() {
        let error = LoopbackEndpoint::bind(Some("http://127.0.0.1/callback")).unwrap_err();
        assert!(error.to_string().contains("explicit loopback port"));
    }

    #[test]
    fn loopback_endpoint_accepts_ephemeral_port_and_reports_bound_uri() {
        let endpoint = LoopbackEndpoint::bind(Some("http://127.0.0.1:0/callback")).unwrap();
        assert!(endpoint.redirect_uri().starts_with("http://127.0.0.1:"));
        assert!(endpoint.redirect_uri().ends_with("/callback"));
    }
    #[test]
    fn loopback_endpoint_preserves_registered_localhost_redirect() {
        let endpoint = LoopbackEndpoint::bind(Some("http://localhost:0/auth/callback")).unwrap();
        assert!(endpoint.redirect_uri().starts_with("http://localhost:"));
        let port = endpoint.listener.local_addr().unwrap().port();
        TcpStream::connect(("localhost", port)).unwrap();
        assert!(endpoint.redirect_uri().ends_with("/auth/callback"));
    }

    #[test]
    fn callback_accepts_exact_path_and_state() {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
        listener.set_nonblocking(true).unwrap();
        let address = listener.local_addr().unwrap();
        let sender = thread::spawn(move || {
            let mut stream = TcpStream::connect(address).unwrap();
            stream
                .write_all(b"GET /callback?code=one-time-code&state=csrf HTTP/1.1\r\nHost: localhost\r\n\r\n")
                .unwrap();
        });
        let code =
            wait_for_callback(listener, "/callback", "csrf", Duration::from_secs(2)).unwrap();
        sender.join().unwrap();
        assert_eq!(code, "one-time-code");
    }

    #[test]
    fn callback_ignores_idle_browser_preconnect() {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
        listener.set_nonblocking(true).unwrap();
        let address = listener.local_addr().unwrap();
        let sender = thread::spawn(move || {
            let idle = TcpStream::connect(address).unwrap();
            thread::sleep(Duration::from_millis(50));
            drop(idle);
            let mut stream = TcpStream::connect(address).unwrap();
            stream
                .write_all(
                    b"GET /callback?code=after-preconnect&state=csrf HTTP/1.1\r\nHost: localhost\r\n\r\n",
                )
                .unwrap();
        });
        let code =
            wait_for_callback(listener, "/callback", "csrf", Duration::from_secs(3)).unwrap();
        sender.join().unwrap();
        assert_eq!(code, "after-preconnect");
    }

    #[test]
    fn callback_rejects_wrong_path() {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
        listener.set_nonblocking(true).unwrap();
        let address = listener.local_addr().unwrap();
        let sender = thread::spawn(move || {
            let mut stream = TcpStream::connect(address).unwrap();
            stream
                .write_all(
                    b"GET /wrong?code=one-time-code&state=csrf HTTP/1.1\r\nHost: localhost\r\n\r\n",
                )
                .unwrap();
        });
        let error =
            wait_for_callback(listener, "/callback", "csrf", Duration::from_secs(2)).unwrap_err();
        sender.join().unwrap();
        assert!(error.to_string().contains("path did not match"));
    }

    #[test]
    fn callback_rejects_wrong_state() {
        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
        listener.set_nonblocking(true).unwrap();
        let address = listener.local_addr().unwrap();
        let sender = thread::spawn(move || {
            let mut stream = TcpStream::connect(address).unwrap();
            stream
                .write_all(b"GET /callback?code=one-time-code&state=wrong HTTP/1.1\r\nHost: localhost\r\n\r\n")
                .unwrap();
        });
        let error =
            wait_for_callback(listener, "/callback", "csrf", Duration::from_secs(2)).unwrap_err();
        sender.join().unwrap();
        assert!(error.to_string().contains("state did not match"));
    }
}