Add initial yaak-proxy crate

crates/yaak-proxy/src/request.rs

@@ -0,0 +1,390 @@
+use std::convert::Infallible;
+use std::sync::Arc;
+use std::sync::atomic::Ordering;
+use std::sync::mpsc as std_mpsc;
+use std::time::Instant;
+
+use bytes::Bytes;
+use http_body_util::{BodyExt, Full};
+use hyper::body::Incoming;
+use hyper::header::HeaderMap;
+use hyper::service::service_fn;
+use hyper::{Method, Request, Response, StatusCode, Uri};
+use hyper_util::client::legacy::Client;
+use hyper_util::rt::TokioExecutor;
+use hyper_util::server::conn::auto;
+use rustls::ClientConfig;
+use rustls::pki_types::ServerName;
+use tokio::net::TcpStream;
+use tokio_rustls::TlsAcceptor;
+
+use crate::body::MeasuredBody;
+use crate::cert::CertificateAuthority;
+use crate::{ProxyEvent, REQUEST_ID};
+
+type BoxBody = http_body_util::combinators::BoxBody<Bytes, hyper::Error>;
+
+fn full_body(bytes: Bytes) -> BoxBody {
+    Full::new(bytes).map_err(|never| match never {}).boxed()
+}
+
+fn measured_incoming(
+    incoming: Incoming,
+    id: u64,
+    start: Instant,
+    tx: std_mpsc::Sender<ProxyEvent>,
+) -> BoxBody {
+    MeasuredBody::new(incoming, id, start, tx).boxed()
+}
+
+fn version_str(v: hyper::Version) -> String {
+    match v {
+        hyper::Version::HTTP_09 => "HTTP/0.9",
+        hyper::Version::HTTP_10 => "HTTP/1.0",
+        hyper::Version::HTTP_11 => "HTTP/1.1",
+        hyper::Version::HTTP_2 => "HTTP/2",
+        hyper::Version::HTTP_3 => "HTTP/3",
+        _ => "unknown",
+    }
+    .to_string()
+}
+
+fn emit_request_events(
+    tx: &std_mpsc::Sender<ProxyEvent>,
+    id: u64,
+    headers: &HeaderMap,
+    body: &Option<Vec<u8>>,
+) {
+    for (name, value) in headers.iter() {
+        let _ = tx.send(ProxyEvent::RequestHeader {
+            id,
+            name: name.to_string(),
+            value: value.to_str().unwrap_or("<binary>").to_string(),
+        });
+    }
+    if let Some(body) = body {
+        let _ = tx.send(ProxyEvent::RequestBody {
+            id,
+            body: body.clone(),
+        });
+    }
+}
+
+fn emit_response_events(
+    tx: &std_mpsc::Sender<ProxyEvent>,
+    id: u64,
+    resp: &Response<Incoming>,
+    start: &Instant,
+) {
+    let _ = tx.send(ProxyEvent::ResponseStart {
+        id,
+        status: resp.status().as_u16(),
+        http_version: version_str(resp.version()),
+        elapsed_ms: start.elapsed().as_millis() as u64,
+    });
+    for (name, value) in resp.headers().iter() {
+        let _ = tx.send(ProxyEvent::ResponseHeader {
+            id,
+            name: name.to_string(),
+            value: value.to_str().unwrap_or("<binary>").to_string(),
+        });
+    }
+}
+
+pub(crate) async fn handle_request(
+    req: Request<Incoming>,
+    event_tx: std_mpsc::Sender<ProxyEvent>,
+    ca: Arc<CertificateAuthority>,
+) -> Result<Response<BoxBody>, Infallible> {
+    let result = if req.method() == Method::CONNECT {
+        handle_connect(req, event_tx, ca).await
+    } else {
+        handle_http(req, event_tx).await
+    };
+    match result {
+        Ok(resp) => Ok(resp),
+        Err(e) => {
+            eprintln!("Proxy error: {e}");
+            Ok(Response::builder()
+                .status(StatusCode::BAD_GATEWAY)
+                .body(full_body(Bytes::from(format!("Proxy error: {e}"))))
+                .unwrap())
+        }
+    }
+}
+
+async fn handle_http(
+    req: Request<Incoming>,
+    event_tx: std_mpsc::Sender<ProxyEvent>,
+) -> Result<Response<BoxBody>, Box<dyn std::error::Error + Send + Sync>> {
+    let id = REQUEST_ID.fetch_add(1, Ordering::Relaxed);
+    let method = req.method().to_string();
+    let uri = req.uri().to_string();
+    let http_version = version_str(req.version());
+    let start = Instant::now();
+
+    let _ = event_tx.send(ProxyEvent::RequestStart {
+        id,
+        method,
+        url: uri.clone(),
+        http_version,
+    });
+
+    let client: Client<_, Full<Bytes>> = Client::builder(TokioExecutor::new()).build_http();
+
+    let (parts, body) = req.into_parts();
+    let body_bytes = body.collect().await?.to_bytes();
+    let request_body = if body_bytes.is_empty() {
+        None
+    } else {
+        Some(body_bytes.to_vec())
+    };
+    emit_request_events(&event_tx, id, &parts.headers, &request_body);
+
+    let outgoing_req = Request::from_parts(parts, Full::new(body_bytes));
+
+    match client.request(outgoing_req).await {
+        Ok(resp) => {
+            emit_response_events(&event_tx, id, &resp, &start);
+
+            let (parts, body) = resp.into_parts();
+            Ok(Response::from_parts(
+                parts,
+                measured_incoming(body, id, start, event_tx),
+            ))
+        }
+        Err(e) => {
+            let _ = event_tx.send(ProxyEvent::Error {
+                id,
+                error: e.to_string(),
+            });
+            Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
+        }
+    }
+}
+
+async fn handle_connect(
+    req: Request<Incoming>,
+    event_tx: std_mpsc::Sender<ProxyEvent>,
+    ca: Arc<CertificateAuthority>,
+) -> Result<Response<BoxBody>, Box<dyn std::error::Error + Send + Sync>> {
+    let authority = req
+        .uri()
+        .authority()
+        .map(|a| a.to_string())
+        .unwrap_or_default();
+    let (host, port) = parse_host_port(&authority);
+
+    let server_config = ca.server_config(&host)?;
+    let acceptor = TlsAcceptor::from(server_config);
+
+    let target_addr = format!("{host}:{port}");
+
+    tokio::spawn(async move {
+        let upgraded = match hyper::upgrade::on(req).await {
+            Ok(u) => u,
+            Err(e) => {
+                eprintln!("CONNECT upgrade failed: {e}");
+                return;
+            }
+        };
+
+        let tls_stream = match acceptor
+            .accept(hyper_util::rt::TokioIo::new(upgraded))
+            .await
+        {
+            Ok(s) => s,
+            Err(e) => {
+                eprintln!("TLS accept failed for {host}: {e}");
+                return;
+            }
+        };
+
+        let tx = event_tx.clone();
+        let host_for_requests = host.clone();
+        let mut builder = auto::Builder::new(TokioExecutor::new());
+        builder
+            .http1()
+            .preserve_header_case(true)
+            .title_case_headers(true);
+        if let Err(e) = builder
+            .serve_connection_with_upgrades(
+                hyper_util::rt::TokioIo::new(tls_stream),
+                service_fn(move |req| {
+                    let tx = tx.clone();
+                    let host = host_for_requests.clone();
+                    let target_addr = target_addr.clone();
+                    async move { handle_tunneled_request(req, tx, &host, &target_addr).await }
+                }),
+            )
+            .await
+        {
+            eprintln!("MITM connection error for {host}: {e}");
+        }
+    });
+
+    Ok(Response::new(full_body(Bytes::new())))
+}
+
+async fn handle_tunneled_request(
+    req: Request<Incoming>,
+    event_tx: std_mpsc::Sender<ProxyEvent>,
+    host: &str,
+    target_addr: &str,
+) -> Result<Response<BoxBody>, Infallible> {
+    let result = forward_https(req, event_tx, host, target_addr).await;
+    match result {
+        Ok(resp) => Ok(resp),
+        Err(e) => {
+            eprintln!("HTTPS forward error: {e:?}");
+            Ok(Response::builder()
+                .status(StatusCode::BAD_GATEWAY)
+                .body(full_body(Bytes::from(format!("Proxy error: {e}"))))
+                .unwrap())
+        }
+    }
+}
+
+enum HttpSender {
+    H1(hyper::client::conn::http1::SendRequest<Full<Bytes>>),
+    H2(hyper::client::conn::http2::SendRequest<Full<Bytes>>),
+}
+
+impl HttpSender {
+    async fn send_request(
+        &mut self,
+        req: Request<Full<Bytes>>,
+    ) -> Result<Response<Incoming>, hyper::Error> {
+        match self {
+            HttpSender::H1(s) => s.send_request(req).await,
+            HttpSender::H2(s) => s.send_request(req).await,
+        }
+    }
+}
+
+async fn forward_https(
+    req: Request<Incoming>,
+    event_tx: std_mpsc::Sender<ProxyEvent>,
+    host: &str,
+    target_addr: &str,
+) -> Result<Response<BoxBody>, Box<dyn std::error::Error + Send + Sync>> {
+    let id = REQUEST_ID.fetch_add(1, Ordering::Relaxed);
+    let method = req.method().to_string();
+    let http_version = version_str(req.version());
+    let path = req
+        .uri()
+        .path_and_query()
+        .map(|pq| pq.to_string())
+        .unwrap_or_else(|| "/".into());
+    let uri_str = format!("https://{host}{path}");
+    let start = Instant::now();
+
+    let _ = event_tx.send(ProxyEvent::RequestStart {
+        id,
+        method,
+        url: uri_str.clone(),
+        http_version,
+    });
+
+    // Connect to upstream with TLS
+    let tcp_stream = TcpStream::connect(target_addr).await?;
+
+    let mut root_store = rustls::RootCertStore::empty();
+    for cert in rustls_native_certs::load_native_certs().certs {
+        let _ = root_store.add(cert);
+    }
+
+    let mut tls_config =
+        ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
+            .with_safe_default_protocol_versions()?
+            .with_root_certificates(root_store)
+            .with_no_client_auth();
+    tls_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
+
+    let connector = tokio_rustls::TlsConnector::from(Arc::new(tls_config));
+    let server_name = ServerName::try_from(host.to_string())?;
+    let tls_stream = connector.connect(server_name, tcp_stream).await?;
+
+    let negotiated_h2 = tls_stream
+        .get_ref()
+        .1
+        .alpn_protocol()
+        .map_or(false, |p| p == b"h2");
+
+    let io = hyper_util::rt::TokioIo::new(tls_stream);
+
+    let mut sender = if negotiated_h2 {
+        let (sender, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor::new())
+            .handshake(io)
+            .await?;
+        tokio::spawn(async move {
+            if let Err(e) = conn.await {
+                eprintln!("Upstream h2 connection error: {e}");
+            }
+        });
+        HttpSender::H2(sender)
+    } else {
+        let (sender, conn) = hyper::client::conn::http1::Builder::new()
+            .preserve_header_case(true)
+            .title_case_headers(true)
+            .handshake(io)
+            .await?;
+        tokio::spawn(async move {
+            if let Err(e) = conn.await {
+                eprintln!("Upstream h1 connection error: {e}");
+            }
+        });
+        HttpSender::H1(sender)
+    };
+
+    // Capture request metadata
+    let (mut parts, body) = req.into_parts();
+    let body_bytes = body.collect().await?.to_bytes();
+    let request_body = if body_bytes.is_empty() {
+        None
+    } else {
+        Some(body_bytes.to_vec())
+    };
+    emit_request_events(&event_tx, id, &parts.headers, &request_body);
+
+    if negotiated_h2 {
+        // HTTP/2 requires absolute-form URI with scheme + authority
+        parts.uri = uri_str.parse::<Uri>()?;
+    } else {
+        parts.uri = path.parse::<Uri>()?;
+    }
+
+    if !parts.headers.contains_key(hyper::header::HOST) {
+        parts.headers.insert(hyper::header::HOST, host.parse()?);
+    }
+
+    let outgoing = Request::from_parts(parts, Full::new(body_bytes));
+
+    match sender.send_request(outgoing).await {
+        Ok(resp) => {
+            emit_response_events(&event_tx, id, &resp, &start);
+
+            let (parts, body) = resp.into_parts();
+            Ok(Response::from_parts(
+                parts,
+                measured_incoming(body, id, start, event_tx),
+            ))
+        }
+        Err(e) => {
+            let _ = event_tx.send(ProxyEvent::Error {
+                id,
+                error: e.to_string(),
+            });
+            Err(Box::new(e) as Box<dyn std::error::Error + Send + Sync>)
+        }
+    }
+}
+
+fn parse_host_port(authority: &str) -> (String, u16) {
+    if let Some((host, port_str)) = authority.rsplit_once(':') {
+        if let Ok(port) = port_str.parse::<u16>() {
+            return (host.to_string(), port);
+        }
+    }
+    (authority.to_string(), 443)
+}