1 files changed, 337 insertions, 0 deletions

diff --git a/src/metrics/connection.rs b/src/metrics/connection.rs
new file mode 100644
index 0000000..6a7f406
--- /dev/null
+++ b/src/metrics/connection.rs
@@ -0,0 +1,337 @@
+//! Connection tracking with privacy-preserving abuse detection.
+//!
+//! This module tracks WebSocket connections per IP address internally for abuse
+//! detection, but NEVER exposes IP addresses in Prometheus metrics. Only aggregate
+//! counts are exposed.
+//!
+//! # Privacy Model
+//!
+//! | Data | Location | Exposed? |
+//! |------|----------|----------|
+//! | Total connections | Prometheus | ✅ Yes |
+//! | Unique IP count | Prometheus | ✅ Yes |
+//! | Flagged abuser count | Prometheus | ✅ Yes |
+//! | Actual IP addresses | Internal HashMap | ❌ No |
+//! | IP + abuse flag | Logs (when flagged) | ⚠️ Logs only |
+
+use std::net::IpAddr;
+use std::time::Instant;
+
+use dashmap::DashMap;
+use prometheus::{IntGauge, Opts, Registry};
+use tracing::warn;
+
+/// Information about connections from a specific IP address.
+struct ConnectionInfo {
+    /// Number of active connections from this IP
+    count: u32,
+    /// When the first connection from this IP was established
+    first_seen: Instant,
+    /// Whether this IP has been flagged as potentially abusive
+    flagged_as_abuse: bool,
+}
+
+/// Tracks WebSocket connections per IP with abuse detection.
+///
+/// # Thread Safety
+///
+/// Uses `DashMap` for lock-free concurrent access, as connection tracking
+/// happens across multiple tokio tasks.
+///
+/// # Privacy
+///
+/// IP addresses are stored internally only for abuse detection and are
+/// NEVER exposed in Prometheus metrics. Only aggregate counts are exposed:
+/// - Total active connections
+/// - Number of unique IPs
+/// - Number of IPs flagged as potential abusers
+pub struct ConnectionTracker {
+    /// Active connections per IP (INTERNAL ONLY - never exposed to metrics)
+    connections: DashMap<IpAddr, ConnectionInfo>,
+    
+    /// Threshold for abuse flagging (connections per IP)
+    abuse_threshold: u32,
+    
+    /// Prometheus gauge: total active connections
+    active_connections: IntGauge,
+    
+    /// Prometheus gauge: number of unique IPs connected
+    unique_ips: IntGauge,
+    
+    /// Prometheus gauge: number of IPs flagged as potential abusers
+    flagged_abusers: IntGauge,
+}
+
+impl ConnectionTracker {
+    /// Creates a new ConnectionTracker and registers metrics with Prometheus.
+    ///
+    /// # Arguments
+    ///
+    /// * `abuse_threshold` - Number of connections from a single IP before flagging
+    /// * `registry` - Prometheus registry to register metrics with
+    pub fn new(abuse_threshold: u32, registry: &Registry) -> Self {
+        let active_connections = IntGauge::with_opts(
+            Opts::new(
+                "ngit_websocket_connections_active",
+                "Current active WebSocket connections",
+            )
+        ).unwrap();
+        registry.register(Box::new(active_connections.clone())).unwrap();
+
+        let unique_ips = IntGauge::with_opts(
+            Opts::new(
+                "ngit_websocket_unique_ips",
+                "Number of unique IP addresses connected (NOT the IPs themselves)",
+            )
+        ).unwrap();
+        registry.register(Box::new(unique_ips.clone())).unwrap();
+
+        let flagged_abusers = IntGauge::with_opts(
+            Opts::new(
+                "ngit_websocket_flagged_abusers",
+                "Number of IPs exceeding connection threshold",
+            )
+        ).unwrap();
+        registry.register(Box::new(flagged_abusers.clone())).unwrap();
+
+        Self {
+            connections: DashMap::new(),
+            abuse_threshold,
+            active_connections,
+            unique_ips,
+            flagged_abusers,
+        }
+    }
+
+    /// Called when a new WebSocket connection is established.
+    ///
+    /// This method:
+    /// 1. Increments the connection count for this IP
+    /// 2. Checks if the IP has exceeded the abuse threshold
+    /// 3. Logs a warning if abuse is detected (IP is logged here only)
+    /// 4. Updates Prometheus metrics (aggregate counts only)
+    ///
+    /// # Privacy
+    ///
+    /// The IP address is logged only when abuse is detected. It is NEVER
+    /// exposed in Prometheus metrics.
+    pub fn on_connect(&self, ip: IpAddr) {
+        let mut is_new_ip = false;
+        let mut newly_flagged = false;
+
+        self.connections
+            .entry(ip)
+            .and_modify(|info| {
+                info.count += 1;
+                // Check if this connection pushes us over the threshold
+                if !info.flagged_as_abuse && info.count >= self.abuse_threshold {
+                    info.flagged_as_abuse = true;
+                    newly_flagged = true;
+                }
+            })
+            .or_insert_with(|| {
+                is_new_ip = true;
+                ConnectionInfo {
+                    count: 1,
+                    first_seen: Instant::now(),
+                    flagged_as_abuse: false,
+                }
+            });
+
+        // Update Prometheus metrics (aggregate counts only)
+        self.active_connections.inc();
+        
+        if is_new_ip {
+            self.unique_ips.inc();
+        }
+
+        if newly_flagged {
+            self.flagged_abusers.inc();
+            // Log the abuse detection - IP is only exposed in logs, not metrics
+            warn!(
+                ip = %ip,
+                threshold = self.abuse_threshold,
+                "Potential abuse detected: IP exceeded connection threshold"
+            );
+        }
+    }
+
+    /// Called when a WebSocket connection is closed.
+    ///
+    /// This method:
+    /// 1. Decrements the connection count for this IP
+    /// 2. Removes the IP from tracking if count reaches 0
+    /// 3. Updates the abuse flag count if the IP was flagged
+    /// 4. Updates Prometheus metrics (aggregate counts only)
+    pub fn on_disconnect(&self, ip: IpAddr) {
+        let mut remove_entry = false;
+        let mut was_flagged = false;
+        let mut had_connection = false;
+
+        if let Some(mut entry) = self.connections.get_mut(&ip) {
+            had_connection = true;
+            entry.count = entry.count.saturating_sub(1);
+            if entry.count == 0 {
+                remove_entry = true;
+                was_flagged = entry.flagged_as_abuse;
+            }
+        }
+
+        // Remove the entry if count is 0
+        if remove_entry {
+            self.connections.remove(&ip);
+            self.unique_ips.dec();
+            if was_flagged {
+                self.flagged_abusers.dec();
+            }
+        }
+
+        // Update total connections only if this IP had a tracked connection
+        if had_connection {
+            self.active_connections.dec();
+        }
+    }
+
+    /// Returns the current number of active connections.
+    pub fn active_connections(&self) -> u64 {
+        self.active_connections.get() as u64
+    }
+
+    /// Returns the current number of unique IPs.
+    pub fn unique_ip_count(&self) -> u64 {
+        self.unique_ips.get() as u64
+    }
+
+    /// Returns the current number of flagged abusers.
+    pub fn flagged_abuser_count(&self) -> u64 {
+        self.flagged_abusers.get() as u64
+    }
+
+    /// Returns the connection count for a specific IP (for internal use only).
+    ///
+    /// # Privacy
+    ///
+    /// This is an internal method. The returned data should NEVER be exposed
+    /// in metrics or logs without privacy consideration.
+    #[cfg(test)]
+    pub(crate) fn connection_count(&self, ip: &IpAddr) -> Option<u32> {
+        self.connections.get(ip).map(|info| info.count)
+    }
+
+    /// Returns whether an IP is flagged as abusive (for internal use only).
+    #[cfg(test)]
+    pub(crate) fn is_flagged(&self, ip: &IpAddr) -> bool {
+        self.connections
+            .get(ip)
+            .map(|info| info.flagged_as_abuse)
+            .unwrap_or(false)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::net::{Ipv4Addr, Ipv6Addr};
+
+    fn test_registry() -> Registry {
+        Registry::new()
+    }
+
+    #[test]
+    fn test_connection_tracking() {
+        let registry = test_registry();
+        let tracker = ConnectionTracker::new(5, &registry);
+        let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
+
+        // Connect
+        tracker.on_connect(ip);
+        assert_eq!(tracker.active_connections(), 1);
+        assert_eq!(tracker.unique_ip_count(), 1);
+        assert_eq!(tracker.connection_count(&ip), Some(1));
+
+        // Connect again from same IP
+        tracker.on_connect(ip);
+        assert_eq!(tracker.active_connections(), 2);
+        assert_eq!(tracker.unique_ip_count(), 1); // Still 1 unique IP
+        assert_eq!(tracker.connection_count(&ip), Some(2));
+
+        // Disconnect one
+        tracker.on_disconnect(ip);
+        assert_eq!(tracker.active_connections(), 1);
+        assert_eq!(tracker.unique_ip_count(), 1);
+        assert_eq!(tracker.connection_count(&ip), Some(1));
+
+        // Disconnect last
+        tracker.on_disconnect(ip);
+        assert_eq!(tracker.active_connections(), 0);
+        assert_eq!(tracker.unique_ip_count(), 0);
+        assert_eq!(tracker.connection_count(&ip), None);
+    }
+
+    #[test]
+    fn test_multiple_ips() {
+        let registry = test_registry();
+        let tracker = ConnectionTracker::new(5, &registry);
+        let ip1 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
+        let ip2 = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 2));
+        let ip3 = IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1));
+
+        tracker.on_connect(ip1);
+        tracker.on_connect(ip2);
+        tracker.on_connect(ip3);
+
+        assert_eq!(tracker.active_connections(), 3);
+        assert_eq!(tracker.unique_ip_count(), 3);
+
+        tracker.on_disconnect(ip2);
+        assert_eq!(tracker.active_connections(), 2);
+        assert_eq!(tracker.unique_ip_count(), 2);
+    }
+
+    #[test]
+    fn test_abuse_detection() {
+        let registry = test_registry();
+        let threshold = 3;
+        let tracker = ConnectionTracker::new(threshold, &registry);
+        let abuser_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1));
+        let normal_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
+
+        // Normal user with 1 connection
+        tracker.on_connect(normal_ip);
+        assert!(!tracker.is_flagged(&normal_ip));
+        assert_eq!(tracker.flagged_abuser_count(), 0);
+
+        // Abuser approaching threshold
+        tracker.on_connect(abuser_ip);
+        tracker.on_connect(abuser_ip);
+        assert!(!tracker.is_flagged(&abuser_ip));
+        assert_eq!(tracker.flagged_abuser_count(), 0);
+
+        // Abuser hits threshold
+        tracker.on_connect(abuser_ip);
+        assert!(tracker.is_flagged(&abuser_ip));
+        assert_eq!(tracker.flagged_abuser_count(), 1);
+
+        // Normal user still not flagged
+        assert!(!tracker.is_flagged(&normal_ip));
+
+        // Abuser disconnects all - should be removed from flagged count
+        tracker.on_disconnect(abuser_ip);
+        tracker.on_disconnect(abuser_ip);
+        tracker.on_disconnect(abuser_ip);
+        assert_eq!(tracker.flagged_abuser_count(), 0);
+        assert_eq!(tracker.active_connections(), 1); // Only normal user remains
+    }
+
+    #[test]
+    fn test_disconnect_without_connect() {
+        let registry = test_registry();
+        let tracker = ConnectionTracker::new(5, &registry);
+        let ip = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1));
+
+        // Disconnect without connect should not panic or go negative
+        tracker.on_disconnect(ip);
+        assert_eq!(tracker.active_connections(), 0);
+        assert_eq!(tracker.unique_ip_count(), 0);
+    }
+}
\ No newline at end of file