feat(udp,http3): add UDP datagram handler relay support and stream HTTP/3 request bodies to backends

rust/crates/rustproxy-http/src/h3_service.rs

@@ -4,11 +4,14 @@
 //! and forwards them to backends using the same routing and pool infrastructure
 //! as the HTTP/1+2 proxy.
 
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::Duration;
 
 use arc_swap::ArcSwap;
 use bytes::{Buf, Bytes};
+use http_body::Frame;
 use tracing::{debug, warn};
 
 use rustproxy_config::{RouteConfig, TransportProtocol};
@@ -165,15 +168,6 @@ async fn handle_h3_request(
     let backend_port = target.port.resolve(port);
     let backend_addr = format!("{}:{}", backend_host, backend_port);
 
-    // Read request body
-    let mut body_data = Vec::new();
-    while let Some(mut chunk) = stream.recv_data().await
-        .map_err(|e| anyhow::anyhow!("Failed to read H3 request body: {}", e))?
-    {
-        body_data.extend_from_slice(chunk.chunk());
-        chunk.advance(chunk.remaining());
-    }
-
     // Connect to backend via TCP HTTP/1.1 with timeout
     let tcp_stream = tokio::time::timeout(
         connect_timeout,
@@ -194,11 +188,37 @@ async fn handle_h3_request(
         }
     });
 
-    let body = http_body_util::Full::new(Bytes::from(body_data));
+    // Stream request body from H3 client to backend via an mpsc channel.
+    // This avoids buffering the entire request body in memory.
+    let (body_tx, body_rx) = tokio::sync::mpsc::channel::<Bytes>(4);
+    let total_bytes_in = Arc::new(std::sync::atomic::AtomicU64::new(0));
+    let total_bytes_in_writer = Arc::clone(&total_bytes_in);
+
+    // Spawn the H3 body reader task
+    let body_reader = tokio::spawn(async move {
+        while let Ok(Some(mut chunk)) = stream.recv_data().await {
+            let data = Bytes::copy_from_slice(chunk.chunk());
+            total_bytes_in_writer.fetch_add(data.len() as u64, std::sync::atomic::Ordering::Relaxed);
+            chunk.advance(chunk.remaining());
+            if body_tx.send(data).await.is_err() {
+                break;
+            }
+        }
+        stream
+    });
+
+    // Create a body that polls from the mpsc receiver
+    let body = H3RequestBody { receiver: body_rx };
     let backend_req = build_backend_request(&method, &backend_addr, &path, &host, &request, body)?;
+
     let response = sender.send_request(backend_req).await
         .map_err(|e| anyhow::anyhow!("Backend request failed: {}", e))?;
 
+    // Await the body reader to get the stream back
+    let mut stream = body_reader.await
+        .map_err(|e| anyhow::anyhow!("Body reader task failed: {}", e))?;
+    let total_bytes_in = total_bytes_in.load(std::sync::atomic::Ordering::Relaxed);
+
     // Build H3 response
     let status = response.status();
     let mut h3_response = hyper::Response::builder().status(status);
@@ -252,7 +272,7 @@ async fn handle_h3_request(
 
     // Record metrics
     let route_id = route.name.as_deref().or(route.id.as_deref());
-    metrics.record_bytes(0, total_bytes_out, route_id, Some(client_ip));
+    metrics.record_bytes(total_bytes_in, total_bytes_out, route_id, Some(client_ip));
 
     // Finish the stream
     stream.finish().await
@@ -262,14 +282,14 @@ async fn handle_h3_request(
 }
 
 /// Build an HTTP/1.1 backend request from the H3 frontend request.
-fn build_backend_request(
+fn build_backend_request<B>(
     method: &hyper::Method,
     backend_addr: &str,
     path: &str,
     host: &str,
     original_request: &hyper::Request<()>,
-    body: http_body_util::Full<Bytes>,
-) -> anyhow::Result<hyper::Request<http_body_util::Full<Bytes>>> {
+    body: B,
+) -> anyhow::Result<hyper::Request<B>> {
     let mut req = hyper::Request::builder()
         .method(method)
         .uri(format!("http://{}{}", backend_addr, path))
@@ -286,3 +306,27 @@ fn build_backend_request(
     req.body(body)
         .map_err(|e| anyhow::anyhow!("Failed to build backend request: {}", e))
 }
+
+/// A streaming request body backed by an mpsc channel receiver.
+///
+/// Implements `http_body::Body` so hyper can poll chunks as they arrive
+/// from the H3 client, avoiding buffering the entire request body in memory.
+struct H3RequestBody {
+    receiver: tokio::sync::mpsc::Receiver<Bytes>,
+}
+
+impl http_body::Body for H3RequestBody {
+    type Data = Bytes;
+    type Error = hyper::Error;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        match self.receiver.poll_recv(cx) {
+            Poll::Ready(Some(data)) => Poll::Ready(Some(Ok(Frame::data(data)))),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}