lobe_core/engine/

capture.rs

use std::net::SocketAddr;
use std::sync::mpsc::Sender;
use std::sync::Arc;
use std::time::{Instant, SystemTime, UNIX_EPOCH};

use serde::{Deserialize, Serialize};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::{lookup_host, TcpListener, TcpStream};
use tokio_rustls::rustls::pki_types::ServerName;
use tokio_rustls::rustls::{ClientConfig, RootCertStore};
use tokio_rustls::TlsConnector;
use url::Url;
use webpki_roots::TLS_SERVER_ROOTS;

use crate::engine::pool::ConnectionPool;
use crate::engine::proxy::ProxyMeasurement;
use crate::engine::tls::{report_for_http, report_for_https_failure, report_for_https_success};
use crate::error::{Result, TloxError};
use crate::models::{TimingReport, TlsReport};

/// On-disk format for `lobe capture` session exports (also read by `lobe explain`).
/// Bumping `version` is a breaking change to the file format.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CaptureSessionExport {
    pub version: u8,
    pub listen_addr: String,
    pub upstream: String,
    pub exported_at_ms: i64,
    pub event_count: usize,
    pub events: Vec<CapturedExchange>,
    /// Optional git/CI context attached when the CLI is invoked from CI
    /// with `--pr` / `--repo` / `--branch` / `--commit`. Enables the PR bot
    /// to correlate uploads with the PR that produced them.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub git_context: Option<GitContext>,
}

/// PR + repo metadata passed from CI so uploaded sessions can be tied back
/// to the PR that produced them. All fields optional so partial context still
/// helps (e.g. baseline runs on main have a branch + commit but no PR number).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GitContext {
    /// Full repo identifier — "owner/repo" style (e.g. "kpwithcode/lobe").
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub repo: Option<String>,
    /// GitHub PR number, when the session was captured in a PR CI job.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub pr_number: Option<u32>,
    /// Branch the CI ran against.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub branch: Option<String>,
    /// Full git SHA of the commit under test.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub commit: Option<String>,
    /// GitHub Actions run identifier for linking back from PR comments.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub ci_run_id: Option<String>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CapturedExchange {
    pub target: String,
    pub request_host: String,
    pub request_method: String,
    pub request_path: String,
    pub request_headers: Vec<(String, String)>,
    pub response_headers: Vec<(String, String)>,
    pub protocol_version: Option<String>,
    pub redirect_location: Option<String>,
    pub response_status_code: Option<u16>,
    pub response_status_text: Option<String>,
    pub response_bytes: u64,
    pub report: TimingReport,
    pub tls: Option<TlsReport>,
    pub error_message: Option<String>,
    pub created_at_ms: i64,
    /// Whether the underlying TCP/TLS/HTTP-2 connection was newly opened for
    /// this request or reused from a previous one. Downstream anomaly
    /// detection uses this to skip DNS/TCP/TLS baseline checks on warm
    /// requests (those phases should be ~0 on reuse). Currently always
    /// `"new"` because Lobe opens a fresh connection per request — the real
    /// pool lands as a follow-up "Production Ready" milestone.
    #[serde(default = "default_connection_reuse")]
    pub connection_reuse: String,
}

fn default_connection_reuse() -> String {
    "new".to_string()
}

