Fixed routing across endpoint refresh

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

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<RuntimeGeneration> {
-    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<bool>,
+) -> Arc<RuntimeGeneration> {
+    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,

src/transport/middle_proxy/send/selection.rs

@@ -16,19 +16,37 @@ impl MePool {
         include_warm: bool,
     ) -> Vec<usize> {
         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);
             }
         }

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<MePool>, Arc<SecureRandom>) {
+    make_pool_with_decision(NetworkDecision::default()).await
+}
+
+async fn make_pool_with_decision(
+    decision: NetworkDecision,
+) -> (Arc<MePool>, Arc<SecureRandom>) {
     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<MePool>, Arc<SecureRandom>) {
         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;

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);

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<WebSession>,
     manager: Arc<WebProcessRuntime>,
     generation: Arc<RuntimeGeneration>,
+    admission_tx: watch::Sender<bool>,
 }
 
 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);