konduit-public/konduit-platform/src/stats.rs

// Connection statistics tracking for konduit
// Provides real-time metrics and historical data for GUI display

use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Instant;

/// Ring buffer size for historical data (300 seconds = 5 minutes at 1 sample/sec)
const HISTORY_SIZE: usize = 300;

/// Connection statistics snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConnectionStats {
    /// Connection status: "connected", "disconnected", "connecting"
    pub status: String,

    /// Duration in seconds since connection started
    pub duration: u64,

    /// Total bytes uploaded
    pub upload_bytes: u64,

    /// Total bytes downloaded
    pub download_bytes: u64,

    /// Current upload speed (bytes/sec)
    pub upload_speed: u64,

    /// Current download speed (bytes/sec)
    pub download_speed: u64,

    /// Server address
    pub server: String,

    /// Peer ID
    pub peer_id: String,

    /// Local TUN interface IP (e.g., "10.0.0.2")
    pub local_ip: Option<String>,

    /// Upload speed history (last 300 seconds)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub upload_history: Option<Vec<u64>>,

    /// Download speed history (last 300 seconds)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub download_history: Option<Vec<u64>>,
}

/// Statistics tracker for a VPN connection
pub struct StatsTracker {
    /// Connection start time
    start_time: Instant,

    /// Total bytes uploaded (atomic for thread-safe updates)
    upload_bytes: Arc<AtomicU64>,

    /// Total bytes downloaded (atomic for thread-safe updates)
    download_bytes: Arc<AtomicU64>,

    /// Previous upload bytes (for speed calculation)
    prev_upload: AtomicU64,

    /// Previous download bytes (for speed calculation)
    prev_download: AtomicU64,

    /// Upload speed history ring buffer
    upload_history: parking_lot::Mutex<RingBuffer>,

    /// Download speed history ring buffer
    download_history: parking_lot::Mutex<RingBuffer>,

    /// Server address
    server: String,

    /// Peer ID
    peer_id: String,

    /// Local TUN IP
    local_ip: parking_lot::Mutex<Option<String>>,
}

/// Ring buffer for historical speed data
struct RingBuffer {
    data: Vec<u64>,
    index: usize,
    filled: bool,
}

impl RingBuffer {
    fn new() -> Self {
        Self {
            data: vec![0; HISTORY_SIZE],
            index: 0,
            filled: false,
        }
    }

    fn push(&mut self, value: u64) {
        self.data[self.index] = value;
        self.index = (self.index + 1) % HISTORY_SIZE;
        if self.index == 0 {
            self.filled = true;
        }
    }

    fn to_vec(&self) -> Vec<u64> {
        if !self.filled {
            // Not filled yet, return only valid data (newest first)
            self.data[0..self.index].iter().rev().copied().collect()
        } else {
            // Buffer is full, return in correct order (newest first)
            let mut result = Vec::with_capacity(HISTORY_SIZE);

            // Start from current position - 1 (newest) and go backwards
            for i in 0..HISTORY_SIZE {
                let idx = (self.index + HISTORY_SIZE - 1 - i) % HISTORY_SIZE;
                result.push(self.data[idx]);
            }

            result
        }
    }
}

impl StatsTracker {
    /// Create a new stats tracker
    pub fn new(server: String, peer_id: String) -> Self {
        Self {
            start_time: Instant::now(),
            upload_bytes: Arc::new(AtomicU64::new(0)),
            download_bytes: Arc::new(AtomicU64::new(0)),
            prev_upload: AtomicU64::new(0),
            prev_download: AtomicU64::new(0),
            upload_history: parking_lot::Mutex::new(RingBuffer::new()),
            download_history: parking_lot::Mutex::new(RingBuffer::new()),
            server,
            peer_id,
            local_ip: parking_lot::Mutex::new(None),
        }
    }

    /// Get atomic upload counter (for incrementing from packet handlers)
    pub fn upload_counter(&self) -> Arc<AtomicU64> {
        Arc::clone(&self.upload_bytes)
    }

    /// Get atomic download counter (for incrementing from packet handlers)
    pub fn download_counter(&self) -> Arc<AtomicU64> {
        Arc::clone(&self.download_bytes)
    }

    /// Record bytes uploaded (convenience method)
    pub fn record_upload(&self, bytes: usize) {
        self.upload_bytes.fetch_add(bytes as u64, Ordering::Relaxed);
    }

    /// Record bytes downloaded (convenience method)
    pub fn record_download(&self, bytes: usize) {
        self.download_bytes
            .fetch_add(bytes as u64, Ordering::Relaxed);
    }

    /// Set local TUN interface IP
    pub fn set_local_ip(&self, ip: String) {
        *self.local_ip.lock() = Some(ip);
    }

