From a8d28e6b37a39242f3535a2005108ece728b4254 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Sun, 11 Oct 2026 13:55:01 +0300 Subject: [PATCH] Fix stalled ME relays by observing child task completion #879 --- src/proxy/middle_relay/session.rs | 41 ++-- src/proxy/middle_relay/session/children.rs | 51 ++++ src/transport/middle_proxy/mod.rs | 4 + .../tests/relay_completion_tests.rs | 224 ++++++++++++++++++ 4 files changed, 293 insertions(+), 27 deletions(-) create mode 100644 src/transport/middle_proxy/tests/relay_completion_tests.rs diff --git a/src/proxy/middle_relay/session.rs b/src/proxy/middle_relay/session.rs index 3531cb7..827cb65 100644 --- a/src/proxy/middle_relay/session.rs +++ b/src/proxy/middle_relay/session.rs @@ -239,6 +239,8 @@ where }; let mut main_result: Result<()> = Ok(()); + let mut c2me_completed = None; + let mut writer_completed = None; let mut client_closed = false; let mut frame_counter: u64 = 0; let mut route_watch_open = true; @@ -303,6 +305,15 @@ where } tokio::select! { + // A terminal child result must wake an otherwise idle client relay. + joined = &mut child_tasks.c2me_sender => { + c2me_completed = Some(joined); + break; + } + joined = &mut child_tasks.me_writer => { + writer_completed = Some(joined); + break; + } _ = session_cancel.cancelled() => { warn!( user = %user, @@ -464,33 +475,9 @@ where } drop(c2me_tx); - let c2me_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.c2me_sender).await { - Ok(joined) => { - joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME sender join error: {e}")))) - } - Err(_) => { - stats.increment_me_child_join_timeout_total(); - stats.increment_me_child_abort_total(); - child_tasks.c2me_sender.abort(); - Err(ProxyError::Proxy("ME sender join timeout".into())) - } - }; - - flow_cancel.cancel(); - if let Some(stop_tx) = child_tasks.stop_tx.take() { - let _ = stop_tx.send(()); - } - let mut writer_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.me_writer).await { - Ok(joined) => { - joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}")))) - } - Err(_) => { - stats.increment_me_child_join_timeout_total(); - stats.increment_me_child_abort_total(); - child_tasks.me_writer.abort(); - Err(ProxyError::Proxy("ME writer join timeout".into())) - } - }; + let (c2me_result, mut writer_result) = child_tasks + .finish(c2me_completed, writer_completed, stats.as_ref()) + .await; // A client-initiated close can unregister the ME channel before its writer exits. if client_closed && matches!(writer_result, Err(ProxyError::MiddleConnectionLost)) { diff --git a/src/proxy/middle_relay/session/children.rs b/src/proxy/middle_relay/session/children.rs index 1b1cd8c..39aed52 100644 --- a/src/proxy/middle_relay/session/children.rs +++ b/src/proxy/middle_relay/session/children.rs @@ -1,12 +1,63 @@ use super::*; +/// Owns relay tasks and their cancellation signals through bounded cleanup. pub(super) struct RelayChildTasks { + /// Serializes client frames into the ME route. pub(super) c2me_sender: AbortOnDropHandle>, + /// Delivers ME responses to the client. pub(super) me_writer: AbortOnDropHandle>, + /// Interrupts downstream quota and transport waits during cleanup. pub(super) flow_cancel: CancellationToken, + /// Stops the downstream task after the sender has drained or timed out. pub(super) stop_tx: Option>, } +impl RelayChildTasks { + /// Reuses consumed join results and preserves sender-first bounded shutdown. + pub(super) async fn finish( + &mut self, + c2me_completed: Option, tokio::task::JoinError>>, + writer_completed: Option, tokio::task::JoinError>>, + stats: &Stats, + ) -> (Result<()>, Result<()>) { + // A JoinHandle whose output was selected must never be polled again. + let c2me_join = match c2me_completed { + Some(joined) => Ok(joined), + None => timeout(ME_CHILD_JOIN_TIMEOUT, &mut self.c2me_sender).await, + }; + let c2me_result = match c2me_join { + Ok(joined) => joined + .unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME sender join error: {e}")))), + Err(_) => { + stats.increment_me_child_join_timeout_total(); + stats.increment_me_child_abort_total(); + self.c2me_sender.abort(); + Err(ProxyError::Proxy("ME sender join timeout".into())) + } + }; + + self.flow_cancel.cancel(); + if let Some(stop_tx) = self.stop_tx.take() { + let _ = stop_tx.send(()); + } + let writer_join = match writer_completed { + Some(joined) => Ok(joined), + None => timeout(ME_CHILD_JOIN_TIMEOUT, &mut self.me_writer).await, + }; + let writer_result = match writer_join { + Ok(joined) => joined + .unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}")))), + Err(_) => { + stats.increment_me_child_join_timeout_total(); + stats.increment_me_child_abort_total(); + self.me_writer.abort(); + Err(ProxyError::Proxy("ME writer join timeout".into())) + } + }; + (c2me_result, writer_result) + } +} + impl Drop for RelayChildTasks { fn drop(&mut self) { self.flow_cancel.cancel(); diff --git a/src/transport/middle_proxy/mod.rs b/src/transport/middle_proxy/mod.rs index 968aa2d..66b17ff 100644 --- a/src/transport/middle_proxy/mod.rs +++ b/src/transport/middle_proxy/mod.rs @@ -47,6 +47,10 @@ mod send; #[cfg(test)] #[path = "tests/send_adversarial_tests.rs"] mod send_adversarial_tests; +// Full relay lifecycle checks exercise propagation of ME task completion. +#[cfg(test)] +#[path = "tests/relay_completion_tests.rs"] +mod relay_completion_tests; mod wire; /// Network-free writer fixtures shared by admission and recovery regressions. diff --git a/src/transport/middle_proxy/tests/relay_completion_tests.rs b/src/transport/middle_proxy/tests/relay_completion_tests.rs new file mode 100644 index 0000000..1251e98 --- /dev/null +++ b/src/transport/middle_proxy/tests/relay_completion_tests.rs @@ -0,0 +1,224 @@ +use super::pool_writer_security_tests::make_pool; +use super::registry::RouteResult; +use super::{MePool, MeResponse}; +use crate::config::{MeRouteNoWriterMode, ProxyConfig}; +use crate::crypto::AesCtr; +use crate::error::{ProxyError, Result}; +use crate::protocol::constants::ProtoTag; +use crate::proxy::handshake::HandshakeSuccess; +use crate::proxy::middle_relay::handle_via_middle_proxy; +use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; +use crate::proxy::shared_state::ProxySharedState; +use crate::stats::Stats; +use crate::stream::{BufferPool, CryptoReader, CryptoWriter}; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream, duplex}; +use tokio::time::timeout; +use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; + +struct RelayFixture { + pool: Arc, + stats: Arc, + client: DuplexStream, + task: AbortOnDropHandle>, + cancel: CancellationToken, + routes: RouteRuntimeController, + conn_id: u64, +} + +impl RelayFixture { + async fn new(mode: MeRouteNoWriterMode) -> Self { + let mut pool = make_pool().await; + let route = Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().route_runtime).unwrap(); + route + .me_route_no_writer_mode + .store(mode.as_u8(), Ordering::Relaxed); + route.me_route_no_writer_wait = Duration::from_millis(20); + route.me_route_hybrid_max_wait = Duration::from_millis(50); + let stats = pool.stats.clone(); + let mut config = ProxyConfig::default(); + // The observation deadline must expire well before ordinary client idle handling. + config.timeouts.relay_client_idle_soft_secs = 60; + config.timeouts.relay_client_idle_hard_secs = 120; + let routes = RouteRuntimeController::new(RelayRouteMode::Middle); + let cancel = CancellationToken::new(); + // This isolated pool has no other registrants; reserve the preceding ID. + let (preceding_id, _) = pool.registry.register().await; + pool.registry.unregister(preceding_id).await; + let conn_id = preceding_id.wrapping_add(1); + let (server, client) = duplex(16 * 1024); + let (reader, writer) = tokio::io::split(server); + let success = HandshakeSuccess { + user: "completion-test".into(), + dc_idx: 2, + proto_tag: ProtoTag::Intermediate, + dec_key: [0; 32], + dec_iv: 0, + enc_key: [0; 32], + enc_iv: 0, + peer: "127.0.0.1:50000".parse().unwrap(), + is_tls: false, + }; + let quota = stats.current_user_quota_handle(&success.user); + let task = AbortOnDropHandle::new(tokio::spawn(handle_via_middle_proxy( + CryptoReader::new(reader, AesCtr::new(&[0; 32], 0)), + CryptoWriter::new(writer, AesCtr::new(&[0; 32], 0), 64 * 1024), + success, + pool.clone(), + stats.clone(), + Arc::new(config), + Arc::new(BufferPool::new()), + "127.0.0.1:443".parse().unwrap(), + pool.rng.clone(), + routes.subscribe(), + routes.snapshot(), + 0xB5, + cancel.clone(), + ProxySharedState::new(), + quota, + ))); + timeout(Duration::from_secs(1), async { + loop { + if stats.get_current_connections_me() == 1 { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("relay must register its connection"); + Self { + pool, + stats, + client, + task, + cancel, + routes, + conn_id, + } + } + + async fn send_frame(&mut self) { + let mut frame = 2048u32.to_le_bytes().to_vec(); + frame.extend_from_slice(&[0x42; 2048]); + let encrypted = AesCtr::new(&[0; 32], 0).encrypt(&frame); + self.client.write_all(&encrypted).await.unwrap(); + } + + async fn finish(&mut self) -> Result<()> { + self.finish_with_deadline(Duration::from_secs(1)).await + } + + async fn finish_with_deadline(&mut self, deadline: Duration) -> Result<()> { + let result = timeout(deadline, &mut self.task) + .await + .expect("completed ME child must terminate relay while client remains open") + .expect("relay must not panic or poll a completed child twice"); + assert_eq!(self.stats.get_current_connections_me(), 0); + assert!(self.pool.registry.active_conn_ids().await.is_empty()); + assert_eq!( + self.pool + .registry + .route(self.conn_id, MeResponse::Close) + .await, + RouteResult::NoConn, + "the unbound route and its byte budget must be unregistered" + ); + let mut byte = [0u8; 1]; + assert_eq!( + timeout(Duration::from_secs(1), self.client.read(&mut byte)) + .await + .unwrap() + .unwrap(), + 0 + ); + assert!(self.pool.shutdown_until(Duration::from_secs(1)).await); + result + } +} + +#[tokio::test] +async fn no_writer_error_closes_silent_client_after_one_frame() { + for mode in [ + MeRouteNoWriterMode::AsyncRecoveryFailfast, + MeRouteNoWriterMode::HybridAsyncPersistent, + ] { + let mut fixture = RelayFixture::new(mode).await; + fixture.send_frame().await; + // Prove the sender has already failed, independently of parent relay completion. + timeout(Duration::from_secs(1), async { + while fixture.stats.get_me_no_writer_failfast_total() + + fixture.stats.get_me_hybrid_timeout_total() + == 0 + { + tokio::task::yield_now().await; + } + }) + .await + .expect("empty pool must reach its configured no-writer deadline"); + let error = fixture.finish().await.unwrap_err(); + assert!( + matches!(error, ProxyError::Proxy(ref message) + if message.contains("failfast window") || message.contains("hybrid timeout")), + "original routing error must reach the caller: {error}" + ); + } +} + +#[tokio::test] +async fn downstream_close_terminates_silent_client() { + let mut fixture = RelayFixture::new(MeRouteNoWriterMode::AsyncRecoveryFailfast).await; + fixture + .pool + .registry + .route(fixture.conn_id, MeResponse::Close) + .await; + fixture.finish().await.unwrap(); +} + +#[tokio::test] +async fn lost_me_route_terminates_silent_client_with_original_error() { + let mut fixture = RelayFixture::new(MeRouteNoWriterMode::AsyncRecoveryFailfast).await; + fixture.pool.registry.unregister(fixture.conn_id).await; + assert!(matches!( + fixture.finish().await, + Err(ProxyError::MiddleConnectionLost) + )); +} + +#[tokio::test] +async fn healthy_idle_remains_open_and_client_eof_is_clean() { + let mut fixture = RelayFixture::new(MeRouteNoWriterMode::AsyncRecoveryFailfast).await; + assert!( + timeout(Duration::from_millis(100), &mut fixture.task) + .await + .is_err() + ); + assert_eq!(fixture.stats.get_current_connections_me(), 1); + fixture.client.shutdown().await.unwrap(); + fixture.finish().await.unwrap(); +} + +#[tokio::test] +async fn user_cancellation_still_releases_relay() { + let mut fixture = RelayFixture::new(MeRouteNoWriterMode::AsyncRecoveryFailfast).await; + fixture.cancel.cancel(); + assert!(matches!( + fixture.finish().await, + Err(ProxyError::UserDisabled { .. }) + )); +} + +#[tokio::test] +async fn route_cutover_still_releases_relay() { + let mut fixture = RelayFixture::new(MeRouteNoWriterMode::AsyncRecoveryFailfast).await; + fixture.routes.set_mode(RelayRouteMode::Direct); + // The existing cutover stagger intentionally waits between one and two seconds. + assert!(matches!( + fixture.finish_with_deadline(Duration::from_secs(3)).await, + Err(ProxyError::RouteSwitched) + )); +}