#[derive(Debug)]
pub struct CaptureProxy {
    listen_addr: String,
    upstream: Url,
    sender: Sender<CapturedExchange>,
    /// Shared across all incoming connections so requests to the same
    /// upstream benefit from connection reuse — HTTP/2 stream multiplexing
    /// on top of a single TCP+TLS session.
    pool: ConnectionPool,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct IncomingRequest {
    method: String,
    path: String,
    version: String,
    headers: Vec<(String, String)>,
    body: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ForwardResponse {
    raw_response: Vec<u8>,
    metadata: ResponseMetadata,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct ResponseMetadata {
    total_ms: u64,
    /// Upstream round-trip duration in microseconds, captured at the same
    /// instant as `total_ms`. Used by `handle_connection` to derive the
    /// proxy's own added latency with sub-millisecond precision.
    upstream_us: u64,
    status_code: Option<u16>,
    status_text: Option<String>,
    bytes: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct UpstreamResponse {
    status_code: u16,
    status_text: String,
    protocol_version: String,
    response_headers: Vec<(String, String)>,
    redirect_location: Option<String>,
    body: Vec<u8>,
}

impl CaptureProxy {
    pub fn new(listen_addr: String, upstream: &str, sender: Sender<CapturedExchange>) -> Result<Self> {
        let upstream = Url::parse(upstream)
            .map_err(|_| TloxError::InvalidTarget(upstream.to_string()))?;

        Ok(Self {
            listen_addr,
            upstream,
            sender,
            pool: ConnectionPool::new(),
        })
    }

    pub fn listen_addr(&self) -> &str {
        &self.listen_addr
    }

    pub fn upstream(&self) -> &Url {
        &self.upstream
    }

    pub async fn run(self) -> Result<()> {
        let listener = TcpListener::bind(&self.listen_addr).await?;
        let upstream = self.upstream.clone();
        let sender = self.sender.clone();
        let pool = self.pool.clone();

        loop {
            let (stream, _) = listener.accept().await?;
            let upstream = upstream.clone();
            let sender = sender.clone();
            let pool = pool.clone();

            tokio::spawn(async move {
                let error_target = upstream.to_string();
                if let Err(error) =
                    handle_connection(stream, upstream, sender.clone(), pool).await
                {
                    let exchange = CapturedExchange {
                        target: error_target,
                        request_host: "n/a".to_string(),
                        request_method: "UNKNOWN".to_string(),
                        request_path: "/".to_string(),
                        request_headers: Vec::new(),
                        response_headers: Vec::new(),
                        protocol_version: None,
                        redirect_location: None,
                        response_status_code: Some(502),
                        response_status_text: Some("Bad Gateway".to_string()),
                        response_bytes: 0,
                        report: TimingReport::default(),
                        tls: None,
                        error_message: Some(error.to_string()),
                        created_at_ms: current_time_ms().unwrap_or_default(),
                        connection_reuse: "new".to_string(),
                    };
                    let _ = sender.send(exchange);
                }
            });
        }
    }
}

async fn handle_connection(
    mut inbound: TcpStream,
    upstream: Url,
    sender: Sender<CapturedExchange>,
    pool: ConnectionPool,
) -> Result<()> {
    let Some(request) = read_http_request(&mut inbound).await? else {
        return Ok(());
    };
    // Start of Lobe's own serving work — the client request is now in hand.
    // Everything from here until the response is flushed back, minus the
    // upstream round trip, is proxy overhead.
    let received_at = Instant::now();

    let result = forward_request(&upstream, &request, &pool).await;

    match result {
        Ok((response, mut exchange)) => {
            inbound.write_all(&response.raw_response).await?;
            inbound.flush().await?;
            let client_facing_us = received_at.elapsed().as_micros() as u64;
            exchange.report.proxy_overhead_us =
                client_facing_us.saturating_sub(response.metadata.upstream_us);
            let _ = sender.send(exchange);
        }
        Err((exchange, response_bytes)) => {
            // Failed upstream request — overhead isn't meaningful, leave it 0.
            inbound.write_all(&response_bytes).await?;
            inbound.flush().await?;
            let _ = sender.send(exchange);
        }
    }

    Ok(())
}

async fn forward_request(
    upstream: &Url,
    request: &IncomingRequest,
    pool: &ConnectionPool,
) -> std::result::Result<(ForwardResponse, CapturedExchange), (CapturedExchange, Vec<u8>)> {
    let is_https = upstream.scheme() == "https";
    let target_url = match build_upstream_url(upstream, &request.path) {
        Ok(url) => url,
        Err(error) => {
            let exchange = error_exchange(
                upstream,
                request,
                TimingReport::default(),
                tls_report_for_failure(is_https, error.to_string()),
                error.to_string(),
            );
            return Err((exchange, bad_gateway_response("invalid upstream url")));
        }
    };

    let host = match target_url.host_str() {
        Some(host) => host.to_string(),
        None => {
            let message = "upstream is missing a host".to_string();
            let exchange = error_exchange(
                upstream,
                request,
                TimingReport::default(),
                tls_report_for_failure(is_https, message.clone()),
                message,
            );
            return Err((exchange, bad_gateway_response("upstream missing host")));
        }
    };

    let port = match target_url.port_or_known_default() {
        Some(port) => port,
        None => {
            let message = "upstream is missing a port".to_string();
            let exchange = error_exchange(
                upstream,
                request,
                TimingReport::default(),
                tls_report_for_failure(is_https, message.clone()),
                message,
            );
            return Err((exchange, bad_gateway_response("upstream missing port")));
        }
    };

    let mut measurement = ProxyMeasurement::new();
    let total_start = Instant::now();

    // Warm path — if we already have an HTTP/2 sender cached for this host,
    // skip DNS+TCP+TLS+h2-handshake entirely and multiplex a new stream over
    // the existing connection. All connection-setup phases stay at 0ms
    // because the underlying transport is already established.
    if is_https {
        if let Some(cached_sender) = pool.get_h2(host.as_str(), port).await {
            let request_target = build_request_target(&target_url);
            let warm_outcome = send_over_h2(
                cached_sender,
                &host,
                &request_target,
                request,
                &mut measurement,
            )
            .await;

            match warm_outcome {
                Ok(upstream_response) => {
                    let elapsed = total_start.elapsed();
                    let total_ms = elapsed.as_millis() as u64;
                    let mut report = measurement.finish_report();
                    report.total_ms = total_ms;

                    let raw_response = build_raw_response(
                        upstream_response.status_code,
                        &upstream_response.status_text,
                        &upstream_response.response_headers,
                        &upstream_response.body,
                    );
                    let metadata = ResponseMetadata {
                        total_ms,
                        upstream_us: elapsed.as_micros() as u64,
                        status_code: Some(upstream_response.status_code),
                        status_text: Some(upstream_response.status_text.clone()),
                        bytes: upstream_response.body.len() as u64,
                    };
                    // Reused connection has no TLS handshake data to report
                    // — the TLS metadata from the original handshake is
                    // still valid but we don't re-record it per request.
                    let tls_report = report_for_https_success(None, None, None);
                    let mut exchange = success_exchange(
                        &target_url,
                        request,
                        report,
                        request.headers.clone(),
                        upstream_response.response_headers.clone(),
                        Some(upstream_response.protocol_version.clone()),
                        upstream_response.redirect_location.clone(),
                        Some(tls_report),
                        metadata.clone(),
                    );
                    exchange.connection_reuse = "reused".to_string();

                    return Ok((ForwardResponse { raw_response, metadata }, exchange));
                }
                Err(_error) => {
                    // Cached sender is dead — evict and fall through to a
                    // fresh handshake below. Don't fail the request just
                    // because the pool went stale.
                    pool.invalidate_h2(host.as_str(), port).await;
                }
            }
        }
    }

    let dns_start = Instant::now();
    let addr: SocketAddr = match lookup_host((host.as_str(), port)).await {
        Ok(mut addresses) => match addresses.next() {
            Some(addr) => addr,
            None => {
                let message = "no addresses resolved".to_string();
                let report = finish_report(&measurement, total_start);
                let tls = tls_report_for_failure(is_https, message.clone());
                let exchange = error_exchange(upstream, request, report, tls, message);
                return Err((exchange, bad_gateway_response("upstream unresolved")));
            }
        },
        Err(error) => {
            let report = finish_report(&measurement, total_start);
            let tls = tls_report_for_failure(is_https, error.to_string());
            let exchange = error_exchange(upstream, request, report, tls, error.to_string());
            return Err((exchange, bad_gateway_response("dns resolution failed")));
        }
    };
    measurement.record_dns_ms(dns_start.elapsed().as_millis() as u64);

    let tcp_start = Instant::now();
    let stream = match TcpStream::connect(addr).await {
        Ok(stream) => stream,
        Err(error) => {
            let report = finish_report(&measurement, total_start);
            let tls = tls_report_for_failure(is_https, error.to_string());
            let exchange = error_exchange(upstream, request, report, tls, error.to_string());
            return Err((exchange, bad_gateway_response("tcp connect failed")));
        }
    };
    measurement.record_tcp_ms(tcp_start.elapsed().as_millis() as u64);

    let request_target = build_request_target(&target_url);

    let outcome = if is_https {
        exchange_https(
            stream,
            &host,
            port,
            &request_target,
            request,
            &mut measurement,
            pool,
        )
        .await
    } else {
        exchange_http(stream, &host, &request_target, request, &mut measurement).await
    };

    match outcome {
        Ok((upstream_response, tls_report)) => {
            let elapsed = total_start.elapsed();
            let total_ms = elapsed.as_millis() as u64;
            let mut report = measurement.finish_report();
            report.total_ms = total_ms;

            let raw_response = build_raw_response(
                upstream_response.status_code,
                &upstream_response.status_text,
                &upstream_response.response_headers,
                &upstream_response.body,
            );

            let metadata = ResponseMetadata {
                total_ms,
                upstream_us: elapsed.as_micros() as u64,
                status_code: Some(upstream_response.status_code),
                status_text: Some(upstream_response.status_text.clone()),
                bytes: upstream_response.body.len() as u64,
            };

            let exchange = success_exchange(
                &target_url,
                request,
                report,
                request.headers.clone(),
                upstream_response.response_headers.clone(),
                Some(upstream_response.protocol_version.clone()),
                upstream_response.redirect_location.clone(),
                Some(tls_report),
                metadata.clone(),
            );

            Ok((
                ForwardResponse {
                    raw_response,
                    metadata,
                },
                exchange,
            ))
        }
        Err(error) => {
            let report = finish_report(&measurement, total_start);
            let tls = tls_report_for_failure(is_https, error.to_string());
            let exchange = error_exchange(upstream, request, report, tls, error.to_string());
            Err((exchange, bad_gateway_response("upstream exchange failed")))
        }
    }
}

async fn exchange_http(
    mut stream: TcpStream,
    host: &str,
    request_target: &str,
    request: &IncomingRequest,
    measurement: &mut ProxyMeasurement,
) -> std::io::Result<(UpstreamResponse, TlsReport)> {
    let response = exchange_over_stream(&mut stream, host, request_target, request, measurement).await?;
    Ok((response, report_for_http()))
}

async fn exchange_https(
    stream: TcpStream,
    host: &str,
    port: u16,
    request_target: &str,
    request: &IncomingRequest,
    measurement: &mut ProxyMeasurement,
    pool: &ConnectionPool,
) -> std::io::Result<(UpstreamResponse, TlsReport)> {
    let tls_start = Instant::now();
    let connector = build_tls_connector();
    let server_name = ServerName::try_from(host.to_string())
        .map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidInput, error.to_string()))?;

    let mut tls_stream = connector
        .connect(server_name, stream)
        .await
        .map_err(|error| std::io::Error::other(error))?;
    measurement.record_tls_ms(tls_start.elapsed().as_millis() as u64);

    // ALPN-branch: honor whichever protocol the server negotiated. HTTP/1.1
    // stays on the proven raw-tokio path; HTTP/2 goes through the hyper-based
    // h2 helper. The h1 path below is byte-for-byte unchanged from the
    // pre-h2 shipping code.
    let negotiated_h2 = {
        let (_, connection) = tls_stream.get_ref();
        connection.alpn_protocol() == Some(b"h2".as_slice())
    };

    // Capture TLS metadata BEFORE handing the stream off — after h2 handshake
    // the stream is consumed by hyper and we can't inspect it anymore.
    let tls_metadata = {
        let (_, connection) = tls_stream.get_ref();
        (
            connection.protocol_version().map(|value| format!("{value:?}")),
            connection
                .negotiated_cipher_suite()
                .map(|value| format!("{:?}", value.suite())),
        )
    };

    let response = if negotiated_h2 {
        exchange_over_h2(
            tls_stream,
            host,
            port,
            request_target,
            request,
            measurement,
            pool,
        )
        .await?
    } else {
        exchange_over_stream(&mut tls_stream, host, request_target, request, measurement).await?
    };

    let (tls_version, cipher_suite) = tls_metadata;
    let tls_report = report_for_https_success(tls_version, cipher_suite, None);

    Ok((response, tls_report))
}

async fn exchange_over_stream<S>(
    stream: &mut S,
    host: &str,
    request_target: &str,
    request: &IncomingRequest,
    measurement: &mut ProxyMeasurement,
) -> std::io::Result<UpstreamResponse>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let request_bytes = serialize_request(&request.method, request_target, host, &request.headers, &request.body);

    let ttfb_start = Instant::now();
    stream.write_all(&request_bytes).await?;
    stream.flush().await?;

    let mut buffer = Vec::new();
    let mut chunk = [0_u8; 4096];
    let mut download_start: Option<Instant> = None;

    loop {
        let bytes_read = stream.read(&mut chunk).await?;
        if bytes_read == 0 {
            break;
        }

        if download_start.is_none() {
            measurement.record_ttfb_ms(ttfb_start.elapsed().as_millis() as u64);
            download_start = Some(Instant::now());
        }

        buffer.extend_from_slice(&chunk[..bytes_read]);
    }

    if let Some(start) = download_start {
        measurement.record_download_ms(start.elapsed().as_millis() as u64);
    }

    parse_upstream_response(&buffer)
}

/// HTTP/2 exchange over an already-TLS-handshaken stream. Called when ALPN
/// negotiated `h2`. Uses `hyper` for the HTTP/2 protocol so we get proper stream
/// multiplexing, HPACK, and flow control for free.
///
/// This path is entirely separate from the HTTP/1.1 raw-tokio implementation
/// (`exchange_over_stream`). If h1 is negotiated, this function is never called
/// and the h1 path is byte-for-byte unchanged.
async fn exchange_over_h2<S>(
    stream: S,
    host: &str,
    port: u16,
    request_target: &str,
    request: &IncomingRequest,
    measurement: &mut ProxyMeasurement,
    pool: &ConnectionPool,
) -> std::io::Result<UpstreamResponse>
where
    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + Unpin + 'static,
{
    use http_body_util::Full;
    use hyper::body::Bytes;
    use hyper::client::conn::http2;
    use hyper_util::rt::{TokioExecutor, TokioIo};

    let io = TokioIo::new(stream);
    let (sender, connection) = http2::handshake::<_, _, Full<Bytes>>(TokioExecutor::new(), io)
        .await
        .map_err(|error| std::io::Error::other(error))?;

    // Drive the connection I/O in the background. Because `sender` is
    // `Clone` (it's a channel handle), we can cache one clone in the pool
    // and use another to send this request. The connection driver stays
    // alive as long as ANY clone of the sender is live — so once we cache
    // it in the pool, the connection persists across many requests until
    // the pool evicts it.
    tokio::spawn(async move {
        let _ = connection.await;
    });

    // Cache the sender BEFORE sending this request. Even if this request
    // fails, the sender is still valid for future retries.
    pool.store_h2(host, port, sender.clone()).await;

    send_over_h2(sender, host, request_target, request, measurement).await
}

/// Send a single HTTP/2 request over an already-established `SendRequest`.
/// Used both on the warm path (cached sender from the pool) and internally
/// from `exchange_over_h2` (fresh sender we just handshaked).
async fn send_over_h2(
    mut sender: crate::engine::pool::PooledH2Sender,
    host: &str,
    request_target: &str,
    request: &IncomingRequest,
    measurement: &mut ProxyMeasurement,
) -> std::io::Result<UpstreamResponse> {
    use http::Request as HttpRequest;
    use http_body_util::{BodyExt, Full};
    use hyper::body::Bytes;

    // Build the h2 request. For HTTP/2 the URI must include scheme + authority
    // so hyper can populate the `:scheme`/`:authority` pseudo-headers; we do NOT
    // forward the client's `Host` header (h2 has no Host header — it's the
    // `:authority` pseudo-header instead).
    let uri: http::Uri = format!("https://{host}{request_target}")
        .parse()
        .map_err(|error: http::uri::InvalidUri| std::io::Error::other(error))?;

    let method = http::Method::from_bytes(request.method.as_bytes())
        .map_err(|error| std::io::Error::other(error))?;

    let mut builder = HttpRequest::builder().method(method).uri(uri);

    for (name, value) in &request.headers {
        if header_is_hop_by_hop_or_h2_incompatible(name) {
            continue;
        }
        builder = builder.header(name.as_str(), value.as_str());
    }

    let body = Full::new(Bytes::from(request.body.clone()));
    let req = builder
        .body(body)
        .map_err(|error| std::io::Error::other(error))?;

    let ttfb_start = Instant::now();
    let response = sender
        .send_request(req)
        .await
        .map_err(|error| std::io::Error::other(error))?;
    measurement.record_ttfb_ms(ttfb_start.elapsed().as_millis() as u64);

    let status_code = response.status().as_u16();
    let status_text = response
        .status()
        .canonical_reason()
        .unwrap_or("")
        .to_string();

    let mut response_headers: Vec<(String, String)> = Vec::new();
    let mut redirect_location: Option<String> = None;
    for (name, value) in response.headers() {
        let value_str = match value.to_str() {
            Ok(text) => text.to_string(),
            Err(_) => continue,
        };
        if name.as_str().eq_ignore_ascii_case("location") {
            redirect_location = Some(value_str.clone());
        }
        response_headers.push((name.as_str().to_string(), value_str));
    }

    let download_start = Instant::now();
    let body_bytes = response
        .into_body()
        .collect()
        .await
        .map_err(|error| std::io::Error::other(error))?
        .to_bytes();
    measurement.record_download_ms(download_start.elapsed().as_millis() as u64);

    Ok(UpstreamResponse {
        status_code,
        status_text,
        protocol_version: "HTTP/2".to_string(),
        response_headers,
        redirect_location,
        body: body_bytes.to_vec(),
    })
}

/// Headers we should not forward across HTTP/2. Some are hop-by-hop and
/// per-RFC 7230 not meaningful end-to-end; some (like `host`) are replaced by
/// h2 pseudo-headers; some (like `transfer-encoding`) are outright forbidden
/// by RFC 7540 §8.1.2.2.
fn header_is_hop_by_hop_or_h2_incompatible(name: &str) -> bool {
    matches!(
        name.to_ascii_lowercase().as_str(),
        "host"
            | "connection"
            | "proxy-connection"
            | "keep-alive"
            | "transfer-encoding"
            | "te"
            | "trailer"
            | "upgrade"
            | "content-length"
    )
}

fn serialize_request(
    method: &str,
    request_target: &str,
    host: &str,
    headers: &[(String, String)],
    body: &[u8],
) -> Vec<u8> {
    let mut out = format!("{method} {request_target} HTTP/1.1\r\n").into_bytes();
    out.extend_from_slice(format!("Host: {host}\r\n").as_bytes());
    out.extend_from_slice(b"Connection: close\r\n");

    for (name, value) in headers {
        if name.eq_ignore_ascii_case("host")
            || name.eq_ignore_ascii_case("connection")
            || name.eq_ignore_ascii_case("proxy-connection")
            || name.eq_ignore_ascii_case("content-length")
        {
            continue;
        }
        out.extend_from_slice(name.as_bytes());
        out.extend_from_slice(b": ");
        out.extend_from_slice(value.as_bytes());
        out.extend_from_slice(b"\r\n");
    }

    if !body.is_empty() {
        out.extend_from_slice(format!("Content-Length: {}\r\n", body.len()).as_bytes());
    }

    out.extend_from_slice(b"\r\n");
    out.extend_from_slice(body);
    out
}

fn parse_upstream_response(buffer: &[u8]) -> std::io::Result<UpstreamResponse> {
    let header_end = find_header_end(buffer).ok_or_else(|| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, "missing response headers")
    })?;

    let header_bytes = &buffer[..header_end];
    let body_bytes = &buffer[header_end + 4..];

    let header_text = std::str::from_utf8(header_bytes)
        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "response headers not utf8"))?;

    let mut lines = header_text.split("\r\n");
    let status_line = lines
        .next()
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing status line"))?;

    let mut parts = status_line.splitn(3, ' ');
    let protocol_version = parts.next().unwrap_or("HTTP/1.1").to_string();
    let status_code = parts
        .next()
        .and_then(|value| value.parse::<u16>().ok())
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing status code"))?;
    let status_text = parts.next().unwrap_or("").trim().to_string();

    let mut response_headers: Vec<(String, String)> = Vec::new();
    let mut is_chunked = false;
    let mut redirect_location: Option<String> = None;

    for line in lines {
        let Some((name, value)) = line.split_once(':') else {
            continue;
        };
        let name = name.trim();
        let value = value.trim();
        if name.eq_ignore_ascii_case("transfer-encoding")
            && value.to_ascii_lowercase().contains("chunked")
        {
            is_chunked = true;
        }
        if name.eq_ignore_ascii_case("location") {
            redirect_location = Some(value.to_string());
        }
        response_headers.push((name.to_string(), value.to_string()));
    }

    let body = if is_chunked {
        decode_chunked(body_bytes)?
    } else {
        body_bytes.to_vec()
    };

    Ok(UpstreamResponse {
        status_code,
        status_text,
        protocol_version,
        response_headers,
        redirect_location,
        body,
    })
}