    /// Update speed calculations and history (call this every second)
    pub fn update(&self) {
        let current_upload = self.upload_bytes.load(Ordering::Relaxed);
        let current_download = self.download_bytes.load(Ordering::Relaxed);

        let prev_upload = self.prev_upload.swap(current_upload, Ordering::Relaxed);
        let prev_download = self.prev_download.swap(current_download, Ordering::Relaxed);

        let upload_speed = current_upload.saturating_sub(prev_upload);
        let download_speed = current_download.saturating_sub(prev_download);

        self.upload_history.lock().push(upload_speed);
        self.download_history.lock().push(download_speed);
    }

    /// Get current statistics snapshot
    pub fn get_stats(&self, include_history: bool) -> ConnectionStats {
        let current_upload = self.upload_bytes.load(Ordering::Relaxed);
        let current_download = self.download_bytes.load(Ordering::Relaxed);
        let prev_upload = self.prev_upload.load(Ordering::Relaxed);
        let prev_download = self.prev_download.load(Ordering::Relaxed);

        ConnectionStats {
            status: "connected".to_string(),
            duration: self.start_time.elapsed().as_secs(),
            upload_bytes: current_upload,
            download_bytes: current_download,
            upload_speed: current_upload.saturating_sub(prev_upload),
            download_speed: current_download.saturating_sub(prev_download),
            server: self.server.clone(),
            peer_id: self.peer_id.clone(),
            local_ip: self.local_ip.lock().clone(),
            upload_history: if include_history {
                Some(self.upload_history.lock().to_vec())
            } else {
                None
            },
            download_history: if include_history {
                Some(self.download_history.lock().to_vec())
            } else {
                None
            },
        }
    }

    /// Get stats as JSON string
    pub fn get_stats_json(&self, include_history: bool) -> Result<String, serde_json::Error> {
        serde_json::to_string(&self.get_stats(include_history))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread;

    #[test]
    fn test_stats_tracker_creation() {
        let tracker = StatsTracker::new("utilbox.eu:8443".to_string(), "test".to_string());

        let stats = tracker.get_stats(false);
        assert_eq!(stats.status, "connected");
        assert_eq!(stats.upload_bytes, 0);
        assert_eq!(stats.download_bytes, 0);
        assert_eq!(stats.server, "utilbox.eu:8443");
        assert_eq!(stats.peer_id, "test");
    }

    #[test]
    fn test_record_traffic() {
        let tracker = StatsTracker::new("utilbox.eu:8443".to_string(), "test".to_string());

        tracker.record_upload(1024);
        tracker.record_download(2048);

        let stats = tracker.get_stats(false);
        assert_eq!(stats.upload_bytes, 1024);
        assert_eq!(stats.download_bytes, 2048);
    }

    #[test]
    fn test_ring_buffer() {
        let mut buffer = RingBuffer::new();

        // Add some values
        for i in 0..10 {
            buffer.push(i);
        }

        let vec = buffer.to_vec();
        assert_eq!(vec.len(), 10);
        assert_eq!(vec[0], 9); // Newest first
        assert_eq!(vec[9], 0); // Oldest last
    }

    #[test]
    fn test_ring_buffer_wraparound() {
        let mut buffer = RingBuffer::new();

        // Fill buffer and wrap around
        for i in 0..350 {
            buffer.push(i);
        }

        let vec = buffer.to_vec();
        assert_eq!(vec.len(), HISTORY_SIZE);
        assert_eq!(vec[0], 349); // Newest
        assert_eq!(vec[HISTORY_SIZE - 1], 50); // Oldest (350 - 300)
    }

    #[test]
    fn test_json_serialization() {
        let tracker = StatsTracker::new("utilbox.eu:8443".to_string(), "test".to_string());

        tracker.record_upload(1024);
        tracker.record_download(2048);
        tracker.set_local_ip("10.0.0.2".to_string());

        let json = tracker.get_stats_json(false).unwrap();
        assert!(json.contains("\"status\":\"connected\""));
        assert!(json.contains("\"upload_bytes\":1024"));
        assert!(json.contains("\"download_bytes\":2048"));
        assert!(json.contains("\"local_ip\":\"10.0.0.2\""));
    }

    #[test]
    fn test_atomic_counters() {
        let tracker = StatsTracker::new("utilbox.eu:8443".to_string(), "test".to_string());

        let upload_counter = tracker.upload_counter();
        let download_counter = tracker.download_counter();

        // Simulate concurrent updates
        let handles: Vec<_> = (0..10)
            .map(|_| {
                let up = Arc::clone(&upload_counter);
                let down = Arc::clone(&download_counter);
                thread::spawn(move || {
                    for _ in 0..100 {
                        up.fetch_add(1, Ordering::Relaxed);
                        down.fetch_add(2, Ordering::Relaxed);
                    }
                })
            })
            .collect();

        for handle in handles {
            handle.join().unwrap();
        }

        let stats = tracker.get_stats(false);
        assert_eq!(stats.upload_bytes, 1000); // 10 threads * 100 increments
        assert_eq!(stats.download_bytes, 2000); // 10 threads * 100 * 2
    }
}