Handle source connections concurrently

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/mod.rs

@@ -0,0 +1,276 @@
+//! Per-request forwarding, deadlines, and response lifetime tracking.
+pub(crate) mod response_idle;
+#[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 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,
+}
+
+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(),
+        }
+    }
+}
+
+pub(crate) type ProxyResponse =
+    Response<http_body_util::combinators::BoxBody<bytes::Bytes, hyper::Error>>;
+
+pub(crate) struct PendingRequest {
+    pub request: http::Request<Incoming>,
+    pub response_tx: oneshot::Sender<ProxyResponse>,
+    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,
+}
+
+/// 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<HttpSender>) -> Option<HttpSender> {
+    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<HttpSender>,
+    result: Result<ForwardResult, tokio::task::JoinError>,
+) -> 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,
+    measurements: Option<MultiMeasurements>,
+    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();
+    // Never forward measurements supplied by the target service.
+    headers.remove(MEASUREMENT_HEADER);
+    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<Arc<OwnedSemaphorePermit>>,
+        finished: Option<oneshot::Sender<()>>,
+    }
+}
+
+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<Option<Result<Frame<Self::Data>, 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()
+    }
+}