fn decode_chunked(input: &[u8]) -> std::io::Result<Vec<u8>> {
    let mut out = Vec::with_capacity(input.len());
    let mut cursor = 0usize;

    loop {
        let line_end = match find_crlf(&input[cursor..]) {
            Some(pos) => cursor + pos,
            None => {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    "unterminated chunk size line",
                ));
            }
        };

        let size_line = std::str::from_utf8(&input[cursor..line_end])
            .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "chunk size not utf8"))?;
        let size_hex = size_line.split(';').next().unwrap_or("").trim();
        let chunk_size = usize::from_str_radix(size_hex, 16)
            .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "invalid chunk size"))?;

        cursor = line_end + 2;

        if chunk_size == 0 {
            break;
        }

        if cursor + chunk_size > input.len() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "chunk body exceeds buffer",
            ));
        }

        out.extend_from_slice(&input[cursor..cursor + chunk_size]);
        cursor += chunk_size;

        if input.get(cursor..cursor + 2) != Some(b"\r\n") {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "missing CRLF after chunk",
            ));
        }
        cursor += 2;
    }

    Ok(out)
}

fn find_crlf(buffer: &[u8]) -> Option<usize> {
    buffer.windows(2).position(|window| window == b"\r\n")
}

fn build_tls_config() -> ClientConfig {
    let roots = RootCertStore::from_iter(TLS_SERVER_ROOTS.iter().cloned());
    let mut config = ClientConfig::builder()
        .with_root_certificates(roots)
        .with_no_client_auth();
    config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
    config
}

