From 1a18f2452fc1863ce52b9d013c1a0fc1e33720cb Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Sun, 23 Aug 2026 20:43:22 +0300 Subject: [PATCH] Fixed routing across endpoint refresh Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/maestro/generation.rs | 12 +++- src/transport/middle_proxy/send/selection.rs | 38 ++++++++--- .../tests/send_adversarial_tests.rs | 63 ++++++++++++++++++- src/web/session/backend.rs | 4 ++ src/web/session/backend_tests.rs | 43 ++++++++++++- 5 files changed, 146 insertions(+), 14 deletions(-) diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs index de29b7e..8a8c5de 100644 --- a/src/maestro/generation.rs +++ b/src/maestro/generation.rs @@ -305,8 +305,18 @@ impl RuntimeGeneration { #[cfg(test)] /// Builds a lightweight runtime generation without network startup tasks. pub(crate) fn test_runtime_generation(id: u64, config: ProxyConfig) -> Arc { - let (config_tx, config_rx) = watch::channel(Arc::new(config.clone())); let (_admission_tx, admission_rx) = watch::channel(true); + test_runtime_generation_with_admission(id, config, admission_rx) +} + +#[cfg(test)] +/// Builds a lightweight runtime generation with a controllable admission gate. +pub(crate) fn test_runtime_generation_with_admission( + id: u64, + config: ProxyConfig, + admission_rx: watch::Receiver, +) -> Arc { + let (config_tx, config_rx) = watch::channel(Arc::new(config.clone())); let stats = Arc::new(Stats::new()); let upstream_manager = Arc::new(UpstreamManager::new( config.upstreams, diff --git a/src/transport/middle_proxy/send/selection.rs b/src/transport/middle_proxy/send/selection.rs index 834e0c0..837062c 100644 --- a/src/transport/middle_proxy/send/selection.rs +++ b/src/transport/middle_proxy/send/selection.rs @@ -16,19 +16,37 @@ impl MePool { include_warm: bool, ) -> Vec { let preferred_snapshot = self.preferred_endpoints_by_dc.load(); - let Some(preferred) = preferred_snapshot.get(&routed_dc) else { - return Vec::new(); - }; - if preferred.is_empty() { - return Vec::new(); + let mut out = Vec::new(); + if let Some(preferred) = preferred_snapshot + .get(&routed_dc) + .filter(|preferred| !preferred.is_empty()) + { + for (idx, w) in writers.iter().enumerate() { + if !self.writer_eligible_for_selection(w, include_warm) { + continue; + } + if w.writer_dc == routed_dc && preferred.binary_search(&w.addr).is_ok() { + out.push(idx); + } + } + } + if !out.is_empty() || !include_warm { + return out; } - let mut out = Vec::new(); + // A map update publishes desired endpoints before replacement coverage is + // guaranteed. Preserve the existing same-DC writer as the final tier so + // the data plane remains available while the pool-owned reinit converges. for (idx, w) in writers.iter().enumerate() { - if !self.writer_eligible_for_selection(w, include_warm) { - continue; - } - if w.writer_dc == routed_dc && preferred.binary_search(&w.addr).is_ok() { + let family_enabled = if w.addr.is_ipv4() { + self.decision.ipv4_me + } else { + self.decision.ipv6_me + }; + if family_enabled + && w.writer_dc == routed_dc + && self.writer_eligible_for_selection(w, true) + { out.push(idx); } } diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 3b6e718..eac095f 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -15,6 +15,12 @@ use crate::network::probe::NetworkDecision; use crate::stats::Stats; async fn make_pool() -> (Arc, Arc) { + make_pool_with_decision(NetworkDecision::default()).await +} + +async fn make_pool_with_decision( + decision: NetworkDecision, +) -> (Arc, Arc) { let general = GeneralConfig { me_route_no_writer_mode: MeRouteNoWriterMode::AsyncRecoveryFailfast, me_route_no_writer_wait_ms: 50, @@ -40,7 +46,7 @@ async fn make_pool() -> (Arc, Arc) { HashMap::new(), HashMap::new(), None, - NetworkDecision::default(), + decision, None, rng.clone(), Arc::new(Stats::default()), @@ -213,6 +219,61 @@ fn proxy_req_our_addr_from_payload(payload: &[u8]) -> SocketAddr { ) } +#[tokio::test] +async fn send_proxy_req_uses_live_same_dc_writer_while_preferred_endpoint_refills() { + let decision = NetworkDecision { + ipv4_dc: true, + ipv4_me: true, + effective_prefer: 4, + ..NetworkDecision::default() + }; + let (pool, _rng) = make_pool_with_decision(decision).await; + let old_positive_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 3, 2)), 443); + let old_negative_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 3, 3)), 443); + let mut old_positive_rx = insert_writer(&pool, 41, 2, old_positive_addr, true).await; + let _old_negative_rx = insert_writer(&pool, 42, -2, old_negative_addr, true).await; + + let new_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 0, 2, 42)), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(new_addr.ip(), new_addr.port())])]), + None, + ) + .await; + + assert!(pool.admission_ready_conditional_cast().await); + assert_eq!( + pool.preferred_endpoints_by_dc + .load() + .get(&2) + .cloned() + .unwrap_or_default(), + vec![new_addr] + ); + + let (conn_id, _rx) = pool.registry.register().await; + let result = pool + .send_proxy_req( + conn_id, + 2, + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30004), + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443), + b"cold-route", + 0, + None, + None, + ) + .await; + + assert!( + result.is_ok(), + "a live same-DC writer must bridge preferred-endpoint refill: {result:?}" + ); + assert_eq!( + recv_data_count(&mut old_positive_rx, Duration::from_millis(50)).await, + 1 + ); +} + #[tokio::test] async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() { let (pool, _rng) = make_pool().await; diff --git a/src/web/session/backend.rs b/src/web/session/backend.rs index 293ccdf..69acef6 100644 --- a/src/web/session/backend.rs +++ b/src/web/session/backend.rs @@ -21,6 +21,10 @@ impl WebSession { return; }; let generation = manager.active_generation(); + if !*generation.admission_rx.borrow() { + self.stream_finished(stream_id, peer_port); + return; + } let Ok(connection_permit) = generation.max_connections.clone().try_acquire_owned() else { manager.record_stream_rejected(); self.stream_finished(stream_id, peer_port); diff --git a/src/web/session/backend_tests.rs b/src/web/session/backend_tests.rs index 245444c..b09989d 100644 --- a/src/web/session/backend_tests.rs +++ b/src/web/session/backend_tests.rs @@ -5,6 +5,7 @@ use std::time::Duration; use arc_swap::ArcSwap; use tokio::net::TcpListener; +use tokio::sync::watch; use super::*; use crate::config::{ @@ -12,7 +13,9 @@ use crate::config::{ WebSecretMode, }; use crate::crypto::{AesCtr, sha256}; -use crate::maestro::generation::{RuntimeGeneration, test_runtime_generation}; +use crate::maestro::generation::{ + RuntimeGeneration, test_runtime_generation_with_admission, +}; use crate::protocol::constants::{ DC_IDX_POS, HANDSHAKE_LEN, IV_LEN, PREKEY_LEN, PROTO_TAG_POS, ProtoTag, SKIP_LEN, }; @@ -23,6 +26,7 @@ struct TestRuntime { session: Arc, manager: Arc, generation: Arc, + admission_tx: watch::Sender, } impl TestRuntime { @@ -106,7 +110,8 @@ fn test_runtime_with_dc( config.rebuild_runtime_user_auth().unwrap(); let limits = config.web.limits.clone(); let timeouts = config.web.timeouts.clone(); - let generation = test_runtime_generation(1, config); + let (admission_tx, admission_rx) = watch::channel(true); + let generation = test_runtime_generation_with_admission(1, config, admission_rx); let manager = WebProcessRuntime::start(Arc::new(ArcSwap::from(Arc::clone(&generation)))); let session = WebSession::new( Arc::downgrade(&manager), @@ -121,6 +126,7 @@ fn test_runtime_with_dc( session, manager, generation, + admission_tx, } } @@ -264,6 +270,39 @@ async fn delayed_valid_handshake_reaches_the_authenticated_relay() { drop(upstream); } +#[tokio::test] +async fn closed_generation_admission_rejects_web_stream_before_backend() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let runtime = test_runtime_with_dc(WebCarrier::Https, 1, Some(listener.local_addr().unwrap())); + runtime.admission_tx.send_replace(false); + + assert_eq!(runtime.process_frame(1, 1, FrameType::Open, &[]), Ok(1)); + settle_tasks().await; + assert!( + tokio::time::timeout(Duration::from_millis(50), listener.accept()) + .await + .is_err(), + "WEB stream bypassed the closed generation admission gate" + ); + assert!(!runtime.session.state.lock().streams.contains_key(&1)); + assert_eq!(runtime.generation.max_connections.available_permits(), 64); + + runtime.admission_tx.send_replace(true); + assert_eq!(runtime.process_frame(2, 2, FrameType::Open, &[]), Ok(2)); + assert_eq!( + runtime.process_frame(2, 3, FrameType::Data, &valid_plain_handshake()), + Ok(3) + ); + let (upstream, _) = tokio::time::timeout(Duration::from_secs(1), listener.accept()) + .await + .expect("a new WEB stream did not start after admission reopened") + .unwrap(); + assert_eq!(runtime.generation.max_connections.available_permits(), 63); + + runtime.shutdown().await; + drop(upstream); +} + #[tokio::test(start_paused = true)] async fn silent_streams_do_not_consume_active_handshake_capacity() { let runtime = test_runtime(WebCarrier::Https, 1);