From 93fd7260931dc0ea9a08e572e49afe65ad27b4c4 Mon Sep 17 00:00:00 2001 From: fannnzhang Date: Sat, 11 Jul 2026 16:57:33 +0800 Subject: [PATCH 1/2] fix: harden connection lifecycle, connect budgets, and follow-up paths Close production-readiness gaps from the deep review: pool eviction now aborts owned hyper tasks and clears bindings, connect_timeout covers TLS and protocol binding, HTTP/2 temporary unreadiness no longer parks on the wrong wait signal, follow-ups drain intermediate bodies and strip body headers on method-changing redirects, decompression failures force connection discard, and WebSocket engines honor size limits without request-path panics. --- crates/openwire-core/src/body.rs | 2 +- crates/openwire-core/src/context.rs | 13 + crates/openwire-tungstenite/src/lib.rs | 5 +- crates/openwire/src/bridge.rs | 3 + crates/openwire/src/compression.rs | 33 ++- .../openwire/src/connection/fast_fallback.rs | 56 +++- crates/openwire/src/connection/mod.rs | 3 +- crates/openwire/src/connection/pool.rs | 245 ++++++++++++------ crates/openwire/src/policy/follow_up.rs | 71 ++++- crates/openwire/src/transport/bindings.rs | 73 ++++-- crates/openwire/src/transport/body.rs | 39 ++- crates/openwire/src/transport/connect.rs | 14 +- crates/openwire/src/transport/service.rs | 138 ++++++++-- crates/openwire/src/transport/tests.rs | 4 +- crates/openwire/src/websocket/handshake.rs | 19 +- .../openwire/src/websocket/native/engine.rs | 2 +- crates/openwire/src/websocket/native/mask.rs | 15 +- crates/openwire/tests/integration.rs | 5 +- docs/ARCHITECTURE.md | 16 ++ 19 files changed, 579 insertions(+), 177 deletions(-) diff --git a/crates/openwire-core/src/body.rs b/crates/openwire-core/src/body.rs index 776092a..a42bac2 100644 --- a/crates/openwire-core/src/body.rs +++ b/crates/openwire-core/src/body.rs @@ -141,7 +141,7 @@ impl Body for RequestBody { fn is_end_stream(&self) -> bool { match &self.inner { RequestBodyInner::Empty => true, - RequestBodyInner::Replayable { emitted, .. } => *emitted, + RequestBodyInner::Replayable { bytes, emitted } => *emitted || bytes.is_empty(), RequestBodyInner::Streaming { inner } => inner.is_end_stream(), } } diff --git a/crates/openwire-core/src/context.rs b/crates/openwire-core/src/context.rs index 4b8a3f5..d228e91 100644 --- a/crates/openwire-core/src/context.rs +++ b/crates/openwire-core/src/context.rs @@ -53,6 +53,8 @@ struct CallContextInner { created_at: Instant, deadline: Option, connection_established: AtomicBool, + /// When set, response-body Drop discards the connection (decode/body errors). + body_force_discard: AtomicBool, tls_alpn_preference: TlsAlpnPreference, } @@ -86,6 +88,7 @@ impl CallContext { created_at, deadline, connection_established: AtomicBool::new(false), + body_force_discard: AtomicBool::new(false), tls_alpn_preference, }), } @@ -113,6 +116,16 @@ impl CallContext { &self.inner.listener } + pub fn mark_body_force_discard(&self) { + self.inner + .body_force_discard + .store(true, Ordering::Release); + } + + pub fn body_force_discard(&self) -> bool { + self.inner.body_force_discard.load(Ordering::Acquire) + } + pub fn created_at(&self) -> Instant { self.inner.created_at } diff --git a/crates/openwire-tungstenite/src/lib.rs b/crates/openwire-tungstenite/src/lib.rs index 72247c6..93b4130 100644 --- a/crates/openwire-tungstenite/src/lib.rs +++ b/crates/openwire-tungstenite/src/lib.rs @@ -73,10 +73,13 @@ impl WebSocketEngine for TungsteniteEngine { // BoxConnection (hyper::rt::Read+Write) → tokio AsyncRead+Write. let tokio_io = TokioIo::new(io); + let mut ws_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(); + ws_config.max_message_size = Some(config.max_message_size); + ws_config.max_frame_size = Some(config.max_frame_size); let stream = WebSocketStream::from_raw_socket( tokio_io, tokio_tungstenite::tungstenite::protocol::Role::Client, - None, + Some(ws_config), ) .await; diff --git a/crates/openwire/src/bridge.rs b/crates/openwire/src/bridge.rs index bbe0ac7..a82cab7 100644 --- a/crates/openwire/src/bridge.rs +++ b/crates/openwire/src/bridge.rs @@ -42,6 +42,8 @@ impl Interceptor for BridgeInterceptor { let transparent_compression = normalization.as_ref().copied().unwrap_or(false); Box::pin(async move { normalization?; + #[cfg(feature = "compression")] + let ctx = exchange.context().clone(); let response = next.run(exchange).await?; #[cfg(feature = "compression")] { @@ -50,6 +52,7 @@ impl Interceptor for BridgeInterceptor { response, &request_method, max_decompressed_body_bytes, + ctx, )); } } diff --git a/crates/openwire/src/compression.rs b/crates/openwire/src/compression.rs index 2b12771..06cb63b 100644 --- a/crates/openwire/src/compression.rs +++ b/crates/openwire/src/compression.rs @@ -12,7 +12,7 @@ use http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH, RANGE}; use http::{HeaderMap, HeaderValue, Method, Response, StatusCode}; use http_body::{Body, Frame, SizeHint}; use http_body_util::BodyExt; -use openwire_core::{RequestBody, ResponseBody, WireError}; +use openwire_core::{CallContext, RequestBody, ResponseBody, WireError}; use pin_project_lite::pin_project; const ACCEPTED_ENCODINGS: HeaderValue = HeaderValue::from_static("br, gzip, deflate, zstd"); @@ -42,6 +42,7 @@ pub(crate) fn decode_response( response: Response, request_method: &Method, max_decompressed_body_bytes: usize, + ctx: CallContext, ) -> Response { if !response_can_have_body(request_method, response.status()) { return response; @@ -62,7 +63,7 @@ pub(crate) fn decode_response( .map(|encoding| encoding.as_str()) .collect::>() .join(", "); - let body = DecodedResponseBody::new(body, encodings, label, max_decompressed_body_bytes); + let body = DecodedResponseBody::new(body, encodings, label, max_decompressed_body_bytes, ctx); Response::from_parts(parts, ResponseBody::new(body.boxed())) } @@ -150,6 +151,7 @@ pin_project! { label: String, max_decompressed_body_bytes: usize, decoded_bytes: usize, + ctx: CallContext, } } @@ -159,6 +161,7 @@ impl DecodedResponseBody { encodings: Vec, label: String, max_decompressed_body_bytes: usize, + ctx: CallContext, ) -> Self { let stream = body.into_data_stream().map_err(wire_error_to_io); let reader = stream.into_async_read(); @@ -173,6 +176,7 @@ impl DecodedResponseBody { label, max_decompressed_body_bytes, decoded_bytes: 0, + ctx, } } } @@ -191,6 +195,7 @@ impl Body for DecodedResponseBody { Poll::Ready(Ok(0)) => Poll::Ready(None), Poll::Ready(Ok(read)) => { let Some(total) = this.decoded_bytes.checked_add(read) else { + this.ctx.mark_body_force_discard(); return Poll::Ready(Some(Err(WireError::body( format!( "decompressed {} response exceeded size limit {}", @@ -200,6 +205,7 @@ impl Body for DecodedResponseBody { )))); }; if total > *this.max_decompressed_body_bytes { + this.ctx.mark_body_force_discard(); return Poll::Ready(Some(Err(WireError::body( format!( "decompressed {} response exceeded size limit {}", @@ -214,6 +220,7 @@ impl Body for DecodedResponseBody { ))))) } Poll::Ready(Err(error)) => { + this.ctx.mark_body_force_discard(); Poll::Ready(Some(Err(io_error_to_wire(error, this.label.as_str())))) } Poll::Pending => Poll::Pending, @@ -479,7 +486,16 @@ mod tests { .body(ResponseBody::empty()) .expect("response"); - let response = decode_response(response, &Method::GET, DEFAULT_MAX_DECOMPRESSED_BODY_BYTES); + let response = decode_response( + response, + &Method::GET, + DEFAULT_MAX_DECOMPRESSED_BODY_BYTES, + openwire_core::CallContext::new( + std::sync::Arc::new(openwire_core::NoopEventListener) + as openwire_core::SharedEventListener, + None, + ), + ); assert!(response.headers().get(CONTENT_ENCODING).is_none()); assert!(response.headers().get(CONTENT_LENGTH).is_none()); @@ -493,7 +509,16 @@ mod tests { .body(ResponseBody::empty()) .expect("response"); - let response = decode_response(response, &Method::GET, DEFAULT_MAX_DECOMPRESSED_BODY_BYTES); + let response = decode_response( + response, + &Method::GET, + DEFAULT_MAX_DECOMPRESSED_BODY_BYTES, + openwire_core::CallContext::new( + std::sync::Arc::new(openwire_core::NoopEventListener) + as openwire_core::SharedEventListener, + None, + ), + ); assert_eq!(response.headers().get(CONTENT_ENCODING).unwrap(), "made-up"); assert_eq!(response.headers().get(CONTENT_LENGTH).unwrap(), "20"); diff --git a/crates/openwire/src/connection/fast_fallback.rs b/crates/openwire/src/connection/fast_fallback.rs index ef42dd9..27e5c4a 100644 --- a/crates/openwire/src/connection/fast_fallback.rs +++ b/crates/openwire/src/connection/fast_fallback.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use std::time::Duration; use futures_channel::mpsc; -use futures_util::future::{AbortHandle, Abortable}; +use futures_util::future::{select, AbortHandle, Abortable, Either}; use futures_util::stream::StreamExt; use hyper::rt::Timer; use hyper::Uri; @@ -376,6 +376,8 @@ impl FastFallbackDialer { route_plan: RoutePlan, deps: DirectDialDeps, ) -> Result<(BoxConnection, FastFallbackOutcome), WireError> { + let connect_timeout = deps.connect_timeout; + let timer = deps.runtime.timer.clone(); self.dial_route_plan( ctx, uri, @@ -398,7 +400,18 @@ impl FastFallbackDialer { }, move |ctx, uri, _route, stream| { let tls_connector = deps.tls_connector.clone(); - async move { finalize_direct_connection(ctx, uri, stream, tls_connector).await } + let timer = timer.clone(); + async move { + finalize_direct_connection( + ctx, + uri, + stream, + tls_connector, + timer, + connect_timeout, + ) + .await + } }, ) .await @@ -410,6 +423,8 @@ async fn finalize_direct_connection( uri: Uri, stream: BoxConnection, tls_connector: Option>, + timer: SharedTimer, + connect_timeout: Option, ) -> Result { if !uri .scheme_str() @@ -425,7 +440,42 @@ async fn finalize_direct_connection( )); }; - tls_connector.connect(ctx, uri, stream).await + with_connect_timeout( + timer, + connect_timeout, + tls_connector.connect(ctx, uri, stream), + "TLS handshake", + ) + .await +} + +/// Applies `connect_timeout` to a connect-stage future (TLS or protocol bind). +pub(crate) async fn with_connect_timeout( + timer: SharedTimer, + connect_timeout: Option, + future: F, + stage: &'static str, +) -> Result +where + F: Future>, +{ + let Some(timeout) = connect_timeout else { + return future.await; + }; + if timeout.is_zero() { + return Err(WireError::connect_timeout(format!( + "connect timed out before {stage}" + ))); + } + + let future = Box::pin(future); + let sleep = timer.sleep(timeout); + match select(future, sleep).await { + Either::Left((result, _sleep)) => result, + Either::Right((_ready, _future)) => Err(WireError::connect_timeout(format!( + "connect timed out after {timeout:?} during {stage}" + ))), + } } fn failure_stage(error: &WireError) -> ConnectFailureStage { diff --git a/crates/openwire/src/connection/mod.rs b/crates/openwire/src/connection/mod.rs index fd03d5e..d75ab4e 100644 --- a/crates/openwire/src/connection/mod.rs +++ b/crates/openwire/src/connection/mod.rs @@ -12,7 +12,8 @@ pub(crate) use exchange_finder::{ ResolvedAddress, }; pub(crate) use fast_fallback::{ - DirectDialDeps, FastFallbackDialer, FastFallbackOutcome, FastFallbackRuntime, + with_connect_timeout, DirectDialDeps, FastFallbackDialer, FastFallbackOutcome, + FastFallbackRuntime, }; pub(crate) use limits::{ ConnectionAvailability, ConnectionLimiter, ConnectionPermit, RequestAdmissionLimiter, diff --git a/crates/openwire/src/connection/pool.rs b/crates/openwire/src/connection/pool.rs index daa7e2e..bdf0f2c 100644 --- a/crates/openwire/src/connection/pool.rs +++ b/crates/openwire/src/connection/pool.rs @@ -1,6 +1,6 @@ use std::collections::{HashMap, HashSet}; use std::net::SocketAddr; -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; use std::time::Duration; use openwire_core::ConnectionId; @@ -11,6 +11,10 @@ use super::{ }; use crate::sync_util::lock_mutex; +/// Invoked after a connection is removed from pool metadata so transport can +/// abort the owned hyper task and drop bindings. Must not re-enter the pool. +pub(crate) type PoolEvictionHook = Arc; + #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct PoolSettings { pub(crate) idle_timeout: Option, @@ -53,6 +57,7 @@ pub(crate) struct PoolStats { pub(crate) struct ConnectionPool { settings: PoolSettings, state: Mutex, + eviction_hook: Mutex>, } #[derive(Debug, Default)] @@ -67,35 +72,46 @@ impl ConnectionPool { Self { settings, state: Mutex::new(PoolState::default()), + eviction_hook: Mutex::new(None), } } + /// Installs a hook that runs after a connection leaves the pool. + /// + /// Transport uses this to abort hyper connection tasks and clear bindings + /// so idle eviction actually closes sockets. + pub(crate) fn set_eviction_hook(&self, hook: PoolEvictionHook) { + *lock_mutex(&self.eviction_hook) = Some(hook); + } + pub(crate) fn settings(&self) -> &PoolSettings { &self.settings } pub(crate) fn insert(&self, connection: RealConnection) { let address = connection.address().clone(); - { - let mut state = lock_mutex(&self.state); - state - .by_address - .entry(address.clone()) - .or_default() - .push(connection.clone()); - register_connection(&mut state, &connection); - prune_address(&self.settings, &mut state, &address); - } + let mut state = lock_mutex(&self.state); + state + .by_address + .entry(address.clone()) + .or_default() + .push(connection.clone()); + register_connection(&mut state, &connection); + let evicted = prune_address(&self.settings, &mut state, &address); + drop(state); + self.notify_evictions(evicted); } pub(crate) fn acquire(&self, address: &Address) -> Option { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return None; - } - - let connections = state.by_address.get_mut(address)?; - connections.iter().find(|conn| conn.try_acquire()).cloned() + let evicted = prune_address(&self.settings, &mut state, address); + let result = state + .by_address + .get_mut(address) + .and_then(|connections| connections.iter().find(|conn| conn.try_acquire()).cloned()); + drop(state); + self.notify_evictions(evicted); + result } pub(crate) fn acquire_with_in_use_hint( @@ -103,31 +119,33 @@ impl ConnectionPool { address: &Address, ) -> (Option, bool) { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return (None, false); - } - - let Some(connections) = state.by_address.get_mut(address) else { - return (None, false); + let evicted = prune_address(&self.settings, &mut state, address); + let result = match state.by_address.get_mut(address) { + Some(connections) => { + let connection = connections.iter().find(|conn| conn.try_acquire()).cloned(); + let has_in_use = has_in_use_connection_unpruned( + connections, + connection.as_ref().map(RealConnection::id), + ); + (connection, has_in_use) + } + None => (None, false), }; - let connection = connections.iter().find(|conn| conn.try_acquire()).cloned(); - let has_in_use = has_in_use_connection_unpruned( - connections, - connection.as_ref().map(RealConnection::id), - ); - (connection, has_in_use) + drop(state); + self.notify_evictions(evicted); + result } pub(crate) fn has_in_use_connection(&self, address: &Address) -> bool { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return false; - } - - state + let evicted = prune_address(&self.settings, &mut state, address); + let result = state .by_address .get(address) - .is_some_and(|connections| has_in_use_connection_unpruned(connections, None)) + .is_some_and(|connections| has_in_use_connection_unpruned(connections, None)); + drop(state); + self.notify_evictions(evicted); + result } pub(crate) fn acquire_coalesced( @@ -153,8 +171,13 @@ impl ConnectionPool { } } + let mut evicted = Vec::new(); for candidate_address in addresses_to_prune { - prune_address(&self.settings, &mut state, &candidate_address); + evicted.extend(prune_address( + &self.settings, + &mut state, + &candidate_address, + )); } let mut candidates = Vec::new(); @@ -174,6 +197,7 @@ impl ConnectionPool { } drop(state); + self.notify_evictions(evicted); candidates .into_iter() @@ -187,8 +211,9 @@ impl ConnectionPool { let address = connection.address().clone(); let mut state = lock_mutex(&self.state); - prune_address(&self.settings, &mut state, &address); - + let evicted = prune_address(&self.settings, &mut state, &address); + drop(state); + self.notify_evictions(evicted); true } @@ -198,19 +223,20 @@ impl ConnectionPool { connection_id: ConnectionId, ) -> Option { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return None; - } - if state.by_id.get(&connection_id) != Some(address) { - return None; - } - - state.by_address.get(address).and_then(|connections| { - connections - .iter() - .find(|connection| connection.id() == connection_id) - .cloned() - }) + let evicted = prune_address(&self.settings, &mut state, address); + let result = if state.by_id.get(&connection_id) != Some(address) { + None + } else { + state.by_address.get(address).and_then(|connections| { + connections + .iter() + .find(|connection| connection.id() == connection_id) + .cloned() + }) + }; + drop(state); + self.notify_evictions(evicted); + result } pub(crate) fn acquire_by_id( @@ -219,21 +245,38 @@ impl ConnectionPool { connection_id: ConnectionId, ) -> Option { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return None; - } - if state.by_id.get(&connection_id) != Some(address) { - return None; - } - - let connections = state.by_address.get_mut(address)?; - connections.iter().find_map(|connection| { - (connection.id() == connection_id && connection.try_acquire()) - .then(|| connection.clone()) - }) + let evicted = prune_address(&self.settings, &mut state, address); + let result = if state.by_id.get(&connection_id) != Some(address) { + None + } else { + state.by_address.get_mut(address).and_then(|connections| { + connections.iter().find_map(|connection| { + (connection.id() == connection_id && connection.try_acquire()) + .then(|| connection.clone()) + }) + }) + }; + drop(state); + self.notify_evictions(evicted); + result } pub(crate) fn remove(&self, connection_id: ConnectionId) -> Option { + let removed = self.remove_without_hook(connection_id); + if removed.is_some() { + self.notify_evictions(vec![connection_id]); + } + removed + } + + /// Removes pool metadata without running the eviction hook. + /// + /// Used by transport teardown after it has already aborted the connection + /// task, so the hook does not re-enter abort/remove. + pub(crate) fn remove_without_hook( + &self, + connection_id: ConnectionId, + ) -> Option { let mut state = lock_mutex(&self.state); let address = state.by_id.remove(&connection_id)?; let mut removed = None; @@ -264,31 +307,48 @@ impl ConnectionPool { pub(crate) fn stats(&self, address: &Address) -> PoolStats { let mut state = lock_mutex(&self.state); - if !prune_address(&self.settings, &mut state, address) { - return PoolStats::default(); - } - let Some(connections) = state.by_address.get(address) else { - return PoolStats::default(); + let evicted = prune_address(&self.settings, &mut state, address); + let stats = match state.by_address.get(address) { + Some(connections) => connections.iter().fold( + PoolStats::default(), + |mut stats, connection| { + stats.total += 1; + match connection.snapshot().allocation { + ConnectionAllocationState::Idle => stats.idle += 1, + ConnectionAllocationState::InUse { .. } => stats.in_use += 1, + ConnectionAllocationState::Closed => {} + } + stats + }, + ), + None => PoolStats::default(), }; - - connections - .iter() - .fold(PoolStats::default(), |mut stats, connection| { - stats.total += 1; - match connection.snapshot().allocation { - ConnectionAllocationState::Idle => stats.idle += 1, - ConnectionAllocationState::InUse { .. } => stats.in_use += 1, - ConnectionAllocationState::Closed => {} - } - stats - }) + drop(state); + self.notify_evictions(evicted); + stats } pub(crate) fn prune_all(&self) { let mut state = lock_mutex(&self.state); let addresses = state.by_address.keys().cloned().collect::>(); + let mut evicted = Vec::new(); for address in addresses { - prune_address(&self.settings, &mut state, &address); + evicted.extend(prune_address(&self.settings, &mut state, &address)); + } + drop(state); + self.notify_evictions(evicted); + } + + fn notify_evictions(&self, connection_ids: Vec) { + if connection_ids.is_empty() { + return; + } + let hook = lock_mutex(&self.eviction_hook).clone(); + let Some(hook) = hook else { + return; + }; + for connection_id in connection_ids { + hook(connection_id); } } } @@ -301,26 +361,32 @@ impl std::fmt::Debug for ConnectionPool { } } -fn prune_address(settings: &PoolSettings, state: &mut PoolState, address: &Address) -> bool { +/// Returns connection ids closed and removed from the pool. +fn prune_address( + settings: &PoolSettings, + state: &mut PoolState, + address: &Address, +) -> Vec { let (removed, empty) = { let Some(connections) = state.by_address.get_mut(address) else { - return false; + return Vec::new(); }; let removed = prune_connections(settings, connections); (removed, connections.is_empty()) }; + let mut ids = Vec::with_capacity(removed.len()); for connection in &removed { unregister_connection(state, connection); + ids.push(connection.id()); } if empty { state.by_address.remove(address); - return false; } - true + ids } fn prune_connections( @@ -330,6 +396,15 @@ fn prune_connections( let mut removed = Vec::new(); connections.retain(|connection| { if !connection.is_healthy() { + // Keep unhealthy connections that still have live allocations so + // HTTP/2 multiplex bookkeeping can drain cleanly. They are not + // re-acquired because try_acquire requires Healthy. + if matches!( + connection.snapshot().allocation, + ConnectionAllocationState::InUse { .. } + ) { + return true; + } connection.close(); removed.push(connection.clone()); return false; diff --git a/crates/openwire/src/policy/follow_up.rs b/crates/openwire/src/policy/follow_up.rs index 72fef3c..d52941a 100644 --- a/crates/openwire/src/policy/follow_up.rs +++ b/crates/openwire/src/policy/follow_up.rs @@ -2,7 +2,8 @@ use std::task::{Context, Poll}; use std::time::SystemTime; use http::header::{ - AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE, COOKIE, HOST, LOCATION, RETRY_AFTER, SET_COOKIE, + AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, COOKIE, EXPECT, HOST, + LOCATION, RETRY_AFTER, SET_COOKIE, TRANSFER_ENCODING, }; use http::{HeaderMap, Method, Request, Response, StatusCode, Uri, Version}; use openwire_core::{ @@ -143,6 +144,7 @@ impl Service for FollowUpPolicyService { response_status = %response.status(), "following authentication challenge", ); + let _ = drain_intermediate_response(response).await; request = next_request; attempt = next_attempt; continue; @@ -173,7 +175,7 @@ impl Service for FollowUpPolicyService { "retrying request after HTTP/2 misdirected request", ); - drop(response); + let _ = drain_intermediate_response(response).await; request = snapshot.to_retry_request(policy_trace)?; request.extensions_mut().insert(NoCoalescedConnections); attempt = next_attempt; @@ -200,7 +202,7 @@ impl Service for FollowUpPolicyService { "retrying request after retryable response status", ); - drop(response); + let _ = drain_intermediate_response(response).await; request = snapshot.to_retry_request(policy_trace)?; attempt = next_attempt; continue; @@ -257,8 +259,9 @@ impl Service for FollowUpPolicyService { "following redirect", ); + let status = response.status(); let Some(next_request) = snapshot.into_redirect_request( - response.status(), + status, next_uri.clone(), policy_trace, selected_proxy, @@ -266,6 +269,7 @@ impl Service for FollowUpPolicyService { else { return Ok(response); }; + let _ = drain_intermediate_response(response).await; ctx.listener().redirect(&ctx, redirects + 1, &next_uri); request = next_request; redirects += 1; @@ -463,6 +467,29 @@ fn should_retry_misdirected_request( .is_some() } + +const INTERMEDIATE_BODY_DRAIN_LIMIT: u64 = 256 * 1024; + +async fn drain_intermediate_response( + response: Response, +) -> Result<(), WireError> { + let body = response.into_body(); + let mut body = std::pin::pin!(body); + let mut read = 0u64; + use http_body_util::BodyExt; + while let Some(frame) = body.frame().await { + let frame = frame?; + if let Some(data) = frame.data_ref() { + read = read.saturating_add(data.len() as u64); + if read > INTERMEDIATE_BODY_DRAIN_LIMIT { + // Oversized intermediate body: drop the rest and let lease discard. + break; + } + } + } + Ok(()) +} + fn store_response_cookies( response: &Response, request_uri: &Uri, @@ -587,6 +614,9 @@ impl RequestSnapshot { if should_switch_to_get { headers.remove(CONTENT_LENGTH); headers.remove(CONTENT_TYPE); + headers.remove(CONTENT_ENCODING); + headers.remove(TRANSFER_ENCODING); + headers.remove(EXPECT); } let mut request = Request::builder() @@ -922,6 +952,39 @@ mod tests { assert!(next.is_none()); } + + #[test] + fn redirect_to_get_strips_body_describing_headers() { + use http::header::{ + CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, EXPECT, TRANSFER_ENCODING, + }; + let mut request = Request::builder() + .method("POST") + .uri("http://source.test/start") + .header(CONTENT_TYPE, "application/json") + .header(CONTENT_ENCODING, "gzip") + .header(TRANSFER_ENCODING, "chunked") + .header(EXPECT, "100-continue") + .body(RequestBody::from_static(b"{}")) + .expect("request"); + let snapshot = RequestSnapshot::capture(&request); + let next = snapshot + .into_redirect_request( + StatusCode::FOUND, + "http://source.test/next".parse().expect("redirect uri"), + PolicyTraceContext::default(), + None, + ) + .expect("redirect request") + .expect("followable"); + assert_eq!(next.method(), http::Method::GET); + assert!(next.headers().get(CONTENT_TYPE).is_none()); + assert!(next.headers().get(CONTENT_ENCODING).is_none()); + assert!(next.headers().get(TRANSFER_ENCODING).is_none()); + assert!(next.headers().get(EXPECT).is_none()); + assert!(next.headers().get(CONTENT_LENGTH).is_none()); + } + struct ReadinessTrackingService { was_polled: bool, is_clone: bool, diff --git a/crates/openwire/src/transport/bindings.rs b/crates/openwire/src/transport/bindings.rs index 68a4b2d..fbbcef3 100644 --- a/crates/openwire/src/transport/bindings.rs +++ b/crates/openwire/src/transport/bindings.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, Weak}; use hyper::client::conn::{http1, http2}; @@ -104,8 +103,6 @@ impl ConnectionBindings { if binding.sender.is_closed() { remove_stale = true; BindingAcquireResult::Stale - } else if !binding.sender.is_ready() { - BindingAcquireResult::Busy } else { BindingAcquireResult::Acquired(AcquiredBinding::Http2 { info: binding.info.clone(), @@ -166,43 +163,55 @@ pub(super) struct ConnectionTaskRegistry { #[derive(Default)] pub(super) struct ConnectionTaskRegistryInner { - next_id: AtomicU64, - handles: Mutex>>, + handles_by_connection: Mutex>>, } impl ConnectionTaskRegistry { - pub(super) fn reserve(&self) -> (u64, Weak) { - let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed) + 1; - lock_mutex(&self.inner.handles).insert(id, None); - (id, Arc::downgrade(&self.inner)) - } - - pub(super) fn attach(&self, task_id: u64, handle: BoxTaskHandle) { - let mut handles = lock_mutex(&self.inner.handles); - if let Some(slot) = handles.get_mut(&task_id) { - *slot = Some(handle); - return; + pub(super) fn attach_connection(&self, connection_id: ConnectionId, handle: BoxTaskHandle) { + let mut handles = lock_mutex(&self.inner.handles_by_connection); + if let Some(previous) = handles.insert(connection_id, Some(handle)) { + if let Some(previous) = previous { + previous.abort(); + } } - drop(handles); - handle.abort(); } - pub(super) fn cancel(&self, task_id: u64) { - lock_mutex(&self.inner.handles).remove(&task_id); + pub(super) fn abort_connection(&self, connection_id: ConnectionId) { + if let Some(Some(handle)) = + lock_mutex(&self.inner.handles_by_connection).remove(&connection_id) + { + handle.abort(); + } } - pub(super) fn complete_weak(inner: &Weak, task_id: u64) { + pub(super) fn complete_connection_weak( + inner: &Weak, + connection_id: ConnectionId, + ) { let Some(inner) = inner.upgrade() else { return; }; - lock_mutex(&inner.handles).remove(&task_id); + lock_mutex(&inner.handles_by_connection).remove(&connection_id); + } + + pub(super) fn downgrade(&self) -> Weak { + Arc::downgrade(&self.inner) + } + + pub(super) fn teardown_connection( + &self, + bindings: &ConnectionBindings, + connection_id: ConnectionId, + ) { + bindings.remove(connection_id); + self.abort_connection(connection_id); } #[cfg(test)] pub(super) fn poison_handles_for_test(&self) { let _guard = self .inner - .handles + .handles_by_connection .lock() .expect("poison connection task registry lock for test"); panic!("poison connection task registry"); @@ -211,7 +220,7 @@ impl ConnectionTaskRegistry { impl Drop for ConnectionTaskRegistryInner { fn drop(&mut self) { - let handles = lock_mutex(&self.handles); + let handles = lock_mutex(&self.handles_by_connection); for handle in handles.values().filter_map(Option::as_ref) { handle.abort(); } @@ -221,6 +230,7 @@ impl Drop for ConnectionTaskRegistryInner { pub(super) fn release_acquired_connection( exchange_finder: &Arc, bindings: &Arc, + tasks: &ConnectionTaskRegistry, availability: &ConnectionAvailability, connection: RealConnection, binding: AcquiredBinding, @@ -233,8 +243,7 @@ pub(super) fn release_acquired_connection( availability.notify(); return; } - bindings.remove(connection.id()); - let _ = exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection(exchange_finder, bindings, tasks, connection.id()); availability.notify(); } AcquiredBinding::Http2 { .. } => { @@ -242,8 +251,18 @@ pub(super) fn release_acquired_connection( availability.notify(); return; } - let _ = exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection(exchange_finder, bindings, tasks, connection.id()); availability.notify(); } } } + +pub(super) fn teardown_pooled_connection( + exchange_finder: &Arc, + bindings: &ConnectionBindings, + tasks: &ConnectionTaskRegistry, + connection_id: ConnectionId, +) { + tasks.teardown_connection(bindings, connection_id); + let _ = exchange_finder.pool().remove_without_hook(connection_id); +} diff --git a/crates/openwire/src/transport/body.rs b/crates/openwire/src/transport/body.rs index 1595ed3..1f71d9b 100644 --- a/crates/openwire/src/transport/body.rs +++ b/crates/openwire/src/transport/body.rs @@ -18,7 +18,9 @@ use openwire_core::{ use crate::connection::{ConnectionAvailability, ExchangeFinder, RealConnection}; -use super::bindings::{ConnectionBindings, ConnectionTaskRegistry}; +use super::bindings::{ + teardown_pooled_connection, ConnectionBindings, ConnectionTaskRegistry, +}; pub(super) struct BoundResponse { pub(super) response: Response, @@ -188,6 +190,21 @@ impl ObservedIncomingBody { return; } self.finished = true; + if self.ctx.body_force_discard() { + // Decode/body-layer errors set this flag so Drop discards the + // connection instead of treating the body as a clean abandon. + self.ctx + .listener() + .response_body_failed( + &self.ctx, + &WireError::body( + "response body failed after intermediate processing error", + std::io::Error::other("body force discard"), + ), + ); + self.discard_connection(); + return; + } self.ctx .listener() .response_body_end(&self.ctx, self.bytes_read); @@ -254,7 +271,7 @@ fn release_response_lease(state: ResponseLeaseState) { exchange_finder, ctx, availability, - .. + _tasks, }, .. } => { @@ -266,8 +283,12 @@ fn release_response_lease(state: ResponseLeaseState) { ctx.listener().connection_released(&ctx, connection.id()); return; } - bindings.remove(connection.id()); - let _ = exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &exchange_finder, + &bindings, + &_tasks, + connection.id(), + ); availability.notify(); ctx.listener().connection_released(&ctx, connection.id()); } @@ -307,12 +328,16 @@ fn evict_response_lease_state(state: ResponseLeaseState, mark_unhealthy: bool) { exchange_finder, ctx, availability, - .. + _tasks, }, .. } => { - bindings.remove(connection.id()); - let _ = exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &exchange_finder, + &bindings, + &_tasks, + connection.id(), + ); availability.notify(); ctx.listener().connection_released(&ctx, connection.id()); } diff --git a/crates/openwire/src/transport/connect.rs b/crates/openwire/src/transport/connect.rs index dbdb777..7b9a78f 100644 --- a/crates/openwire/src/transport/connect.rs +++ b/crates/openwire/src/transport/connect.rs @@ -376,6 +376,8 @@ async fn connect_via_http_proxy( target_uri, tunneled, deps.tls_connector.clone(), + deps.timer.clone(), + deps.connect_timeout, ) .await } @@ -461,6 +463,8 @@ async fn connect_via_socks_proxy( target_uri, proxied, deps.tls_connector.clone(), + deps.timer.clone(), + deps.connect_timeout, ) .await } @@ -480,6 +484,8 @@ async fn connect_target_tls_if_needed( target_uri: Uri, stream: BoxConnection, tls_connector: Option>, + timer: SharedTimer, + connect_timeout: Option, ) -> Result { if !target_uri .scheme_str() @@ -495,7 +501,13 @@ async fn connect_target_tls_if_needed( ) .with_authority_from_uri(&target_uri) })?; - tls_connector.connect(ctx, target_uri, stream).await + crate::connection::with_connect_timeout( + timer, + connect_timeout, + tls_connector.connect(ctx, target_uri, stream), + "TLS handshake", + ) + .await } pub(super) async fn establish_connect_tunnel( diff --git a/crates/openwire/src/transport/service.rs b/crates/openwire/src/transport/service.rs index 3156b30..72a000b 100644 --- a/crates/openwire/src/transport/service.rs +++ b/crates/openwire/src/transport/service.rs @@ -1,9 +1,12 @@ use std::io; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Duration; +use futures_util::future::{select, Either}; use http::{Request, Response}; use hyper::client::conn::{http1, http2, TrySendError}; +use hyper::rt::Timer; use openwire_core::{ BoxConnection, BoxFuture, CallContext, Connection, ConnectionInfo, Exchange, HyperExecutor, RequestBody, ResponseBody, SharedTimer, WireError, WireErrorKind, WireExecutor, @@ -25,8 +28,8 @@ use crate::proxy::{SelectedProxy, SharedProxySelector}; use crate::trace::PolicyTraceContext; use super::bindings::{ - release_acquired_connection, AcquiredBinding, BindingAcquireResult, ConnectionBindings, - ConnectionTaskRegistry, + release_acquired_connection, teardown_pooled_connection, AcquiredBinding, BindingAcquireResult, + ConnectionBindings, ConnectionTaskRegistry, }; use super::body::{ spawn_body_deadline_signal, BoundResponse, ObservedIncomingBody, ResponseLease, @@ -55,6 +58,7 @@ struct SelectedConnection { coalesced: bool, exchange_finder: Arc, bindings: Arc, + tasks: ConnectionTaskRegistry, availability: ConnectionAvailability, } @@ -68,6 +72,7 @@ struct SelectedConnectionInit { coalesced: bool, exchange_finder: Arc, bindings: Arc, + tasks: ConnectionTaskRegistry, availability: ConnectionAvailability, } @@ -92,6 +97,7 @@ type SelectedConnectionSendParts = ( bool, Arc, Arc, + ConnectionTaskRegistry, ConnectionAvailability, ); @@ -107,6 +113,7 @@ impl SelectedConnection { coalesced: init.coalesced, exchange_finder: init.exchange_finder, bindings: init.bindings, + tasks: init.tasks, availability: init.availability, } } @@ -142,6 +149,7 @@ impl SelectedConnection { self.coalesced, self.exchange_finder.clone(), self.bindings.clone(), + self.tasks.clone(), self.availability.clone(), )) } @@ -159,6 +167,7 @@ impl Drop for SelectedConnection { release_acquired_connection( &self.exchange_finder, &self.bindings, + &self.tasks, &self.availability, connection, binding, @@ -211,6 +220,17 @@ impl TransportService { config.max_connections_per_host, connection_availability.clone(), ); + let bindings = Arc::new(ConnectionBindings::default()); + let connection_tasks = ConnectionTaskRegistry::default(); + { + let bindings = bindings.clone(); + let tasks = connection_tasks.clone(); + exchange_finder + .pool() + .set_eviction_hook(Arc::new(move |connection_id| { + tasks.teardown_connection(&bindings, connection_id); + })); + } Self { connector, config, @@ -222,8 +242,8 @@ impl TransportService { on_pooled_connection_published, connection_limiter, connection_availability, - bindings: Arc::new(ConnectionBindings::default()), - connection_tasks: ConnectionTaskRegistry::default(), + bindings, + connection_tasks, } } @@ -392,6 +412,7 @@ impl TransportService { coalesced: false, exchange_finder: self.exchange_finder.clone(), bindings: self.bindings.clone(), + tasks: self.connection_tasks.clone(), availability: self.connection_availability.clone(), })); } @@ -400,7 +421,12 @@ impl TransportService { waitable_pooled_connection = true; } BindingAcquireResult::Stale => { - let _ = self.exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &self.exchange_finder, + &self.bindings, + &self.connection_tasks, + connection.id(), + ); self.connection_availability.notify(); } } @@ -507,7 +533,12 @@ impl TransportService { let binding = match protocol { ConnectionProtocol::Http1 => { - let (sender, task) = bind_http1(stream).await?; + let (sender, task) = with_connect_stage_timeout( + self.timer.clone(), + connect_timeout, + bind_http1(stream), + ) + .await?; self.bindings.insert_http1(info.id, info, sender); let binding = match self.bindings.acquire(connection.id()) { BindingAcquireResult::Acquired(binding) => binding, @@ -523,8 +554,12 @@ impl TransportService { }; self.exchange_finder.pool().insert(connection.clone()); if let Err(error) = self.spawn_http1_task(connection.clone(), task, span.clone()) { - self.bindings.remove(connection.id()); - let _ = self.exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &self.exchange_finder, + &self.bindings, + &self.connection_tasks, + connection.id(), + ); self.connection_availability.notify(); return Err(error); } @@ -532,11 +567,15 @@ impl TransportService { binding } ConnectionProtocol::Http2 => { - let (sender, task) = bind_http2( - stream, - &self.config, - HyperExecutor(self.executor.clone()), + let (sender, task) = with_connect_stage_timeout( self.timer.clone(), + connect_timeout, + bind_http2( + stream, + &self.config, + HyperExecutor(self.executor.clone()), + self.timer.clone(), + ), ) .await?; self.bindings.insert_http2(info.id, info, sender); @@ -554,8 +593,12 @@ impl TransportService { }; self.exchange_finder.pool().insert(connection.clone()); if let Err(error) = self.spawn_http2_task(connection.clone(), task, span.clone()) { - self.bindings.remove(connection.id()); - let _ = self.exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &self.exchange_finder, + &self.bindings, + &self.connection_tasks, + connection.id(), + ); self.connection_availability.notify(); return Err(error); } @@ -574,6 +617,7 @@ impl TransportService { coalesced: false, exchange_finder: self.exchange_finder.clone(), bindings: self.bindings.clone(), + tasks: self.connection_tasks.clone(), availability: self.connection_availability.clone(), })) } @@ -609,6 +653,7 @@ impl TransportService { coalesced: true, exchange_finder: self.exchange_finder.clone(), bindings: self.bindings.clone(), + tasks: self.connection_tasks.clone(), availability: self.connection_availability.clone(), })); } @@ -616,7 +661,12 @@ impl TransportService { let _ = self.exchange_finder.release(&connection); } BindingAcquireResult::Stale => { - let _ = self.exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection( + &self.exchange_finder, + &self.bindings, + &self.connection_tasks, + connection.id(), + ); self.connection_availability.notify(); } } @@ -640,12 +690,13 @@ impl TransportService { let bindings = self.bindings.clone(); let pool = self.exchange_finder.pool().clone(); let availability = self.connection_availability.clone(); - let (task_id, registry) = self.connection_tasks.reserve(); + let registry = self.connection_tasks.downgrade(); let future = Box::pin( async move { let result = task.await; bindings.remove(connection_id); - let _ = pool.remove(connection_id); + let _ = pool.remove_without_hook(connection_id); + ConnectionTaskRegistry::complete_connection_weak(®istry, connection_id); availability.notify(); if let Err(error) = result { tracing::debug!( @@ -654,17 +705,17 @@ impl TransportService { "owned HTTP/1 connection task failed", ); } - ConnectionTaskRegistry::complete_weak(®istry, task_id); } .instrument(span), ); match self.executor.spawn(future) { Ok(handle) => { - self.connection_tasks.attach(task_id, handle); + self.connection_tasks + .attach_connection(connection_id, handle); Ok(()) } Err(error) => { - self.connection_tasks.cancel(task_id); + self.connection_tasks.abort_connection(connection_id); Err(error) } } @@ -680,12 +731,13 @@ impl TransportService { let bindings = self.bindings.clone(); let pool = self.exchange_finder.pool().clone(); let availability = self.connection_availability.clone(); - let (task_id, registry) = self.connection_tasks.reserve(); + let registry = self.connection_tasks.downgrade(); let future = Box::pin( async move { let result = task.await; bindings.remove(connection_id); - let _ = pool.remove(connection_id); + let _ = pool.remove_without_hook(connection_id); + ConnectionTaskRegistry::complete_connection_weak(®istry, connection_id); availability.notify(); if let Err(error) = result { tracing::debug!( @@ -694,17 +746,17 @@ impl TransportService { "owned HTTP/2 connection task failed", ); } - ConnectionTaskRegistry::complete_weak(®istry, task_id); } .instrument(span), ); match self.executor.spawn(future) { Ok(handle) => { - self.connection_tasks.attach(task_id, handle); + self.connection_tasks + .attach_connection(connection_id, handle); Ok(()) } Err(error) => { - self.connection_tasks.cancel(task_id); + self.connection_tasks.abort_connection(connection_id); Err(error) } } @@ -780,6 +832,7 @@ async fn send_bound_request( coalesced, exchange_finder, bindings, + tasks_reg, availability, ) = selected.into_send_parts()?; let request = prepare_request_for_send( @@ -797,6 +850,7 @@ async fn send_bound_request( &connection, &exchange_finder, &bindings, + &tasks_reg, &availability, &ctx, ); @@ -812,6 +866,7 @@ async fn send_bound_request( &connection, &exchange_finder, &bindings, + &tasks_reg, &availability, &ctx, ); @@ -843,6 +898,7 @@ async fn send_bound_request( &connection, &exchange_finder, &bindings, + &tasks_reg, &availability, &ctx, ); @@ -857,6 +913,7 @@ async fn send_bound_request( &connection, &exchange_finder, &bindings, + &tasks_reg, &availability, &ctx, ); @@ -948,13 +1005,13 @@ fn cleanup_failed_request( connection: &RealConnection, exchange_finder: &Arc, bindings: &Arc, + tasks: &ConnectionTaskRegistry, availability: &ConnectionAvailability, ctx: &CallContext, ) { match connection.protocol() { ConnectionProtocol::Http1 => { - bindings.remove(connection.id()); - let _ = exchange_finder.pool().remove(connection.id()); + teardown_pooled_connection(exchange_finder, bindings, tasks, connection.id()); } ConnectionProtocol::Http2 => { connection.mark_unhealthy(); @@ -964,3 +1021,30 @@ fn cleanup_failed_request( availability.notify(); ctx.listener().connection_released(ctx, connection.id()); } + +async fn with_connect_stage_timeout( + timer: SharedTimer, + connect_timeout: Option, + future: F, +) -> Result +where + F: std::future::Future>, +{ + let Some(timeout) = connect_timeout else { + return future.await; + }; + if timeout.is_zero() { + return Err(WireError::connect_timeout( + "connect timed out before protocol binding", + )); + } + + let future = Box::pin(future); + let sleep = timer.sleep(timeout); + match select(future, sleep).await { + Either::Left((result, _sleep)) => result, + Either::Right((_ready, _future)) => Err(WireError::connect_timeout(format!( + "connect timed out after {timeout:?} during protocol binding" + ))), + } +} diff --git a/crates/openwire/src/transport/tests.rs b/crates/openwire/src/transport/tests.rs index 7bbaa6e..8bd97f0 100644 --- a/crates/openwire/src/transport/tests.rs +++ b/crates/openwire/src/transport/tests.rs @@ -1359,6 +1359,6 @@ fn connection_task_registry_recovers_after_mutex_poisoning() { let _ = panic::catch_unwind(AssertUnwindSafe(|| registry.poison_handles_for_test())); - let (task_id, _weak) = registry.reserve(); - registry.cancel(task_id); + let connection_id = openwire_core::next_connection_id(); + registry.abort_connection(connection_id); } diff --git a/crates/openwire/src/websocket/handshake.rs b/crates/openwire/src/websocket/handshake.rs index 3a80286..370bce2 100644 --- a/crates/openwire/src/websocket/handshake.rs +++ b/crates/openwire/src/websocket/handshake.rs @@ -37,10 +37,15 @@ pub(crate) fn derive_accept(client_key: &str) -> String { base64::engine::general_purpose::STANDARD.encode(hasher.finalize()) } -pub(crate) fn generate_client_key() -> String { +pub(crate) fn generate_client_key() -> Result { let mut bytes = [0u8; 16]; - getrandom::getrandom(&mut bytes).expect("getrandom failed"); - base64::engine::general_purpose::STANDARD.encode(bytes) + getrandom::getrandom(&mut bytes).map_err(|error| { + WireError::internal( + "failed to generate WebSocket client key", + std::io::Error::other(error.to_string()), + ) + })?; + Ok(base64::engine::general_purpose::STANDARD.encode(bytes)) } pub(crate) fn inject_handshake(request: &mut Request) -> Result<(), WireError> { @@ -63,7 +68,7 @@ pub(crate) fn inject_handshake(request: &mut Request) -> Result<(), .headers_mut() .insert("sec-websocket-version", HeaderValue::from_static("13")); - let key = generate_client_key(); + let key = generate_client_key()?; let accept = derive_accept(&key); request.headers_mut().insert( "sec-websocket-key", @@ -546,15 +551,15 @@ mod tests { #[test] fn client_key_is_24_base64_chars() { - let k = generate_client_key(); + let k = generate_client_key().expect("key"); assert_eq!(k.len(), 24); assert!(k.ends_with('='), "16-byte base64 always ends with ="); } #[test] fn client_key_is_random() { - let a = generate_client_key(); - let b = generate_client_key(); + let a = generate_client_key().expect("key"); + let b = generate_client_key().expect("key"); assert_ne!(a, b); } diff --git a/crates/openwire/src/websocket/native/engine.rs b/crates/openwire/src/websocket/native/engine.rs index 40b4f24..df99075 100644 --- a/crates/openwire/src/websocket/native/engine.rs +++ b/crates/openwire/src/websocket/native/engine.rs @@ -133,7 +133,7 @@ impl Sink for NativeSink { validate_outbound_engine_frame(&item)?; let me = self.get_mut(); - let key = random_mask_key(); + let key = random_mask_key()?; let header = match &item { EngineFrame::Text(_) => FrameHeader { fin: true, diff --git a/crates/openwire/src/websocket/native/mask.rs b/crates/openwire/src/websocket/native/mask.rs index 7b91eff..ea2c29b 100644 --- a/crates/openwire/src/websocket/native/mask.rs +++ b/crates/openwire/src/websocket/native/mask.rs @@ -4,10 +4,15 @@ pub(crate) fn mask_in_place(payload: &mut [u8], key: [u8; 4]) { } } -pub(crate) fn random_mask_key() -> [u8; 4] { +pub(crate) fn random_mask_key() -> Result<[u8; 4], openwire_core::websocket::WebSocketEngineError> { let mut key = [0u8; 4]; - getrandom::getrandom(&mut key).expect("getrandom failed"); - key + getrandom::getrandom(&mut key).map_err(|error| { + openwire_core::websocket::WebSocketEngineError::Io(openwire_core::WireError::internal( + "failed to generate WebSocket mask key", + std::io::Error::other(error.to_string()), + )) + })?; + Ok(key) } #[cfg(test)] @@ -36,8 +41,8 @@ mod tests { #[test] fn random_keys_differ_across_calls() { - let a = random_mask_key(); - let b = random_mask_key(); + let a = random_mask_key().expect("mask key"); + let b = random_mask_key().expect("mask key"); // 4 bytes of entropy — collisions theoretically possible but vanishingly rare. assert_ne!(a, b); } diff --git a/crates/openwire/tests/integration.rs b/crates/openwire/tests/integration.rs index e6f342c..5f171c7 100644 --- a/crates/openwire/tests/integration.rs +++ b/crates/openwire/tests/integration.rs @@ -4866,8 +4866,11 @@ async fn retry_and_redirect_events_follow_stable_order_and_trace_fields() { "connect_end ", "connection_acquired ", "response_headers_end 302 Found", - "redirect 1 http://openwire.test:", + // Intermediate redirect body is drained so the connection can be + // reused before the follow-up hop starts. + "response_body_end ", "connection_released ", + "redirect 1 http://openwire.test:", "connection_acquired ", "response_headers_end 200 OK", "response_body_end 20", diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index e9a16b9..57b01a0 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -329,3 +329,19 @@ Optional live-network smoke suite: ```bash cargo test -p openwire --test live_network -- --ignored --test-threads=1 ``` + +## Connection teardown and connect budgets + +- Pool eviction (idle timeout, max idle, max lifetime, explicit remove) aborts the + owned hyper connection task and clears protocol bindings so sockets are closed, + not only removed from reuse indexes. +- `connect_timeout` covers TCP establishment, TLS handshake, and protocol binding + for direct and proxy-tunneled paths. +- HTTP/2 temporary sender unreadiness is handled by awaiting `ready()` on the + acquired sender under the call deadline, not by parking on pool availability. +- Intermediate follow-up responses (auth / redirect / status retry) drain the + body up to a small cap before the next network attempt so HTTP/1 connections + can be reused when possible. +- Transparent decompression failures mark the call for connection discard so + HTTP/2 connections are not returned to the pool as healthy after a body error. + From 9b8fa98bf887d500beae70125c0bb09d81048870 Mon Sep 17 00:00:00 2001 From: fannnzhang Date: Sat, 11 Jul 2026 17:00:36 +0800 Subject: [PATCH 2/2] chore: fix clippy deny findings and rustfmt Resolve workspace clippy issues under `-D warnings` (field reassignment with Default, collapsible match, unused mut) and apply rustfmt across touched crates. --- crates/openwire-core/src/context.rs | 4 +-- crates/openwire-tungstenite/src/lib.rs | 8 ++++-- crates/openwire/src/connection/pool.rs | 25 +++++++++-------- crates/openwire/src/policy/follow_up.rs | 12 +++----- crates/openwire/src/transport/bindings.rs | 6 ++-- crates/openwire/src/transport/body.rs | 34 +++++++---------------- 6 files changed, 35 insertions(+), 54 deletions(-) diff --git a/crates/openwire-core/src/context.rs b/crates/openwire-core/src/context.rs index d228e91..24bfc5b 100644 --- a/crates/openwire-core/src/context.rs +++ b/crates/openwire-core/src/context.rs @@ -117,9 +117,7 @@ impl CallContext { } pub fn mark_body_force_discard(&self) { - self.inner - .body_force_discard - .store(true, Ordering::Release); + self.inner.body_force_discard.store(true, Ordering::Release); } pub fn body_force_discard(&self) -> bool { diff --git a/crates/openwire-tungstenite/src/lib.rs b/crates/openwire-tungstenite/src/lib.rs index 93b4130..41b23bc 100644 --- a/crates/openwire-tungstenite/src/lib.rs +++ b/crates/openwire-tungstenite/src/lib.rs @@ -73,9 +73,11 @@ impl WebSocketEngine for TungsteniteEngine { // BoxConnection (hyper::rt::Read+Write) → tokio AsyncRead+Write. let tokio_io = TokioIo::new(io); - let mut ws_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default(); - ws_config.max_message_size = Some(config.max_message_size); - ws_config.max_frame_size = Some(config.max_frame_size); + let ws_config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig { + max_message_size: Some(config.max_message_size), + max_frame_size: Some(config.max_frame_size), + ..Default::default() + }; let stream = WebSocketStream::from_raw_socket( tokio_io, tokio_tungstenite::tungstenite::protocol::Role::Client, diff --git a/crates/openwire/src/connection/pool.rs b/crates/openwire/src/connection/pool.rs index bdf0f2c..cd489ea 100644 --- a/crates/openwire/src/connection/pool.rs +++ b/crates/openwire/src/connection/pool.rs @@ -309,18 +309,19 @@ impl ConnectionPool { let mut state = lock_mutex(&self.state); let evicted = prune_address(&self.settings, &mut state, address); let stats = match state.by_address.get(address) { - Some(connections) => connections.iter().fold( - PoolStats::default(), - |mut stats, connection| { - stats.total += 1; - match connection.snapshot().allocation { - ConnectionAllocationState::Idle => stats.idle += 1, - ConnectionAllocationState::InUse { .. } => stats.in_use += 1, - ConnectionAllocationState::Closed => {} - } - stats - }, - ), + Some(connections) => { + connections + .iter() + .fold(PoolStats::default(), |mut stats, connection| { + stats.total += 1; + match connection.snapshot().allocation { + ConnectionAllocationState::Idle => stats.idle += 1, + ConnectionAllocationState::InUse { .. } => stats.in_use += 1, + ConnectionAllocationState::Closed => {} + } + stats + }) + } None => PoolStats::default(), }; drop(state); diff --git a/crates/openwire/src/policy/follow_up.rs b/crates/openwire/src/policy/follow_up.rs index d52941a..15a2741 100644 --- a/crates/openwire/src/policy/follow_up.rs +++ b/crates/openwire/src/policy/follow_up.rs @@ -2,8 +2,8 @@ use std::task::{Context, Poll}; use std::time::SystemTime; use http::header::{ - AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, COOKIE, EXPECT, HOST, - LOCATION, RETRY_AFTER, SET_COOKIE, TRANSFER_ENCODING, + AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, COOKIE, EXPECT, HOST, LOCATION, + RETRY_AFTER, SET_COOKIE, TRANSFER_ENCODING, }; use http::{HeaderMap, Method, Request, Response, StatusCode, Uri, Version}; use openwire_core::{ @@ -467,12 +467,9 @@ fn should_retry_misdirected_request( .is_some() } - const INTERMEDIATE_BODY_DRAIN_LIMIT: u64 = 256 * 1024; -async fn drain_intermediate_response( - response: Response, -) -> Result<(), WireError> { +async fn drain_intermediate_response(response: Response) -> Result<(), WireError> { let body = response.into_body(); let mut body = std::pin::pin!(body); let mut read = 0u64; @@ -952,13 +949,12 @@ mod tests { assert!(next.is_none()); } - #[test] fn redirect_to_get_strips_body_describing_headers() { use http::header::{ CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE, EXPECT, TRANSFER_ENCODING, }; - let mut request = Request::builder() + let request = Request::builder() .method("POST") .uri("http://source.test/start") .header(CONTENT_TYPE, "application/json") diff --git a/crates/openwire/src/transport/bindings.rs b/crates/openwire/src/transport/bindings.rs index fbbcef3..985ca8a 100644 --- a/crates/openwire/src/transport/bindings.rs +++ b/crates/openwire/src/transport/bindings.rs @@ -169,10 +169,8 @@ pub(super) struct ConnectionTaskRegistryInner { impl ConnectionTaskRegistry { pub(super) fn attach_connection(&self, connection_id: ConnectionId, handle: BoxTaskHandle) { let mut handles = lock_mutex(&self.inner.handles_by_connection); - if let Some(previous) = handles.insert(connection_id, Some(handle)) { - if let Some(previous) = previous { - previous.abort(); - } + if let Some(Some(previous)) = handles.insert(connection_id, Some(handle)) { + previous.abort(); } } diff --git a/crates/openwire/src/transport/body.rs b/crates/openwire/src/transport/body.rs index 1f71d9b..e63c3a5 100644 --- a/crates/openwire/src/transport/body.rs +++ b/crates/openwire/src/transport/body.rs @@ -18,9 +18,7 @@ use openwire_core::{ use crate::connection::{ConnectionAvailability, ExchangeFinder, RealConnection}; -use super::bindings::{ - teardown_pooled_connection, ConnectionBindings, ConnectionTaskRegistry, -}; +use super::bindings::{teardown_pooled_connection, ConnectionBindings, ConnectionTaskRegistry}; pub(super) struct BoundResponse { pub(super) response: Response, @@ -193,15 +191,13 @@ impl ObservedIncomingBody { if self.ctx.body_force_discard() { // Decode/body-layer errors set this flag so Drop discards the // connection instead of treating the body as a clean abandon. - self.ctx - .listener() - .response_body_failed( - &self.ctx, - &WireError::body( - "response body failed after intermediate processing error", - std::io::Error::other("body force discard"), - ), - ); + self.ctx.listener().response_body_failed( + &self.ctx, + &WireError::body( + "response body failed after intermediate processing error", + std::io::Error::other("body force discard"), + ), + ); self.discard_connection(); return; } @@ -283,12 +279,7 @@ fn release_response_lease(state: ResponseLeaseState) { ctx.listener().connection_released(&ctx, connection.id()); return; } - teardown_pooled_connection( - &exchange_finder, - &bindings, - &_tasks, - connection.id(), - ); + teardown_pooled_connection(&exchange_finder, &bindings, &_tasks, connection.id()); availability.notify(); ctx.listener().connection_released(&ctx, connection.id()); } @@ -332,12 +323,7 @@ fn evict_response_lease_state(state: ResponseLeaseState, mark_unhealthy: bool) { }, .. } => { - teardown_pooled_connection( - &exchange_finder, - &bindings, - &_tasks, - connection.id(), - ); + teardown_pooled_connection(&exchange_finder, &bindings, &_tasks, connection.id()); availability.notify(); ctx.listener().connection_released(&ctx, connection.id()); }