shadow_network_sim/

correlation.rs

//! Traffic correlation analysis — detect relationships between flows.

use serde::{Serialize, Deserialize};

/// A recorded traffic flow for correlation analysis.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FlowRecord {
    /// Unique flow identifier.
    pub flow_id: String,
    /// Packet sizes in order.
    pub packet_sizes: Vec<usize>,
    /// Timestamps in microseconds.
    pub timestamps_us: Vec<u64>,
    /// Direction markers (true = outgoing, false = incoming).
    pub directions: Vec<bool>,
    /// Source label (for testing purposes).
    pub source_label: String,
}

impl FlowRecord {
    pub fn new(flow_id: impl Into<String>) -> Self {
        Self {
            flow_id: flow_id.into(),
            packet_sizes: Vec::new(),
            timestamps_us: Vec::new(),
            directions: Vec::new(),
            source_label: String::new(),
        }
    }

    /// Add a packet to this flow.
    pub fn add_packet(&mut self, size: usize, timestamp_us: u64, outgoing: bool) {
        self.packet_sizes.push(size);
        self.timestamps_us.push(timestamp_us);
        self.directions.push(outgoing);
    }

    /// Duration of the flow in microseconds.
    pub fn duration_us(&self) -> u64 {
        if self.timestamps_us.len() < 2 {
            return 0;
        }
        self.timestamps_us.last().unwrap() - self.timestamps_us.first().unwrap()
    }

    /// Total bytes transferred.
    pub fn total_bytes(&self) -> usize {
        self.packet_sizes.iter().sum()
    }

    /// Inter-packet delays in microseconds.
    pub fn inter_packet_delays(&self) -> Vec<u64> {
        self.timestamps_us
            .windows(2)
            .map(|w| w[1] - w[0])
            .collect()
    }

    /// Average packet size.
    pub fn avg_packet_size(&self) -> f64 {
        if self.packet_sizes.is_empty() {
            return 0.0;
        }
        self.packet_sizes.iter().sum::<usize>() as f64 / self.packet_sizes.len() as f64
    }

    /// Packet count.
    pub fn packet_count(&self) -> usize {
        self.packet_sizes.len()
    }
}

/// Result of correlating two flows.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrelationResult {
    pub flow_a: String,
    pub flow_b: String,
    /// Timing correlation coefficient (0.0–1.0).
    pub timing_correlation: f64,
    /// Size correlation coefficient (0.0–1.0).
    pub size_correlation: f64,
    /// Combined correlation score.
    pub combined_score: f64,
    /// Whether the correlation is significant.
    pub is_correlated: bool,
}

/// Analyzer for detecting traffic correlation patterns.
pub struct CorrelationAnalyzer {
    /// Threshold above which two flows are considered correlated.
    pub correlation_threshold: f64,
    /// Tolerance for timing offset (microseconds).
    pub timing_tolerance_us: u64,
}

impl CorrelationAnalyzer {
    pub fn new() -> Self {
        Self {
            correlation_threshold: 0.7,
            timing_tolerance_us: 10_000, // 10ms
        }
    }

    pub fn with_threshold(mut self, threshold: f64) -> Self {
        self.correlation_threshold = threshold;
        self
    }

    /// Analyze timing correlation between two flows.
    pub fn timing_correlation(&self, a: &FlowRecord, b: &FlowRecord) -> f64 {
        let delays_a: Vec<f64> = a.inter_packet_delays().iter().map(|&d| d as f64).collect();
        let delays_b: Vec<f64> = b.inter_packet_delays().iter().map(|&d| d as f64).collect();

        if delays_a.is_empty() || delays_b.is_empty() {
            return 0.0;
        }

        pearson(&delays_a, &delays_b).abs()
    }

    /// Analyze size-pattern correlation between two flows.
    pub fn size_correlation(&self, a: &FlowRecord, b: &FlowRecord) -> f64 {
        let sizes_a: Vec<f64> = a.packet_sizes.iter().map(|&s| s as f64).collect();
        let sizes_b: Vec<f64> = b.packet_sizes.iter().map(|&s| s as f64).collect();

        if sizes_a.is_empty() || sizes_b.is_empty() {
            return 0.0;
        }

        pearson(&sizes_a, &sizes_b).abs()
    }

    /// Full correlation analysis between two flows.
    pub fn correlate(&self, a: &FlowRecord, b: &FlowRecord) -> CorrelationResult {
        let timing = self.timing_correlation(a, b);
        let size = self.size_correlation(a, b);

        // Weighted combination: timing is more important
        let combined = 0.6 * timing + 0.4 * size;

        CorrelationResult {
            flow_a: a.flow_id.clone(),
            flow_b: b.flow_id.clone(),
            timing_correlation: timing,
            size_correlation: size,
            combined_score: combined,
            is_correlated: combined >= self.correlation_threshold,
        }
    }

