yaak-mountain-loop/crates/yaak-proxy/src/lib.rs

pub mod body;
pub mod cert;
mod connection;
mod request;

use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::sync::mpsc as std_mpsc;

use cert::CertificateAuthority;
use tokio::net::TcpListener;

use connection::handle_connection;

static REQUEST_ID: AtomicU64 = AtomicU64::new(1);

/// Granular events emitted during request/response lifecycle.
/// Each event carries a request `id` so consumers can correlate events.
#[derive(Debug, Clone)]
pub enum ProxyEvent {
    /// A new request has been received from the client.
    RequestStart {
        id: u64,
        method: String,
        url: String,
        http_version: String,
    },
    /// A request header sent to the upstream server.
    RequestHeader {
        id: u64,
        name: String,
        value: String,
    },
    /// The full request body (buffered before forwarding).
    RequestBody { id: u64, body: Vec<u8> },
    /// Response headers received from upstream.
    ResponseStart {
        id: u64,
        status: u16,
        http_version: String,
        elapsed_ms: u64,
    },
    /// A response header received from the upstream server.
    ResponseHeader {
        id: u64,
        name: String,
        value: String,
    },
    /// A chunk of the response body was received (emitted per-frame).
    ResponseBodyChunk { id: u64, bytes: usize },
    /// The response body stream has completed.
    ResponseBodyComplete {
        id: u64,
        body: Option<Vec<u8>>,
        size: u64,
        elapsed_ms: u64,
    },
    /// The upstream request failed.
    Error { id: u64, error: String },
}

/// Accumulated view of a proxied request, built from `ProxyEvent`s.
#[derive(Debug, Clone)]
pub struct CapturedRequest {
    pub id: u64,
    pub method: String,
    pub url: String,
    pub status: Option<u16>,
    pub elapsed_ms: Option<u64>,
    pub http_version: String,
    pub remote_http_version: Option<String>,
    pub request_headers: Vec<(String, String)>,
    pub request_body: Option<Vec<u8>>,
    pub response_headers: Vec<(String, String)>,
    pub response_body: Option<Vec<u8>>,
    pub response_body_size: u64,
    pub state: RequestState,
    pub error: Option<String>,
}

#[derive(Debug, Clone, PartialEq)]
pub enum RequestState {
    Sending,
    Receiving,
    Complete,
    Error,
}

pub struct ProxyHandle {
    shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
    thread_handle: Option<std::thread::JoinHandle<()>>,
    event_rx: Option<std_mpsc::Receiver<ProxyEvent>>,
    pub port: u16,
    pub ca_pem: String,
}

impl ProxyHandle {
    /// Take the event receiver. Can only be called once — returns `None` after the first call.
    pub fn take_event_rx(&mut self) -> Option<std_mpsc::Receiver<ProxyEvent>> {
        self.event_rx.take()
    }
}

impl Drop for ProxyHandle {
    fn drop(&mut self) {
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(());
        }
        if let Some(handle) = self.thread_handle.take() {
            let _ = handle.join();
        }
    }
}

pub fn start_proxy(port: u16) -> Result<ProxyHandle, String> {
    let ca = CertificateAuthority::new().map_err(|e| format!("Failed to create CA: {e}"))?;
    let ca_pem = ca.ca_pem();
    let ca = Arc::new(ca);

    let (event_tx, event_rx) = std_mpsc::channel();
    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
    let (ready_tx, ready_rx) = std_mpsc::channel();

    let thread_handle = std::thread::spawn(move || {
        let rt = match tokio::runtime::Runtime::new() {
            Ok(rt) => rt,
            Err(e) => {
                let _ = ready_tx.send(Err(format!("Failed to create runtime: {e}")));
                return;
            }
        };

        rt.block_on(async move {
            let addr = SocketAddr::from(([127, 0, 0, 1], port));
            let listener = match TcpListener::bind(addr).await {
                Ok(l) => l,
                Err(e) => {
                    let _ = ready_tx.send(Err(format!("Failed to bind: {e}")));
                    return;
                }
            };

            let bound_port = listener.local_addr().unwrap().port();
            let _ = ready_tx.send(Ok(bound_port));

            let mut shutdown_rx = shutdown_rx;
            loop {
                tokio::select! {
                    result = listener.accept() => {
                        match result {
                            Ok((stream, _addr)) => {
                                let tx = event_tx.clone();
                                let ca = ca.clone();
                                tokio::spawn(async move {
                                    if let Err(e) = handle_connection(stream, tx, ca).await {
                                        eprintln!("Connection error: {e}");
                                    }
                                });
                            }
                            Err(e) => eprintln!("Accept error: {e}"),
                        }
                    }
                    _ = &mut shutdown_rx => {
                        break;
                    }
                }
            }
        });
    });

    match ready_rx.recv() {
        Ok(Ok(bound_port)) => Ok(ProxyHandle {
            shutdown_tx: Some(shutdown_tx),
            thread_handle: Some(thread_handle),
            event_rx: Some(event_rx),
            port: bound_port,
            ca_pem,
        }),
        Ok(Err(e)) => Err(e),
        Err(_) => Err("Proxy thread died before binding".into()),
    }
}