From abfcbe2b38636802fa784e72ebf40de15f001b5e Mon Sep 17 00:00:00 2001 From: peg Date: Fri, 11 Sep 2026 14:13:10 +0200 Subject: [PATCH 1/7] Improve concurrency adddressing issue #118 --- README.md | 2 + src/client_request/mod.rs | 235 +++++++++++++++++ src/client_request/tests.rs | 494 +++++++++++++++++++++++++++++++++++ src/client_request/upload.rs | 140 ++++++++++ src/http_version.rs | 25 +- src/lib.rs | 184 ++++++------- src/main.rs | 18 +- 7 files changed, 1002 insertions(+), 96 deletions(-) create mode 100644 src/client_request/mod.rs create mode 100644 src/client_request/tests.rs create mode 100644 src/client_request/upload.rs diff --git a/README.md b/README.md index 7d97492..b187ce6 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,8 @@ These are the attestation type names used in the HTTP headers, and the measureme - `--pccs-url` selects the PCCS used to retrieve collateral when verifying DCAP attestations. It defaults to Intel PCS. - `client`, `get-tls-cert`, and `attested-get` accept `--allow-self-signed` to permit a self-signed remote TLS certificate. - `client` and `server` accept `--listen-addr-healthcheck` to start a separate HTTP health-check listener. +- `client --request-timeout-secs` sets the deadline from receipt of request headers through queueing, upload, and receipt of response headers (default: 60 seconds). Expired requests receive HTTP 504 and are not retried. If a response has already started, an unfinished upload is canceled at the deadline; its response status cannot be changed. Response bodies can continue streaming after that deadline once the upload completes. +- `client --max-in-flight-requests` limits admitted requests, including streaming responses (default: 64). HTTP/2 requests run concurrently; HTTP/1.1 uses one request at a time and reconnects after a timeout or cancellation. Requests waiting for capacity are subject to the same deadline. Library callers can set these limits with `ProxyClient::with_request_options` and `ProxyClientOptions`. - `get-tls-cert --out-measurements ` writes the verified remote measurements as JSON in addition to writing the certificate chain to standard output. - If `server` is started without `--tls-private-key-path` and `--tls-certificate-path`, it generates a self-signed certificate for its listening IP address. diff --git a/src/client_request/mod.rs b/src/client_request/mod.rs new file mode 100644 index 0000000..422f520 --- /dev/null +++ b/src/client_request/mod.rs @@ -0,0 +1,235 @@ +//! Per-request forwarding, deadlines, and response lifetime tracking. +#[cfg(test)] +mod tests; +mod upload; +use std::{ + num::NonZeroUsize, + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Duration, +}; +pub(crate) use upload::RequestBody; + +use http_body_util::BodyExt; +use hyper::{ + Response, + body::{Body, Frame, Incoming, SizeHint}, +}; +use tokio::{ + sync::{OwnedSemaphorePermit, oneshot}, + time::Instant, +}; + +use crate::{ + ATTESTATION_TYPE_HEADER, MEASUREMENT_HEADER, + attestation::{AttestationType, measurements::MultiMeasurements}, + full, + http_version::HttpSender, + update_header, +}; + +/// Limits for requests accepted by a proxy client. +#[derive(Clone, Copy, Debug)] +pub struct ProxyClientOptions { + /// Deadline covering queueing, request upload, and waiting for response headers. + /// Response bodies may continue streaming after this deadline. + pub request_timeout: Duration, + /// Maximum admitted requests, including responses whose bodies are still streaming. + /// HTTP/1.1 forwards one request at a time on its shared connection. + pub max_in_flight_requests: NonZeroUsize, +} + +impl Default for ProxyClientOptions { + fn default() -> Self { + Self { + request_timeout: Duration::from_secs(60), + max_in_flight_requests: NonZeroUsize::new(64).unwrap(), + } + } +} + +pub(crate) type ProxyResponse = + Response>; + +pub(crate) struct PendingRequest { + pub request: http::Request, + pub response_tx: oneshot::Sender, + pub deadline: Instant, + pub permit: OwnedSemaphorePermit, +} + +pub(crate) fn gateway_timeout() -> ProxyResponse { + let mut response = Response::new(full("Request deadline exceeded")); + *response.status_mut() = http::StatusCode::GATEWAY_TIMEOUT; + response +} + +pub(crate) struct ForwardResult { + pub sender: HttpSender, + pub reconnect: bool, +} + +pub(crate) async fn forward( + mut sender: HttpSender, + pending: PendingRequest, + measurements: Option, + attestation_type: AttestationType, +) -> ForwardResult { + let PendingRequest { + request, + mut response_tx, + deadline, + permit, + } = pending; + let http1 = matches!(sender, HttpSender::Http1(_)); + // Expired or canceled queued requests must never be sent to the backend. + if response_tx.is_closed() || Instant::now() >= deadline { + let _ = response_tx.send(gateway_timeout()); + return ForwardResult { + sender, + reconnect: false, + }; + } + + let permit = Arc::new(permit); + let (parts, body) = request.into_parts(); + let (body, upload_guard, mut upload_finished) = RequestBody::new(body, permit.clone()); + let request = http::Request::from_parts(parts, body); + + let response = tokio::select! { + biased; + _ = response_tx.closed() => None, + _ = tokio::time::sleep_until(deadline) => { + let _ = response_tx.send(gateway_timeout()); + return ForwardResult { reconnect: http1 || sender.is_closed(), sender }; + } + result = async { + sender.ready().await?; + sender.send_request(request).await + } => Some(result), + }; + let mut response = match response { + Some(Ok(response)) => response, + failure => { + if let Some(Err(error)) = failure { + tracing::warn!("Failed to send request to proxy-server: {error}"); + let mut response = Response::new(full(format!("Request failed: {error}"))); + *response.status_mut() = http::StatusCode::BAD_GATEWAY; + let _ = response_tx.send(response); + } + // HTTP/2 stream failures/cancellations must not interrupt other streams. + return ForwardResult { + reconnect: http1 || sender.is_closed(), + sender, + }; + } + }; + + // These measurements belong to the connection used for this request. + let headers = response.headers_mut(); + if let Some(measurements) = measurements { + match measurements.to_header_format() { + Ok(value) => { + headers.insert(MEASUREMENT_HEADER, value); + } + Err(error) => tracing::error!("Failed to encode measurement values: {error}"), + } + } + update_header(headers, ATTESTATION_TYPE_HEADER, attestation_type.as_str()); + + let (finished_tx, mut finished_rx) = oneshot::channel(); + let response = response.map(|inner| { + let mut body = TrackedBody { + inner, + permit: Some(permit.clone()), + finished: Some(finished_tx), + }; + if body.inner.is_end_stream() { + body.finish(); + } + body.boxed() + }); + let _ = response_tx.send(response); + + // Early response headers do not mean that the upload has finished. Keep its + // deadline and the shared permit alive until both directions have completed. + let mut uploaded = false; + let mut responded = false; + let mut upload_ok = true; + while !uploaded || !responded { + tokio::select! { + biased; + result = &mut upload_finished, if !uploaded => { + uploaded = true; + upload_ok = result.is_ok(); + } + result = &mut finished_rx, if !responded => { + responded = true; + if result.is_err() { + // Dropping this guard stops an upload still owned by Hyper. + drop(upload_guard); + return ForwardResult { reconnect: http1 || sender.is_closed(), sender }; + } + } + _ = tokio::time::sleep_until(deadline), if !uploaded => { + // Headers may already have reached the caller, so a 504 can no + // longer replace them. Cancel the upload instead. + drop(upload_guard); + return ForwardResult { reconnect: http1 || sender.is_closed(), sender }; + } + } + } + ForwardResult { + sender, + reconnect: http1 && !upload_ok, + } +} + +pin_project_lite::pin_project! { + struct TrackedBody { + #[pin] + inner: Incoming, + permit: Option>, + finished: Option>, + } +} + +impl TrackedBody { + fn finish(&mut self) { + self.permit.take(); + if let Some(finished) = self.finished.take() { + let _ = finished.send(()); + } + } +} + +impl Body for TrackedBody { + type Data = bytes::Bytes; + type Error = hyper::Error; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let mut this = self.project(); + let frame = this.inner.as_mut().poll_frame(cx); + if matches!(frame, Poll::Ready(Some(Err(_)))) { + this.permit.take(); + this.finished.take(); + } else if matches!(frame, Poll::Ready(None)) || this.inner.is_end_stream() { + this.permit.take(); + if let Some(finished) = this.finished.take() { + let _ = finished.send(()); + } + } + frame + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + fn size_hint(&self) -> SizeHint { + self.inner.size_hint() + } +} diff --git a/src/client_request/tests.rs b/src/client_request/tests.rs new file mode 100644 index 0000000..c29fea5 --- /dev/null +++ b/src/client_request/tests.rs @@ -0,0 +1,494 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use axum::{ + Router, + routing::{get, post}, +}; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::Notify, + task::JoinSet, + time::timeout, +}; + +use super::ProxyClientOptions; +use crate::{ + AttestationGenerator, AttestationVerifier, ProxyClient, ProxyServer, + http_version::{ALPN_H2, ALPN_HTTP11}, + test_helpers::{generate_certificate_chain, generate_tls_config}, +}; + +struct Fixture { + url: String, + addr: std::net::SocketAddr, + connections: Arc, + _tasks: JoinSet<()>, +} + +async fn proxy(app: Router, protocol: &[u8], slots: usize, request_timeout: Duration) -> Fixture { + let mut tasks = JoinSet::new(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let target = listener.local_addr().unwrap(); + tasks.spawn(async move { axum::serve(listener, app).await.unwrap() }); + + let (certs, key) = generate_certificate_chain("127.0.0.1".parse().unwrap()); + let (mut server_config, mut client_config) = generate_tls_config(certs.clone(), key); + server_config.alpn_protocols = vec![protocol.to_vec()]; + client_config.alpn_protocols = vec![protocol.to_vec()]; + let server = ProxyServer::new_with_tls_config( + certs, + server_config, + "127.0.0.1:0", + target.to_string(), + AttestationGenerator::with_no_attestation(), + AttestationVerifier::expect_none(), + ) + .await + .unwrap(); + let target = server.local_addr().unwrap(); + let connections = Arc::new(AtomicUsize::new(0)); + let counter = connections.clone(); + tasks.spawn(async move { + loop { + server.accept().await.unwrap(); + counter.fetch_add(1, Ordering::SeqCst); + } + }); + let client = ProxyClient::new_with_tls_config( + client_config, + "127.0.0.1:0", + target.to_string(), + AttestationGenerator::with_no_attestation(), + AttestationVerifier::expect_none(), + None, + ) + .await + .unwrap() + .with_request_options(ProxyClientOptions { + request_timeout, + max_in_flight_requests: slots.try_into().unwrap(), + }); + let addr = client.local_addr().unwrap(); + tasks.spawn(async move { + loop { + client.accept().await.unwrap(); + } + }); + Fixture { + url: format!("http://{addr}"), + addr, + connections, + _tasks: tasks, + } +} + +fn http_client() -> reqwest::Client { + reqwest::Client::builder() + .no_proxy() + .timeout(Duration::from_secs(5)) + .build() + .unwrap() +} + +#[tokio::test] +async fn http2_stalled_request_does_not_block_fast_request() { + let entered = Arc::new(Notify::new()); + let signal = entered.clone(); + let app = Router::new() + .route( + "/slow", + get(move || { + let signal = signal.clone(); + async move { + signal.notify_one(); + std::future::pending::<&'static str>().await + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_H2, 2, Duration::from_millis(800)).await; + let client = http_client(); + let slow = tokio::spawn(client.get(format!("{}/slow", fixture.url)).send()); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + let fast = timeout( + Duration::from_millis(400), + client.get(format!("{}/fast", fixture.url)).send(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(fast.text().await.unwrap(), "fast"); + assert_eq!( + slow.await.unwrap().unwrap().status(), + http::StatusCode::GATEWAY_TIMEOUT + ); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn http1_timeout_reconnects_without_replaying_post() { + let calls = Arc::new(AtomicUsize::new(0)); + let count = calls.clone(); + let app = Router::new() + .route( + "/slow", + post(move || { + let count = count.clone(); + async move { + count.fetch_add(1, Ordering::SeqCst); + std::future::pending::<&'static str>().await + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_HTTP11, 2, Duration::from_millis(300)).await; + let client = http_client(); + let slow = client + .post(format!("{}/slow", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(slow.status(), http::StatusCode::GATEWAY_TIMEOUT); + let fast = client + .get(format!("{}/fast", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(fast.text().await.unwrap(), "fast"); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 2); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn streaming_bodies_hold_capacity_and_expired_requests_are_not_forwarded() { + for (protocol, slots) in [(ALPN_H2, 1), (ALPN_HTTP11, 2)] { + let calls = Arc::new(AtomicUsize::new(0)); + let count = calls.clone(); + let (body_tx, body_rx) = tokio::sync::mpsc::channel::< + Result, std::convert::Infallible>, + >(1); + // A stream which remains open until the test drops body_tx. + let body_rx = Arc::new(tokio::sync::Mutex::new(Some(body_rx))); + let app = Router::new() + .route( + "/stream", + get(move || { + let body_rx = body_rx.clone(); + async move { + let rx = body_rx.lock().await.take().unwrap(); + axum::body::Body::new(TestBody(rx)) + } + }), + ) + .route( + "/fast", + get(move || { + let count = count.clone(); + async move { + count.fetch_add(1, Ordering::SeqCst); + "fast" + } + }), + ); + let fixture = proxy(app, protocol, slots, Duration::from_millis(300)).await; + let client = http_client(); + let stream = client + .get(format!("{}/stream", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(stream.status(), http::StatusCode::OK); + let blocked = client + .get(format!("{}/fast", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(blocked.status(), http::StatusCode::GATEWAY_TIMEOUT); + assert_eq!(calls.load(Ordering::SeqCst), 0); + drop(body_tx); + assert!(stream.bytes().await.unwrap().is_empty()); + let fast = client + .get(format!("{}/fast", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(fast.text().await.unwrap(), "fast"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } +} + +// Avoid a new stream adapter dependency for a controllable streaming response. +struct TestBody( + tokio::sync::mpsc::Receiver, std::convert::Infallible>>, +); + +impl hyper::body::Body for TestBody { + type Data = bytes::Bytes; + type Error = std::convert::Infallible; + fn poll_frame( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll, Self::Error>>> { + self.0.poll_recv(cx) + } +} + +#[tokio::test] +async fn http2_disconnected_caller_releases_slot_without_reconnect() { + let entered = Arc::new(Notify::new()); + let signal = entered.clone(); + let app = Router::new() + .route( + "/slow", + get(move || { + let signal = signal.clone(); + async move { + signal.notify_one(); + std::future::pending::<&'static str>().await + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_H2, 1, Duration::from_secs(5)).await; + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"GET /slow HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + drop(source); + let response = timeout( + Duration::from_secs(1), + http_client().get(format!("{}/fast", fixture.url)).send(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(response.text().await.unwrap(), "fast"); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn stalled_upload_times_out_without_blocking_http2() { + let entered = Arc::new(Notify::new()); + let signal = entered.clone(); + let app = Router::new() + .route( + "/upload", + post(move |request: axum::extract::Request| { + let signal = signal.clone(); + async move { + signal.notify_one(); + let _ = axum::body::to_bytes(request.into_body(), 1024).await; + "upload finished" + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_H2, 2, Duration::from_millis(800)).await; + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"POST /upload HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\n\r\nx") + .await + .unwrap(); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + let fast = timeout( + Duration::from_millis(400), + http_client().get(format!("{}/fast", fixture.url)).send(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(fast.text().await.unwrap(), "fast"); + let mut response = [0; 1024]; + let len = timeout(Duration::from_secs(2), source.read(&mut response)) + .await + .unwrap() + .unwrap(); + assert!(String::from_utf8_lossy(&response[..len]).starts_with("HTTP/1.1 504")); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn dropping_streaming_response_releases_slot() { + for protocol in [ALPN_HTTP11, ALPN_H2] { + let (_body_tx, body_rx) = tokio::sync::mpsc::channel(1); + let body_rx = Arc::new(tokio::sync::Mutex::new(Some(body_rx))); + let app = + Router::new() + .route( + "/stream", + get(move || { + let body_rx = body_rx.clone(); + async move { + axum::body::Body::new(TestBody(body_rx.lock().await.take().unwrap())) + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, protocol, 1, Duration::from_secs(5)).await; + let client = http_client(); + let response = client + .get(format!("{}/stream", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + drop(response); + let fast = timeout( + Duration::from_secs(1), + client.get(format!("{}/fast", fixture.url)).send(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(fast.text().await.unwrap(), "fast"); + let expected_connections = if protocol == ALPN_HTTP11 { 2 } else { 1 }; + assert_eq!( + fixture.connections.load(Ordering::SeqCst), + expected_connections + ); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn http1_clean_close_preserves_response() { + let app = Router::new().route( + "/", + get(|| async { ([(http::header::CONNECTION, "close")], "ok") }), + ); + let fixture = proxy(app, ALPN_HTTP11, 2, Duration::from_secs(3)).await; + let client = http_client(); + for i in 0..30 { + let response = client + .get(format!("{}/", fixture.url)) + .send() + .await + .unwrap(); + let status = response.status(); + let body = response.text().await.unwrap(); + assert_eq!(status, http::StatusCode::OK, "request {i}: {body}"); + assert_eq!(body, "ok"); + } +} + +#[tokio::test] +async fn early_response_keeps_upload_bounded() { + let active = Arc::new(AtomicUsize::new(0)); + let counter = active.clone(); + let app = Router::new().route( + "/", + post(move |request: axum::extract::Request| { + let counter = counter.clone(); + async move { + counter.fetch_add(1, Ordering::SeqCst); + tokio::spawn(async move { + let _ = axum::body::to_bytes(request.into_body(), 1024).await; + counter.fetch_sub(1, Ordering::SeqCst); + }); + http::StatusCode::OK + } + }), + ); + let fixture = proxy(app, ALPN_H2, 1, Duration::from_millis(300)).await; + let mut sources = Vec::new(); + for i in 0..3 { + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"POST / HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100\r\n\r\nx") + .await + .unwrap(); + let mut response = [0; 1024]; + if i > 0 { + // The previous response was empty, but its unfinished upload must + // still occupy the sole slot until its deadline cancels it. + assert!( + timeout(Duration::from_millis(50), source.read(&mut response)) + .await + .is_err() + ); + } + let n = timeout(Duration::from_secs(2), source.read(&mut response)) + .await + .unwrap() + .unwrap(); + assert!(String::from_utf8_lossy(&response[..n]).starts_with("HTTP/1.1 ")); + sources.push(source); + } + tokio::time::sleep(Duration::from_millis(500)).await; + assert_eq!( + active.load(Ordering::SeqCst), + 0, + "uploads survived their request deadline" + ); +} + +#[tokio::test] +async fn early_response_allows_upload_to_finish_before_releasing_slot() { + let (uploaded_tx, uploaded_rx) = tokio::sync::oneshot::channel(); + let uploaded_tx = Arc::new(tokio::sync::Mutex::new(Some(uploaded_tx))); + let app = Router::new() + .route( + "/upload", + post(move |request: axum::extract::Request| { + let uploaded_tx = uploaded_tx.clone(); + async move { + tokio::spawn(async move { + let body = axum::body::to_bytes(request.into_body(), 1024) + .await + .unwrap(); + uploaded_tx.lock().await.take().unwrap().send(body).unwrap(); + }); + http::StatusCode::OK + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_H2, 1, Duration::from_secs(5)).await; + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"POST /upload HTTP/1.1\r\nHost: localhost\r\nContent-Length: 3\r\n\r\na") + .await + .unwrap(); + let mut response = [0; 1024]; + let len = timeout(Duration::from_secs(2), source.read(&mut response)) + .await + .unwrap() + .unwrap(); + assert!(String::from_utf8_lossy(&response[..len]).starts_with("HTTP/1.1 200")); + + let mut fast = tokio::spawn(http_client().get(format!("{}/fast", fixture.url)).send()); + assert!( + timeout(Duration::from_millis(100), &mut fast) + .await + .is_err() + ); + source.write_all(b"bc").await.unwrap(); + assert_eq!( + timeout(Duration::from_secs(2), uploaded_rx) + .await + .unwrap() + .unwrap(), + "abc" + ); + let response = timeout(Duration::from_secs(2), fast) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(response.text().await.unwrap(), "fast"); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); +} diff --git a/src/client_request/upload.rs b/src/client_request/upload.rs new file mode 100644 index 0000000..c32961e --- /dev/null +++ b/src/client_request/upload.rs @@ -0,0 +1,140 @@ +//! Upload cancellation must work even when Hyper is waiting for HTTP/2 capacity +//! and is not polling the body. Shared state lets the forwarding task drop the +//! source body independently. +use hyper::body::{Body, Frame, Incoming, SizeHint}; +use std::{ + io, + pin::Pin, + sync::{Arc, Mutex}, + task::{Context, Poll, Waker}, +}; +use tokio::sync::{OwnedSemaphorePermit, oneshot}; + +struct UploadState { + inner: Option, + finished: Option>, + complete: bool, + waker: Option, +} + +impl UploadState { + fn finish(&mut self) { + self.inner.take(); + self.complete = true; + self.waker.take(); + if let Some(finished) = self.finished.take() { + let _ = finished.send(()); + } + } +} + +fn cancel(state: &Mutex) { + let waker = { + let mut state = state.lock().unwrap(); + state.inner.take(); + state.finished.take(); + state.waker.take() + }; + if let Some(waker) = waker { + waker.wake(); + } +} + +pub(super) struct UploadGuard(Arc>); + +impl Drop for UploadGuard { + fn drop(&mut self) { + cancel(&self.0); + } +} + +/// Hyper sends uploads independently of response futures. Keep the permit until +/// the upload finishes or Hyper drops the canceled stream. +pub(crate) struct RequestBody { + state: Arc>, + permit: Option>, +} + +impl RequestBody { + pub(super) fn new( + inner: Incoming, + permit: Arc, + ) -> (Self, UploadGuard, oneshot::Receiver<()>) { + let (finished, receiver) = oneshot::channel(); + let mut state = UploadState { + inner: Some(inner), + finished: Some(finished), + complete: false, + waker: None, + }; + if state.inner.as_ref().unwrap().is_end_stream() { + state.finish(); + } + let state = Arc::new(Mutex::new(state)); + ( + Self { + state: state.clone(), + permit: Some(permit), + }, + UploadGuard(state), + receiver, + ) + } +} + +impl Drop for RequestBody { + fn drop(&mut self) { + cancel(&self.state); + } +} + +impl Body for RequestBody { + type Data = bytes::Bytes; + type Error = io::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let frame = { + let mut state = self.state.lock().unwrap(); + if state.complete { + Poll::Ready(None) + } else if let Some(inner) = state.inner.as_mut() { + let frame = Pin::new(&mut *inner).poll_frame(cx); + if matches!(frame, Poll::Ready(Some(Err(_)))) { + state.inner.take(); + state.finished.take(); + } else if matches!(frame, Poll::Ready(None)) || inner.is_end_stream() { + state.finish(); + } else { + state.waker = Some(cx.waker().clone()); + } + frame.map(|frame| frame.map(|frame| frame.map_err(io::Error::other))) + } else { + Poll::Ready(Some(Err(io::Error::new( + io::ErrorKind::Interrupted, + "request upload canceled", + )))) + } + }; + if matches!(frame, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) || self.is_end_stream() { + self.permit.take(); + } + frame + } + + fn is_end_stream(&self) -> bool { + self.state.lock().unwrap().complete + } + + fn size_hint(&self) -> SizeHint { + self.state + .lock() + .unwrap() + .inner + .as_ref() + .map(Body::size_hint) + .unwrap_or_default() + } +} diff --git a/src/http_version.rs b/src/http_version.rs index bef817c..90c7421 100644 --- a/src/http_version.rs +++ b/src/http_version.rs @@ -1,4 +1,5 @@ //! HTTP Version support and negotiation +use crate::client_request::RequestBody; use hyper::Response; use hyper_util::rt::TokioIo; use std::pin::Pin; @@ -52,17 +53,17 @@ impl HttpVersion { } } -type Http1Sender = hyper::client::conn::http1::SendRequest; -type Http2Sender = hyper::client::conn::http2::SendRequest; +type Http1Sender = hyper::client::conn::http1::SendRequest; +type Http2Sender = hyper::client::conn::http2::SendRequest; type Http1Connection = hyper::client::conn::http1::Connection< TokioIo>, - hyper::body::Incoming, + RequestBody, >; type Http2Connection = hyper::client::conn::http2::Connection< TokioIo>, - hyper::body::Incoming, + RequestBody, crate::TokioExecutor, >; @@ -85,9 +86,23 @@ impl From for HttpSender { } impl HttpSender { + pub async fn ready(&mut self) -> Result<(), hyper::Error> { + match self { + Self::Http1(sender) => sender.ready().await, + Self::Http2(sender) => sender.ready().await, + } + } + + pub fn is_closed(&self) -> bool { + match self { + Self::Http1(sender) => sender.is_closed(), + Self::Http2(sender) => sender.is_closed(), + } + } + pub async fn send_request( &mut self, - request: http::Request, + request: http::Request, ) -> Result, hyper::Error> { match self { Self::Http1(sender) => sender.send_request(request).await, diff --git a/src/lib.rs b/src/lib.rs index 8b4a833..da54922 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,7 +9,10 @@ pub use attested_tls; pub use attested_tls::attestation; pub use attested_tls::attestation::AttestationGenerator; +mod client_request; mod http_version; +pub use client_request::ProxyClientOptions; +use client_request::{PendingRequest, forward, gateway_timeout}; #[cfg(test)] mod test_helpers; @@ -23,7 +26,7 @@ use std::{net::SocketAddr, num::TryFromIntError, sync::Arc, time::Duration}; use thiserror::Error; use tokio::io; use tokio::net::{TcpListener, TcpStream, ToSocketAddrs}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Semaphore, mpsc, oneshot}; use tokio_rustls::rustls::server::{VerifierBuilderError, WebPkiClientVerifier}; use tokio_rustls::rustls::{ self, ClientConfig, RootCertStore, ServerConfig, pki_types::CertificateDer, @@ -56,11 +59,6 @@ const SERVER_RECONNECT_MAX_BACKOFF_SECS: u64 = 120; const KEEP_ALIVE_INTERVAL: u64 = 30; const KEEP_ALIVE_TIMEOUT: u64 = 10; -type RequestWithResponseSender = ( - http::Request, - oneshot::Sender>, hyper::Error>>, -); - /// Adds HTTP 1 and 2 to the list of allowed protocols fn ensure_proxy_alpn_protocols(alpn_protocols: &mut Vec>) { for protocol in [ALPN_H2, ALPN_HTTP11] { @@ -359,10 +357,19 @@ pub struct ProxyClient { /// The underlying TCP listener listener: TcpListener, /// A channel for sending requests to the connection to the proxy-server - requests_tx: mpsc::Sender, + requests_tx: mpsc::Sender, + options: ProxyClientOptions, + request_slots: Arc, } impl ProxyClient { + /// Configure request limits before accepting source connections. + pub fn with_request_options(mut self, options: ProxyClientOptions) -> Self { + self.request_slots = Arc::new(Semaphore::new(options.max_in_flight_requests.get())); + self.options = options; + self + } + /// Start with optional TLS client auth pub async fn new( cert_and_key: Option, @@ -438,12 +445,7 @@ impl ProxyClient { let target = host_to_host_with_port(target_name); // Channel for getting incoming requests from the source client - let (requests_tx, mut requests_rx) = mpsc::channel::<( - http::Request, - oneshot::Sender< - Result>, hyper::Error>, - >, - )>(1024); + let (requests_tx, mut requests_rx) = mpsc::channel::(1024); // used only to signal "initial connect succeeded" or "failed with error" let (ready_tx, ready_rx) = oneshot::channel::>(); @@ -451,8 +453,11 @@ impl ProxyClient { tokio::spawn(async move { let mut first = true; let mut ready_tx = Some(ready_tx); + // Retired connections may still have complete responses waiting to be + // delivered. Drain their workers without blocking a fresh connection. + let mut draining = tokio::task::JoinSet::new(); 'reconnect: loop { - let (mut sender, conn, measurements, remote_attestation_type) = + let (sender, conn, measurements, remote_attestation_type) = // Connect to the proxy server and provide / verify attestation match Self::setup_connection_with_backoff(&target, &attested_tls_client, first) .await @@ -479,79 +484,48 @@ impl ProxyClient { } }; - let (conn_done_tx, mut conn_done_rx) = - tokio::sync::watch::channel::>(None); - - tokio::spawn(async move { - let res = conn.await; - let _ = conn_done_tx.send(res.err()); - }); + // The connection driver is stopped on reconnect. Request workers + // retain their own deadlines and connection-specific measurements. + let mut connection = tokio::task::JoinSet::new(); + connection.spawn(conn); + let mut in_flight = tokio::task::JoinSet::new(); + let mut sender = Some(sender); loop { tokio::select! { - // Read an incoming request from the channel (from the source client) - incoming_req_option = requests_rx.recv() => { - if let Some((req, response_tx)) = incoming_req_option { - debug!("[proxy-client] Read incoming request from source client: {req:?}"); - // Attempt to forward it to the proxy server - let (response, should_reconnect) = match sender.send_request(req).await { - Ok(mut resp) => { - debug!("[proxy-client] Read response from proxy-server: {resp:?}"); - // If we have measurements from the proxy-server, inject them into the - // response header - let headers = resp.headers_mut(); - if let Some(measurements) = measurements.clone() { - match measurements.to_header_format() { - Ok(header_value) => { - headers.insert(MEASUREMENT_HEADER, header_value); - } - Err(e) => { - // This error is highly unlikely - that the measurement values fail to - // encode to JSON or fit in an HTTP header - error!("Failed to encode measurement values: {e}"); - } - } - } - - update_header( - headers, - ATTESTATION_TYPE_HEADER, - remote_attestation_type.as_str(), - ); - (Ok(resp.map(|b| b.boxed())), false) - } - Err(e) => { - warn!("Failed to send request to proxy-server: {e}"); - let mut resp = Response::new(full(format!("Request failed: {e}"))); - *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY; - - (Ok(resp), true) + incoming = requests_rx.recv(), if sender.is_some() => { + let Some(pending) = incoming else { + break 'reconnect; + }; + let request_sender = match sender.as_ref().unwrap() { + HttpSender::Http2(inner) => HttpSender::Http2(inner.clone()), + HttpSender::Http1(_) => sender.take().unwrap(), + }; + in_flight.spawn(forward( + request_sender, + pending, + measurements.clone(), + remote_attestation_type, + )); + } + result = in_flight.join_next(), if !in_flight.is_empty() => { + match result { + Some(Ok(result)) if !result.reconnect => { + if matches!(result.sender, HttpSender::Http1(_)) { + sender = Some(result.sender); } - }; - - // Send the response back to the source client - if response_tx.send(response).is_err() { - warn!("Failed to forward response to source client, probably they dropped the connection"); - } - - if should_reconnect { - // Leave the inner loop and continue on the reconnect loop - warn!("Reconnecting to proxy-server due to failed request"); - break; } - } else { - // The request sender was dropped - so no more incoming requests - debug!("Request sender dropped - leaving connection handler loop"); - break 'reconnect; + _ => break, } } - - // Connection closed - _ = conn_done_rx.changed() => { - // Leave the inner loop and continue on the reconnect loop + _ = connection.join_next() => { warn!("Connection dropped - reconnecting..."); break; } - }; + _ = draining.join_next(), if !draining.is_empty() => {} + } + } + if !in_flight.is_empty() { + draining.spawn(async move { while in_flight.join_next().await.is_some() {} }); } } }); @@ -560,6 +534,10 @@ impl ProxyClient { Ok(Ok(())) => Ok(Self { listener, requests_tx, + options: ProxyClientOptions::default(), + request_slots: Arc::new(Semaphore::new( + ProxyClientOptions::default().max_in_flight_requests.get(), + )), }), Ok(Err(e)) => Err(e), Err(e) => Err(e.into()), @@ -576,9 +554,13 @@ impl ProxyClient { let (inbound, _client_addr) = self.listener.accept().await?; let requests_tx = self.requests_tx.clone(); + let options = self.options; + let request_slots = self.request_slots.clone(); let handle = tokio::spawn(async move { - if let Err(err) = Self::handle_connection(inbound, requests_tx).await { + if let Err(err) = + Self::handle_connection(inbound, requests_tx, options, request_slots).await + { warn!("Failed to handle connection from source client: {err}"); } }); @@ -589,7 +571,9 @@ impl ProxyClient { /// Handle an incoming connection from the source client async fn handle_connection( inbound: TcpStream, - requests_tx: mpsc::Sender, + requests_tx: mpsc::Sender, + options: ProxyClientOptions, + request_slots: Arc, ) -> Result<(), ProxyError> { tracing::debug!("proxy-client accepted connection"); @@ -597,8 +581,9 @@ impl ProxyClient { let http = hyper::server::conn::http1::Builder::new(); let service = service_fn(move |req| { let requests_tx = requests_tx.clone(); + let request_slots = request_slots.clone(); async move { - match Self::handle_http_request(req, requests_tx).await { + match Self::handle_http_request(req, requests_tx, options, request_slots).await { Ok(res) => { Ok::>, hyper::Error>(res) } @@ -684,13 +669,13 @@ impl ProxyClient { .keep_alive_interval(Some(Duration::from_secs(KEEP_ALIVE_INTERVAL))) .keep_alive_timeout(Duration::from_secs(KEEP_ALIVE_TIMEOUT)) .keep_alive_while_idle(true) - .handshake::<_, hyper::body::Incoming>(outbound_io) + .handshake::<_, client_request::RequestBody>(outbound_io) .await?; (sender.into(), conn.into()) } HttpVersion::Http1 => { let (sender, conn) = hyper::client::conn::http1::Builder::new() - .handshake::<_, hyper::body::Incoming>(outbound_io) + .handshake::<_, client_request::RequestBody>(outbound_io) .await?; (sender.into(), conn.into()) } @@ -703,11 +688,32 @@ impl ProxyClient { // Handle a request from the source client to the proxy server async fn handle_http_request( req: hyper::Request, - requests_tx: mpsc::Sender, + requests_tx: mpsc::Sender, + options: ProxyClientOptions, + request_slots: Arc, ) -> Result>, ProxyError> { - let (response_tx, response_rx) = oneshot::channel(); - requests_tx.send((req, response_tx)).await?; - Ok(response_rx.await??) + let deadline = tokio::time::Instant::now() + options.request_timeout; + let result = tokio::time::timeout_at(deadline, async { + let permit = request_slots + .acquire_owned() + .await + .expect("request semaphore is never closed"); + let (response_tx, response_rx) = oneshot::channel(); + requests_tx + .send(PendingRequest { + request: req, + response_tx, + deadline, + permit, + }) + .await?; + Ok::<_, ProxyError>(response_rx.await?) + }) + .await; + match result { + Ok(result) => result, + Err(_) => Ok(gateway_timeout()), + } } } @@ -759,8 +765,8 @@ pub enum ProxyError { AttestedTls(#[from] AttestedTlsError), } -impl From> for ProxyError { - fn from(_err: mpsc::error::SendError) -> Self { +impl From> for ProxyError { + fn from(_err: mpsc::error::SendError) -> Self { Self::MpscSend } } diff --git a/src/main.rs b/src/main.rs index 7cef2f0..dbdac8d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,14 +4,16 @@ use clap::{Parser, Subcommand}; use std::{ fs::File, net::{IpAddr, SocketAddr}, + num::{NonZeroU64, NonZeroUsize}, path::PathBuf, + time::Duration, }; use tokio::io::AsyncWriteExt; use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer}; use tracing::level_filters::LevelFilter; use attested_tls_proxy::{ - AttestationGenerator, ProxyClient, ProxyServer, + AttestationGenerator, ProxyClient, ProxyClientOptions, ProxyServer, attested_get::{attested_get, split_target_and_path}, attested_tls::{ TlsCertAndKey, @@ -66,6 +68,12 @@ enum CliCommand { listen_addr: SocketAddr, /// The hostname:port or ip:port of the proxy server (port defaults to 443) target_addr: String, + /// Request deadline in seconds, including queueing and waiting for response headers + #[arg(long, default_value = "60")] + request_timeout_secs: NonZeroU64, + /// Maximum in-flight requests, including streaming responses + #[arg(long, default_value = "64")] + max_in_flight_requests: NonZeroUsize, /// Type of attestation to present (dafaults to 'auto' for automatic detection) /// If other than None, a TLS key and certicate must also be given #[arg(long, env = "CLIENT_ATTESTATION_TYPE")] @@ -245,6 +253,8 @@ async fn main() -> anyhow::Result<()> { CliCommand::Client { listen_addr, target_addr, + request_timeout_secs, + max_in_flight_requests, client_attestation_type, tls_private_key_path, tls_certificate_path, @@ -311,7 +321,11 @@ async fn main() -> anyhow::Result<()> { remote_tls_cert, ) .await? - }; + } + .with_request_options(ProxyClientOptions { + request_timeout: Duration::from_secs(request_timeout_secs.get()), + max_in_flight_requests, + }); loop { if let Err(err) = client.accept().await { From 13ab49c0b7d47c30f4c77c9a486495c8bb740583 Mon Sep 17 00:00:00 2001 From: peg Date: Fri, 11 Sep 2026 14:52:27 +0200 Subject: [PATCH 2/7] Cover edge cases --- README.md | 3 + src/client_request/mod.rs | 39 +++++ src/client_request/response_idle.rs | 171 +++++++++++++++++++ src/client_request/tests.rs | 248 ++++++++++++++++++++++++++++ src/lib.rs | 63 ++++--- src/main.rs | 7 + 6 files changed, 507 insertions(+), 24 deletions(-) create mode 100644 src/client_request/response_idle.rs diff --git a/README.md b/README.md index b187ce6..9d32e34 100644 --- a/README.md +++ b/README.md @@ -78,10 +78,13 @@ These are the attestation type names used in the HTTP headers, and the measureme - `client`, `get-tls-cert`, and `attested-get` accept `--allow-self-signed` to permit a self-signed remote TLS certificate. - `client` and `server` accept `--listen-addr-healthcheck` to start a separate HTTP health-check listener. - `client --request-timeout-secs` sets the deadline from receipt of request headers through queueing, upload, and receipt of response headers (default: 60 seconds). Expired requests receive HTTP 504 and are not retried. If a response has already started, an unfinished upload is canceled at the deadline; its response status cannot be changed. Response bodies can continue streaming after that deadline once the upload completes. +- `client --response-body-idle-timeout-secs` closes a source connection if its active response body makes no socket write progress for this interval (default: 60 seconds). This covers silent backends and clients that stop reading, even when body polling is blocked. The affected request releases its capacity; HTTP/1.1 reconnects upstream, while other HTTP/2 streams remain usable. Responses already started are truncated rather than replaced with a 504. Streams that keep making progress may continue indefinitely. - `client --max-in-flight-requests` limits admitted requests, including streaming responses (default: 64). HTTP/2 requests run concurrently; HTTP/1.1 uses one request at a time and reconnects after a timeout or cancellation. Requests waiting for capacity are subject to the same deadline. Library callers can set these limits with `ProxyClient::with_request_options` and `ProxyClientOptions`. - `get-tls-cert --out-measurements ` writes the verified remote measurements as JSON in addition to writing the certificate chain to standard output. - If `server` is started without `--tls-private-key-path` and `--tls-certificate-path`, it generates a self-signed certificate for its listening IP address. +These limits also apply to library callers: `ProxyClient::new*` defaults to a 60-second request deadline, a 60-second response-body idle timeout, and 64 in-flight requests. Set `ProxyClientOptions` with `with_request_options` before accepting connections to adjust them, including for long-polling services or streams with long gaps between messages. The idle timeout is inactive while waiting for response headers or between requests on a keep-alive connection. + ## Protocol Specification A proxy-client will immediately attempt to connect to the given proxy-server. diff --git a/src/client_request/mod.rs b/src/client_request/mod.rs index 422f520..0cc3b67 100644 --- a/src/client_request/mod.rs +++ b/src/client_request/mod.rs @@ -1,4 +1,5 @@ //! Per-request forwarding, deadlines, and response lifetime tracking. +pub(crate) mod response_idle; #[cfg(test)] mod tests; mod upload; @@ -35,6 +36,9 @@ pub struct ProxyClientOptions { /// Deadline covering queueing, request upload, and waiting for response headers. /// Response bodies may continue streaming after this deadline. pub request_timeout: Duration, + /// Maximum time without response bytes being written to the source while a + /// response body is active. Expiry closes the source connection. + pub response_body_idle_timeout: Duration, /// Maximum admitted requests, including responses whose bodies are still streaming. /// HTTP/1.1 forwards one request at a time on its shared connection. pub max_in_flight_requests: NonZeroUsize, @@ -44,6 +48,7 @@ impl Default for ProxyClientOptions { fn default() -> Self { Self { request_timeout: Duration::from_secs(60), + response_body_idle_timeout: Duration::from_secs(60), max_in_flight_requests: NonZeroUsize::new(64).unwrap(), } } @@ -70,6 +75,40 @@ pub(crate) struct ForwardResult { pub reconnect: bool, } +/// Borrow the shared HTTP/2 sender or take the exclusive HTTP/1 sender. +/// A closed connection must be replaced before dispatching the queued request. +pub(crate) fn take_sender(sender: &mut Option) -> Option { + match sender.as_ref()? { + inner if inner.is_closed() => None, + HttpSender::Http2(inner) => Some(HttpSender::Http2(inner.clone())), + HttpSender::Http1(_) => sender.take(), + } +} + +/// Restore an exclusive sender, returning whether the connection must be replaced. +pub(crate) fn worker_finished( + sender: &mut Option, + result: Result, +) -> bool { + match result { + Ok(result) => { + if result.reconnect { + return true; + } + if matches!(result.sender, HttpSender::Http1(_)) { + *sender = Some(result.sender); + } + false + } + Err(error) => { + tracing::error!(%error, "Request worker failed"); + // HTTP/1 lost its exclusive sender. HTTP/2 retains a shared sender + // and other streams can continue after this worker unwinds. + sender.as_ref().is_none_or(HttpSender::is_closed) + } + } +} + pub(crate) async fn forward( mut sender: HttpSender, pending: PendingRequest, diff --git a/src/client_request/response_idle.rs b/src/client_request/response_idle.rs new file mode 100644 index 0000000..aede7f1 --- /dev/null +++ b/src/client_request/response_idle.rs @@ -0,0 +1,171 @@ +//! Watch response progress outside Hyper's body polling. A blocked source socket +//! must still time out even when Hyper has stopped asking for body frames. +use super::ProxyResponse; +use http_body_util::BodyExt; +use hyper::body::{Body, Frame, SizeHint}; +use std::{ + io, + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; +use tokio::{ + io::{AsyncRead, AsyncWrite, ReadBuf}, + net::TcpStream, + sync::watch, + time::Instant, +}; + +#[derive(Clone)] +pub(crate) struct ResponseActivity(watch::Sender>); + +pub(crate) fn new( + stream: TcpStream, + timeout: Duration, +) -> (IdleIo, ResponseActivity, impl Future) { + let (tx, rx) = watch::channel(None); + let activity = ResponseActivity(tx); + ( + IdleIo { + stream, + activity: activity.clone(), + }, + activity, + wait_for_idle(rx, timeout), + ) +} + +async fn wait_for_idle(mut activity: watch::Receiver>, timeout: Duration) { + loop { + let last_write = *activity.borrow_and_update(); + if let Some(last_write) = last_write { + tokio::select! { + result = activity.changed() => { if result.is_err() { return; } } + _ = tokio::time::sleep_until(last_write + timeout) => { + // A write may race with timer expiry. Check the latest value. + if activity.borrow().is_some_and(|at| Instant::now() >= at + timeout) { + return; + } + } + } + } else if activity.changed().await.is_err() { + return; + } + } +} + +impl ResponseActivity { + pub(crate) fn track(&self, response: ProxyResponse) -> ProxyResponse { + self.0.send_replace(Some(Instant::now())); + response.map(|inner| { + let mut body = IdleBody { + inner, + activity: Some(self.clone()), + }; + if body.inner.is_end_stream() { + body.finish(); + } + body.boxed() + }) + } + + fn wrote_bytes(&self, result: &Poll>) { + if matches!(result, Poll::Ready(Ok(n)) if *n > 0) { + self.0.send_if_modified(|at| { + if at.is_some() { + *at = Some(Instant::now()); + true + } else { + false + } + }); + } + } +} + +pub(crate) struct IdleIo { + stream: TcpStream, + activity: ResponseActivity, +} + +impl AsyncRead for IdleIo { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.stream).poll_read(cx, buf) + } +} + +impl AsyncWrite for IdleIo { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let result = Pin::new(&mut self.stream).poll_write(cx, buf); + self.activity.wrote_bytes(&result); + result + } + fn poll_write_vectored( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll> { + let result = Pin::new(&mut self.stream).poll_write_vectored(cx, bufs); + self.activity.wrote_bytes(&result); + result + } + fn is_write_vectored(&self) -> bool { + self.stream.is_write_vectored() + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_flush(cx) + } + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.stream).poll_shutdown(cx) + } +} + +struct IdleBody { + inner: http_body_util::combinators::BoxBody, + activity: Option, +} + +impl IdleBody { + fn finish(&mut self) { + if let Some(activity) = self.activity.take() { + activity.0.send_replace(None); + } + } +} + +impl Drop for IdleBody { + fn drop(&mut self) { + self.finish(); + } +} + +impl Body for IdleBody { + type Data = bytes::Bytes; + type Error = hyper::Error; + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let result = Pin::new(&mut self.inner).poll_frame(cx); + if matches!(result, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) + || self.inner.is_end_stream() + { + self.finish(); + } + result + } + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + fn size_hint(&self) -> SizeHint { + self.inner.size_hint() + } +} diff --git a/src/client_request/tests.rs b/src/client_request/tests.rs index c29fea5..8aaf00b 100644 --- a/src/client_request/tests.rs +++ b/src/client_request/tests.rs @@ -33,6 +33,23 @@ struct Fixture { } async fn proxy(app: Router, protocol: &[u8], slots: usize, request_timeout: Duration) -> Fixture { + proxy_with_idle_timeout( + app, + protocol, + slots, + request_timeout, + Duration::from_secs(60), + ) + .await +} + +async fn proxy_with_idle_timeout( + app: Router, + protocol: &[u8], + slots: usize, + request_timeout: Duration, + response_body_idle_timeout: Duration, +) -> Fixture { let mut tasks = JoinSet::new(); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let target = listener.local_addr().unwrap(); @@ -74,6 +91,7 @@ async fn proxy(app: Router, protocol: &[u8], slots: usize, request_timeout: Dura .with_request_options(ProxyClientOptions { request_timeout, max_in_flight_requests: slots.try_into().unwrap(), + response_body_idle_timeout, }); let addr = client.local_addr().unwrap(); tasks.spawn(async move { @@ -97,6 +115,83 @@ fn http_client() -> reqwest::Client { .unwrap() } +// Real Hyper senders over an in-memory connection make worker failures and +// connection closure deterministic, without racing TCP shutdown against dispatch. +async fn sender_for_test(http2: bool) -> (crate::http_version::HttpSender, JoinSet<()>) { + use hyper_util::rt::TokioIo; + let (client, server) = tokio::io::duplex(4096); + let mut tasks = JoinSet::new(); + let service = hyper::service::service_fn(|_| async { + Ok::<_, std::convert::Infallible>(hyper::Response::new(crate::full("ok"))) + }); + let sender = if http2 { + tasks.spawn(async move { + let _ = hyper::server::conn::http2::Builder::new(crate::TokioExecutor) + .serve_connection(TokioIo::new(server), service) + .await; + }); + let (sender, connection) = + hyper::client::conn::http2::handshake(crate::TokioExecutor, TokioIo::new(client)) + .await + .unwrap(); + tasks.spawn(async move { + let _ = connection.await; + }); + sender.into() + } else { + tasks.spawn(async move { + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(TokioIo::new(server), service) + .await; + }); + let (sender, connection) = hyper::client::conn::http1::handshake(TokioIo::new(client)) + .await + .unwrap(); + tasks.spawn(async move { + let _ = connection.await; + }); + sender.into() + }; + (sender, tasks) +} + +#[tokio::test] +async fn worker_panic_preserves_http2_sender_but_reconnects_http1() { + for http2 in [false, true] { + let (sender, _tasks) = sender_for_test(http2).await; + let mut sender = Some(sender); + let worker_sender = super::take_sender(&mut sender).unwrap(); + let failure: Result = tokio::spawn(async move { + let _sender = worker_sender; + panic!("simulated forwarding worker panic"); + }) + .await; + assert!(matches!(&failure, Err(error) if error.is_panic())); + assert_eq!(super::worker_finished(&mut sender, failure), !http2); + if http2 { + let mut next = super::take_sender(&mut sender).unwrap(); + timeout(Duration::from_secs(1), next.ready()) + .await + .unwrap() + .unwrap(); + assert!(!next.is_closed()); + } else { + assert!(sender.is_none()); + } + } +} + +#[tokio::test] +async fn closed_sender_is_not_dispatched() { + for http2 in [false, true] { + let (sender, mut tasks) = sender_for_test(http2).await; + tasks.shutdown().await; + let mut sender = Some(sender); + assert!(sender.as_ref().unwrap().is_closed()); + assert!(super::take_sender(&mut sender).is_none()); + } +} + #[tokio::test] async fn http2_stalled_request_does_not_block_fast_request() { let entered = Arc::new(Notify::new()); @@ -492,3 +587,156 @@ async fn early_response_allows_upload_to_finish_before_releasing_slot() { assert_eq!(response.text().await.unwrap(), "fast"); assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); } + +#[tokio::test] +async fn non_reading_source_times_out_and_releases_capacity() { + for protocol in [ALPN_HTTP11, ALPN_H2] { + let started = Arc::new(Notify::new()); + let signal = started.clone(); + let app = Router::new() + .route( + "/stream", + get(move || { + let signal = signal.clone(); + async move { + let (tx, rx) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + let chunk = bytes::Bytes::from(vec![b'x'; 64 * 1024]); + // An endless body eventually fills the source TCP window. + while tx + .send(Ok(hyper::body::Frame::data(chunk.clone()))) + .await + .is_ok() + { + signal.notify_one(); + } + }); + axum::body::Body::new(TestBody(rx)) + } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy_with_idle_timeout( + app, + protocol, + 1, + Duration::from_secs(5), + Duration::from_millis(250), + ) + .await; + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"GET /stream HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + timeout(Duration::from_secs(2), started.notified()) + .await + .unwrap(); + // Keep the connection open without reading any response bytes. The only + // slot must become available well before the five-second request deadline. + let response = timeout( + Duration::from_secs(3), + http_client().get(format!("{}/fast", fixture.url)).send(), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(response.text().await.unwrap(), "fast"); + assert_eq!( + fixture.connections.load(Ordering::SeqCst), + if protocol == ALPN_HTTP11 { 2 } else { 1 } + ); + drop(source); + } +} + +#[tokio::test] +async fn silent_response_body_times_out_and_releases_capacity() { + for protocol in [ALPN_HTTP11, ALPN_H2] { + let (_tx, rx) = tokio::sync::mpsc::channel(1); + let rx = Arc::new(tokio::sync::Mutex::new(Some(rx))); + let app = Router::new() + .route( + "/silent", + get(move || { + let rx = rx.clone(); + async move { axum::body::Body::new(TestBody(rx.lock().await.take().unwrap())) } + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy_with_idle_timeout( + app, + protocol, + 1, + Duration::from_secs(5), + Duration::from_millis(200), + ) + .await; + let client = http_client(); + let response = client + .get(format!("{}/silent", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + assert!( + timeout(Duration::from_secs(2), response.bytes()) + .await + .unwrap() + .is_err() + ); + let response = client + .get(format!("{}/fast", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.text().await.unwrap(), "fast"); + } +} + +#[tokio::test] +async fn progressing_response_outlives_idle_and_request_deadlines() { + for protocol in [ALPN_HTTP11, ALPN_H2] { + let app = Router::new().route( + "/", + get(|| async { + let (tx, rx) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + for _ in 0..10 { + tx.send(Ok(hyper::body::Frame::data(bytes::Bytes::from_static( + b"x", + )))) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + } + }); + axum::body::Body::new(TestBody(rx)) + }), + ); + let fixture = proxy_with_idle_timeout( + app, + protocol, + 1, + Duration::from_millis(200), + Duration::from_millis(200), + ) + .await; + let client = http_client(); + let response = client + .get(format!("{}/", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.text().await.unwrap(), "xxxxxxxxxx"); + // An idle keep-alive connection has no active body and must not time out. + tokio::time::sleep(Duration::from_millis(300)).await; + let response = client + .get(format!("{}/", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.text().await.unwrap(), "xxxxxxxxxx"); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); + } +} diff --git a/src/lib.rs b/src/lib.rs index da54922..032b7f7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,7 +12,7 @@ pub use attested_tls::attestation::AttestationGenerator; mod client_request; mod http_version; pub use client_request::ProxyClientOptions; -use client_request::{PendingRequest, forward, gateway_timeout}; +use client_request::{PendingRequest, forward, gateway_timeout, take_sender, worker_finished}; #[cfg(test)] mod test_helpers; @@ -456,6 +456,7 @@ impl ProxyClient { // Retired connections may still have complete responses waiting to be // delivered. Drain their workers without blocking a fresh connection. let mut draining = tokio::task::JoinSet::new(); + let mut deferred = None; 'reconnect: loop { let (sender, conn, measurements, remote_attestation_type) = // Connect to the proxy server and provide / verify attestation @@ -492,13 +493,21 @@ impl ProxyClient { let mut sender = Some(sender); loop { tokio::select! { - incoming = requests_rx.recv(), if sender.is_some() => { + incoming = async { + match deferred.take() { + Some(pending) => Some(pending), + None => requests_rx.recv().await, + } + }, if sender.is_some() => { let Some(pending) = incoming else { break 'reconnect; }; - let request_sender = match sender.as_ref().unwrap() { - HttpSender::Http2(inner) => HttpSender::Http2(inner.clone()), - HttpSender::Http1(_) => sender.take().unwrap(), + let Some(request_sender) = take_sender(&mut sender) else { + // This request has not been dispatched. Preserve its + // deadline and permit across reconnect; never replay + // a request already handed to a worker. + deferred = Some(pending); + break; }; in_flight.spawn(forward( request_sender, @@ -508,13 +517,9 @@ impl ProxyClient { )); } result = in_flight.join_next(), if !in_flight.is_empty() => { - match result { - Some(Ok(result)) if !result.reconnect => { - if matches!(result.sender, HttpSender::Http1(_)) { - sender = Some(result.sender); - } - } - _ => break, + if let Some(result) = result + && worker_finished(&mut sender, result) { + break; } } _ = connection.join_next() => { @@ -577,28 +582,38 @@ impl ProxyClient { ) -> Result<(), ProxyError> { tracing::debug!("proxy-client accepted connection"); + let (inbound, activity, idle_timeout) = + client_request::response_idle::new(inbound, options.response_body_idle_timeout); + // Setup http server and handler let http = hyper::server::conn::http1::Builder::new(); let service = service_fn(move |req| { let requests_tx = requests_tx.clone(); let request_slots = request_slots.clone(); + let activity = activity.clone(); async move { - match Self::handle_http_request(req, requests_tx, options, request_slots).await { - Ok(res) => { - Ok::>, hyper::Error>(res) - } - Err(e) => { - warn!("send_request error: {e}"); - let mut resp = Response::new(full(format!("Request failed: {e}"))); - *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY; - Ok(resp) - } - } + let response = + match Self::handle_http_request(req, requests_tx, options, request_slots).await + { + Ok(res) => res, + Err(e) => { + warn!("send_request error: {e}"); + let mut resp = Response::new(full(format!("Request failed: {e}"))); + *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY; + resp + } + }; + Ok::<_, hyper::Error>(activity.track(response)) } }); let io = TokioIo::new(inbound); - http.serve_connection(io, service).await?; + tokio::select! { + result = http.serve_connection(io, service) => result?, + _ = idle_timeout => { + warn!("Closing source connection after response-body idle timeout"); + } + } Ok(()) } diff --git a/src/main.rs b/src/main.rs index dbdac8d..b25e89f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -71,6 +71,9 @@ enum CliCommand { /// Request deadline in seconds, including queueing and waiting for response headers #[arg(long, default_value = "60")] request_timeout_secs: NonZeroU64, + /// Close a source connection if its response body makes no write progress for this many seconds + #[arg(long, default_value = "60")] + response_body_idle_timeout_secs: NonZeroU64, /// Maximum in-flight requests, including streaming responses #[arg(long, default_value = "64")] max_in_flight_requests: NonZeroUsize, @@ -254,6 +257,7 @@ async fn main() -> anyhow::Result<()> { listen_addr, target_addr, request_timeout_secs, + response_body_idle_timeout_secs, max_in_flight_requests, client_attestation_type, tls_private_key_path, @@ -324,6 +328,9 @@ async fn main() -> anyhow::Result<()> { } .with_request_options(ProxyClientOptions { request_timeout: Duration::from_secs(request_timeout_secs.get()), + response_body_idle_timeout: Duration::from_secs( + response_body_idle_timeout_secs.get(), + ), max_in_flight_requests, }); From 89148b984216e3c22b7287a312d211e02c9094c0 Mon Sep 17 00:00:00 2001 From: peg Date: Mon, 14 Sep 2026 09:26:22 +0200 Subject: [PATCH 3/7] Cover edge case for upload cancellation --- Cargo.lock | 1 + Cargo.toml | 1 + src/client_request/http2.rs | 191 +++++++++++++++++++ src/client_request/mod.rs | 20 +- src/client_request/response_idle.rs | 4 +- src/client_request/tests.rs | 278 +++++++++++++++++++++++++++- src/client_request/upload.rs | 118 +++++++++++- src/http_version.rs | 30 +-- src/lib.rs | 19 +- 9 files changed, 622 insertions(+), 40 deletions(-) create mode 100644 src/client_request/http2.rs diff --git a/Cargo.lock b/Cargo.lock index b14146a..f06f474 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -680,6 +680,7 @@ dependencies = [ "axum", "bytes", "clap", + "h2", "http", "http-body-util", "hyper", diff --git a/Cargo.toml b/Cargo.toml index 545bf7c..95317a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ rustls-pemfile = "2.2.0" anyhow = "1.0.100" pem-rfc7468 = { version = "0.7.0", features = ["std"] } hyper = { version = "1.7.0", features = ["server", "http2"] } +h2 = "0.4.12" hyper-util = { version = "0.1.17", features = ["tokio"] } http-body-util = "0.1.3" bytes = "1.11.1" diff --git a/src/client_request/http2.rs b/src/client_request/http2.rs new file mode 100644 index 0000000..9a3be8c --- /dev/null +++ b/src/client_request/http2.rs @@ -0,0 +1,191 @@ +//! HTTP/2 forwarding with an explicit reset handle for each request upload. +//! Hyper's client hides this handle after returning response headers, so a body +//! waiting for send capacity cannot otherwise be canceled through its Body API. +use std::{ + future::poll_fn, + pin::Pin, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll}, + time::Duration, +}; + +use bytes::Bytes; +use http_body_util::BodyExt; +use hyper::body::{Body, Frame, SizeHint}; +use tokio::io::{AsyncRead, AsyncWrite}; + +use super::{BoxError, ProxyResponse, RequestBody}; +use crate::ProxyError; + +pub(crate) type Connection = Pin> + Send>>; + +#[derive(Clone)] +pub(crate) struct Sender { + inner: h2::client::SendRequest, + closed: Arc, +} + +pub(crate) async fn handshake(io: T) -> Result<(Sender, Connection), ProxyError> +where + T: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let (inner, mut connection) = h2::client::Builder::new() + .initial_window_size(2 * 1024 * 1024) + .initial_connection_window_size(5 * 1024 * 1024) + .max_header_list_size(16 * 1024) + .enable_push(false) + .handshake(io) + .await?; + let ping = connection.ping_pong().expect("ping handle available once"); + let closed = Arc::new(AtomicBool::new(false)); + let sender = Sender { + inner, + closed: closed.clone(), + }; + // Mark the sender closed on completion, cancellation, or a keep-alive failure. + struct Closed(Arc); + impl Drop for Closed { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + let closed = Closed(closed); + let connection = Box::pin(async move { + let _closed = closed; + tokio::select! { + result = &mut connection => result.map_err(Into::into), + result = keep_alive(ping) => result, + } + }); + Ok((sender, connection)) +} + +async fn keep_alive(mut ping: h2::PingPong) -> Result<(), ProxyError> { + loop { + tokio::time::sleep(Duration::from_secs(crate::KEEP_ALIVE_INTERVAL)).await; + tokio::time::timeout( + Duration::from_secs(crate::KEEP_ALIVE_TIMEOUT), + ping.ping(h2::Ping::opaque()), + ) + .await + .map_err(|_| { + std::io::Error::new(std::io::ErrorKind::TimedOut, "HTTP/2 keep-alive timed out") + })??; + } +} + +impl Sender { + pub(crate) fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } + + pub(crate) async fn ready(&mut self) -> Result<(), ProxyError> { + poll_fn(|cx| self.inner.poll_ready(cx)) + .await + .map_err(|error| { + self.closed.store(true, Ordering::Release); + error.into() + }) + } + + pub(crate) async fn send_request( + &mut self, + request: http::Request, + ) -> Result { + let (mut parts, body) = request.into_parts(); + strip_connection_headers(&mut parts.headers); + if let Some(length) = body.size_hint().exact() + && (length != 0 + || matches!( + parts.method, + http::Method::POST | http::Method::PUT | http::Method::PATCH + )) + { + parts + .headers + .entry(http::header::CONTENT_LENGTH) + .or_insert(length.into()); + } + let end = body.is_end_stream(); + let (response, stream) = self + .inner + .send_request(http::Request::from_parts(parts, ()), end)?; + if !end { + body.send_http2(stream); + } + let response = response.await?; + Ok(response.map(|inner| { + ResponseBody { + inner, + data_done: false, + } + .boxed() + })) + } +} + +fn strip_connection_headers(headers: &mut http::HeaderMap) { + if let Some(connection) = headers.remove(http::header::CONNECTION) + && let Ok(names) = connection.to_str() + { + for name in names.split(',') { + headers.remove(name.trim()); + } + } + for name in [ + "keep-alive", + "proxy-connection", + "transfer-encoding", + "upgrade", + ] { + headers.remove(name); + } + if headers + .get(http::header::TE) + .is_some_and(|value| value != "trailers") + { + headers.remove(http::header::TE); + } +} + +struct ResponseBody { + inner: h2::RecvStream, + data_done: bool, +} + +impl Body for ResponseBody { + type Data = Bytes; + type Error = BoxError; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, BoxError>>> { + if !self.data_done { + match std::task::ready!(self.inner.poll_data(cx)) { + Some(Ok(data)) => { + if let Err(error) = self.inner.flow_control().release_capacity(data.len()) { + return Poll::Ready(Some(Err(error.into()))); + } + return Poll::Ready(Some(Ok(Frame::data(data)))); + } + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), + None => self.data_done = true, + } + } + self.inner.poll_trailers(cx).map(|result| match result { + Ok(trailers) => trailers.map(|trailers| Ok(Frame::trailers(trailers))), + Err(error) => Some(Err(error.into())), + }) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + fn size_hint(&self) -> SizeHint { + SizeHint::default() + } +} diff --git a/src/client_request/mod.rs b/src/client_request/mod.rs index 0cc3b67..79a19f5 100644 --- a/src/client_request/mod.rs +++ b/src/client_request/mod.rs @@ -1,4 +1,5 @@ //! Per-request forwarding, deadlines, and response lifetime tracking. +pub(crate) mod http2; pub(crate) mod response_idle; #[cfg(test)] mod tests; @@ -54,8 +55,9 @@ impl Default for ProxyClientOptions { } } +pub(crate) type BoxError = Box; pub(crate) type ProxyResponse = - Response>; + Response>; pub(crate) struct PendingRequest { pub request: http::Request, @@ -65,7 +67,11 @@ pub(crate) struct PendingRequest { } pub(crate) fn gateway_timeout() -> ProxyResponse { - let mut response = Response::new(full("Request deadline exceeded")); + let mut response = Response::new( + full("Request deadline exceeded") + .map_err(Into::into) + .boxed(), + ); *response.status_mut() = http::StatusCode::GATEWAY_TIMEOUT; response } @@ -153,7 +159,11 @@ pub(crate) async fn forward( failure => { if let Some(Err(error)) = failure { tracing::warn!("Failed to send request to proxy-server: {error}"); - let mut response = Response::new(full(format!("Request failed: {error}"))); + let mut response = Response::new( + full(format!("Request failed: {error}")) + .map_err(Into::into) + .boxed(), + ); *response.status_mut() = http::StatusCode::BAD_GATEWAY; let _ = response_tx.send(response); } @@ -228,7 +238,7 @@ pub(crate) async fn forward( pin_project_lite::pin_project! { struct TrackedBody { #[pin] - inner: Incoming, + inner: http_body_util::combinators::BoxBody, permit: Option>, finished: Option>, } @@ -245,7 +255,7 @@ impl TrackedBody { impl Body for TrackedBody { type Data = bytes::Bytes; - type Error = hyper::Error; + type Error = BoxError; fn poll_frame( self: Pin<&mut Self>, diff --git a/src/client_request/response_idle.rs b/src/client_request/response_idle.rs index aede7f1..bda256c 100644 --- a/src/client_request/response_idle.rs +++ b/src/client_request/response_idle.rs @@ -129,7 +129,7 @@ impl AsyncWrite for IdleIo { } struct IdleBody { - inner: http_body_util::combinators::BoxBody, + inner: http_body_util::combinators::BoxBody, activity: Option, } @@ -149,7 +149,7 @@ impl Drop for IdleBody { impl Body for IdleBody { type Data = bytes::Bytes; - type Error = hyper::Error; + type Error = super::BoxError; fn poll_frame( mut self: Pin<&mut Self>, cx: &mut Context<'_>, diff --git a/src/client_request/tests.rs b/src/client_request/tests.rs index 8aaf00b..a31978f 100644 --- a/src/client_request/tests.rs +++ b/src/client_request/tests.rs @@ -130,10 +130,7 @@ async fn sender_for_test(http2: bool) -> (crate::http_version::HttpSender, JoinS .serve_connection(TokioIo::new(server), service) .await; }); - let (sender, connection) = - hyper::client::conn::http2::handshake(crate::TokioExecutor, TokioIo::new(client)) - .await - .unwrap(); + let (sender, connection) = super::http2::handshake(client).await.unwrap(); tasks.spawn(async move { let _ = connection.await; }); @@ -740,3 +737,276 @@ async fn progressing_response_outlives_idle_and_request_deadlines() { assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); } } + +#[tokio::test] +async fn flow_control_blocked_upload_deadline_releases_capacity() { + let app = Router::new() + .route( + "/upload", + post(|request: axum::extract::Request| async move { + // Keep the body alive without consuming it. A sufficiently large + // upload fills both the target socket and the HTTP/2 stream window. + tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(3)).await; + drop(request); + }); + http::StatusCode::OK + }), + ) + .route("/fast", get(|| async { "fast" })); + let fixture = proxy(app, ALPN_H2, 1, Duration::from_millis(400)).await; + let mut source = TcpStream::connect(fixture.addr).await.unwrap(); + source + .write_all(b"POST /upload HTTP/1.1\r\nHost: localhost\r\nContent-Length: 100000000\r\n\r\n") + .await + .unwrap(); + let (mut reader, mut writer) = source.into_split(); + let mut tasks = JoinSet::new(); + tasks.spawn(async move { + let _ = writer.write_all(&vec![b'x'; 16 * 1024 * 1024]).await; + std::future::pending::<()>().await; + }); + let mut headers = [0; 1024]; + let n = timeout(Duration::from_secs(2), reader.read(&mut headers)) + .await + .unwrap() + .unwrap(); + assert!(String::from_utf8_lossy(&headers[..n]).starts_with("HTTP/1.1 200")); + tokio::time::sleep(Duration::from_millis(600)).await; + let response = http_client() + .get(format!("{}/fast", fixture.url)) + .send() + .await + .unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + assert_eq!(response.text().await.unwrap(), "fast"); + assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); +} + +// Obtain a real Incoming body without a TCP socket or TLS handshake. The source +// connection stays driven while the forwarding worker consumes the request. +async fn incoming_request( + raw: Vec, + tasks: &mut JoinSet<()>, +) -> http::Request { + let (mut source, server) = tokio::io::duplex(4096); + let (tx, rx) = tokio::sync::oneshot::channel(); + let tx = Arc::new(std::sync::Mutex::new(Some(tx))); + tasks.spawn(async move { + let service = hyper::service::service_fn(move |request| { + tx.lock().unwrap().take().unwrap().send(request).unwrap(); + std::future::pending::< + Result< + http::Response>, + std::convert::Infallible, + >, + >() + }); + let _ = hyper::server::conn::http1::Builder::new() + .serve_connection(hyper_util::rt::TokioIo::new(server), service) + .await; + }); + tasks.spawn(async move { + let _ = source.write_all(&raw).await; + std::future::pending::<()>().await; + }); + rx.await.unwrap() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn blocked_uploads_reset_streams_and_preserve_other_http2_responses() { + use http_body_util::BodyExt; + use std::future::poll_fn; + let mut tasks = JoinSet::new(); + let (client, server) = tokio::io::duplex(4096); + let (reset_tx, mut reset_rx) = tokio::sync::mpsc::channel(8); + let finish_other = Arc::new(Notify::new()); + let finish = finish_other.clone(); + tasks.spawn(async move { + // No upload can make progress, regardless of socket buffer sizes. + let mut connection = h2::server::Builder::new() + .initial_window_size(0) + .handshake::<_, bytes::Bytes>(server) + .await + .unwrap(); + let mut streams = JoinSet::new(); + while let Some(request) = connection.accept().await { + let (request, mut respond) = request.unwrap(); + let other = request.uri().path() == "/other"; + let mut response = if request.uri().path() == "/late" { + None + } else { + Some( + respond + .send_response(http::Response::new(()), !other) + .unwrap(), + ) + }; + let reset_tx = reset_tx.clone(); + let finish = finish.clone(); + streams.spawn(async move { + if other { + finish.notified().await; + response + .as_mut() + .unwrap() + .send_data(bytes::Bytes::from_static(b"still streaming"), true) + .unwrap(); + } else { + let reason = match response.as_mut() { + Some(response) => poll_fn(|cx| response.poll_reset(cx)).await.unwrap(), + None => poll_fn(|cx| respond.poll_reset(cx)).await.unwrap(), + }; + assert_eq!(reason, h2::Reason::CANCEL); + reset_tx.send(()).await.unwrap(); + } + drop(request); + }); + } + }); + let (sender, connection) = super::http2::handshake(client).await.unwrap(); + tasks.spawn(async move { + let _ = connection.await; + }); + let mut sender = Some(crate::http_version::HttpSender::Http2(sender)); + let slots = Arc::new(tokio::sync::Semaphore::new(2)); + let (tx, rx) = tokio::sync::oneshot::channel(); + let other_request = incoming_request( + b"GET /other HTTP/1.1\r\nHost: localhost\r\n\r\n".to_vec(), + &mut tasks, + ) + .await; + let other = tokio::spawn(super::forward( + super::take_sender(&mut sender).unwrap(), + super::PendingRequest { + request: other_request, + response_tx: tx, + deadline: tokio::time::Instant::now() + Duration::from_secs(5), + permit: slots.clone().acquire_owned().await.unwrap(), + }, + None, + crate::attestation::AttestationType::None, + )); + let response = timeout(Duration::from_secs(1), rx).await.unwrap().unwrap(); + // Cover both completed and partial source bodies, with and without early + // response headers. Even a fully read source body must wait for send capacity. + for (path, length, supplied) in [ + ("/upload", 1, 1), + ("/upload", 1_000_000, 128 * 1024), + ("/late", 1, 1), + ("/late", 1_000_000, 128 * 1024), + ] { + let mut raw = + format!("POST {path} HTTP/1.1\r\nHost: localhost\r\nContent-Length: {length}\r\n\r\n") + .into_bytes(); + raw.extend(vec![b'x'; supplied]); + let request = incoming_request(raw, &mut tasks).await; + let (tx, rx) = tokio::sync::oneshot::channel(); + let worker = tokio::spawn(super::forward( + super::take_sender(&mut sender).unwrap(), + super::PendingRequest { + request, + response_tx: tx, + deadline: tokio::time::Instant::now() + Duration::from_millis(150), + permit: slots.clone().acquire_owned().await.unwrap(), + }, + None, + crate::attestation::AttestationType::None, + )); + assert_eq!( + rx.await.unwrap().status(), + if path == "/late" { + http::StatusCode::GATEWAY_TIMEOUT + } else { + http::StatusCode::OK + }, + ); + let result = timeout(Duration::from_secs(1), worker) + .await + .unwrap() + .unwrap(); + assert!(!result.reconnect); + timeout(Duration::from_secs(1), reset_rx.recv()) + .await + .unwrap() + .unwrap(); + let permit = timeout(Duration::from_secs(1), slots.clone().acquire_owned()) + .await + .unwrap() + .unwrap(); + drop(permit); + } + finish_other.notify_one(); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "still streaming" + ); + assert!(!other.await.unwrap().reconnect); +} + +#[tokio::test] +async fn http2_upload_and_response_preserve_trailers() { + use http_body_util::BodyExt; + let mut tasks = JoinSet::new(); + let (client, server) = tokio::io::duplex(4096); + tasks.spawn(async move { + let mut connection = h2::server::handshake(server).await.unwrap(); + let mut streams = JoinSet::new(); + while let Some(request) = connection.accept().await { + let (request, mut respond) = request.unwrap(); + streams.spawn(async move { + assert!(!request.headers().contains_key("connection")); + assert!(!request.headers().contains_key("x-hop")); + assert!(!request.headers().contains_key("transfer-encoding")); + assert_eq!(request.headers()["te"], "trailers"); + let mut response = respond + .send_response(http::Response::new(()), false) + .unwrap(); + let mut body = request.into_body(); + let mut received = Vec::new(); + while let Some(data) = body.data().await { + let data = data.unwrap(); + body.flow_control().release_capacity(data.len()).unwrap(); + received.extend_from_slice(&data); + } + assert_eq!(received, b"abcdef"); + let trailers = body.trailers().await.unwrap().unwrap(); + assert_eq!(trailers["x-checksum"], "valid"); + response + .send_data(bytes::Bytes::from(received), false) + .unwrap(); + response.send_trailers(trailers).unwrap(); + }); + } + }); + let (sender, connection) = super::http2::handshake(client).await.unwrap(); + tasks.spawn(async move { + let _ = connection.await; + }); + let request = incoming_request( + b"POST / HTTP/1.1\r\nHost: localhost\r\nTransfer-Encoding: chunked\r\nTE: trailers\r\nTrailer: x-checksum\r\nConnection: x-hop\r\nx-hop: remove\r\n\r\n3\r\nabc\r\n3\r\ndef\r\n0\r\nx-checksum: valid\r\n\r\n".to_vec(), + &mut tasks, + ).await; + let slots = Arc::new(tokio::sync::Semaphore::new(1)); + let (tx, rx) = tokio::sync::oneshot::channel(); + let worker = tokio::spawn(super::forward( + crate::http_version::HttpSender::Http2(sender), + super::PendingRequest { + request, + response_tx: tx, + deadline: tokio::time::Instant::now() + Duration::from_secs(2), + permit: slots.clone().acquire_owned().await.unwrap(), + }, + None, + crate::attestation::AttestationType::None, + )); + let response = timeout(Duration::from_secs(2), rx).await.unwrap().unwrap(); + let body = timeout(Duration::from_secs(2), response.into_body().collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(body.trailers().unwrap()["x-checksum"], "valid"); + assert_eq!(body.to_bytes(), "abcdef"); + assert!(!worker.await.unwrap().reconnect); + assert_eq!(slots.available_permits(), 1); +} diff --git a/src/client_request/upload.rs b/src/client_request/upload.rs index c32961e..b2a6588 100644 --- a/src/client_request/upload.rs +++ b/src/client_request/upload.rs @@ -1,6 +1,6 @@ -//! Upload cancellation must work even when Hyper is waiting for HTTP/2 capacity -//! and is not polling the body. Shared state lets the forwarding task drop the -//! source body independently. +//! Upload cancellation must work while waiting for HTTP/2 capacity, without +//! polling the body. Shared state lets the forwarding task reset the +//! HTTP/2 stream and drop the source body independently. use hyper::body::{Body, Frame, Incoming, SizeHint}; use std::{ io, @@ -15,6 +15,11 @@ struct UploadState { finished: Option>, complete: bool, waker: Option, + // Retain both the wire reset and local task cancellation handles because + // Body::poll_frame cannot interrupt a flow-control capacity wait. + http2: Option>>>, + upload_task: Option, + canceled: bool, } impl UploadState { @@ -22,7 +27,9 @@ impl UploadState { self.inner.take(); self.complete = true; self.waker.take(); - if let Some(finished) = self.finished.take() { + if self.http2.is_none() + && let Some(finished) = self.finished.take() + { let _ = finished.send(()); } } @@ -31,6 +38,15 @@ impl UploadState { fn cancel(state: &Mutex) { let waker = { let mut state = state.lock().unwrap(); + state.canceled = true; + if let Some(stream) = state.http2.take() { + stream.lock().unwrap().send_reset(h2::Reason::CANCEL); + } + // A local reset need not wake our own capacity waiter. Abort the task + // as well, so it drops its body and permit without any peer progress. + if let Some(task) = state.upload_task.take() { + task.abort(); + } state.inner.take(); state.finished.take(); state.waker.take() @@ -48,8 +64,8 @@ impl Drop for UploadGuard { } } -/// Hyper sends uploads independently of response futures. Keep the permit until -/// the upload finishes or Hyper drops the canceled stream. +/// Uploads run independently of response futures. Keep the permit until the +/// upload finishes or its canceled task drops the body. pub(crate) struct RequestBody { state: Arc>, permit: Option>, @@ -66,6 +82,9 @@ impl RequestBody { finished: Some(finished), complete: false, waker: None, + http2: None, + upload_task: None, + canceled: false, }; if state.inner.as_ref().unwrap().is_end_stream() { state.finish(); @@ -80,6 +99,87 @@ impl RequestBody { receiver, ) } + + pub(crate) fn send_http2(mut self, stream: h2::SendStream) { + use http_body_util::BodyExt; + use std::future::poll_fn; + + let stream = Arc::new(Mutex::new(stream)); + { + let mut state = self.state.lock().unwrap(); + if state.canceled { + stream.lock().unwrap().send_reset(h2::Reason::CANCEL); + return; + } + state.http2 = Some(stream.clone()); + } + let state = self.state.clone(); + let task = tokio::spawn(async move { + let result: Result<(), super::BoxError> = async { + loop { + let frame = tokio::select! { + biased; + reset = poll_fn(|cx| stream.lock().unwrap().poll_reset(cx)) => { + return Err(reset.map(h2::Error::from).unwrap_or_else(|error| error).into()); + } + frame = self.frame() => frame, + }; + let Some(frame) = frame else { break; }; + let frame = frame?; + match frame.into_data() { + Ok(mut data) => { + while !data.is_empty() { + stream.lock().unwrap().reserve_capacity(data.len()); + let capacity = + poll_fn(|cx| stream.lock().unwrap().poll_capacity(cx)) + .await + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::BrokenPipe, + "HTTP/2 upload closed", + ) + })??; + if capacity == 0 { + continue; + } + let chunk = data.split_to(capacity.min(data.len())); + stream.lock().unwrap().send_data(chunk, false)?; + } + } + Err(frame) => { + if let Ok(trailers) = frame.into_trailers() { + stream.lock().unwrap().send_trailers(trailers)?; + return Ok(()); + } + } + } + } + stream + .lock() + .unwrap() + .send_data(bytes::Bytes::new(), true)?; + Ok(()) + } + .await; + if let Err(error) = result { + tracing::debug!(%error, "HTTP/2 request upload failed"); + } else { + let mut state = self.state.lock().unwrap(); + state.http2.take(); + state.upload_task.take(); + if let Some(finished) = state.finished.take() { + let _ = finished.send(()); + } + } + // On failure, Drop resets the stream and closes upload_finished. + }); + let mut state = state.lock().unwrap(); + if state.canceled { + task.abort(); + } else if state.http2.is_some() { + state.upload_task = Some(task.abort_handle()); + } + } } impl Drop for RequestBody { @@ -118,7 +218,11 @@ impl Body for RequestBody { )))) } }; - if matches!(frame, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) || self.is_end_stream() { + let http2 = self.state.lock().unwrap().http2.is_some(); + if !http2 + && (matches!(frame, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) + || self.is_end_stream()) + { self.permit.take(); } frame diff --git a/src/http_version.rs b/src/http_version.rs index 90c7421..eb408bc 100644 --- a/src/http_version.rs +++ b/src/http_version.rs @@ -1,6 +1,10 @@ //! HTTP Version support and negotiation use crate::client_request::RequestBody; -use hyper::Response; +use crate::{ + ProxyError, + client_request::{ProxyResponse, http2}, +}; +use http_body_util::BodyExt; use hyper_util::rt::TokioIo; use std::pin::Pin; use std::task::{Context, Poll}; @@ -54,18 +58,14 @@ impl HttpVersion { } type Http1Sender = hyper::client::conn::http1::SendRequest; -type Http2Sender = hyper::client::conn::http2::SendRequest; +type Http2Sender = http2::Sender; type Http1Connection = hyper::client::conn::http1::Connection< TokioIo>, RequestBody, >; -type Http2Connection = hyper::client::conn::http2::Connection< - TokioIo>, - RequestBody, - crate::TokioExecutor, ->; +type Http2Connection = http2::Connection; /// A protocol version agnostic HTTP sender pub enum HttpSender { @@ -86,9 +86,9 @@ impl From for HttpSender { } impl HttpSender { - pub async fn ready(&mut self) -> Result<(), hyper::Error> { + pub async fn ready(&mut self) -> Result<(), ProxyError> { match self { - Self::Http1(sender) => sender.ready().await, + Self::Http1(sender) => sender.ready().await.map_err(Into::into), Self::Http2(sender) => sender.ready().await, } } @@ -103,9 +103,13 @@ impl HttpSender { pub async fn send_request( &mut self, request: http::Request, - ) -> Result, hyper::Error> { + ) -> Result { match self { - Self::Http1(sender) => sender.send_request(request).await, + Self::Http1(sender) => sender + .send_request(request) + .await + .map(|response| response.map(|body| body.map_err(Into::into).boxed())) + .map_err(Into::into), Self::Http2(sender) => sender.send_request(request).await, } } @@ -133,11 +137,11 @@ impl From for HttpConnection { } impl Future for HttpConnection { - type Output = Result<(), hyper::Error>; + type Output = Result<(), ProxyError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { match self.project() { - HttpConnectionProj::Http1 { inner } => inner.poll(cx), + HttpConnectionProj::Http1 { inner } => inner.poll(cx).map_err(Into::into), HttpConnectionProj::Http2 { inner } => inner.poll(cx), } } diff --git a/src/lib.rs b/src/lib.rs index 032b7f7..57f9ff5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -598,7 +598,11 @@ impl ProxyClient { Ok(res) => res, Err(e) => { warn!("send_request error: {e}"); - let mut resp = Response::new(full(format!("Request failed: {e}"))); + let mut resp = Response::new( + full(format!("Request failed: {e}")) + .map_err(Into::into) + .boxed(), + ); *resp.status_mut() = hyper::StatusCode::BAD_GATEWAY; resp } @@ -679,13 +683,8 @@ impl ProxyClient { let outbound_io = TokioIo::new(tls_stream); let (sender, conn) = match http_version { HttpVersion::Http2 => { - let (sender, conn) = hyper::client::conn::http2::Builder::new(TokioExecutor) - .timer(hyper_util::rt::tokio::TokioTimer::new()) - .keep_alive_interval(Some(Duration::from_secs(KEEP_ALIVE_INTERVAL))) - .keep_alive_timeout(Duration::from_secs(KEEP_ALIVE_TIMEOUT)) - .keep_alive_while_idle(true) - .handshake::<_, client_request::RequestBody>(outbound_io) - .await?; + let (sender, conn) = + client_request::http2::handshake(outbound_io.into_inner()).await?; (sender.into(), conn.into()) } HttpVersion::Http1 => { @@ -706,7 +705,7 @@ impl ProxyClient { requests_tx: mpsc::Sender, options: ProxyClientOptions, request_slots: Arc, - ) -> Result>, ProxyError> { + ) -> Result { let deadline = tokio::time::Instant::now() + options.request_timeout; let result = tokio::time::timeout_at(deadline, async { let permit = request_slots @@ -770,6 +769,8 @@ pub enum ProxyError { BadDnsName(#[from] tokio_rustls::rustls::pki_types::InvalidDnsNameError), #[error("HTTP: {0}")] Hyper(#[from] hyper::Error), + #[error("HTTP/2: {0}")] + Http2(#[from] h2::Error), #[error("JSON: {0}")] Json(#[from] serde_json::Error), #[error("Could not forward response - sender was dropped")] From 9db6a7a83027239881dbf6ddfd4f6220933b49c7 Mon Sep 17 00:00:00 2001 From: peg Date: Mon, 14 Sep 2026 10:48:21 +0200 Subject: [PATCH 4/7] Improve idle timer --- src/client_request/response_idle.rs | 58 +++++++++++++++++++---- src/client_request/tests.rs | 73 +++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 10 deletions(-) diff --git a/src/client_request/response_idle.rs b/src/client_request/response_idle.rs index aede7f1..b88db0b 100644 --- a/src/client_request/response_idle.rs +++ b/src/client_request/response_idle.rs @@ -16,9 +16,18 @@ use tokio::{ time::Instant, }; +/// Shares response progress between the body, socket, and idle timer. #[derive(Clone)] -pub(crate) struct ResponseActivity(watch::Sender>); +pub(crate) struct ResponseActivity(watch::Sender>); +/// Tracks write progress and body completion until the response is flushed. +#[derive(Clone, Copy)] +struct ActiveResponse { + last_write: Instant, + body_finished: bool, +} + +/// Wraps a source socket with response tracking and an idle timeout future. pub(crate) fn new( stream: TcpStream, timeout: Duration, @@ -35,15 +44,16 @@ pub(crate) fn new( ) } -async fn wait_for_idle(mut activity: watch::Receiver>, timeout: Duration) { +/// Waits until an active response becomes idle or its activity channel closes. +async fn wait_for_idle(mut activity: watch::Receiver>, timeout: Duration) { loop { - let last_write = *activity.borrow_and_update(); + let last_write = activity.borrow_and_update().map(|active| active.last_write); if let Some(last_write) = last_write { tokio::select! { result = activity.changed() => { if result.is_err() { return; } } _ = tokio::time::sleep_until(last_write + timeout) => { // A write may race with timer expiry. Check the latest value. - if activity.borrow().is_some_and(|at| Instant::now() >= at + timeout) { + if activity.borrow().is_some_and(|active| Instant::now() >= active.last_write + timeout) { return; } } @@ -55,8 +65,12 @@ async fn wait_for_idle(mut activity: watch::Receiver>, timeout: } impl ResponseActivity { + /// Starts idle tracking and wraps the response body to observe completion. pub(crate) fn track(&self, response: ProxyResponse) -> ProxyResponse { - self.0.send_replace(Some(Instant::now())); + self.0.send_replace(Some(ActiveResponse { + last_write: Instant::now(), + body_finished: false, + })); response.map(|inner| { let mut body = IdleBody { inner, @@ -69,11 +83,12 @@ impl ResponseActivity { }) } + /// Refreshes the active response's timestamp after a successful nonempty write. fn wrote_bytes(&self, result: &Poll>) { if matches!(result, Poll::Ready(Ok(n)) if *n > 0) { - self.0.send_if_modified(|at| { - if at.is_some() { - *at = Some(Instant::now()); + self.0.send_if_modified(|active| { + if let Some(active) = active { + active.last_write = Instant::now(); true } else { false @@ -83,6 +98,7 @@ impl ResponseActivity { } } +/// Records socket write progress and clears idle tracking after the final flush. pub(crate) struct IdleIo { stream: TcpStream, activity: ResponseActivity, @@ -120,23 +136,45 @@ impl AsyncWrite for IdleIo { fn is_write_vectored(&self) -> bool { self.stream.is_write_vectored() } + /// Disarms the idle timer once a completed response has been flushed. fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - Pin::new(&mut self.stream).poll_flush(cx) + let result = Pin::new(&mut self.stream).poll_flush(cx); + if matches!(result, Poll::Ready(Ok(()))) { + // Hyper flushes its write buffer before flushing the underlying IO. + // Only then are the final body bytes no longer waiting to be written. + self.activity.0.send_if_modified(|active| { + if active.is_some_and(|active| active.body_finished) { + *active = None; + true + } else { + false + } + }); + } + result } fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { Pin::new(&mut self.stream).poll_shutdown(cx) } } +/// Reports body completion while leaving buffered writes covered by the idle timer. struct IdleBody { inner: http_body_util::combinators::BoxBody, activity: Option, } impl IdleBody { + /// Marks the body finished without disarming the timer before the final flush. fn finish(&mut self) { if let Some(activity) = self.activity.take() { - activity.0.send_replace(None); + // Hyper may still have the final frame buffered. Keep watching for + // write progress until IdleIo observes a successful flush. + activity.0.send_modify(|active| { + if let Some(active) = active { + active.body_finished = true; + } + }); } } } diff --git a/src/client_request/tests.rs b/src/client_request/tests.rs index 8aaf00b..0d073cb 100644 --- a/src/client_request/tests.rs +++ b/src/client_request/tests.rs @@ -740,3 +740,76 @@ async fn progressing_response_outlives_idle_and_request_deadlines() { assert_eq!(fixture.connections.load(Ordering::SeqCst), 1); } } + +/// Serves a single-frame response through the idle wrapper on a reusable connection. +async fn finite_idle_response(body: bytes::Bytes) -> (TcpStream, JoinSet<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let source = TcpStream::connect(listener.local_addr().unwrap()) + .await + .unwrap(); + let (inbound, _) = listener.accept().await.unwrap(); + let (io, activity, idle) = super::response_idle::new(inbound, Duration::from_millis(200)); + let mut tasks = JoinSet::new(); + tasks.spawn(async move { + let service = hyper::service::service_fn(move |_| { + let response = activity.track(hyper::Response::new(crate::full(body.clone()))); + async { Ok::<_, std::convert::Infallible>(response) } + }); + tokio::select! { + result = hyper::server::conn::http1::Builder::new() + .serve_connection(hyper_util::rt::TokioIo::new(io), service) => result.unwrap(), + _ = idle => (), + } + }); + (source, tasks) +} + +/// Checks that a blocked final frame remains subject to the idle timeout. +#[tokio::test] +async fn idle_timeout_covers_final_buffered_frame() { + let (mut source, mut tasks) = + finite_idle_response(bytes::Bytes::from(vec![b'x'; 16 * 1024 * 1024])).await; + source + .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n") + .await + .unwrap(); + // Keep the socket open without reading. Consuming the final body frame must + // not disable the timeout while its bytes are still buffered by Hyper. + timeout(Duration::from_secs(2), tasks.join_next()) + .await + .expect("final buffered frame escaped the idle timeout") + .unwrap() + .unwrap(); +} + +/// Checks that flushed empty and nonempty responses leave keep-alive connections usable. +#[tokio::test] +async fn flushed_responses_leave_source_keep_alive() { + use http_body_util::BodyExt; + // Empty responses must also disarm the timer after their headers are flushed. + for body in [bytes::Bytes::new(), bytes::Bytes::from_static(b"ok")] { + let (source, mut tasks) = finite_idle_response(body.clone()).await; + let (mut sender, connection) = + hyper::client::conn::http1::handshake(hyper_util::rt::TokioIo::new(source)) + .await + .unwrap(); + tasks.spawn(async move { + connection.await.unwrap(); + }); + for _ in 0..2 { + let response = sender + .send_request(http::Request::new(crate::full(""))) + .await + .unwrap(); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + body + ); + tokio::time::sleep(Duration::from_millis(350)).await; + assert!( + !sender.is_closed(), + "flushed response timed out on keep-alive connection" + ); + } + } +} From d47d29116e6670c9f48ed343ae6137bf452155c4 Mon Sep 17 00:00:00 2001 From: peg Date: Mon, 14 Sep 2026 11:21:11 +0200 Subject: [PATCH 5/7] Avoid panicing when using mutex --- src/client_request/upload.rs | 58 +++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/src/client_request/upload.rs b/src/client_request/upload.rs index c32961e..0f70828 100644 --- a/src/client_request/upload.rs +++ b/src/client_request/upload.rs @@ -30,7 +30,11 @@ impl UploadState { fn cancel(state: &Mutex) { let waker = { - let mut state = state.lock().unwrap(); + // Cleanup must also work during unwinding after a panic under this lock. + // Recover the guard only to discard the upload, never to resume it. + let mut state = state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); state.inner.take(); state.finished.take(); state.waker.take() @@ -96,28 +100,33 @@ impl Body for RequestBody { mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll, Self::Error>>> { - let frame = { - let mut state = self.state.lock().unwrap(); - if state.complete { - Poll::Ready(None) - } else if let Some(inner) = state.inner.as_mut() { - let frame = Pin::new(&mut *inner).poll_frame(cx); - if matches!(frame, Poll::Ready(Some(Err(_)))) { - state.inner.take(); - state.finished.take(); - } else if matches!(frame, Poll::Ready(None)) || inner.is_end_stream() { - state.finish(); + let frame = match self.state.lock() { + Err(_) => Poll::Ready(Some(Err(io::Error::other("request upload state poisoned")))), + Ok(mut state) => { + if state.complete { + Poll::Ready(None) + } else if let Some(inner) = state.inner.as_mut() { + let frame = Pin::new(&mut *inner).poll_frame(cx); + if matches!(frame, Poll::Ready(Some(Err(_)))) { + state.inner.take(); + state.finished.take(); + } else if matches!(frame, Poll::Ready(None)) || inner.is_end_stream() { + state.finish(); + } else { + state.waker = Some(cx.waker().clone()); + } + frame.map(|frame| frame.map(|frame| frame.map_err(io::Error::other))) } else { - state.waker = Some(cx.waker().clone()); + Poll::Ready(Some(Err(io::Error::new( + io::ErrorKind::Interrupted, + "request upload canceled", + )))) } - frame.map(|frame| frame.map(|frame| frame.map_err(io::Error::other))) - } else { - Poll::Ready(Some(Err(io::Error::new( - io::ErrorKind::Interrupted, - "request upload canceled", - )))) } }; + if matches!(frame, Poll::Ready(Some(Err(_)))) { + cancel(&self.state); + } if matches!(frame, Poll::Ready(None) | Poll::Ready(Some(Err(_)))) || self.is_end_stream() { self.permit.take(); } @@ -125,16 +134,17 @@ impl Body for RequestBody { } fn is_end_stream(&self) -> bool { - self.state.lock().unwrap().complete + self.state + .lock() + .map(|state| state.complete) + .unwrap_or(false) } fn size_hint(&self) -> SizeHint { self.state .lock() - .unwrap() - .inner - .as_ref() - .map(Body::size_hint) + .ok() + .and_then(|state| state.inner.as_ref().map(Body::size_hint)) .unwrap_or_default() } } From 29fd1322ad307c6ba407f17e53975ff7c58287a4 Mon Sep 17 00:00:00 2001 From: peg Date: Tue, 15 Sep 2026 08:39:57 +0200 Subject: [PATCH 6/7] Validate max inflight requests value following review --- README.md | 2 +- src/main.rs | 55 ++++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 1b843a0..0d03647 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ These are the attestation type names used in the HTTP headers, and the measureme - `client` and `server` accept `--listen-addr-healthcheck` to start a separate HTTP health-check listener. - `client --request-timeout-secs` sets the deadline from receipt of request headers through queueing, upload, and receipt of response headers (default: 60 seconds). Expired requests receive HTTP 504 and are not retried. If a response has already started, an unfinished upload is canceled at the deadline; its response status cannot be changed. Response bodies can continue streaming after that deadline once the upload completes. - `client --response-body-idle-timeout-secs` closes a source connection if its active response body makes no socket write progress for this interval (default: 60 seconds). This covers silent backends and clients that stop reading, even when body polling is blocked. The affected request releases its capacity; HTTP/1.1 reconnects upstream, while other HTTP/2 streams remain usable. Responses already started are truncated rather than replaced with a 504. Streams that keep making progress may continue indefinitely. -- `client --max-in-flight-requests` limits admitted requests, including streaming responses (default: 64). HTTP/2 requests run concurrently; HTTP/1.1 uses one request at a time and reconnects after a timeout or cancellation. Requests waiting for capacity are subject to the same deadline. Library callers can set these limits with `ProxyClient::with_request_options` and `ProxyClientOptions`. +- `client --max-in-flight-requests` limits admitted requests, including streaming responses (default: 64). HTTP/2 requests run concurrently; HTTP/1.1 uses one request at a time and reconnects after a timeout or cancellation. Requests waiting for capacity are subject to the same deadline. - `get-tls-cert --out-measurements ` writes the verified remote measurements as JSON in addition to writing the certificate chain to standard output. - `attested-get` does not follow HTTP redirects and exits with an error on 3xx responses. Its library API returns the original response for inspection. The loopback request ignores environment proxy settings. - If `server` is started without `--tls-private-key-path` and `--tls-certificate-path`, it generates a self-signed certificate for its listening IP address. diff --git a/src/main.rs b/src/main.rs index 89538ad..f4d5283 100644 --- a/src/main.rs +++ b/src/main.rs @@ -75,7 +75,7 @@ enum CliCommand { #[arg(long, default_value = "60")] response_body_idle_timeout_secs: NonZeroU64, /// Maximum in-flight requests, including streaming responses - #[arg(long, default_value = "64")] + #[arg(long, default_value = "64", value_parser = parse_max_in_flight_requests)] max_in_flight_requests: NonZeroUsize, /// Type of attestation to present (dafaults to 'auto' for automatic detection) /// If other than None, a TLS key and certicate must also be given @@ -544,3 +544,56 @@ fn certs_to_pem_string(certs: &[CertificateDer<'_>]) -> Result Result { + let count = value + .parse::() + .map_err(|error| error.to_string())?; + if count.get() > tokio::sync::Semaphore::MAX_PERMITS { + return Err(format!( + "must not exceed {}", + tokio::sync::Semaphore::MAX_PERMITS, + )); + } + Ok(count) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Checks that CLI parsing rejects invalid limits before starting the proxy. + #[test] + fn max_in_flight_requests_validates_semaphore_limit() { + for count in [ + 0, + 1, + tokio::sync::Semaphore::MAX_PERMITS, + tokio::sync::Semaphore::MAX_PERMITS + 1, + ] { + let result = Cli::try_parse_from([ + "attested-tls-proxy", + "client", + "localhost:443", + "--max-in-flight-requests", + &count.to_string(), + ]); + if (1..=tokio::sync::Semaphore::MAX_PERMITS).contains(&count) { + let CliCommand::Client { + max_in_flight_requests, + .. + } = result.unwrap().command + else { + panic!("expected client command"); + }; + assert_eq!(max_in_flight_requests.get(), count); + } else { + assert_eq!( + result.unwrap_err().kind(), + clap::error::ErrorKind::ValueValidation + ); + } + } + } +} From f25feef911d82a3cf5ef3aeb41f85e6b9a369849 Mon Sep 17 00:00:00 2001 From: peg Date: Wed, 23 Sep 2026 10:28:19 +0200 Subject: [PATCH 7/7] On graceful shutdown of HTTP2 connection, drain existing requests --- src/client_request/tests.rs | 187 ++++++++++++++++++++++++++++++++++++ src/lib.rs | 22 +++-- 2 files changed, 203 insertions(+), 6 deletions(-) diff --git a/src/client_request/tests.rs b/src/client_request/tests.rs index 31df35c..868a503 100644 --- a/src/client_request/tests.rs +++ b/src/client_request/tests.rs @@ -1085,3 +1085,190 @@ async fn flushed_responses_leave_source_keep_alive() { } } } + +/// Selects how the accepted response ends after its connection is retired. +#[derive(Clone, Copy)] +enum DrainEnd { + Complete, + Idle, + Disconnect, + Empty, +} + +/// Exercises real proxy reconnection while an upstream HTTP/2 connection drains. +async fn check_http2_drain(end: DrainEnd) { + let mut tasks = JoinSet::new(); + let (certs, key) = generate_certificate_chain("127.0.0.1".parse().unwrap()); + let (mut server_config, mut client_config) = generate_tls_config(certs.clone(), key); + server_config.alpn_protocols = vec![ALPN_H2.to_vec()]; + client_config.alpn_protocols = vec![ALPN_H2.to_vec()]; + let server = crate::AttestedTlsServer::new_with_tls_config( + certs, + server_config, + AttestationGenerator::with_no_attestation(), + AttestationVerifier::expect_none(), + ) + .unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let target = listener.local_addr().unwrap(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let (retired_tx, retired_rx) = tokio::sync::oneshot::channel(); + let (finish_tx, finish_rx) = tokio::sync::oneshot::channel(); + let (closed_tx, closed_rx) = tokio::sync::oneshot::channel(); + tasks.spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let (tls, _, _) = server.handle_connection(socket).await.unwrap(); + let mut connection = h2::server::handshake(tls).await.unwrap(); + let mut ping = connection.ping_pong().unwrap(); + let (request, mut respond) = connection.accept().await.unwrap().unwrap(); + assert_eq!(request.uri().path(), "/old"); + let response = respond + .send_response(http::Response::new(()), matches!(end, DrainEnd::Empty)) + .unwrap(); + drop(respond); + drop(request); + let mut streams = JoinSet::new(); + streams.spawn(async move { + let mut response = response; + finish_rx.await.unwrap(); + match end { + DrainEnd::Complete => response + .send_data(bytes::Bytes::from_static(b"complete"), true) + .unwrap(), + DrainEnd::Idle => { + let _ = std::future::poll_fn(|cx| response.poll_reset(cx)).await; + } + DrainEnd::Empty | DrainEnd::Disconnect => {} + } + }); + tokio::select! { + _ = shutdown_rx => {}, + _ = connection.accept() => panic!("unexpected request before GOAWAY"), + } + connection.graceful_shutdown(); + let mut drivers = JoinSet::new(); + drivers.spawn(async move { + while let Some(Ok(request)) = connection.accept().await { + panic!("unexpected request on retired connection: {request:?}"); + } + }); + // A pong confirms the client processed preceding GOAWAY frames. + // An empty connection can close before replying, which also retires it. + let _ = ping.ping(h2::Ping::opaque()).await; + retired_tx.send(()).unwrap(); + let (socket, _) = listener.accept().await.unwrap(); + let (tls, _, _) = server.handle_connection(socket).await.unwrap(); + let mut replacement = h2::server::handshake(tls).await.unwrap(); + let (request, mut respond) = replacement.accept().await.unwrap().unwrap(); + assert_eq!(request.uri().path(), "/fresh"); + respond + .send_response(http::Response::new(()), true) + .unwrap(); + drop(respond); + drop(request); + let mut replacements = JoinSet::new(); + replacements.spawn(async move { while replacement.accept().await.is_some() {} }); + while streams.join_next().await.is_some() {} + if matches!(end, DrainEnd::Disconnect) { + drivers.abort_all(); + } + drivers.join_next().await.unwrap().unwrap_or_else(|error| { + assert!(matches!(end, DrainEnd::Disconnect) && error.is_cancelled()); + }); + closed_tx.send(()).unwrap(); + while replacements.join_next().await.is_some() {} + }); + let client = Arc::new( + ProxyClient::new_with_tls_config( + client_config, + "127.0.0.1:0", + target.to_string(), + AttestationGenerator::with_no_attestation(), + AttestationVerifier::expect_none(), + None, + ) + .await + .unwrap() + .with_request_options(ProxyClientOptions { + request_timeout: Duration::from_secs(5), + response_body_idle_timeout: Duration::from_millis(800), + max_in_flight_requests: 2.try_into().unwrap(), + }), + ); + let url = format!("http://{}", client.local_addr().unwrap()); + let acceptor = client.clone(); + tasks.spawn(async move { + loop { + acceptor.accept().await.unwrap(); + } + }); + let old = http_client() + .get(format!("{url}/old")) + .send() + .await + .unwrap(); + assert_eq!(old.status(), http::StatusCode::OK); + shutdown_tx.send(()).unwrap(); + timeout(Duration::from_secs(2), retired_rx) + .await + .unwrap() + .unwrap(); + // Discovering GOAWAY may fail this request; it must never reach the peer or be replayed. + let probe = http_client() + .get(format!("{url}/fresh")) + .send() + .await + .unwrap(); + if probe.status() == http::StatusCode::BAD_GATEWAY { + let fresh = http_client() + .get(format!("{url}/fresh")) + .send() + .await + .unwrap(); + assert_eq!(fresh.status(), http::StatusCode::OK); + } else { + assert_eq!(probe.status(), http::StatusCode::OK); + } + finish_tx.send(()).unwrap(); + let result = old.text().await; + match end { + DrainEnd::Complete => assert_eq!(result.unwrap(), "complete"), + DrainEnd::Empty => assert_eq!(result.unwrap(), ""), + DrainEnd::Idle | DrainEnd::Disconnect => assert!(result.is_err()), + } + timeout(Duration::from_secs(2), closed_rx) + .await + .unwrap() + .unwrap(); + timeout(Duration::from_secs(2), async { + while client.request_slots.available_permits() != 2 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); +} + +/// Checks that a fresh connection serves requests before an accepted response finishes. +#[tokio::test] +async fn http2_goaway_drains_accepted_response() { + check_http2_drain(DrainEnd::Complete).await; +} + +/// Checks that retiring a connection does not disable response idle timeouts. +#[tokio::test] +async fn http2_goaway_drain_respects_idle_timeout() { + check_http2_drain(DrainEnd::Idle).await; +} + +/// Checks that abrupt failure of a draining connection releases request capacity. +#[tokio::test] +async fn http2_goaway_drain_disconnect_releases_capacity() { + check_http2_drain(DrainEnd::Disconnect).await; +} + +/// Checks that retirement without an active response releases the old connection. +#[tokio::test] +async fn http2_goaway_empty_connection_finishes() { + check_http2_drain(DrainEnd::Empty).await; +} diff --git a/src/lib.rs b/src/lib.rs index 574614a..9e2db67 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -454,8 +454,8 @@ impl ProxyClient { tokio::spawn(async move { let mut first = true; let mut ready_tx = Some(ready_tx); - // Retired connections may still have complete responses waiting to be - // delivered. Drain their workers without blocking a fresh connection. + // Retired HTTP/2 connections and their workers drain in the background + // without blocking a fresh connection. let mut draining = tokio::task::JoinSet::new(); let mut deferred = None; 'reconnect: loop { @@ -486,8 +486,9 @@ impl ProxyClient { } }; - // The connection driver is stopped on reconnect. Request workers - // retain their own deadlines and connection-specific measurements. + // HTTP/2 drivers survive retirement so accepted streams can finish. + // Workers retain their deadlines and connection-specific measurements. + let http2 = matches!(sender, HttpSender::Http2(_)); let mut connection = tokio::task::JoinSet::new(); connection.spawn(conn); let mut in_flight = tokio::task::JoinSet::new(); @@ -530,8 +531,17 @@ impl ProxyClient { _ = draining.join_next(), if !draining.is_empty() => {} } } - if !in_flight.is_empty() { - draining.spawn(async move { while in_flight.join_next().await.is_some() {} }); + drop(sender); + if !http2 { + connection.abort_all(); + } + if !in_flight.is_empty() || !connection.is_empty() { + draining.spawn(async move { + tokio::join!( + async { while in_flight.join_next().await.is_some() {} }, + async { while connection.join_next().await.is_some() {} }, + ); + }); } } });