fn build_tls_connector() -> TlsConnector {
    TlsConnector::from(Arc::new(build_tls_config()))
}

fn finish_report(measurement: &ProxyMeasurement, total_start: Instant) -> TimingReport {
    let mut report = measurement.finish_report();
    report.total_ms = total_start.elapsed().as_millis() as u64;
    report
}

async fn read_http_request(stream: &mut TcpStream) -> std::io::Result<Option<IncomingRequest>> {
    let mut buffer = Vec::new();
    let mut chunk = [0_u8; 2048];

    loop {
        let bytes_read = stream.read(&mut chunk).await?;
        if bytes_read == 0 {
            if buffer.is_empty() {
                return Ok(None);
            }
            break;
        }

        buffer.extend_from_slice(&chunk[..bytes_read]);

        if let Some(header_end) = find_header_end(&buffer) {
            let content_length = parse_content_length(&buffer[..header_end]).unwrap_or(0);
            let total_length = header_end + 4 + content_length;

            while buffer.len() < total_length {
                let bytes_read = stream.read(&mut chunk).await?;
                if bytes_read == 0 {
                    break;
                }
                buffer.extend_from_slice(&chunk[..bytes_read]);
            }

            return parse_http_request(&buffer).map(Some);
        }
    }

    parse_http_request(&buffer).map(Some)
}