    /// Analyze all pairs from a set of flows.
    pub fn analyze_all(&self, flows: &[FlowRecord]) -> Vec<CorrelationResult> {
        let mut results = Vec::new();
        for i in 0..flows.len() {
            for j in (i + 1)..flows.len() {
                results.push(self.correlate(&flows[i], &flows[j]));
            }
        }
        results
    }

    /// Find flows correlated with a target flow.
    pub fn find_correlated(
        &self,
        target: &FlowRecord,
        candidates: &[FlowRecord],
    ) -> Vec<CorrelationResult> {
        candidates
            .iter()
            .filter(|c| c.flow_id != target.flow_id)
            .map(|c| self.correlate(target, c))
            .filter(|r| r.is_correlated)
            .collect()
    }

    /// Evaluate how well traffic shaping resists correlation.
    pub fn resistance_score(&self, original: &FlowRecord, shaped: &FlowRecord) -> f64 {
        let result = self.correlate(original, shaped);
        // Lower correlation = better resistance
        1.0 - result.combined_score
    }
}

impl Default for CorrelationAnalyzer {
    fn default() -> Self {
        Self::new()
    }
}

/// Pearson correlation coefficient.
fn pearson(a: &[f64], b: &[f64]) -> f64 {
    let n = a.len().min(b.len());
    if n < 2 {
        return 0.0;
    }

    let mean_a = a[..n].iter().sum::<f64>() / n as f64;
    let mean_b = b[..n].iter().sum::<f64>() / n as f64;

    let mut cov = 0.0;
    let mut var_a = 0.0;
    let mut var_b = 0.0;

    for i in 0..n {
        let da = a[i] - mean_a;
        let db = b[i] - mean_b;
        cov += da * db;
        var_a += da * da;
        var_b += db * db;
    }

    let denom = (var_a * var_b).sqrt();
    if denom < 1e-12 {
        0.0
    } else {
        cov / denom
    }
}

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

    fn make_flow(id: &str, sizes: Vec<usize>, timestamps: Vec<u64>) -> FlowRecord {
        let mut flow = FlowRecord::new(id);
        for i in 0..sizes.len() {
            flow.add_packet(sizes[i], timestamps[i], true);
        }
        flow
    }

    #[test]
    fn test_identical_flows_correlated() {
        let analyzer = CorrelationAnalyzer::new().with_threshold(0.9);
        let a = make_flow("a", vec![100, 200, 150, 300], vec![0, 1000, 2500, 4000]);
        let b = make_flow("b", vec![100, 200, 150, 300], vec![50, 1050, 2550, 4050]);

        let result = analyzer.correlate(&a, &b);
        assert!(result.combined_score > 0.9);
        assert!(result.is_correlated);
    }

    #[test]
    fn test_different_flows_uncorrelated() {
        let analyzer = CorrelationAnalyzer::new().with_threshold(0.7);
        let a = make_flow("a", vec![100, 200, 100, 200], vec![0, 1000, 2000, 3000]);
        let b = make_flow("b", vec![500, 50, 500, 50], vec![0, 5000, 5500, 20000]);

        let result = analyzer.correlate(&a, &b);
        assert!(result.combined_score < 0.7, "Different patterns should not correlate: {}", result.combined_score);
    }

    #[test]
    fn test_resistance_score() {
        let analyzer = CorrelationAnalyzer::new();
        let original = make_flow("orig", vec![100, 200, 150], vec![0, 1000, 2500]);
        let shaped = make_flow("shaped", vec![256, 256, 256], vec![0, 1000, 2000]);

        let resistance = analyzer.resistance_score(&original, &shaped);
        // Constant-size traffic should reduce correlation
        assert!(resistance > 0.0, "Shaped traffic should show some resistance");
    }

    #[test]
    fn test_analyze_all_pairs() {
        let analyzer = CorrelationAnalyzer::new();
        let flows = vec![
            make_flow("a", vec![100, 200, 300], vec![0, 1000, 2000]),
            make_flow("b", vec![100, 200, 300], vec![50, 1050, 2050]),
            make_flow("c", vec![400, 50, 400], vec![0, 5000, 10000]),
        ];

        let results = analyzer.analyze_all(&flows);
        assert_eq!(results.len(), 3); // C(3,2) = 3 pairs
    }

    #[test]
    fn test_flow_record_stats() {
        let flow = make_flow("test", vec![100, 200, 300], vec![0, 1000, 3000]);
        assert_eq!(flow.packet_count(), 3);
        assert_eq!(flow.total_bytes(), 600);
        assert_eq!(flow.duration_us(), 3000);
        assert!((flow.avg_packet_size() - 200.0).abs() < 0.01);

        let delays = flow.inter_packet_delays();
        assert_eq!(delays, vec![1000, 2000]);
    }
}