Add: Leak-D for dns leak detection

crates/wtfnet-dnsleak/src/lib.rs

@@ -7,7 +7,7 @@ mod rules;
 mod sensor;
 
 use crate::classify::ClassifiedEvent;
-use crate::sensor::capture_events;
+use crate::sensor::{capture_events, SensorEvent, TcpEvent};
 use std::time::Instant;
 use thiserror::Error;
 use tracing::debug;
@@ -50,13 +50,30 @@ pub async fn watch(
     let start = Instant::now();
     let events = capture_events(&options).await?;
     let mut leak_events = Vec::new();
+    let mut dns_cache: std::collections::HashMap<std::net::IpAddr, DnsCacheEntry> =
+        std::collections::HashMap::new();
 
     for event in events {
-        let enriched = enrich_event(event, flow_owner).await;
-        if let Some(decision) = rules::evaluate(&enriched, &options.policy) {
-            let mut leak_event = report::LeakEvent::from_decision(enriched, decision);
-            privacy::apply_privacy(&mut leak_event, options.privacy);
-            leak_events.push(leak_event);
+        match event {
+            SensorEvent::Dns(event) => {
+                let enriched = enrich_event(event, flow_owner).await;
+                if enriched.is_response {
+                    update_dns_cache(&mut dns_cache, &enriched);
+                    continue;
+                }
+                if let Some(decision) = rules::evaluate(&enriched, &options.policy) {
+                    let mut leak_event = report::LeakEvent::from_decision(enriched, decision);
+                    privacy::apply_privacy(&mut leak_event, options.privacy);
+                    leak_events.push(leak_event);
+                }
+            }
+            SensorEvent::Tcp(event) => {
+                if let Some(leak_event) =
+                    evaluate_mismatch(event, flow_owner, &mut dns_cache, options.privacy).await
+                {
+                    leak_events.push(leak_event);
+                }
+            }
         }
     }
 
@@ -100,3 +117,106 @@ async fn enrich_event(
     }
     enriched
 }
+
+struct DnsCacheEntry {
+    qname: String,
+    route_class: RouteClass,
+    timestamp_ms: u128,
+}
+
+const DNS_CACHE_TTL_MS: u128 = 60_000;
+
+fn update_dns_cache(cache: &mut std::collections::HashMap<std::net::IpAddr, DnsCacheEntry>, event: &report::EnrichedEvent) {
+    let Some(qname) = event.qname.as_ref() else { return };
+    let now = event.timestamp_ms;
+    prune_dns_cache(cache, now);
+    for ip in event.answer_ips.iter() {
+        debug!(
+            "dns leak cache insert ip={} qname={} route={:?}",
+            ip, qname, event.route_class
+        );
+        cache.insert(
+            *ip,
+            DnsCacheEntry {
+                qname: qname.clone(),
+                route_class: event.route_class,
+                timestamp_ms: now,
+            },
+        );
+    }
+}
+
+fn prune_dns_cache(
+    cache: &mut std::collections::HashMap<std::net::IpAddr, DnsCacheEntry>,
+    now_ms: u128,
+) {
+    cache.retain(|_, entry| now_ms.saturating_sub(entry.timestamp_ms) <= DNS_CACHE_TTL_MS);
+}
+
+async fn evaluate_mismatch(
+    event: TcpEvent,
+    flow_owner: Option<&dyn FlowOwnerProvider>,
+    cache: &mut std::collections::HashMap<std::net::IpAddr, DnsCacheEntry>,
+    privacy: PrivacyMode,
+) -> Option<LeakEvent> {
+    prune_dns_cache(cache, event.timestamp_ms);
+    debug!(
+        "dns leak tcp syn dst_ip={} dst_port={} cache_size={}",
+        event.dst_ip,
+        event.dst_port,
+        cache.len()
+    );
+    let entry = cache.get(&event.dst_ip)?;
+    let tcp_route = route::route_class_for(event.src_ip, event.dst_ip, event.iface_name.as_deref());
+    if tcp_route == entry.route_class {
+        debug!(
+            "dns leak mismatch skip dst_ip={} tcp_route={:?} dns_route={:?}",
+            event.dst_ip, tcp_route, entry.route_class
+        );
+        return None;
+    }
+
+    let mut enriched = report::EnrichedEvent {
+        timestamp_ms: event.timestamp_ms,
+        proto: wtfnet_platform::FlowProtocol::Tcp,
+        src_ip: event.src_ip,
+        src_port: event.src_port,
+        dst_ip: event.dst_ip,
+        dst_port: event.dst_port,
+        iface_name: event.iface_name.clone(),
+        transport: LeakTransport::Unknown,
+        qname: Some(entry.qname.clone()),
+        qtype: None,
+        rcode: None,
+        is_response: false,
+        answer_ips: Vec::new(),
+        route_class: tcp_route,
+        owner: None,
+        owner_confidence: wtfnet_platform::FlowOwnerConfidence::None,
+        owner_failure: None,
+    };
+
+    if let Some(provider) = flow_owner {
+        let flow = FlowTuple {
+            proto: wtfnet_platform::FlowProtocol::Tcp,
+            src_ip: event.src_ip,
+            src_port: event.src_port,
+            dst_ip: event.dst_ip,
+            dst_port: event.dst_port,
+        };
+        if let Ok(result) = provider.owner_of(flow).await {
+            enriched.owner = result.owner;
+            enriched.owner_confidence = result.confidence;
+            enriched.owner_failure = result.failure_reason;
+        }
+    }
+
+    let decision = rules::LeakDecision {
+        leak_type: report::LeakType::D,
+        severity: Severity::P2,
+        policy_rule_id: "LEAK_D_MISMATCH".to_string(),
+    };
+    let mut leak_event = report::LeakEvent::from_decision(enriched, decision);
+    privacy::apply_privacy(&mut leak_event, privacy);
+    Some(leak_event)
+}