fn parse_http_request(buffer: &[u8]) -> std::io::Result<IncomingRequest> {
    let header_end = find_header_end(buffer)
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing request headers"))?;
    let header_bytes = &buffer[..header_end];
    let body = buffer[header_end + 4..].to_vec();
    let header_text = std::str::from_utf8(header_bytes)
        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidData, "request headers not utf8"))?;
    let mut lines = header_text.split("\r\n");
    let request_line = lines
        .next()
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing request line"))?;
    let mut parts = request_line.split_whitespace();
    let method = parts
        .next()
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing method"))?;
    let path = parts
        .next()
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "missing path"))?;
    let version = parts.next().unwrap_or("HTTP/1.1");

    let headers = lines
        .filter_map(|line| {
            let (name, value) = line.split_once(':')?;
            Some((name.trim().to_string(), value.trim().to_string()))
        })
        .collect::<Vec<_>>();

    Ok(IncomingRequest {
        method: method.to_string(),
        path: normalize_request_path(path),
        version: version.to_string(),
        headers,
        body,
    })
}

fn normalize_request_path(path: &str) -> String {
    if let Ok(url) = Url::parse(path) {
        let mut normalized = url.path().to_string();
        if normalized.is_empty() {
            normalized.push('/');
        }
        if let Some(query) = url.query() {
            normalized.push('?');
            normalized.push_str(query);
        }
        normalized
    } else if path.is_empty() {
        "/".to_string()
    } else {
        path.to_string()
    }
}

