Conntrack Control Method

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

src/proxy/client.rs

@@ -80,7 +80,7 @@ use crate::transport::middle_proxy::MePool;
 use crate::transport::socket::normalize_ip;
 use crate::transport::{UpstreamManager, configure_client_socket, parse_proxy_protocol};
 
-use crate::proxy::direct_relay::handle_via_direct;
+use crate::proxy::direct_relay::handle_via_direct_with_shared;
 use crate::proxy::handshake::{
     HandshakeSuccess, handle_mtproto_handshake_with_shared, handle_tls_handshake_with_shared,
 };
@@ -191,6 +191,24 @@ fn handshake_timeout_with_mask_grace(config: &ProxyConfig) -> Duration {
     }
 }
 
+fn effective_client_first_byte_idle_secs(config: &ProxyConfig, shared: &ProxySharedState) -> u64 {
+    let idle_secs = config.timeouts.client_first_byte_idle_secs;
+    if idle_secs == 0 {
+        return 0;
+    }
+    if shared.conntrack_pressure_active() {
+        idle_secs.min(
+            config
+                .server
+                .conntrack_control
+                .profile
+                .client_first_byte_idle_cap_secs(),
+        )
+    } else {
+        idle_secs
+    }
+}
+
 const MASK_CLASSIFIER_PREFETCH_WINDOW: usize = 16;
 #[cfg(test)]
 const MASK_CLASSIFIER_PREFETCH_TIMEOUT: Duration = Duration::from_millis(5);
@@ -463,10 +481,11 @@ where
 
     debug!(peer = %real_peer, "New connection (generic stream)");
 
-    let first_byte = if config.timeouts.client_first_byte_idle_secs == 0 {
+    let first_byte_idle_secs = effective_client_first_byte_idle_secs(&config, shared.as_ref());
+    let first_byte = if first_byte_idle_secs == 0 {
         None
     } else {
-        let idle_timeout = Duration::from_secs(config.timeouts.client_first_byte_idle_secs);
+        let idle_timeout = Duration::from_secs(first_byte_idle_secs);
         let mut first_byte = [0u8; 1];
         match timeout(idle_timeout, stream.read(&mut first_byte)).await {
             Ok(Ok(0)) => {
@@ -499,15 +518,15 @@ where
                 );
                 return Err(ProxyError::Io(e));
             }
-            Err(_) => {
-                debug!(
-                    peer = %real_peer,
-                    idle_secs = config.timeouts.client_first_byte_idle_secs,
-                    "Closing idle pooled connection before first client byte"
-                );
-                return Ok(());
+                Err(_) => {
+                    debug!(
+                        peer = %real_peer,
+                        idle_secs = first_byte_idle_secs,
+                        "Closing idle pooled connection before first client byte"
+                    );
+                    return Ok(());
+                }
             }
-        }
     };
 
     let handshake_timeout = handshake_timeout_with_mask_grace(&config);
@@ -968,11 +987,12 @@ impl RunningClientHandler {
             }
         }
 
