add handshake failure stage accounting

docs/Architecture/API/API.md

@@ -334,6 +334,7 @@ Returned by `PATCH /v1/config` on success (`200`, or `202` when a reload was acc
 | `connections_bad_total` | `u64` | Failed/invalid client connections. |
 | `connections_bad_by_class` | `ClassCount[]` | Failed/invalid connections grouped by class. |
 | `handshake_failures_by_class` | `ClassCount[]` | Handshake failures grouped by class. |
+| `handshake_failures_by_stage` | `StageCount[]` | Handshake failures grouped by state-machine stage. |
 | `handshake_timeouts_total` | `u64` | Handshake timeout count. |
 | `configured_users` | `usize` | Number of configured users in config. |
 
@@ -343,6 +344,12 @@ Returned by `PATCH /v1/config` on success (`200`, or `202` when a reload was acc
 | `class` | `string` | Failure class label. |
 | `total` | `u64` | Counter value for this class. |
 
+#### `StageCount`
+| Field | Type | Description |
+| --- | --- | --- |
+| `stage` | `string` | State-machine stage label. |
+| `total` | `u64` | Counter value for this stage. |
+
 ### `SystemInfoData`
 | Field | Type | Description |
 | --- | --- | --- |
@@ -927,6 +934,7 @@ JA3 follows the Salesforce ClientHello field order. JA4 follows the FoxIO TLS-cl
 | `connections_bad_total` | `u64` | Failed/invalid connections. |
 | `connections_bad_by_class` | `ClassCount[]` | Failed/invalid connections grouped by class. |
 | `handshake_failures_by_class` | `ClassCount[]` | Handshake failures grouped by class. |
+| `handshake_failures_by_stage` | `StageCount[]` | Handshake failures grouped by state-machine stage. |
 | `handshake_timeouts_total` | `u64` | Handshake timeouts. |
 | `accept_permit_timeout_total` | `u64` | Listener admission permit acquisition timeouts. |
 | `configured_users` | `usize` | Configured user count. |

src/api/mod.rs

@@ -54,8 +54,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, SummaryData, UserActiveIps,
+    PatchUserRequest, ResetUserQuotaResponse, RotateSecretRequest, StageCount, SummaryData,
-    is_valid_username,
+    UserActiveIps, is_valid_username,
 };
 use patch::Patch;
 use runtime_edge::{
@@ -540,12 +540,19 @@ 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,6 +89,12 @@ 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,
@@ -96,6 +102,7 @@ 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,
 }
@@ -113,6 +120,7 @@ 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, UpstreamDcStatus, UpstreamStatus, UpstreamSummaryData, UpstreamsData,
+    MinimalQuarantineData, StageCount, UpstreamDcStatus, UpstreamStatus, UpstreamSummaryData,
-    ZeroAllData, ZeroCodeCount, ZeroCoreData, ZeroDesyncData, ZeroMiddleProxyData, ZeroPoolData,
+    UpstreamsData, ZeroAllData, ZeroCodeCount, ZeroCoreData, ZeroDesyncData, ZeroMiddleProxyData,
-    ZeroUpstreamData,
+    ZeroPoolData, ZeroUpstreamData,
 };
 
 const FEATURE_DISABLED_REASON: &str = "feature_disabled";
@@ -36,6 +36,11 @@ 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()
@@ -50,6 +55,7 @@ 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,6 +739,24 @@ 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"
@@ -3886,6 +3904,7 @@ 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
@@ -3950,6 +3969,9 @@ 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"));
@@ -4156,6 +4178,7 @@ 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, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
 use tokio::net::TcpStream;
@@ -26,6 +26,57 @@ 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>,
@@ -682,10 +733,15 @@ 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?;
@@ -726,6 +782,10 @@ 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) => {
@@ -770,6 +830,10 @@ 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(),
@@ -815,6 +879,10 @@ 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()))?;
@@ -886,6 +954,10 @@ 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);
@@ -937,6 +1009,7 @@ 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,
@@ -948,6 +1021,7 @@ 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,
@@ -1267,9 +1341,14 @@ 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?;
@@ -1283,9 +1362,11 @@ impl RunningClientHandler {
             debug!(peer = %peer, is_tls = is_tls, "Handshake type detected");
 
             if is_tls {
-                self.handle_tls_client(first_bytes, local_addr).await
+                self.handle_tls_client(first_bytes, local_addr, &handshake_stage)
+                    .await
             } else {
-                self.handle_direct_client(first_bytes, local_addr).await
+                self.handle_direct_client(first_bytes, local_addr, &handshake_stage)
+                    .await
             }
         })
         .await
@@ -1294,6 +1375,7 @@ 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,
@@ -1305,6 +1387,7 @@ 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,
@@ -1323,6 +1406,7 @@ impl RunningClientHandler {
         mut self,
         first_bytes: [u8; 5],
         local_addr: SocketAddr,
+        handshake_stage: &AtomicU8,
     ) -> Result<HandshakeOutcome> {
         let peer = self.peer;
 
@@ -1358,6 +1442,7 @@ 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) => {
@@ -1412,6 +1497,7 @@ 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,
@@ -1465,6 +1551,10 @@ 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()
@@ -1541,6 +1631,7 @@ impl RunningClientHandler {
         mut self,
         first_bytes: [u8; 5],
         local_addr: SocketAddr,
+        handshake_stage: &AtomicU8,
     ) -> Result<HandshakeOutcome> {
         let peer = self.peer;
 
@@ -1564,6 +1655,7 @@ 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,6 +50,18 @@ 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,6 +38,21 @@ 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,6 +133,7 @@ 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,