Handle source connections concurrently

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"

README.md

@@ -77,10 +77,15 @@ 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 --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.
 - `get-tls-cert --out-measurements <PATH>` 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.
 
+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.

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<Box<dyn Future<Output = Result<(), ProxyError>> + Send>>;
+
+#[derive(Clone)]
+pub(crate) struct Sender {
+    inner: h2::client::SendRequest<Bytes>,
+    closed: Arc<AtomicBool>,
+}
+
+pub(crate) async fn handshake<T>(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<AtomicBool>);
+    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<RequestBody>,
+    ) -> Result<ProxyResponse, ProxyError> {
+        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<Option<Result<Frame<Bytes>, 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()
+    }
+}