ME/DC Reroute + ME Upper-limit tuning

src/proxy/client.rs

@@ -39,6 +39,7 @@ use crate::proxy::direct_relay::handle_via_direct;
 use crate::proxy::handshake::{HandshakeSuccess, handle_mtproto_handshake, handle_tls_handshake};
 use crate::proxy::masking::handle_bad_client;
 use crate::proxy::middle_relay::handle_via_middle_proxy;
+use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
 
 fn beobachten_ttl(config: &ProxyConfig) -> Duration {
     Duration::from_secs(config.general.beobachten_minutes.saturating_mul(60))
@@ -80,6 +81,7 @@ pub async fn handle_client_stream<S>(
     buffer_pool: Arc<BufferPool>,
     rng: Arc<SecureRandom>,
     me_pool: Option<Arc<MePool>>,
+    route_runtime: Arc<RouteRuntimeController>,
     tls_cache: Option<Arc<TlsFrontCache>>,
     ip_tracker: Arc<UserIpTracker>,
     beobachten: Arc<BeobachtenStore>,
@@ -214,6 +216,7 @@ where
                 RunningClientHandler::handle_authenticated_static(
                     crypto_reader, crypto_writer, success,
                     upstream_manager, stats, config, buffer_pool, rng, me_pool,
+                    route_runtime.clone(),
                     local_addr, real_peer, ip_tracker.clone(),
                 ),
             )))
@@ -274,6 +277,7 @@ where
                     buffer_pool,
                     rng,
                     me_pool,
+                    route_runtime.clone(),
                     local_addr,
                     real_peer,
                     ip_tracker.clone(),
