Fix handshake stage attribution and bound counters

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

src/api/mod.rs

@@ -56,8 +56,8 @@ use events::ApiEventStore;
 use http_utils::{error_response, read_json, read_optional_json, success_response};
 use model::{
     ApiFailure, ClassCount, CreateUserRequest, DeleteUserResponse, HealthData, HealthReadyData,
-    PatchUserRequest, ResetUserQuotaResponse, RotateSecretRequest, StageCount, SummaryData,
+    PatchUserRequest, ResetUserQuotaResponse, RotateSecretRequest, SummaryData, UserActiveIps,
-    UserActiveIps, is_valid_username,
+    is_valid_username,
 };
 use patch::Patch;
 use runtime_edge::{
@@ -571,19 +571,12 @@ async fn handle(
                     .into_iter()
                     .map(|(class, total)| ClassCount { class, total })
                     .collect();
-                let handshake_failures_by_stage = shared
-                    .stats
-                    .get_handshake_failure_stage_counts()
-                    .into_iter()
-                    .map(|(stage, total)| StageCount { stage, total })
-                    .collect();
                 let data = SummaryData {
                     uptime_seconds: shared.stats.uptime_secs(),
                     connections_total: shared.stats.get_connects_all(),
                     connections_bad_total: shared.stats.get_connects_bad(),
                     connections_bad_by_class,
                     handshake_failures_by_class,
-                    handshake_failures_by_stage,
                     handshake_timeouts_total: shared.stats.get_handshake_timeouts(),
                     configured_users: cfg.access.users.len(),
                 };

src/api/model.rs

@@ -89,12 +89,6 @@ pub(super) struct ClassCount {
     pub(super) total: u64,
 }
 
-#[derive(Serialize, Clone)]
-pub(super) struct StageCount {
-    pub(super) stage: String,
-    pub(super) total: u64,
-}
-
 #[derive(Serialize)]
 pub(super) struct SummaryData {
     pub(super) uptime_seconds: f64,
@@ -102,7 +96,6 @@ pub(super) struct SummaryData {
     pub(super) connections_bad_total: u64,
     pub(super) connections_bad_by_class: Vec<ClassCount>,
     pub(super) handshake_failures_by_class: Vec<ClassCount>,
-    pub(super) handshake_failures_by_stage: Vec<StageCount>,
     pub(super) handshake_timeouts_total: u64,
     pub(super) configured_users: usize,
 }
@@ -120,7 +113,6 @@ pub(super) struct ZeroCoreData {
     pub(super) connections_bad_total: u64,
     pub(super) connections_bad_by_class: Vec<ClassCount>,
     pub(super) handshake_failures_by_class: Vec<ClassCount>,
-    pub(super) handshake_failures_by_stage: Vec<StageCount>,
     pub(super) handshake_timeouts_total: u64,
     pub(super) accept_permit_timeout_total: u64,
     pub(super) configured_users: usize,

src/api/runtime_stats.rs

@@ -9,9 +9,9 @@ use super::ApiShared;
 use super::model::{
     ClassCount, DcEndpointWriters, DcStatus, DcStatusData, MeWriterStatus, MeWritersData,
     MeWritersSummary, MinimalAllData, MinimalAllPayload, MinimalDcPathData, MinimalMeRuntimeData,
-    MinimalQuarantineData, StageCount, UpstreamDcStatus, UpstreamStatus, UpstreamSummaryData,
+    MinimalQuarantineData, UpstreamDcStatus, UpstreamStatus, UpstreamSummaryData, UpstreamsData,
-    UpstreamsData, ZeroAllData, ZeroCodeCount, ZeroCoreData, ZeroDesyncData, ZeroMiddleProxyData,
+    ZeroAllData, ZeroCodeCount, ZeroCoreData, ZeroDesyncData, ZeroMiddleProxyData, ZeroPoolData,
-    ZeroPoolData, ZeroUpstreamData,
+    ZeroUpstreamData,
 };
 
 const FEATURE_DISABLED_REASON: &str = "feature_disabled";
@@ -36,11 +36,6 @@ pub(super) fn build_zero_all_data(stats: &Stats, configured_users: usize) -> Zer
         .into_iter()
         .map(|(class, total)| ClassCount { class, total })
         .collect();
-    let handshake_failure_stages = stats
-        .get_handshake_failure_stage_counts()
-        .into_iter()
-        .map(|(stage, total)| StageCount { stage, total })
-        .collect();
     let handshake_error_codes = stats
         .get_me_handshake_error_code_counts()
         .into_iter()
@@ -55,7 +50,6 @@ pub(super) fn build_zero_all_data(stats: &Stats, configured_users: usize) -> Zer
             connections_bad_total: stats.get_connects_bad(),
             connections_bad_by_class: bad_connection_classes,
             handshake_failures_by_class: handshake_failure_classes,
-            handshake_failures_by_stage: handshake_failure_stages,
             handshake_timeouts_total: stats.get_handshake_timeouts(),
             accept_permit_timeout_total: stats.get_accept_permit_timeout_total(),
             configured_users,

src/metrics.rs

@@ -739,24 +739,6 @@ async fn render_metrics(
         }
     }
 
-    let _ = writeln!(
-        out,
-        "# HELP telemt_handshake_failures_by_stage_total Handshake failures by state-machine stage"
-    );
-    let _ = writeln!(
-        out,
-        "# TYPE telemt_handshake_failures_by_stage_total counter"
-    );
-    if core_enabled {
-        for (stage, total) in stats.get_handshake_failure_stage_counts() {
-            let _ = writeln!(
-                out,
-                "telemt_handshake_failures_by_stage_total{{stage=\"{}\"}} {}",
-                stage, total
-            );
-        }
-    }
-
     let _ = writeln!(
         out,
         "# HELP telemt_auth_expensive_checks_total Expensive authentication candidate checks executed during handshake validation"
@@ -3904,7 +3886,6 @@ mod tests {
         stats.increment_connects_bad_with_class("tls_handshake_bad_client");
         stats.increment_handshake_timeouts();
         stats.increment_handshake_failure_class("timeout");
-        stats.increment_handshake_failure_stage("tls_post_serverhello_mtproto");
         shared_state
             .handshake
             .auth_expensive_checks_total
@@ -3969,9 +3950,6 @@ mod tests {
         ));
         assert!(output.contains("telemt_handshake_timeouts_total 1"));
         assert!(output.contains("telemt_handshake_failures_by_class_total{class=\"timeout\"} 1"));
-        assert!(output.contains(
-            "telemt_handshake_failures_by_stage_total{stage=\"tls_post_serverhello_mtproto\"} 1"
-        ));
         assert!(output.contains("telemt_auth_expensive_checks_total 9"));
         assert!(output.contains("telemt_auth_budget_exhausted_total 2"));
         assert!(output.contains("telemt_upstream_connect_attempt_total 2"));
@@ -4178,7 +4156,6 @@ mod tests {
         assert!(output.contains("# TYPE telemt_connections_bad_by_class_total counter"));
         assert!(output.contains("# TYPE telemt_handshake_timeouts_total counter"));
         assert!(output.contains("# TYPE telemt_handshake_failures_by_class_total counter"));
-        assert!(output.contains("# TYPE telemt_handshake_failures_by_stage_total counter"));
         assert!(output.contains("# TYPE telemt_auth_expensive_checks_total counter"));
         assert!(output.contains("# TYPE telemt_auth_budget_exhausted_total counter"));
         assert!(output.contains("# TYPE telemt_upstream_connect_attempt_total counter"));

src/proxy/client.rs

@@ -7,7 +7,7 @@ use std::net::{IpAddr, SocketAddr};
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::OnceLock;
-use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
 use tokio::net::TcpStream;
@@ -26,57 +26,6 @@ enum HandshakeOutcome {
     NeedsMasking(PostHandshakeFuture),
 }
 
-#[derive(Clone, Copy)]
-enum HandshakeFailureStage {
-    FirstPacketPrelude,
-    TlsClientHelloBody,
-    TlsCore,
-    TlsPostServerHelloMtproto,
-    DirectMtproto,
-}
-
-impl HandshakeFailureStage {
-    fn as_label(self) -> &'static str {
-        match self {
-            Self::FirstPacketPrelude => "first_packet_prelude",
-            Self::TlsClientHelloBody => "tls_clienthello_body",
-            Self::TlsCore => "tls_core",
-            Self::TlsPostServerHelloMtproto => "tls_post_serverhello_mtproto",
-            Self::DirectMtproto => "direct_mtproto",
-        }
-    }
-
-    fn as_code(self) -> u8 {
-        match self {
-            Self::FirstPacketPrelude => 0,
-            Self::TlsClientHelloBody => 1,
-            Self::TlsCore => 2,
-            Self::TlsPostServerHelloMtproto => 3,
-            Self::DirectMtproto => 4,
-        }
-    }
-
-    fn from_code(code: u8) -> Self {
-        match code {
-            1 => Self::TlsClientHelloBody,
-            2 => Self::TlsCore,
-            3 => Self::TlsPostServerHelloMtproto,
-            4 => Self::DirectMtproto,
-            _ => Self::FirstPacketPrelude,
-        }
-    }
-}
-
-fn set_handshake_failure_stage(stage: &AtomicU8, value: HandshakeFailureStage) {
-    stage.store(value.as_code(), Ordering::Relaxed);
-}
-
-fn increment_handshake_failure_stage(stats: &Stats, stage: &AtomicU8) {
-    stats.increment_handshake_failure_stage(
-        HandshakeFailureStage::from_code(stage.load(Ordering::Relaxed)).as_label(),
-    );
-}
-
 #[must_use = "UserConnectionReservation must be kept alive to retain user/IP reservation until release or drop"]
 struct UserConnectionReservation {
     stats: Arc<Stats>,
@@ -733,15 +682,10 @@ where
     let config_for_timeout = config.clone();
     let beobachten_for_timeout = beobachten.clone();
     let peer_for_timeout = real_peer.ip();
-    let handshake_stage = AtomicU8::new(HandshakeFailureStage::FirstPacketPrelude.as_code());
 
     // Phase 2: active handshake (with timeout after the first client byte)
     let outcome = match timeout(handshake_timeout, async {
         let mut first_bytes = [0u8; 5];
-        set_handshake_failure_stage(
-            &handshake_stage,
-            HandshakeFailureStage::FirstPacketPrelude,
-        );
         if let Some(first_byte) = first_byte {
             first_bytes[0] = first_byte;
             stream.read_exact(&mut first_bytes[1..]).await?;
@@ -782,10 +726,6 @@ where
 
             let mut handshake = vec![0u8; 5 + tls_len];
             handshake[..5].copy_from_slice(&first_bytes);
-            set_handshake_failure_stage(
-                &handshake_stage,
-                HandshakeFailureStage::TlsClientHelloBody,
-            );
             let body_read = match read_with_progress(&mut stream, &mut handshake[5..]).await {
                 Ok(n) => n,
                 Err(e) => {
@@ -830,10 +770,6 @@ where
 
             let (read_half, write_half) = tokio::io::split(stream);
 
-            set_handshake_failure_stage(
-                &handshake_stage,
-                HandshakeFailureStage::TlsCore,
-            );
             let (mut tls_reader, tls_writer, tls_user) = match handle_tls_handshake_with_shared(
                 &handshake, read_half, write_half, real_peer,
                 &config, &replay_checker, &rng, tls_cache.clone(),
@@ -879,10 +815,6 @@ where
             );
 
             debug!(peer = %peer, "Reading MTProto handshake through TLS");
-            set_handshake_failure_stage(
-                &handshake_stage,
-                HandshakeFailureStage::TlsPostServerHelloMtproto,
-            );
             let mtproto_data = tls_reader.read_exact(HANDSHAKE_LEN).await?;
             let mtproto_handshake: [u8; HANDSHAKE_LEN] = mtproto_data[..].try_into()
                 .map_err(|_| ProxyError::InvalidHandshake("Short MTProto handshake".into()))?;
@@ -954,10 +886,6 @@ where
 
             let mut handshake = [0u8; HANDSHAKE_LEN];
             handshake[..5].copy_from_slice(&first_bytes);
-            set_handshake_failure_stage(
-                &handshake_stage,
-                HandshakeFailureStage::DirectMtproto,
-            );
             stream.read_exact(&mut handshake[5..]).await?;
 
             let (read_half, write_half) = tokio::io::split(stream);
@@ -1009,7 +937,6 @@ where
         Ok(Err(e)) => {
             debug!(peer = %peer, error = %e, "Handshake failed");
             stats_for_timeout.increment_handshake_failure_class(classify_handshake_failure_class(&e));
-            increment_handshake_failure_stage(stats_for_timeout.as_ref(), &handshake_stage);
             record_handshake_failure_class(
                 &beobachten_for_timeout,
                 &config_for_timeout,
@@ -1021,7 +948,6 @@ where
         Err(_) => {
             stats_for_timeout.increment_handshake_timeouts();
             stats_for_timeout.increment_handshake_failure_class("timeout");
-            increment_handshake_failure_stage(stats_for_timeout.as_ref(), &handshake_stage);
             debug!(peer = %peer, "Handshake timeout");
             record_beobachten_class(
                 &beobachten_for_timeout,
@@ -1341,14 +1267,9 @@ impl RunningClientHandler {
         let beobachten_for_timeout = self.beobachten.clone();
         let peer_for_timeout = self.peer.ip();
         let peer_for_log = self.peer;
-        let handshake_stage = AtomicU8::new(HandshakeFailureStage::FirstPacketPrelude.as_code());
 
         let outcome = match timeout(handshake_timeout, async {
             let mut first_bytes = [0u8; 5];
-            set_handshake_failure_stage(
-                &handshake_stage,
-                HandshakeFailureStage::FirstPacketPrelude,
-            );
             if let Some(first_byte) = first_byte {
                 first_bytes[0] = first_byte;
                 self.stream.read_exact(&mut first_bytes[1..]).await?;
@@ -1362,11 +1283,9 @@ impl RunningClientHandler {
             debug!(peer = %peer, is_tls = is_tls, "Handshake type detected");
 
             if is_tls {
-                self.handle_tls_client(first_bytes, local_addr, &handshake_stage)
+                self.handle_tls_client(first_bytes, local_addr).await
-                    .await
             } else {
-                self.handle_direct_client(first_bytes, local_addr, &handshake_stage)
+                self.handle_direct_client(first_bytes, local_addr).await
-                    .await
             }
         })
         .await
@@ -1375,7 +1294,6 @@ impl RunningClientHandler {
             Ok(Err(e)) => {
                 debug!(peer = %peer_for_log, error = %e, "Handshake failed");
                 stats.increment_handshake_failure_class(classify_handshake_failure_class(&e));
-                increment_handshake_failure_stage(stats.as_ref(), &handshake_stage);
                 record_handshake_failure_class(
                     &beobachten_for_timeout,
                     &config_for_timeout,
@@ -1387,7 +1305,6 @@ impl RunningClientHandler {
             Err(_) => {
                 stats.increment_handshake_timeouts();
                 stats.increment_handshake_failure_class("timeout");
-                increment_handshake_failure_stage(stats.as_ref(), &handshake_stage);
                 debug!(peer = %peer_for_log, "Handshake timeout");
                 record_beobachten_class(
                     &beobachten_for_timeout,
@@ -1406,7 +1323,6 @@ impl RunningClientHandler {
         mut self,
         first_bytes: [u8; 5],
         local_addr: SocketAddr,
-        handshake_stage: &AtomicU8,
     ) -> Result<HandshakeOutcome> {
         let peer = self.peer;
 
@@ -1442,7 +1358,6 @@ impl RunningClientHandler {
 
         let mut handshake = vec![0u8; 5 + tls_len];
         handshake[..5].copy_from_slice(&first_bytes);
-        set_handshake_failure_stage(handshake_stage, HandshakeFailureStage::TlsClientHelloBody);
         let body_read = match read_with_progress(&mut self.stream, &mut handshake[5..]).await {
             Ok(n) => n,
             Err(e) => {
@@ -1497,7 +1412,6 @@ impl RunningClientHandler {
 
         let (read_half, write_half) = self.stream.into_split();
 
-        set_handshake_failure_stage(handshake_stage, HandshakeFailureStage::TlsCore);
         let (mut tls_reader, tls_writer, tls_user) = match handle_tls_handshake_with_shared(
             &handshake,
             read_half,
@@ -1551,10 +1465,6 @@ impl RunningClientHandler {
         );
 
         debug!(peer = %peer, "Reading MTProto handshake through TLS");
-        set_handshake_failure_stage(
-            handshake_stage,
-            HandshakeFailureStage::TlsPostServerHelloMtproto,
-        );
         let mtproto_data = tls_reader.read_exact(HANDSHAKE_LEN).await?;
         let mtproto_handshake: [u8; HANDSHAKE_LEN] = mtproto_data[..]
             .try_into()
@@ -1631,7 +1541,6 @@ impl RunningClientHandler {
         mut self,
         first_bytes: [u8; 5],
         local_addr: SocketAddr,
-        handshake_stage: &AtomicU8,
     ) -> Result<HandshakeOutcome> {
         let peer = self.peer;
 
@@ -1655,7 +1564,6 @@ impl RunningClientHandler {
 
         let mut handshake = [0u8; HANDSHAKE_LEN];
         handshake[..5].copy_from_slice(&first_bytes);
-        set_handshake_failure_stage(handshake_stage, HandshakeFailureStage::DirectMtproto);
         self.stream.read_exact(&mut handshake[5..]).await?;
 
         let config = self.config.clone();

src/stats/core_counters.rs

@@ -50,18 +50,6 @@ impl Stats {
             .or_insert_with(|| AtomicU64::new(0));
         entry.fetch_add(1, Ordering::Relaxed);
     }
-
-    pub fn increment_handshake_failure_stage(&self, stage: &'static str) {
-        if !self.telemetry_core_enabled() {
-            return;
-        }
-        let entry = self
-            .handshake_failure_stages
-            .entry(stage)
-            .or_insert_with(|| AtomicU64::new(0));
-        entry.fetch_add(1, Ordering::Relaxed);
-    }
-
     pub fn increment_current_connections_direct(&self) {
         self.current_connections_direct
             .fetch_add(1, Ordering::Relaxed);

src/stats/core_getters.rs

@@ -38,21 +38,6 @@ impl Stats {
         out
     }
 
-    pub fn get_handshake_failure_stage_counts(&self) -> Vec<(String, u64)> {
-        let mut out: Vec<(String, u64)> = self
-            .handshake_failure_stages
-            .iter()
-            .map(|entry| {
-                (
-                    entry.key().to_string(),
-                    entry.value().load(Ordering::Relaxed),
-                )
-            })
-            .collect();
-        out.sort_by(|a, b| a.0.cmp(&b.0));
-        out
-    }
-
     pub fn get_accept_permit_timeout_total(&self) -> u64 {
         self.accept_permit_timeout_total.load(Ordering::Relaxed)
     }

src/stats/mod.rs

@@ -133,7 +133,6 @@ pub struct Stats {
     connects_bad: AtomicU64,
     connects_bad_classes: DashMap<&'static str, AtomicU64>,
     handshake_failure_classes: DashMap<&'static str, AtomicU64>,
-    handshake_failure_stages: DashMap<&'static str, AtomicU64>,
     current_connections_direct: AtomicU64,
     current_connections_me: AtomicU64,
     route_cutover_parked_direct_current: AtomicU64,