diff --git a/README.md b/README.md index 562ea79..0d03647 100644 --- a/README.md +++ b/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 ` 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. diff --git a/src/client_request/mod.rs b/src/client_request/mod.rs new file mode 100644 index 0000000..7ad5fef --- /dev/null +++ b/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>; + +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, +} + +/// 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, + 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(); + // 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>, + 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/response_idle.rs b/src/client_request/response_idle.rs new file mode 100644 index 0000000..b88db0b --- /dev/null +++ b/src/client_request/response_idle.rs @@ -0,0 +1,209 @@ +//! 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, +}; + +/// Shares response progress between the body, socket, and idle timer. +#[derive(Clone)] +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, +) -> (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), + ) +} + +/// 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().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(|active| Instant::now() >= active.last_write + timeout) { + return; + } + } + } + } else if activity.changed().await.is_err() { + return; + } + } +} + +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(ActiveResponse { + last_write: Instant::now(), + body_finished: false, + })); + response.map(|inner| { + let mut body = IdleBody { + inner, + activity: Some(self.clone()), + }; + if body.inner.is_end_stream() { + body.finish(); + } + body.boxed() + }) + } + + /// 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(|active| { + if let Some(active) = active { + active.last_write = Instant::now(); + true + } else { + false + } + }); + } + } +} + +/// Records socket write progress and clears idle tracking after the final flush. +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() + } + /// Disarms the idle timer once a completed response has been flushed. + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + 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() { + // 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; + } + }); + } + } +} + +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 new file mode 100644 index 0000000..0d073cb --- /dev/null +++ b/src/client_request/tests.rs @@ -0,0 +1,815 @@ +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 { + 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(); + 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(), + response_body_idle_timeout, + }); + 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() +} + +// 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()); + 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); +} + +#[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); + } +} + +/// 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" + ); + } + } +} diff --git a/src/client_request/upload.rs b/src/client_request/upload.rs new file mode 100644 index 0000000..0f70828 --- /dev/null +++ b/src/client_request/upload.rs @@ -0,0 +1,150 @@ +//! 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 = { + // 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() + }; + 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 = 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 { + 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(); + } + frame + } + + fn is_end_stream(&self) -> bool { + self.state + .lock() + .map(|state| state.complete) + .unwrap_or(false) + } + + fn size_hint(&self) -> SizeHint { + self.state + .lock() + .ok() + .and_then(|state| state.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 219610a..f85f99f 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, take_sender, worker_finished}; #[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] { @@ -360,10 +358,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, @@ -439,12 +446,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::>(); @@ -452,8 +454,12 @@ 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(); + let mut deferred = None; '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 @@ -480,81 +486,52 @@ 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(); - // Never forward measurements supplied by the target service. - headers.remove(MEASUREMENT_HEADER); - 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) - } - }; - - // 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"); + 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 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, + pending, + measurements.clone(), + remote_attestation_type, + )); + } + result = in_flight.join_next(), if !in_flight.is_empty() => { + if let Some(result) = result + && worker_finished(&mut sender, result) { + 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() {} }); } } }); @@ -563,6 +540,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()), @@ -579,9 +560,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}"); } }); @@ -592,31 +577,44 @@ 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"); + 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).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(()) } @@ -687,13 +685,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()) } @@ -706,11 +704,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()), + } } } @@ -764,8 +783,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 884ca23..f4d5283 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,15 @@ 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, + /// 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", 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 #[arg(long, env = "CLIENT_ATTESTATION_TYPE")] @@ -245,6 +256,9 @@ async fn main() -> anyhow::Result<()> { CliCommand::Client { listen_addr, target_addr, + request_timeout_secs, + response_body_idle_timeout_secs, + max_in_flight_requests, client_attestation_type, tls_private_key_path, tls_certificate_path, @@ -311,7 +325,14 @@ async fn main() -> anyhow::Result<()> { remote_tls_cert, ) .await? - }; + } + .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, + }); loop { if let Err(err) = client.accept().await { @@ -523,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 + ); + } + } + } +}