Optimize throughput: 128K tunnel buffers, streaming bodies, connection pooling

- Increase CONNECT tunnel buffer from 8 KiB to 128 KiB (copy_bidirectional_with_sizes)
- Stream HTTP forward response bodies via Either<Full, Incoming> instead of buffering
- Add pooled hyper_util legacy Client for HTTP forward (OnceLock, 90s idle, 32/host)
- Add throughput tests: 64 MB CONNECT, 32 MB forward, 500-req rate test

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: madeye

src/lib.rs

@@ -22,11 +22,14 @@ pub mod tls;
 
 use std::sync::Arc;
 
-use http_body_util::Full;
+use http_body_util::{Either, Full};
 use hyper::body::{Bytes, Incoming};
 use hyper::service::service_fn;
 use hyper::{Method, Request, Response};
 use hyper_util::rt::TokioIo;
+
+/// Response body type: either a buffered `Full<Bytes>` or a streaming `Incoming`.
+pub type ProxyBody = Either<Full<Bytes>, Incoming>;
 use tokio_rustls::TlsAcceptor;
 use tokio_util::sync::CancellationToken;
 use tracing::error;
@@ -37,7 +40,7 @@ use crate::config::Config;
 pub async fn handle_request(
     req: Request<Incoming>,
     config: &Config,
-) -> Result<Response<Full<Bytes>>, anyhow::Error> {
+) -> Result<Response<ProxyBody>, anyhow::Error> {
     if !stealth::is_proxy_request(&req) {
         return Ok(stealth::fake_404(&config.stealth.server_name));
     }

src/proxy.rs

@@ -2,17 +2,35 @@
 //!
 //! - [`handle_connect`]: Upgrades the client connection and tunnels bytes
 //!   bidirectionally to the target via [`tokio::io::copy_bidirectional`].
-//! - [`handle_forward`]: Rewrites the absolute URI to path-only form, strips
-//!   proxy headers, and forwards the request via hyper's HTTP/1.1 client.
+//! - [`handle_forward`]: Strips proxy headers and forwards the request via
+//!   a pooled HTTP client (or manual connection with TCP Fast Open).
 
 use anyhow::Context;
-use http_body_util::{BodyExt, Full};
+use http_body_util::{Either, Full};
 use hyper::body::{Bytes, Incoming};
 use hyper::{Request, Response, StatusCode};
-use hyper_util::rt::TokioIo;
+use hyper_util::client::legacy::Client;
+use hyper_util::rt::{TokioExecutor, TokioIo};
 use tracing::{error, info};
 
 use crate::net;
+use crate::ProxyBody;
+
+/// Buffer size per direction for CONNECT tunnels (128 KiB).
+/// 16x the default 8 KiB — matches TLS record size and reduces syscall count.
+const TUNNEL_BUF_SIZE: usize = 128 * 1024;
+
+/// Global pooled HTTP/1.1 client for forward proxying (non-TFO path).
+static POOLED_CLIENT: std::sync::OnceLock<Client<hyper_util::client::legacy::connect::HttpConnector, Incoming>> = std::sync::OnceLock::new();
+
+fn get_pooled_client() -> &'static Client<hyper_util::client::legacy::connect::HttpConnector, Incoming> {
+    POOLED_CLIENT.get_or_init(|| {
+        Client::builder(TokioExecutor::new())
+            .pool_idle_timeout(std::time::Duration::from_secs(90))
+            .pool_max_idle_per_host(32)
+            .build_http()
+    })
+}
 
 /// Handle an HTTP `CONNECT` request by establishing a TCP tunnel.
 ///
