WTFnet/crates/wtfnet-dnsleak/src/lib.rs

mod classify;
mod policy;
mod privacy;
mod report;
mod route;
mod rules;
mod sensor;

use crate::classify::ClassifiedEvent;
use crate::sensor::{capture_events, SensorEvent, TcpEvent};
use std::time::Instant;
use thiserror::Error;
use tracing::debug;
use wtfnet_platform::{FlowOwnerProvider, FlowTuple};

pub use crate::policy::{LeakPolicy, LeakPolicyProfile, PolicySummary};
pub use crate::privacy::{apply_privacy, PrivacyMode};
pub use crate::report::{LeakEvent, LeakReport, LeakSummary, LeakTransport, RouteClass, Severity};
pub use crate::sensor::{iface_diagnostics, IfaceDiag};

#[derive(Debug, Error)]
pub enum DnsLeakError {
    #[error("not supported: {0}")]
    NotSupported(String),
    #[error("io error: {0}")]
    Io(String),
    #[error("policy error: {0}")]
    Policy(String),
}

#[derive(Debug, Clone)]
pub struct LeakWatchOptions {
    pub duration_ms: u64,
    pub iface: Option<String>,
    pub policy: LeakPolicy,
    pub privacy: PrivacyMode,
    pub include_events: bool,
}

pub async fn watch(
    options: LeakWatchOptions,
    flow_owner: Option<&dyn FlowOwnerProvider>,
) -> Result<LeakReport, DnsLeakError> {
    debug!(
        duration_ms = options.duration_ms,
        iface = ?options.iface,
        include_events = options.include_events,
        "dns leak watch start"
    );
    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 {
        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);
                }
            }
        }
    }

    let summary = LeakSummary::from_events(&leak_events);
    let report = LeakReport {
        duration_ms: start.elapsed().as_millis() as u64,
        policy: options.policy.summary(),
        summary,
        events: if options.include_events {
            leak_events
        } else {
            Vec::new()
        },
    };
    Ok(report)
}

async fn enrich_event(
    event: ClassifiedEvent,
    flow_owner: Option<&dyn FlowOwnerProvider>,
) -> report::EnrichedEvent {
    let mut enriched = route::enrich_route(event);
    if let Some(provider) = flow_owner {
        let flow = FlowTuple {
            proto: enriched.proto,
            src_ip: enriched.src_ip,
            src_port: enriched.src_port,
            dst_ip: enriched.dst_ip,
            dst_port: enriched.dst_port,
        };
        match provider.owner_of(flow).await {
            Ok(result) => {
                enriched.owner = result.owner;
                enriched.owner_confidence = result.confidence;
                enriched.owner_failure = result.failure_reason;
            }
            Err(err) => {
                enriched.owner_failure = Some(err.message);
            }
        }
    }
    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)
}