-        let first_byte = if self.config.timeouts.client_first_byte_idle_secs == 0 {
+        let first_byte_idle_secs =
+            effective_client_first_byte_idle_secs(&self.config, self.shared.as_ref());
+        let first_byte = if first_byte_idle_secs == 0 {
             None
         } else {
-            let idle_timeout =
-                Duration::from_secs(self.config.timeouts.client_first_byte_idle_secs);
+            let idle_timeout = Duration::from_secs(first_byte_idle_secs);
             let mut first_byte = [0u8; 1];
             match timeout(idle_timeout, self.stream.read(&mut first_byte)).await {
                 Ok(Ok(0)) => {
@@ -1008,7 +1028,7 @@ impl RunningClientHandler {
                 Err(_) => {
                     debug!(
                         peer = %self.peer,
-                        idle_secs = self.config.timeouts.client_first_byte_idle_secs,
+                        idle_secs = first_byte_idle_secs,
                         "Closing idle pooled connection before first client byte"
                     );
                     return Ok(None);
@@ -1395,7 +1415,7 @@ impl RunningClientHandler {
         local_addr: SocketAddr,
         peer_addr: SocketAddr,
         ip_tracker: Arc<UserIpTracker>,
-        _shared: Arc<ProxySharedState>,
+        shared: Arc<ProxySharedState>,
     ) -> Result<()>
     where
         R: AsyncRead + Unpin + Send + 'static,
@@ -1438,12 +1458,12 @@ impl RunningClientHandler {
                     route_runtime.subscribe(),
                     route_snapshot,
                     session_id,
-                    _shared,
+                    shared.clone(),
                 )
                 .await
             } else {
                 warn!("use_middle_proxy=true but MePool not initialized, falling back to direct");
-                handle_via_direct(
+                handle_via_direct_with_shared(
                     client_reader,
                     client_writer,
                     success,
@@ -1455,12 +1475,14 @@ impl RunningClientHandler {
                     route_runtime.subscribe(),
                     route_snapshot,
                     session_id,
+                    local_addr,
+                    shared.clone(),
                 )
                 .await
             }
         } else {
             // Direct mode (original behavior)
-            handle_via_direct(
+            handle_via_direct_with_shared(
                 client_reader,
                 client_writer,
                 success,
@@ -1472,6 +1494,8 @@ impl RunningClientHandler {
                 route_runtime.subscribe(),
                 route_snapshot,
                 session_id,
+                local_addr,
+                shared.clone(),
             )
             .await
         };

src/proxy/direct_relay.rs

@@ -6,6 +6,7 @@ use std::net::SocketAddr;
 use std::path::{Component, Path, PathBuf};
 use std::sync::Arc;
 use std::sync::{Mutex, OnceLock};
+use std::time::Duration;
 
 use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, ReadHalf, WriteHalf, split};
 use tokio::sync::watch;
@@ -16,7 +17,9 @@ use crate::crypto::SecureRandom;
 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::shared_state::{
+    ConntrackCloseEvent, ConntrackClosePublishResult, ConntrackCloseReason, ProxySharedState,
+};
 use crate::proxy::route_mode::{
     ROUTE_SWITCH_ERROR_MSG, RelayRouteMode, RouteCutoverState, affected_cutover_state,
     cutover_stagger_delay,
@@ -225,7 +228,43 @@ fn unknown_dc_test_lock() -> &'static Mutex<()> {
     TEST_LOCK.get_or_init(|| Mutex::new(()))
 }
 
+#[allow(dead_code)]
 pub(crate) async fn handle_via_direct<R, W>(
+    client_reader: CryptoReader<R>,
+    client_writer: CryptoWriter<W>,
+    success: HandshakeSuccess,
+    upstream_manager: Arc<UpstreamManager>,
+    stats: Arc<Stats>,
+    config: Arc<ProxyConfig>,
+    buffer_pool: Arc<BufferPool>,
+    rng: Arc<SecureRandom>,
+    route_rx: watch::Receiver<RouteCutoverState>,
+    route_snapshot: RouteCutoverState,
+    session_id: u64,
+) -> Result<()>
+where
+    R: AsyncRead + Unpin + Send + 'static,
+    W: AsyncWrite + Unpin + Send + 'static,
+{
+    handle_via_direct_with_shared(
+        client_reader,
+        client_writer,
+        success,
+        upstream_manager,
+        stats,
+        config.clone(),
+        buffer_pool,
+        rng,
+        route_rx,
+        route_snapshot,
+        session_id,
+        SocketAddr::from(([0, 0, 0, 0], config.server.port)),
+        ProxySharedState::new(),
+    )
+    .await
+}
+
+pub(crate) async fn handle_via_direct_with_shared<R, W>(
     client_reader: CryptoReader<R>,
     client_writer: CryptoWriter<W>,
     success: HandshakeSuccess,
@@ -237,6 +276,8 @@ pub(crate) async fn handle_via_direct<R, W>(
     mut route_rx: watch::Receiver<RouteCutoverState>,
     route_snapshot: RouteCutoverState,
     session_id: u64,
+    local_addr: SocketAddr,
+    shared: Arc<ProxySharedState>,
 ) -> Result<()>
 where
     R: AsyncRead + Unpin + Send + 'static,
@@ -277,7 +318,18 @@ where
     let _direct_connection_lease = stats.acquire_direct_connection_lease();
 
     let buffer_pool_trim = Arc::clone(&buffer_pool);
-    let relay_result = relay_bidirectional(
+    let relay_activity_timeout = if shared.conntrack_pressure_active() {
+        Duration::from_secs(
+            config
+                .server
+                .conntrack_control
+                .profile
+                .direct_activity_timeout_secs(),
+        )
+    } else {
+        Duration::from_secs(1800)
+    };
+    let relay_result = crate::proxy::relay::relay_bidirectional_with_activity_timeout(
         client_reader,
         client_writer,
         tg_reader,
@@ -288,6 +340,7 @@ where
         Arc::clone(&stats),
         config.access.user_data_quota.get(user).copied(),
         buffer_pool,
+        relay_activity_timeout,
     );
     tokio::pin!(relay_result);
     let relay_result = loop {
@@ -329,9 +382,52 @@ where
         pool_snapshot.allocated,
         pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
     );
+
+    let close_reason = classify_conntrack_close_reason(&relay_result);
+    let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
+        src: success.peer,
+        dst: local_addr,
+        reason: close_reason,
+    });
+    if !matches!(
+        publish_result,
+        ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
+    ) {
+        stats.increment_conntrack_close_event_drop_total();
+    }
+
     relay_result
 }
 
+fn classify_conntrack_close_reason(result: &Result<()>) -> ConntrackCloseReason {
+    match result {
+        Ok(()) => ConntrackCloseReason::NormalEof,
+        Err(crate::error::ProxyError::Io(error))
+            if matches!(error.kind(), std::io::ErrorKind::TimedOut) =>
+        {
+            ConntrackCloseReason::Timeout
+        }
+        Err(crate::error::ProxyError::Io(error))
+            if matches!(
+                error.kind(),
+                std::io::ErrorKind::ConnectionReset
+                    | std::io::ErrorKind::ConnectionAborted
+                    | std::io::ErrorKind::BrokenPipe
+                    | std::io::ErrorKind::NotConnected
+                    | std::io::ErrorKind::UnexpectedEof
+            ) =>
+        {
+            ConntrackCloseReason::Reset
+        }
+        Err(crate::error::ProxyError::Proxy(message))
+            if message.contains("pressure") || message.contains("evicted") =>
+        {
+            ConntrackCloseReason::Pressure
+        }
+        Err(_) => ConntrackCloseReason::Other,
+    }
+}
+
 fn get_dc_addr_static(dc_idx: i16, config: &ProxyConfig) -> Result<SocketAddr> {
     let prefer_v6 = config.network.prefer == 6 && config.network.ipv6.unwrap_or(true);
     let datacenters = if prefer_v6 {

src/proxy/middle_relay.rs

@@ -16,12 +16,14 @@ use tokio::sync::{mpsc, oneshot, watch};
 use tokio::time::timeout;
 use tracing::{debug, info, trace, warn};
 
-use crate::config::ProxyConfig;
+use crate::config::{ConntrackPressureProfile, ProxyConfig};
 use crate::crypto::SecureRandom;
 use crate::error::{ProxyError, Result};
 use crate::protocol::constants::{secure_padding_len, *};
 use crate::proxy::handshake::HandshakeSuccess;
-use crate::proxy::shared_state::ProxySharedState;
+use crate::proxy::shared_state::{
+    ConntrackCloseEvent, ConntrackClosePublishResult, ConntrackCloseReason, ProxySharedState,
+};
 use crate::proxy::route_mode::{
     ROUTE_SWITCH_ERROR_MSG, RelayRouteMode, RouteCutoverState, affected_cutover_state,
     cutover_stagger_delay,
@@ -135,6 +137,10 @@ fn note_relay_pressure_event_in(shared: &ProxySharedState) {
     guard.pressure_event_seq = guard.pressure_event_seq.wrapping_add(1);
 }
 
+pub(crate) fn note_global_relay_pressure(shared: &ProxySharedState) {
+    note_relay_pressure_event_in(shared);
+}
+
 fn relay_pressure_event_seq_in(shared: &ProxySharedState) -> u64 {
     let guard = relay_idle_candidate_registry_lock_in(shared);
     guard.pressure_event_seq
@@ -241,6 +247,23 @@ impl RelayClientIdlePolicy {
             legacy_frame_read_timeout: frame_read_timeout,
         }
     }
+
+    fn apply_pressure_caps(&mut self, profile: ConntrackPressureProfile) {
+        let pressure_soft_idle_cap = Duration::from_secs(profile.middle_soft_idle_cap_secs());
+        let pressure_hard_idle_cap = Duration::from_secs(profile.middle_hard_idle_cap_secs());
+
+        self.soft_idle = self.soft_idle.min(pressure_soft_idle_cap);
+        self.hard_idle = self.hard_idle.min(pressure_hard_idle_cap);
+        if self.soft_idle > self.hard_idle {
+            self.soft_idle = self.hard_idle;
+        }
+        self.legacy_frame_read_timeout = self
+            .legacy_frame_read_timeout
+            .min(pressure_hard_idle_cap);
+        if self.grace_after_downstream_activity > self.hard_idle {
+            self.grace_after_downstream_activity = self.hard_idle;
+        }
+    }
 }
 
 #[derive(Clone, Copy)]
@@ -1027,7 +1050,12 @@ where
     let translated_local_addr = me_pool.translate_our_addr(local_addr);
 
     let frame_limit = config.general.max_client_frame;
-    let relay_idle_policy = RelayClientIdlePolicy::from_config(&config);
+    let mut relay_idle_policy = RelayClientIdlePolicy::from_config(&config);
+    let mut pressure_caps_applied = false;
+    if shared.conntrack_pressure_active() {
+        relay_idle_policy.apply_pressure_caps(config.server.conntrack_control.profile);
+        pressure_caps_applied = true;
+    }
     let session_started_at = forensics.started_at;
     let mut relay_idle_state = RelayClientIdleState::new(session_started_at);
     let last_downstream_activity_ms = Arc::new(AtomicU64::new(0));
@@ -1421,6 +1449,11 @@ where
     let mut route_watch_open = true;
     let mut seen_pressure_seq = relay_pressure_event_seq_in(shared.as_ref());
     loop {
+        if shared.conntrack_pressure_active() && !pressure_caps_applied {
+            relay_idle_policy.apply_pressure_caps(config.server.conntrack_control.profile);
+            pressure_caps_applied = true;
+        }
+
         if relay_idle_policy.enabled
             && maybe_evict_idle_candidate_on_pressure_in(
                 shared.as_ref(),
@@ -1600,6 +1633,20 @@ where
         frames_ok = frame_counter,
         "ME relay cleanup"
     );
+
+    let close_reason = classify_conntrack_close_reason(&result);
+    let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
+        src: peer,
+        dst: local_addr,
+        reason: close_reason,
+    });
+    if !matches!(
+        publish_result,
+        ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
+    ) {
+        stats.increment_conntrack_close_event_drop_total();
+    }
+
     clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
     me_pool.registry().unregister(conn_id).await;
     buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
@@ -1612,6 +1659,33 @@ where
     result
 }
 
+fn classify_conntrack_close_reason(result: &Result<()>) -> ConntrackCloseReason {
+    match result {
+        Ok(()) => ConntrackCloseReason::NormalEof,
+        Err(ProxyError::Io(error)) if matches!(error.kind(), std::io::ErrorKind::TimedOut) => {
+            ConntrackCloseReason::Timeout
+        }
+        Err(ProxyError::Io(error))
+            if matches!(
+                error.kind(),
+                std::io::ErrorKind::ConnectionReset
+                    | std::io::ErrorKind::ConnectionAborted
+                    | std::io::ErrorKind::BrokenPipe
+                    | std::io::ErrorKind::NotConnected
+                    | std::io::ErrorKind::UnexpectedEof
+            ) =>
+        {
+            ConntrackCloseReason::Reset
+        }
+        Err(ProxyError::Proxy(message))
+            if message.contains("pressure") || message.contains("evicted") =>
+        {
+            ConntrackCloseReason::Pressure
+        }
+        Err(_) => ConntrackCloseReason::Other,
+    }
+}
+
 async fn read_client_payload_with_idle_policy_in<R>(
     client_reader: &mut CryptoReader<R>,
     proto_tag: ProtoTag,

src/proxy/relay.rs

@@ -70,6 +70,7 @@ use tracing::{debug, trace, warn};
 ///
 /// iOS keeps Telegram connections alive in background for up to 30 minutes.
 /// Closing earlier causes unnecessary reconnects and handshake overhead.
+#[allow(dead_code)]
 const ACTIVITY_TIMEOUT: Duration = Duration::from_secs(1800);
 
 /// Watchdog check interval — also used for periodic rate logging.
@@ -453,6 +454,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
 /// - Clean shutdown: both write sides are shut down on exit
 /// - Error propagation: quota exits return `ProxyError::DataQuotaExceeded`,
 ///   other I/O failures are returned as `ProxyError::Io`
+#[allow(dead_code)]
 pub async fn relay_bidirectional<CR, CW, SR, SW>(
     client_reader: CR,
     client_writer: CW,
@@ -471,6 +473,42 @@ where
     SR: AsyncRead + Unpin + Send + 'static,
     SW: AsyncWrite + Unpin + Send + 'static,
 {
+    relay_bidirectional_with_activity_timeout(
+        client_reader,
+        client_writer,
+        server_reader,
+        server_writer,
+        c2s_buf_size,
+        s2c_buf_size,
+        user,
+        stats,
+        quota_limit,
+        _buffer_pool,
+        ACTIVITY_TIMEOUT,
+    )
+    .await
+}
+
+pub async fn relay_bidirectional_with_activity_timeout<CR, CW, SR, SW>(
+    client_reader: CR,
+    client_writer: CW,
+    server_reader: SR,
+    server_writer: SW,
+    c2s_buf_size: usize,
+    s2c_buf_size: usize,
+    user: &str,
+    stats: Arc<Stats>,
+    quota_limit: Option<u64>,
+    _buffer_pool: Arc<BufferPool>,
+    activity_timeout: Duration,
+) -> Result<()>
+where
+    CR: AsyncRead + Unpin + Send + 'static,
+    CW: AsyncWrite + Unpin + Send + 'static,
+    SR: AsyncRead + Unpin + Send + 'static,
+    SW: AsyncWrite + Unpin + Send + 'static,
+{
+    let activity_timeout = activity_timeout.max(Duration::from_secs(1));
     let epoch = Instant::now();
     let counters = Arc::new(SharedCounters::new());
     let quota_exceeded = Arc::new(AtomicBool::new(false));
@@ -512,7 +550,7 @@ where
             }
 
             // ── Activity timeout ────────────────────────────────────
-            if idle >= ACTIVITY_TIMEOUT {
+            if idle >= activity_timeout {
                 let c2s = wd_counters.c2s_bytes.load(Ordering::Relaxed);
                 let s2c = wd_counters.s2c_bytes.load(Ordering::Relaxed);
                 warn!(

src/proxy/shared_state.rs

@@ -1,15 +1,40 @@
 use std::collections::HashSet;
 use std::collections::hash_map::RandomState;
-use std::net::IpAddr;
-use std::sync::atomic::AtomicU64;
+use std::net::{IpAddr, SocketAddr};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
 use dashmap::DashMap;
+use tokio::sync::mpsc;
 
 use crate::proxy::handshake::{AuthProbeState, AuthProbeSaturationState};
 use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ConntrackCloseReason {
+    NormalEof,
+    Timeout,
+    Pressure,
+    Reset,
+    Other,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct ConntrackCloseEvent {
+    pub(crate) src: SocketAddr,
+    pub(crate) dst: SocketAddr,
+    pub(crate) reason: ConntrackCloseReason,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ConntrackClosePublishResult {
+    Sent,
+    Disabled,
+    QueueFull,
+    QueueClosed,
+}
+
 pub(crate) struct HandshakeSharedState {
     pub(crate) auth_probe: DashMap<IpAddr, AuthProbeState>,
     pub(crate) auth_probe_saturation: Mutex<Option<AuthProbeSaturationState>>,
@@ -31,6 +56,8 @@ pub(crate) struct MiddleRelaySharedState {
 pub(crate) struct ProxySharedState {
     pub(crate) handshake: HandshakeSharedState,
     pub(crate) middle_relay: MiddleRelaySharedState,
+    pub(crate) conntrack_pressure_active: AtomicBool,
+    pub(crate) conntrack_close_tx: Mutex<Option<mpsc::Sender<ConntrackCloseEvent>>>,
 }
 
 impl ProxySharedState {
@@ -52,6 +79,67 @@ impl ProxySharedState {
                 relay_idle_registry: Mutex::new(RelayIdleCandidateRegistry::default()),
                 relay_idle_mark_seq: AtomicU64::new(0),
             },
+            conntrack_pressure_active: AtomicBool::new(false),
+            conntrack_close_tx: Mutex::new(None),
         })
     }
+
+    pub(crate) fn set_conntrack_close_sender(&self, tx: mpsc::Sender<ConntrackCloseEvent>) {
+        match self.conntrack_close_tx.lock() {
+            Ok(mut guard) => {
+                *guard = Some(tx);
+            }
+            Err(poisoned) => {
+                let mut guard = poisoned.into_inner();
+                *guard = Some(tx);
+                self.conntrack_close_tx.clear_poison();
+            }
+        }
+    }
+
+    pub(crate) fn disable_conntrack_close_sender(&self) {
+        match self.conntrack_close_tx.lock() {
+            Ok(mut guard) => {
+                *guard = None;
+            }
+            Err(poisoned) => {
+                let mut guard = poisoned.into_inner();
+                *guard = None;
+                self.conntrack_close_tx.clear_poison();
+            }
+        }
+    }
+
+    pub(crate) fn publish_conntrack_close_event(
+        &self,
+        event: ConntrackCloseEvent,
+    ) -> ConntrackClosePublishResult {
+        let tx = match self.conntrack_close_tx.lock() {
+            Ok(guard) => guard.clone(),
+            Err(poisoned) => {
+                let guard = poisoned.into_inner();
+                let cloned = guard.clone();
+                self.conntrack_close_tx.clear_poison();
+                cloned
+            }
+        };
+
+        let Some(tx) = tx else {
+            return ConntrackClosePublishResult::Disabled;
+        };
+
+        match tx.try_send(event) {
+            Ok(()) => ConntrackClosePublishResult::Sent,
+            Err(mpsc::error::TrySendError::Full(_)) => ConntrackClosePublishResult::QueueFull,
+            Err(mpsc::error::TrySendError::Closed(_)) => ConntrackClosePublishResult::QueueClosed,
+        }
+    }
+
+    pub(crate) fn set_conntrack_pressure_active(&self, active: bool) {
+        self.conntrack_pressure_active.store(active, Ordering::Relaxed);
+    }
+
+    pub(crate) fn conntrack_pressure_active(&self) -> bool {
+        self.conntrack_pressure_active.load(Ordering::Relaxed)
+    }
 }