@@ -23,7 +41,7 @@ use crate::net;
 pub async fn handle_connect(
     req: Request<Incoming>,
     fast_open: bool,
-) -> anyhow::Result<Response<Full<Bytes>>> {
+) -> anyhow::Result<Response<ProxyBody>> {
     let authority = req
         .uri()
         .authority()
@@ -49,7 +67,7 @@ pub async fn handle_connect(
                 match net::connect(&addr, fast_open).await {
                     Ok(mut target) => {
                         if let Err(e) =
-                            tokio::io::copy_bidirectional(&mut client, &mut target).await
+                            tokio::io::copy_bidirectional_with_sizes(&mut client, &mut target, TUNNEL_BUF_SIZE, TUNNEL_BUF_SIZE).await
                         {
                             error!("tunnel {addr} io error: {e}");
                         }
@@ -68,79 +86,87 @@ pub async fn handle_connect(
     // Return 200 to signal the client that the tunnel is established.
     Ok(Response::builder()
         .status(StatusCode::OK)
-        .body(Full::new(Bytes::new()))
+        .body(Either::Left(Full::new(Bytes::new())))
         .unwrap())
 }
 
 /// Handle a plain HTTP forward proxy request with an absolute URI.
 ///
-/// Rewrites the request URI from absolute form (`http://host/path`) to
-/// path-only (`/path`), removes `Proxy-Authorization` and `Proxy-Connection`
-/// headers, connects to the upstream server, and relays the response back
-/// to the client.
+/// For the common case (`fast_open = false`), uses a pooled HTTP client that
+/// reuses connections across requests. For `fast_open = true`, uses manual
+/// connection setup with TCP Fast Open.
+///
+/// The response body is streamed directly from the upstream server without
+/// buffering the entire body in memory.
 pub async fn handle_forward(
     mut req: Request<Incoming>,
     fast_open: bool,
-) -> anyhow::Result<Response<Full<Bytes>>> {
+) -> anyhow::Result<Response<ProxyBody>> {
     let uri = req.uri().clone();
-    let host = uri
-        .authority()
-        .context("missing authority in forward request")?
-        .to_string();
-
-    let port = uri.port_u16().unwrap_or(match uri.scheme_str() {
-        Some("https") => 443,
-        _ => 80,
-    });
-
-    let addr = if host.contains(':') {
-        host.clone()
-    } else {
-        format!("{host}:{port}")
-    };
-
-    info!("forward {} {} -> {addr}", req.method(), uri);
 
-    // Rewrite the URI to path-only form for the upstream request.
-    let path_and_query = uri
-        .path_and_query()
-        .map(|pq| pq.to_string())
-        .unwrap_or_else(|| "/".to_string());
-    *req.uri_mut() = path_and_query.parse()?;
+    info!("forward {} {}", req.method(), uri);
 
     // Strip hop-by-hop / proxy headers.
     let headers = req.headers_mut();
     headers.remove("proxy-authorization");
     headers.remove("proxy-connection");
 
-    // Connect to the upstream server.
-    let stream = net::connect(&addr, fast_open)
-        .await
-        .with_context(|| format!("connect to {addr}"))?;
-    let io = TokioIo::new(stream);
+    if fast_open {
+        // TFO path: manual connection (pooled client doesn't support custom connectors yet)
+        let host = uri
+            .authority()
+            .context("missing authority in forward request")?
+            .to_string();
 
-    let (mut sender, conn) = hyper::client::conn::http1::handshake(io)
-        .await
-        .context("upstream handshake")?;
-
-    tokio::spawn(async move {
-        if let Err(e) = conn.await {
-            error!("upstream connection error: {e}");
-        }
-    });
+        let port = uri.port_u16().unwrap_or(match uri.scheme_str() {
+            Some("https") => 443,
+            _ => 80,
+        });
 
-    let resp = sender
-        .send_request(req)
-        .await
-        .context("upstream send_request")?;
+        let addr = if host.contains(':') {
+            host.clone()
+        } else {
+            format!("{host}:{port}")
+        };
+
+        // Rewrite the URI to path-only form for the upstream request.
+        let path_and_query = uri
+            .path_and_query()
+            .map(|pq| pq.to_string())
+            .unwrap_or_else(|| "/".to_string());
+        *req.uri_mut() = path_and_query.parse()?;
+
+        let stream = net::connect(&addr, true)
+            .await
+            .with_context(|| format!("connect to {addr}"))?;
+        let io = TokioIo::new(stream);
+
+        let (mut sender, conn) = hyper::client::conn::http1::handshake(io)
+            .await
+            .context("upstream handshake")?;
+
+        tokio::spawn(async move {
+            if let Err(e) = conn.await {
+                error!("upstream connection error: {e}");
+            }
+        });
 
-    // Collect the upstream response body.
-    let (parts, body) = resp.into_parts();
-    let body_bytes = body
-        .collect()
-        .await
-        .context("read upstream body")?
-        .to_bytes();
+        let resp = sender
+            .send_request(req)
+            .await
+            .context("upstream send_request")?;
 
-    Ok(Response::from_parts(parts, Full::new(body_bytes)))
+        let (parts, body) = resp.into_parts();
+        Ok(Response::from_parts(parts, Either::Right(body)))
+    } else {
+        // Pooled client path: connection reuse, automatic URI handling
+        let client = get_pooled_client();
+        let resp = client
+            .request(req)
+            .await
+            .context("pooled client request")?;
+
+        let (parts, body) = resp.into_parts();
+        Ok(Response::from_parts(parts, Either::Right(body)))
+    }
 }

src/stealth.rs

@@ -5,10 +5,12 @@
 //! server. Proxy requests with missing/invalid auth get a `407` so that
 //! real clients (e.g. Chrome) can send credentials.
 
-use http_body_util::Full;
+use http_body_util::{Either, Full};
 use hyper::body::{Bytes, Incoming};
 use hyper::{Method, Request, Response, StatusCode, Version};
 
+use crate::ProxyBody;
+
 /// Returns `true` if the request is a proxy request.
 ///
 /// A request is considered a proxy request if it uses the `CONNECT` method
@@ -35,7 +37,7 @@ pub fn is_proxy_request(req: &Request<Incoming>) -> bool {
 ///
 /// The response includes the configured `Server` header and an HTML body
 /// identical to what nginx produces for a missing page.
-pub fn fake_404(server_name: &str) -> Response<Full<Bytes>> {
+pub fn fake_404(server_name: &str) -> Response<ProxyBody> {
     let body = concat!(
         "<html>\r\n",
         "<head><title>404 Not Found</title></head>\r\n",
@@ -51,7 +53,7 @@ pub fn fake_404(server_name: &str) -> Response<Full<Bytes>> {
         .header("Server", server_name)
         .header("Content-Type", "text/html")
         .header("Content-Length", body.len().to_string())
-        .body(Full::new(Bytes::from(body)))
+        .body(Either::Left(Full::new(Bytes::from(body))))
         .unwrap()
 }
 
@@ -60,12 +62,12 @@ pub fn fake_404(server_name: &str) -> Response<Full<Bytes>> {
 /// Sent when a proxy request (CONNECT or absolute URI) arrives without
 /// valid credentials. The `Proxy-Authenticate` header tells clients like
 /// Chrome to prompt for or resend credentials.
-pub fn proxy_auth_required(server_name: &str) -> Response<Full<Bytes>> {
+pub fn proxy_auth_required(server_name: &str) -> Response<ProxyBody> {
     Response::builder()
         .status(StatusCode::PROXY_AUTHENTICATION_REQUIRED)
         .header("Server", server_name)
         .header("Proxy-Authenticate", "Basic realm=\"Restricted\"")
         .header("Content-Length", "0")
-        .body(Full::new(Bytes::new()))
+        .body(Either::Left(Full::new(Bytes::new())))
         .unwrap()
 }

tests/common/data_server.rs

@@ -0,0 +1,77 @@
+#![allow(dead_code)]
+
+use std::net::SocketAddr;
+
+use http_body_util::Full;
+use hyper::body::{Bytes, Incoming};
+use hyper::server::conn::http1;
+use hyper::service::service_fn;
+use hyper::{Request, Response};
+use hyper_util::rt::TokioIo;
+use tokio::net::TcpListener;
+use tokio_util::sync::CancellationToken;
+
+/// A simple HTTP server that returns a configurable number of zero-filled bytes.
+/// Used for throughput testing.
+pub struct DataServer {
+    pub addr: SocketAddr,
+    shutdown: CancellationToken,
+}
+
+impl DataServer {
+    /// Start a data server that returns `size` bytes of zeros for any request.
+    pub async fn start(size: usize) -> Self {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        let shutdown = CancellationToken::new();
+        let token = shutdown.clone();
+
+        tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                    _ = token.cancelled() => break,
+                    result = listener.accept() => {
+                        let (stream, _) = match result {
+                            Ok(v) => v,
+                            Err(_) => continue,
+                        };
+                        let token = token.clone();
+                        tokio::spawn(async move {
+                            let io = TokioIo::new(stream);
+                            let service = service_fn(move |_req: Request<Incoming>| {
+                                let data = vec![0u8; size];
+                                async move {
+                                    Ok::<_, hyper::Error>(
+                                        Response::builder()
+                                            .status(200)
+                                            .header("Content-Type", "application/octet-stream")
+                                            .header("Content-Length", size.to_string())
+                                            .body(Full::new(Bytes::from(data)))
+                                            .unwrap(),
+                                    )
+                                }
+                            });
+                            let conn = http1::Builder::new().serve_connection(io, service);
+                            tokio::select! {
+                                _ = token.cancelled() => {}
+                                result = conn => {
+                                    if let Err(e) = result {
+                                        eprintln!("data server error: {e}");
+                                    }
+                                }
+                            }
+                        });
+                    }
+                }
+            }
+        });
+
+        DataServer { addr, shutdown }
+    }
+}
+
+impl Drop for DataServer {
+    fn drop(&mut self) {
+        self.shutdown.cancel();
+    }
+}

tests/common/mod.rs

@@ -1,3 +1,4 @@
+pub mod data_server;
 pub mod echo_server;
 pub mod test_server;
 pub mod tls_fixture;