@@ -324,6 +328,7 @@ pub struct RunningClientHandler {
     buffer_pool: Arc<BufferPool>,
     rng: Arc<SecureRandom>,
     me_pool: Option<Arc<MePool>>,
+    route_runtime: Arc<RouteRuntimeController>,
     tls_cache: Option<Arc<TlsFrontCache>>,
     ip_tracker: Arc<UserIpTracker>,
     beobachten: Arc<BeobachtenStore>,
@@ -341,6 +346,7 @@ impl ClientHandler {
         buffer_pool: Arc<BufferPool>,
         rng: Arc<SecureRandom>,
         me_pool: Option<Arc<MePool>>,
+        route_runtime: Arc<RouteRuntimeController>,
         tls_cache: Option<Arc<TlsFrontCache>>,
         ip_tracker: Arc<UserIpTracker>,
         beobachten: Arc<BeobachtenStore>,
@@ -356,6 +362,7 @@ impl ClientHandler {
             buffer_pool,
             rng,
             me_pool,
+            route_runtime,
             tls_cache,
             ip_tracker,
             beobachten,
@@ -597,6 +604,7 @@ impl RunningClientHandler {
                 buffer_pool,
                 self.rng,
                 self.me_pool,
+                self.route_runtime.clone(),
                 local_addr,
                 peer,
                 self.ip_tracker,
@@ -677,6 +685,7 @@ impl RunningClientHandler {
                 buffer_pool,
                 self.rng,
                 self.me_pool,
+                self.route_runtime.clone(),
                 local_addr,
                 peer,
                 self.ip_tracker,
@@ -698,6 +707,7 @@ impl RunningClientHandler {
         buffer_pool: Arc<BufferPool>,
         rng: Arc<SecureRandom>,
         me_pool: Option<Arc<MePool>>,
+        route_runtime: Arc<RouteRuntimeController>,
         local_addr: SocketAddr,
         peer_addr: SocketAddr,
         ip_tracker: Arc<UserIpTracker>,
@@ -713,7 +723,11 @@ impl RunningClientHandler {
             return Err(e);
         }
 
-        let relay_result = if config.general.use_middle_proxy {
+        let route_snapshot = route_runtime.snapshot();
+        let session_id = rng.u64();
+        let relay_result = if config.general.use_middle_proxy
+            && matches!(route_snapshot.mode, RelayRouteMode::Middle)
+        {
             if let Some(ref pool) = me_pool {
                 handle_via_middle_proxy(
                     client_reader,
@@ -725,6 +739,9 @@ impl RunningClientHandler {
                     buffer_pool,
                     local_addr,
                     rng,
+                    route_runtime.subscribe(),
+                    route_snapshot,
+                    session_id,
                 )
                 .await
             } else {
@@ -738,6 +755,9 @@ impl RunningClientHandler {
                     config,
                     buffer_pool,
                     rng,
+                    route_runtime.subscribe(),
+                    route_snapshot,
+                    session_id,
                 )
                 .await
             }
@@ -752,6 +772,9 @@ impl RunningClientHandler {
                 config,
                 buffer_pool,
                 rng,
+                route_runtime.subscribe(),
+                route_snapshot,
+                session_id,
             )
             .await
         };

src/proxy/direct_relay.rs

@@ -5,14 +5,19 @@ use std::sync::Arc;
 
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
 use tokio::net::TcpStream;
+use tokio::sync::watch;
 use tracing::{debug, info, warn};
 
 use crate::config::ProxyConfig;
 use crate::crypto::SecureRandom;
-use crate::error::Result;
+use crate::error::{ProxyError, Result};
 use crate::protocol::constants::*;
 use crate::proxy::handshake::{HandshakeSuccess, encrypt_tg_nonce_with_ciphers, generate_tg_nonce};
 use crate::proxy::relay::relay_bidirectional;
+use crate::proxy::route_mode::{
+    RelayRouteMode, RouteCutoverState, ROUTE_SWITCH_ERROR_MSG, affected_cutover_state,
+    cutover_stagger_delay,
+};
 use crate::stats::Stats;
 use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
 use crate::transport::UpstreamManager;
@@ -26,6 +31,9 @@ pub(crate) async fn handle_via_direct<R, W>(
     config: Arc<ProxyConfig>,
     buffer_pool: Arc<BufferPool>,
     rng: Arc<SecureRandom>,
+    mut route_rx: watch::Receiver<RouteCutoverState>,
+    route_snapshot: RouteCutoverState,
+    session_id: u64,
 ) -> Result<()>
 where
     R: AsyncRead + Unpin + Send + 'static,
@@ -69,8 +77,36 @@ where
         user,
         Arc::clone(&stats),
         buffer_pool,
-    )
-    .await;
+    );
+    tokio::pin!(relay_result);
+    let relay_result = loop {
+        if let Some(cutover) = affected_cutover_state(
+            &route_rx,
+            RelayRouteMode::Direct,
+            route_snapshot.generation,
+        ) {
+            let delay = cutover_stagger_delay(session_id, cutover.generation);
+            warn!(
+                user = %user,
+                target_mode = cutover.mode.as_str(),
+                cutover_generation = cutover.generation,
+                delay_ms = delay.as_millis() as u64,
+                "Cutover affected direct session, closing client connection"
+            );
+            tokio::time::sleep(delay).await;
+            break Err(ProxyError::Proxy(ROUTE_SWITCH_ERROR_MSG.to_string()));
+        }
+        tokio::select! {
+            result = &mut relay_result => {
+                break result;
+            }
+            changed = route_rx.changed() => {
+                if changed.is_err() {
+                    break relay_result.await;
+                }
+            }
+        }
+    };
 
     stats.decrement_current_connections_direct();
     stats.decrement_user_curr_connects(user);

src/proxy/middle_relay.rs

@@ -8,7 +8,7 @@ use std::time::{Duration, Instant};
 
 use bytes::Bytes;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::{mpsc, oneshot, watch};
 use tracing::{debug, trace, warn};
 
 use crate::config::ProxyConfig;
@@ -16,6 +16,10 @@ use crate::crypto::SecureRandom;
 use crate::error::{ProxyError, Result};
 use crate::protocol::constants::{*, secure_padding_len};
 use crate::proxy::handshake::HandshakeSuccess;
+use crate::proxy::route_mode::{
+    RelayRouteMode, RouteCutoverState, ROUTE_SWITCH_ERROR_MSG, affected_cutover_state,
+    cutover_stagger_delay,
+};
 use crate::stats::Stats;
 use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
 use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag};
@@ -228,6 +232,9 @@ pub(crate) async fn handle_via_middle_proxy<R, W>(
     _buffer_pool: Arc<BufferPool>,
     local_addr: SocketAddr,
     rng: Arc<SecureRandom>,
+    mut route_rx: watch::Receiver<RouteCutoverState>,
+    route_snapshot: RouteCutoverState,
+    session_id: u64,
 ) -> Result<()>
 where
     R: AsyncRead + Unpin + Send + 'static,
@@ -267,6 +274,27 @@ where
     stats.increment_user_curr_connects(&user);
     stats.increment_current_connections_me();
 
+    if let Some(cutover) = affected_cutover_state(
+        &route_rx,
+        RelayRouteMode::Middle,
+        route_snapshot.generation,
+    ) {
+        let delay = cutover_stagger_delay(session_id, cutover.generation);
+        warn!(
+            conn_id,
+            target_mode = cutover.mode.as_str(),
+            cutover_generation = cutover.generation,
+            delay_ms = delay.as_millis() as u64,
+            "Cutover affected middle session before relay start, closing client connection"
+        );
+        tokio::time::sleep(delay).await;
+        let _ = me_pool.send_close(conn_id).await;
+        me_pool.registry().unregister(conn_id).await;
+        stats.decrement_current_connections_me();
+        stats.decrement_user_curr_connects(&user);
+        return Err(ProxyError::Proxy(ROUTE_SWITCH_ERROR_MSG.to_string()));
+    }
+
     // Per-user ad_tag from access.user_ad_tags; fallback to general.ad_tag (hot-reloadable)
     let user_tag: Option<Vec<u8>> = config
         .access
@@ -498,46 +526,75 @@ where
     let mut main_result: Result<()> = Ok(());
     let mut client_closed = false;
     let mut frame_counter: u64 = 0;
+    let mut route_watch_open = true;
     loop {
-        match read_client_payload(
-            &mut crypto_reader,
-            proto_tag,
-            frame_limit,
-            &forensics,
-            &mut frame_counter,
-            &stats,
-        ).await {
-            Ok(Some((payload, quickack))) => {
-                trace!(conn_id, bytes = payload.len(), "C->ME frame");
-                forensics.bytes_c2me = forensics
-                    .bytes_c2me
-                    .saturating_add(payload.len() as u64);
-                stats.add_user_octets_from(&user, payload.len() as u64);
-                let mut flags = proto_flags;
-                if quickack {
-                    flags |= RPC_FLAG_QUICKACK;
-                }
-                if payload.len() >= 8 && payload[..8].iter().all(|b| *b == 0) {
-                    flags |= RPC_FLAG_NOT_ENCRYPTED;
-                }
-                // Keep client read loop lightweight: route heavy ME send path via a dedicated task.
-                if enqueue_c2me_command(&c2me_tx, C2MeCommand::Data { payload, flags })
-                    .await
-                    .is_err()
-                {
-                    main_result = Err(ProxyError::Proxy("ME sender channel closed".into()));
-                    break;
+        if let Some(cutover) = affected_cutover_state(
+            &route_rx,
+            RelayRouteMode::Middle,
+            route_snapshot.generation,
+        ) {
+            let delay = cutover_stagger_delay(session_id, cutover.generation);
+            warn!(
+                conn_id,
+                target_mode = cutover.mode.as_str(),
+                cutover_generation = cutover.generation,
+                delay_ms = delay.as_millis() as u64,
+                "Cutover affected middle session, closing client connection"
+            );
+            tokio::time::sleep(delay).await;
+            let _ = enqueue_c2me_command(&c2me_tx, C2MeCommand::Close).await;
+            main_result = Err(ProxyError::Proxy(ROUTE_SWITCH_ERROR_MSG.to_string()));
+            break;
+        }
+
+        tokio::select! {
+            changed = route_rx.changed(), if route_watch_open => {
+                if changed.is_err() {
+                    route_watch_open = false;
                 }
             }
-            Ok(None) => {
-                debug!(conn_id, "Client EOF");
-                client_closed = true;
-                let _ = enqueue_c2me_command(&c2me_tx, C2MeCommand::Close).await;
-                break;
-            }
-            Err(e) => {
-                main_result = Err(e);
-                break;
+            payload_result = read_client_payload(
+                &mut crypto_reader,
+                proto_tag,
+                frame_limit,
+                &forensics,
+                &mut frame_counter,
+                &stats,
+            ) => {
+                match payload_result {
+                    Ok(Some((payload, quickack))) => {
+                        trace!(conn_id, bytes = payload.len(), "C->ME frame");
+                        forensics.bytes_c2me = forensics
+                            .bytes_c2me
+                            .saturating_add(payload.len() as u64);
+                        stats.add_user_octets_from(&user, payload.len() as u64);
+                        let mut flags = proto_flags;
+                        if quickack {
+                            flags |= RPC_FLAG_QUICKACK;
+                        }
+                        if payload.len() >= 8 && payload[..8].iter().all(|b| *b == 0) {
+                            flags |= RPC_FLAG_NOT_ENCRYPTED;
+                        }
+                        // Keep client read loop lightweight: route heavy ME send path via a dedicated task.
+                        if enqueue_c2me_command(&c2me_tx, C2MeCommand::Data { payload, flags })
+                            .await
+                            .is_err()
+                        {
+                            main_result = Err(ProxyError::Proxy("ME sender channel closed".into()));
+                            break;
+                        }
+                    }
+                    Ok(None) => {
+                        debug!(conn_id, "Client EOF");
+                        client_closed = true;
+                        let _ = enqueue_c2me_command(&c2me_tx, C2MeCommand::Close).await;
+                        break;
+                    }
+                    Err(e) => {
+                        main_result = Err(e);
+                        break;
+                    }
+                }
             }
         }
     }

src/proxy/mod.rs

@@ -5,6 +5,7 @@ pub mod direct_relay;
 pub mod handshake;
 pub mod masking;
 pub mod middle_relay;
+pub mod route_mode;
 pub mod relay;
 
 pub use client::ClientHandler;

src/proxy/route_mode.rs

@@ -0,0 +1,117 @@
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
+use std::time::Duration;
+
+use tokio::sync::watch;
+
+pub(crate) const ROUTE_SWITCH_ERROR_MSG: &str = "Route mode switched by cutover";
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[repr(u8)]
+pub(crate) enum RelayRouteMode {
+    Direct = 0,
+    Middle = 1,
+}
+
+impl RelayRouteMode {
+    pub(crate) fn as_u8(self) -> u8 {
+        self as u8
+    }
+
+    pub(crate) fn from_u8(value: u8) -> Self {
+        match value {
+            1 => Self::Middle,
+            _ => Self::Direct,
+        }
+    }
+
+    pub(crate) fn as_str(self) -> &'static str {
+        match self {
+            Self::Direct => "direct",
+            Self::Middle => "middle",
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct RouteCutoverState {
+    pub mode: RelayRouteMode,
+    pub generation: u64,
+}
+
+#[derive(Clone)]
+pub(crate) struct RouteRuntimeController {
+    mode: Arc<AtomicU8>,
+    generation: Arc<AtomicU64>,
+    tx: watch::Sender<RouteCutoverState>,
+}
+
+impl RouteRuntimeController {
+    pub(crate) fn new(initial_mode: RelayRouteMode) -> Self {
+        let initial = RouteCutoverState {
+            mode: initial_mode,
+            generation: 0,
+        };
+        let (tx, _rx) = watch::channel(initial);
+        Self {
+            mode: Arc::new(AtomicU8::new(initial_mode.as_u8())),
+            generation: Arc::new(AtomicU64::new(0)),
+            tx,
+        }
+    }
+
+    pub(crate) fn snapshot(&self) -> RouteCutoverState {
+        RouteCutoverState {
+            mode: RelayRouteMode::from_u8(self.mode.load(Ordering::Relaxed)),
+            generation: self.generation.load(Ordering::Relaxed),
+        }
+    }
+
+    pub(crate) fn subscribe(&self) -> watch::Receiver<RouteCutoverState> {
+        self.tx.subscribe()
+    }
+
+    pub(crate) fn set_mode(&self, mode: RelayRouteMode) -> Option<RouteCutoverState> {
+        let previous = self.mode.swap(mode.as_u8(), Ordering::Relaxed);
+        if previous == mode.as_u8() {
+            return None;
+        }
+        let generation = self.generation.fetch_add(1, Ordering::Relaxed) + 1;
+        let next = RouteCutoverState { mode, generation };
+        self.tx.send_replace(next);
+        Some(next)
+    }
+}
+
+pub(crate) fn is_session_affected_by_cutover(
+    current: RouteCutoverState,
+    _session_mode: RelayRouteMode,
+    session_generation: u64,
+) -> bool {
+    current.generation > session_generation
+}
+
+pub(crate) fn affected_cutover_state(
+    rx: &watch::Receiver<RouteCutoverState>,
+    session_mode: RelayRouteMode,
+    session_generation: u64,
+) -> Option<RouteCutoverState> {
+    let current = *rx.borrow();
+    if is_session_affected_by_cutover(current, session_mode, session_generation) {
+        return Some(current);
+    }
+    None
+}
+
+pub(crate) fn cutover_stagger_delay(session_id: u64, generation: u64) -> Duration {
+    let mut value = session_id
+        ^ generation.rotate_left(17)
+        ^ 0x9e37_79b9_7f4a_7c15;
+    value ^= value >> 30;
+    value = value.wrapping_mul(0xbf58_476d_1ce4_e5b9);
+    value ^= value >> 27;
+    value = value.wrapping_mul(0x94d0_49bb_1331_11eb);
+    value ^= value >> 31;
+    let ms = 1000 + (value % 1000);
+    Duration::from_millis(ms)
+}