fn build_upstream_url(upstream: &Url, request_path: &str) -> std::result::Result<Url, url::ParseError> {
    let request_path = normalize_request_path(request_path);
    let (path_only, query) = match request_path.split_once('?') {
        Some((path, query)) => (path, Some(query)),
        None => (request_path.as_str(), None),
    };
    let request_without_root = path_only.trim_start_matches('/');
    let base_without_trailing = upstream.path().trim_end_matches('/');

    let combined_path = if base_without_trailing.is_empty() || base_without_trailing == "/" {
        format!("/{}", request_without_root)
    } else if request_without_root.is_empty() {
        base_without_trailing.to_string()
    } else {
        format!("{base_without_trailing}/{request_without_root}")
    };

    let mut url = upstream.clone();
    url.set_path(if combined_path.is_empty() { "/" } else { &combined_path });
    url.set_query(query);

    Ok(url)
}

fn build_request_target(url: &Url) -> String {
    let mut target = url.path().to_string();
    if target.is_empty() {
        target.push('/');
    }

    if let Some(query) = url.query() {
        target.push('?');
        target.push_str(query);
    }

    target
}

fn build_raw_response(
    status_code: u16,
    status_text: &str,
    headers: &[(String, String)],
    body: &[u8],
) -> Vec<u8> {
    let mut response = format!("HTTP/1.1 {status_code} {status_text}\r\n").into_bytes();

    for (name, value) in headers {
        if name.eq_ignore_ascii_case("connection")
            || name.eq_ignore_ascii_case("transfer-encoding")
            || name.eq_ignore_ascii_case("content-length")
        {
            continue;
        }

        response.extend_from_slice(name.as_bytes());
        response.extend_from_slice(b": ");
        response.extend_from_slice(value.as_bytes());
        response.extend_from_slice(b"\r\n");
    }

    response
        .extend_from_slice(format!("Content-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes());
    response.extend_from_slice(body);
    response
}

fn success_exchange(
    target_url: &Url,
    request: &IncomingRequest,
    report: TimingReport,
    request_headers: Vec<(String, String)>,
    response_headers: Vec<(String, String)>,
    protocol_version: Option<String>,
    redirect_location: Option<String>,
    tls: Option<TlsReport>,
    metadata: ResponseMetadata,
) -> CapturedExchange {
    CapturedExchange {
        target: target_url.to_string(),
        request_host: target_url.host_str().unwrap_or("n/a").to_string(),
        request_method: request.method.clone(),
        request_path: normalize_request_path(&request.path),
        request_headers,
        response_headers,
        protocol_version,
        redirect_location,
        response_status_code: metadata.status_code,
        response_status_text: metadata.status_text,
        response_bytes: metadata.bytes,
        report,
        tls,
        error_message: None,
        created_at_ms: current_time_ms().unwrap_or_default(),
        connection_reuse: "new".to_string(),
    }
}

fn error_exchange(
    upstream: &Url,
    request: &IncomingRequest,
    report: TimingReport,
    tls: Option<TlsReport>,
    error_message: String,
) -> CapturedExchange {
    let body = error_message.clone();
    CapturedExchange {
        target: upstream.to_string(),
        request_host: upstream.host_str().unwrap_or("n/a").to_string(),
        request_method: request.method.clone(),
        request_path: normalize_request_path(&request.path),
        request_headers: request.headers.clone(),
        response_headers: Vec::new(),
        protocol_version: None,
        redirect_location: None,
        response_status_code: Some(502),
        response_status_text: Some("Bad Gateway".to_string()),
        response_bytes: body.len() as u64,
        report,
        tls,
        error_message: Some(error_message),
        created_at_ms: current_time_ms().unwrap_or_default(),
        connection_reuse: "new".to_string(),
    }
}

fn tls_report_for_failure(is_https: bool, error_message: String) -> Option<TlsReport> {
    if is_https {
        Some(report_for_https_failure(error_message))
    } else {
        Some(report_for_http())
    }
}

fn bad_gateway_response(message: &str) -> Vec<u8> {
    let body = message.as_bytes();
    format!(
        "HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
        body.len(),
        message
    )
    .into_bytes()
}

fn parse_content_length(header_bytes: &[u8]) -> Option<usize> {
    let header_text = std::str::from_utf8(header_bytes).ok()?;
    header_text
        .lines()
        .find_map(|line| {
            let (name, value) = line.split_once(':')?;
            if name.eq_ignore_ascii_case("content-length") {
                value.trim().parse::<usize>().ok()
            } else {
                None
            }
        })
}

fn find_header_end(buffer: &[u8]) -> Option<usize> {
    buffer.windows(4).position(|window| window == b"\r\n\r\n")
}

fn current_time_ms() -> Result<i64> {
    let now = SystemTime::now().duration_since(UNIX_EPOCH)?;
    Ok(now.as_millis() as i64)
}

#[cfg(test)]
mod tests {
    use super::{
        build_raw_response, build_tls_config, build_upstream_url, decode_chunked,
        normalize_request_path, parse_http_request, parse_upstream_response, serialize_request,
    };
    use url::Url;

    #[test]
    fn tls_config_advertises_h2_then_http1_alpn() {
        let config = build_tls_config();
        assert_eq!(
            config.alpn_protocols,
            vec![b"h2".to_vec(), b"http/1.1".to_vec()],
            "capture must advertise h2 first so servers pick it (and we then reject cleanly) \
             instead of silently downgrading to http/1.1 and reporting wrong numbers",
        );
    }

    #[test]
    fn normalize_request_path_extracts_path_from_absolute_url() {
        assert_eq!(
            normalize_request_path("https://example.com/api/health?full=1"),
            "/api/health?full=1"
        );
    }

    #[test]
    fn build_upstream_url_preserves_base_prefix() {
        let upstream = Url::parse("https://example.com/api").expect("url should parse");

        assert_eq!(
            build_upstream_url(&upstream, "/health").expect("url should build").as_str(),
            "https://example.com/api/health"
        );
    }

    #[test]
    fn build_upstream_url_preserves_query_without_encoding_it_into_path() {
        let upstream = Url::parse("http://127.0.0.1:8000").expect("url should parse");

        assert_eq!(
            build_upstream_url(&upstream, "/mlb/run-conversion?timeframe=season")
                .expect("url should build")
                .as_str(),
            "http://127.0.0.1:8000/mlb/run-conversion?timeframe=season"
        );
    }

    #[test]
    fn parse_http_request_reads_method_path_and_body() {
        let request = b"POST /users HTTP/1.1\r\nHost: localhost\r\nContent-Length: 7\r\n\r\nhello=1";
        let parsed = parse_http_request(request).expect("request should parse");

        assert_eq!(parsed.method, "POST");
        assert_eq!(parsed.path, "/users");
        assert_eq!(parsed.version, "HTTP/1.1");
        assert_eq!(parsed.body, b"hello=1");
    }

    #[test]
    fn build_raw_response_preserves_content_encoding_and_strips_hop_headers() {
        let headers = vec![
            ("Content-Type".to_string(), "text/html; charset=utf-8".to_string()),
            ("Content-Encoding".to_string(), "gzip".to_string()),
            ("Transfer-Encoding".to_string(), "chunked".to_string()),
            ("Connection".to_string(), "keep-alive".to_string()),
        ];

        let response = build_raw_response(200, "OK", &headers, b"abc");
        let response_text = String::from_utf8_lossy(&response).to_ascii_lowercase();

        assert!(response_text.contains("content-encoding: gzip\r\n"));
        assert!(response_text.contains("content-type: text/html; charset=utf-8\r\n"));
        assert!(response_text.contains("content-length: 3\r\n"));
        assert!(!response_text.contains("transfer-encoding:"));
        assert!(response_text.ends_with("\r\n\r\nabc"));
    }

    #[test]
    fn serialize_request_forces_close_and_strips_hop_headers() {
        let bytes = serialize_request(
            "POST",
            "/api/users",
            "example.com",
            &[
                ("Content-Type".to_string(), "application/json".to_string()),
                ("Host".to_string(), "old-host".to_string()),
                ("Connection".to_string(), "keep-alive".to_string()),
                ("Content-Length".to_string(), "999".to_string()),
            ],
            b"{}",
        );

        let text = String::from_utf8_lossy(&bytes);
        assert!(text.starts_with("POST /api/users HTTP/1.1\r\n"));
        assert!(text.contains("Host: example.com\r\n"));
        assert!(text.contains("Connection: close\r\n"));
        assert!(text.contains("Content-Length: 2\r\n"));
        assert!(!text.to_ascii_lowercase().contains("host: old-host"));
        assert!(text.ends_with("\r\n\r\n{}"));
    }

    #[test]
    fn parse_upstream_response_reads_status_and_body() {
        let raw = b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nhello world";
        let response = parse_upstream_response(raw).expect("should parse");

        assert_eq!(response.status_code, 200);
        assert_eq!(response.status_text, "OK");
        assert_eq!(response.protocol_version, "HTTP/1.1");
        assert_eq!(response.body, b"hello world");
        assert!(response
            .response_headers
            .iter()
            .any(|(name, value)| name.eq_ignore_ascii_case("content-type") && value == "text/plain"));
    }

    #[test]
    fn parse_upstream_response_decodes_chunked_body() {
        let raw = b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n5\r\nhello\r\n6\r\n world\r\n0\r\n\r\n";
        let response = parse_upstream_response(raw).expect("should parse");

        assert_eq!(response.body, b"hello world");
    }

    #[test]
    fn parse_upstream_response_captures_redirect_location() {
        let raw = b"HTTP/1.1 302 Found\r\nLocation: /next\r\nContent-Length: 0\r\n\r\n";
        let response = parse_upstream_response(raw).expect("should parse");

        assert_eq!(response.status_code, 302);
        assert_eq!(response.redirect_location.as_deref(), Some("/next"));
    }

    #[test]
    fn decode_chunked_handles_extension_after_size() {
        let raw = b"3;name=value\r\nabc\r\n0\r\n\r\n";
        assert_eq!(decode_chunked(raw).expect("should decode"), b"abc");
    }

    #[test]
    fn decode_chunked_rejects_missing_crlf() {
        let raw = b"3\r\nabcXX0\r\n\r\n";
        assert!(decode_chunked(raw).is_err());
    }

    #[test]
    fn h2_header_filter_strips_incompatible_and_hop_by_hop_headers() {
        use super::header_is_hop_by_hop_or_h2_incompatible as filter;

        // Forbidden in HTTP/2 per RFC 7540 §8.1.2.2
        assert!(filter("connection"));
        assert!(filter("Connection"));
        assert!(filter("keep-alive"));
        assert!(filter("proxy-connection"));
        assert!(filter("transfer-encoding"));
        assert!(filter("upgrade"));
        assert!(filter("TE"));

        // Host is replaced by :authority pseudo-header in h2
        assert!(filter("host"));

        // Content-length is derived from the body in h2
        assert!(filter("content-length"));

        // These should pass through
        assert!(!filter("content-type"));
        assert!(!filter("accept"));
        assert!(!filter("authorization"));
        assert!(!filter("user-agent"));
        assert!(!filter("x-custom-header"));
    }
}