nx9-wg/crates/nx9-wireguard/src/engine.rs

//! WireGuard interface controller and live state engine.

use crate::error::{Result, WireGuardError};
use chrono::{NaiveDateTime, Utc};
use nx9_wg_core::types::wireguard::{Interface, Peer, PeerState};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Live statistics for a connected WireGuard peer.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LivePeerStats {
    pub public_key: String,
    pub endpoint: Option<String>,
    pub rx_bytes: u64,
    pub tx_bytes: u64,
    pub last_handshake_at: Option<NaiveDateTime>,
    pub allowed_ips: Vec<String>,
    pub persistent_keepalive: Option<u16>,
}

/// Live status and peer metrics for a WireGuard interface.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LiveInterfaceStats {
    pub name: String,
    pub public_key: String,
    pub listen_port: u16,
    pub fwmark: u32,
    pub peers: Vec<LivePeerStats>,
    #[serde(default)]
    pub addresses: Vec<String>,
    #[serde(default)]
    pub mtu: Option<u32>,
    #[serde(default)]
    pub is_up: bool,
}

/// Abstract WireGuard Engine interface for kernel netlink and simulated environments.
#[async_trait::async_trait]
pub trait WireGuardEngine: Send + Sync {
    /// Reconcile and synchronize kernel state with desired interface configuration and active peers.
    async fn sync_interface(&self, interface: &Interface, peers: &[Peer]) -> Result<()>;

    /// Remove a WireGuard interface from the system.
    async fn delete_interface(&self, name: &str) -> Result<()>;

    /// Read live statistics and peer telemetry from the kernel.
    async fn get_interface_stats(&self, name: &str) -> Result<Option<LiveInterfaceStats>>;

    /// List all managed WireGuard interface names.
    async fn list_interfaces(&self) -> Result<Vec<String>>;
}

/// In-memory simulated WireGuard engine for deterministic tests and non-root development.
#[derive(Debug, Clone, Default)]
pub struct SimulatedWireGuardEngine {
    state: Arc<RwLock<HashMap<String, LiveInterfaceStats>>>,
}

impl SimulatedWireGuardEngine {
    pub fn new() -> Self {
        Self {
            state: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Simulate a handshake from a peer with transfer byte increments.
    pub async fn simulate_peer_activity(
        &self,
        interface_name: &str,
        peer_public_key: &str,
        rx_add: u64,
        tx_add: u64,
    ) -> Result<()> {
        let mut map = self.state.write().await;
        if let Some(iface) = map.get_mut(interface_name) {
            for peer in &mut iface.peers {
                if peer.public_key == peer_public_key {
                    peer.rx_bytes += rx_add;
                    peer.tx_bytes += tx_add;
                    peer.last_handshake_at = Some(Utc::now().naive_utc());
                    return Ok(());
                }
            }
        }
        Err(WireGuardError::Interface(format!(
            "Peer '{peer_public_key}' on interface '{interface_name}' not found"
        )))
    }

    /// Directly inject live interface stats (for testing drift and telemetry scenarios).
    pub async fn inject_interface_stats(&self, stats: LiveInterfaceStats) {
        let mut map = self.state.write().await;
        map.insert(stats.name.clone(), stats);
    }
}

#[async_trait::async_trait]
impl WireGuardEngine for SimulatedWireGuardEngine {
    async fn sync_interface(&self, interface: &Interface, peers: &[Peer]) -> Result<()> {
        let mut map = self.state.write().await;

        let live_peers: Vec<LivePeerStats> = peers
            .iter()
            .filter(|p| p.state == PeerState::Active)
            .map(|p| {
                let allowed_ips: Vec<String> = p
                    .server_wireguard_allowed_ips()
                    .split(',')
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty())
                    .collect();

                LivePeerStats {
                    public_key: p.public_key.as_str().to_string(),
                    endpoint: p.endpoint.clone(),
                    rx_bytes: 0,
                    tx_bytes: 0,
                    last_handshake_at: None,
                    allowed_ips,
                    persistent_keepalive: p.persistent_keepalive,
                }
            })
            .collect();

        let mut addresses = vec![interface.address_v4.to_string()];
        if let Some(ref v6) = interface.address_v6 {
            addresses.push(v6.to_string());
        }

        let stats = LiveInterfaceStats {
            name: interface.name.clone(),
            public_key: interface.public_key.as_str().to_string(),
            listen_port: interface.listen_port,
            fwmark: 0,
            peers: live_peers,
            addresses,
            mtu: interface.mtu.map(|m| m as u32),
            is_up: true,
        };

        map.insert(interface.name.clone(), stats);
        tracing::debug!(interface = %interface.name, "Simulated WireGuard interface synchronized");
        Ok(())
    }

    async fn delete_interface(&self, name: &str) -> Result<()> {
        let mut map = self.state.write().await;
        map.remove(name);
        tracing::debug!(interface = %name, "Simulated WireGuard interface deleted");
        Ok(())
    }

    async fn get_interface_stats(&self, name: &str) -> Result<Option<LiveInterfaceStats>> {
        let map = self.state.read().await;
        Ok(map.get(name).cloned())
    }

    async fn list_interfaces(&self) -> Result<Vec<String>> {
        let map = self.state.read().await;
        Ok(map.keys().cloned().collect())
    }
}

// ── Native Linux WireGuard Engine ─────────────────────────────────────────────
//
// On Linux: the real implementation lives in native_linux.rs and uses
// RTNETLINK + WireGuard Generic Netlink to communicate with the kernel.
//
// On non-Linux platforms: a thin simulation wrapper is provided so that
// the workspace remains portable and tests remain functional.

#[cfg(target_os = "linux")]
pub use crate::native_linux::NativeLinuxWireGuardEngine;

/// Non-Linux fallback: NativeLinuxWireGuardEngine delegates to simulation.
#[cfg(not(target_os = "linux"))]
#[derive(Debug, Clone, Default)]
pub struct NativeLinuxWireGuardEngine {
    simulated_fallback: SimulatedWireGuardEngine,
}

#[cfg(not(target_os = "linux"))]
impl NativeLinuxWireGuardEngine {
    pub fn new() -> Self {
        Self {
            simulated_fallback: SimulatedWireGuardEngine::new(),
        }
    }

    /// Check if Linux kernel WireGuard support is available.
    pub fn is_supported() -> bool {
        false
    }
}

#[cfg(not(target_os = "linux"))]
#[async_trait::async_trait]
impl WireGuardEngine for NativeLinuxWireGuardEngine {
    async fn sync_interface(&self, interface: &Interface, peers: &[Peer]) -> Result<()> {
        self.simulated_fallback
            .sync_interface(interface, peers)
            .await
    }

    async fn delete_interface(&self, name: &str) -> Result<()> {
        self.simulated_fallback.delete_interface(name).await
    }

    async fn get_interface_stats(&self, name: &str) -> Result<Option<LiveInterfaceStats>> {
        self.simulated_fallback.get_interface_stats(name).await
    }

    async fn list_interfaces(&self) -> Result<Vec<String>> {
        self.simulated_fallback.list_interfaces().await
    }
}