From d2e275be561b3f6f81eee39e4b046bb577324c18 Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Tue, 22 Apr 2025 11:35:38 +1000
Subject: [PATCH] feat(connlib): classify UDP traffic by protocol (#8886)

It creates a bit of duplication with code that we have in `snownet` but
it is code that is unlikely to change because the protocols are already
standarised. Contrary to recording the port, the cardinality of these
protocols is much fixed to a much smaller range which will allow us to
safely record these metrics in an actual time-series database further
down the line whilst still reasoning about how much traffic we are
sending over TURN, as STUN or as WireGuard.
---
 rust/connlib/snownet/src/lib.rs  |  4 ++++
 rust/connlib/snownet/src/node.rs | 13 +++++++------
 rust/connlib/tunnel/src/io.rs    |  2 +-
 rust/connlib/tunnel/src/lib.rs   |  4 ++--
 rust/connlib/tunnel/src/otel.rs  | 24 ++++++++++++++++--------
 5 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/rust/connlib/snownet/src/lib.rs b/rust/connlib/snownet/src/lib.rs
index a52c140bb..8940a1b19 100644
--- a/rust/connlib/snownet/src/lib.rs
+++ b/rust/connlib/snownet/src/lib.rs
@@ -20,3 +20,7 @@ pub use node::{
     NoTurnServers, Node, Server, ServerNode, Transmit,
 };
 pub use stats::{ConnectionStats, NodeStats};
+
+pub fn is_wireguard(payload: &[u8]) -> bool {
+    boringtun::noise::Tunn::parse_incoming_packet(payload).is_ok()
+}
diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs
index 01db325da..46d83613b 100644
--- a/rust/connlib/snownet/src/node.rs
+++ b/rust/connlib/snownet/src/node.rs
@@ -927,14 +927,15 @@ where
             };
         }
 
-        match Tunn::parse_incoming_packet(packet) {
-            Ok(_) => tracing::trace!(
+        if crate::is_wireguard(packet) {
+            tracing::trace!(
                 "Packet was a WireGuard packet but no connection handled it. Already disconnected?"
-            ),
-            Err(_) => return ControlFlow::Break(Err(Error::UnknownPacketFormat)),
-        };
+            );
 
-        ControlFlow::Break(Ok(()))
+            return ControlFlow::Break(Ok(()));
+        }
+
+        ControlFlow::Break(Err(Error::UnknownPacketFormat))
     }
 
     fn allocations_drain_events(&mut self) {
diff --git a/rust/connlib/tunnel/src/io.rs b/rust/connlib/tunnel/src/io.rs
index ba61d265b..ca6ccff5b 100644
--- a/rust/connlib/tunnel/src/io.rs
+++ b/rust/connlib/tunnel/src/io.rs
@@ -356,7 +356,7 @@ impl Io {
         self.packet_counter.add(
             1,
             &[
-                crate::otel::network_peer_port(dst.port()),
+                crate::otel::network_protocol_name(payload),
                 crate::otel::network_transport_udp(),
                 crate::otel::network_io_direction_transmit(),
             ],
diff --git a/rust/connlib/tunnel/src/lib.rs b/rust/connlib/tunnel/src/lib.rs
index e58f19093..e98be8e75 100644
--- a/rust/connlib/tunnel/src/lib.rs
+++ b/rust/connlib/tunnel/src/lib.rs
@@ -195,7 +195,7 @@ impl ClientTunnel {
                         self.packet_counter.add(
                             1,
                             &[
-                                crate::otel::network_peer_port(received.from.port()),
+                                crate::otel::network_protocol_name(received.packet),
                                 crate::otel::network_transport_udp(),
                                 crate::otel::network_io_direction_receive(),
                             ],
@@ -327,7 +327,7 @@ impl GatewayTunnel {
                         self.packet_counter.add(
                             1,
                             &[
-                                crate::otel::network_peer_port(received.from.port()),
+                                crate::otel::network_protocol_name(received.packet),
                                 crate::otel::network_transport_udp(),
                                 crate::otel::network_io_direction_receive(),
                             ],
diff --git a/rust/connlib/tunnel/src/otel.rs b/rust/connlib/tunnel/src/otel.rs
index d98107071..3b04b3c4b 100644
--- a/rust/connlib/tunnel/src/otel.rs
+++ b/rust/connlib/tunnel/src/otel.rs
@@ -1,14 +1,6 @@
 use ip_packet::IpPacket;
 use opentelemetry::KeyValue;
 
-// Recording discrete values can lead to a cardinality explosion.
-// We only use metrics for local debugging and not in production.
-// Locally, the set of ports will be small so we don't need to worry about this.
-// If this ever changes, we need to be more clever here in classifying the protocol.
-pub fn network_peer_port(p: u16) -> KeyValue {
-    KeyValue::new("network.peer.port", p as i64)
-}
-
 pub fn network_transport_udp() -> KeyValue {
     KeyValue::new("network.transport", "udp")
 }
@@ -20,6 +12,22 @@ pub fn network_type_for_packet(p: &IpPacket) -> KeyValue {
     }
 }
 
+pub fn network_protocol_name(payload: &[u8]) -> KeyValue {
+    const KEY: &str = "network.protocol.name";
+
+    match payload {
+        [0..3, ..] => KeyValue::new(KEY, "stun"),
+        // Channel-data is a 4-byte header so the actual payload starts on the 5th byte
+        [64..=79, _, _, _, 0..3, ..] => KeyValue::new(KEY, "stun-over-turn"),
+        [64..=79, _, _, _, payload @ ..] if snownet::is_wireguard(payload) => {
+            KeyValue::new(KEY, "wireguard-over-turn")
+        }
+        [64..=79, _, _, _, ..] => KeyValue::new(KEY, "unknown-over-turn"),
+        payload if snownet::is_wireguard(payload) => KeyValue::new(KEY, "wireguard"),
+        _ => KeyValue::new(KEY, "unknown"),
+    }
+}
+
 pub fn network_type_ipv4() -> KeyValue {
     KeyValue::new("network.type", "ipv4")
 }