smartvpn/rust/src/server.rs

use anyhow::Result;
use bytes::BytesMut;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::net::Ipv4Addr;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, Mutex, RwLock};
use tracing::{debug, info, error, warn};

use crate::acl;
use crate::client_registry::{ClientEntry, ClientRegistry};
use crate::codec::{Frame, FrameCodec, PacketType};
use crate::crypto;
use crate::mtu::{MtuConfig, TunnelOverhead};
use crate::network::IpPool;
use crate::ratelimit::TokenBucket;
use crate::transport;
use crate::transport_trait::{self, TransportSink, TransportStream};
use crate::quic_transport;
use crate::tunnel::{self, TunConfig};

/// Dead-peer timeout: 3x max keepalive interval (Healthy=60s).
const DEAD_PEER_TIMEOUT: Duration = Duration::from_secs(180);

/// Destination routing policy for VPN client traffic.
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct DestinationPolicyConfig {
    /// Default action: "forceTarget", "block", or "allow".
    pub default: String,
    /// Target IP for "forceTarget" mode (e.g. "127.0.0.1").
    pub target: Option<String>,
    /// Destinations that pass through directly (not rewritten, not blocked).
    pub allow_list: Option<Vec<String>>,
    /// Destinations always blocked (overrides allowList, deny wins).
    pub block_list: Option<Vec<String>>,
}

/// Server configuration (matches TS IVpnServerConfig).
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ServerConfig {
    pub listen_addr: String,
    pub tls_cert: Option<String>,
    pub tls_key: Option<String>,
    pub private_key: String,
    pub public_key: String,
    pub subnet: String,
    pub dns: Option<Vec<String>>,
    pub mtu: Option<u16>,
    pub keepalive_interval_secs: Option<u64>,
    pub enable_nat: Option<bool>,
    /// Forwarding mode: "tun" (kernel TUN, requires root), "socket" (userspace NAT),
    /// or "testing" (monitoring only, no forwarding). Default: "testing".
    pub forwarding_mode: Option<String>,
    /// Default rate limit for new clients (bytes/sec). None = unlimited.
    pub default_rate_limit_bytes_per_sec: Option<u64>,
    /// Default burst size for new clients (bytes). None = unlimited.
    pub default_burst_bytes: Option<u64>,
    /// Transport mode: "websocket" (default), "quic", or "both".
    pub transport_mode: Option<String>,
    /// QUIC listen address (host:port). Defaults to listen_addr.
    pub quic_listen_addr: Option<String>,
    /// QUIC idle timeout in seconds (default: 30).
    pub quic_idle_timeout_secs: Option<u64>,
    /// Pre-registered clients for IK authentication.
    pub clients: Option<Vec<ClientEntry>>,
    /// Enable PROXY protocol v2 parsing on incoming WebSocket connections.
    /// SECURITY: Must be false when accepting direct client connections.
    pub proxy_protocol: Option<bool>,
    /// Server-level IP block list — applied at TCP accept, before Noise handshake.
    pub connection_ip_block_list: Option<Vec<String>>,
    /// When true and forwarding_mode is "socket", the userspace NAT engine prepends
    /// PROXY protocol v2 headers on outbound TCP connections, conveying the VPN client's
    /// tunnel IP as the source address.
    pub socket_forward_proxy_protocol: Option<bool>,
    /// Destination routing policy for VPN client traffic (socket mode).
    pub destination_policy: Option<DestinationPolicyConfig>,
    /// WireGuard: server X25519 private key (base64). Required when transport includes WG.
    pub wg_private_key: Option<String>,
    /// WireGuard: UDP listen port (default: 51820).
    pub wg_listen_port: Option<u16>,
    /// WireGuard: pre-configured peers.
    pub wg_peers: Option<Vec<crate::wireguard::WgPeerConfig>>,
    /// Public endpoint address for generated client configs (e.g. "vpn.example.com:51820").
    /// Used as WireGuard `Endpoint` and SmartVPN `serverUrl` host.
    /// Defaults to listen_addr.
    pub server_endpoint: Option<String>,
    /// AllowedIPs for generated WireGuard client configs.
    /// Defaults to ["0.0.0.0/0"] (full tunnel).
    #[serde(alias = "clientAllowedIPs")]
    pub client_allowed_ips: Option<Vec<String>>,

    // Bridge mode configuration (forwarding_mode: "bridge")

    /// LAN subnet CIDR for bridge mode (e.g. "192.168.1.0/24").
    pub bridge_lan_subnet: Option<String>,
    /// Physical network interface to bridge (e.g. "eth0"). Auto-detected if omitted.
    pub bridge_physical_interface: Option<String>,
    /// Start of VPN client IP range within the LAN subnet (host offset, e.g. 200).
    pub bridge_ip_range_start: Option<u32>,
    /// End of VPN client IP range within the LAN subnet (host offset, e.g. 250).
    pub bridge_ip_range_end: Option<u32>,
}

/// Information about a connected client.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientInfo {
    pub client_id: String,
    pub assigned_ip: String,
    pub connected_since: String,
    pub bytes_sent: u64,
    pub bytes_received: u64,
    pub packets_dropped: u64,
    pub bytes_dropped: u64,
    pub last_keepalive_at: Option<String>,
    pub keepalives_received: u64,
    pub rate_limit_bytes_per_sec: Option<u64>,
    pub burst_bytes: Option<u64>,
    /// Client's authenticated Noise IK public key (base64).
    pub authenticated_key: String,
    /// Registered client ID from the client registry.
    pub registered_client_id: String,
    /// Real client IP:port (from PROXY protocol header or direct TCP connection).
    pub remote_addr: Option<String>,
    /// Transport used for this connection: "websocket", "quic", or "wireguard".
    pub transport_type: String,
}

/// Server statistics.
#[derive(Debug, Clone, Serialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct ServerStatistics {
    pub bytes_sent: u64,
    pub bytes_received: u64,
    pub packets_sent: u64,
    pub packets_received: u64,
    pub keepalives_sent: u64,
    pub keepalives_received: u64,
    pub uptime_seconds: u64,
    pub active_clients: u64,
    pub total_connections: u64,
    /// Per-transport active client counts.
    pub active_clients_websocket: u64,
    pub active_clients_quic: u64,
    pub active_clients_wireguard: u64,
    /// Per-transport total connection counts.
    pub total_connections_websocket: u64,
    pub total_connections_quic: u64,
    pub total_connections_wireguard: u64,
}

/// The forwarding engine determines how decrypted IP packets are routed.
pub enum ForwardingEngine {
    /// Kernel TUN device — packets written to the TUN, kernel handles routing.
    Tun(tokio::io::WriteHalf<tun::AsyncDevice>),
    /// Userspace NAT — packets sent to smoltcp-based NAT engine via channel.
    Socket(mpsc::Sender<Vec<u8>>),
    /// L2 Bridge — packets sent to BridgeEngine via channel, bridged to host LAN.
    Bridge(mpsc::Sender<Vec<u8>>),
    /// Testing/monitoring — packets are counted but not forwarded.
    Testing,
}

/// Shared server state.
pub struct ServerState {
    pub config: ServerConfig,
    pub ip_pool: Mutex<IpPool>,
    pub clients: RwLock<HashMap<String, ClientInfo>>,
    pub stats: RwLock<ServerStatistics>,
    pub rate_limiters: Mutex<HashMap<String, TokenBucket>>,
    pub mtu_config: MtuConfig,
    pub started_at: std::time::Instant,
    pub client_registry: RwLock<ClientRegistry>,
    /// The forwarding engine for decrypted IP packets.
    pub forwarding_engine: Mutex<ForwardingEngine>,
    /// Routing table: assigned VPN IP → channel sender for return packets.
    pub tun_routes: RwLock<HashMap<Ipv4Addr, mpsc::Sender<Vec<u8>>>>,
    /// Shutdown signal for the forwarding background task (TUN reader or NAT engine).
    pub tun_shutdown: mpsc::Sender<()>,
}

/// The VPN server.
pub struct VpnServer {
    state: Option<Arc<ServerState>>,
    shutdown_tx: Option<mpsc::Sender<()>>,
    wg_command_tx: Option<mpsc::Sender<crate::wireguard::WgCommand>>,
}

impl VpnServer {
    pub fn new() -> Self {
        Self {
            state: None,
            shutdown_tx: None,
            wg_command_tx: None,
        }
    }

    pub async fn start(&mut self, config: ServerConfig) -> Result<()> {
        if self.state.is_some() {
            anyhow::bail!("Server is already running");
        }

        let mode = config.forwarding_mode.as_deref().unwrap_or("testing");
        let ip_pool = if mode == "bridge" {
            let lan_subnet = config.bridge_lan_subnet.as_deref().unwrap_or(&config.subnet);
            let range_start = config.bridge_ip_range_start.unwrap_or(200);
            let range_end = config.bridge_ip_range_end.unwrap_or(250);
            IpPool::new_with_range(lan_subnet, range_start, range_end)?
        } else {
            IpPool::new(&config.subnet)?
        };

        if config.enable_nat.unwrap_or(false) {
            if let Err(e) = crate::network::enable_ip_forwarding() {
                warn!("Failed to enable IP forwarding: {}", e);
            }
            if let Ok(iface) = crate::network::get_default_interface() {
                if let Err(e) = crate::network::setup_nat(&config.subnet, &iface).await {
                    warn!("Failed to setup NAT: {}", e);
                }
            }
        }

        let link_mtu = config.mtu.unwrap_or(1420);
        let gateway_ip = ip_pool.gateway_addr();

        // Create forwarding engine based on mode
        enum ForwardingSetup {
            Tun {
                writer: tokio::io::WriteHalf<tun::AsyncDevice>,
                reader: tokio::io::ReadHalf<tun::AsyncDevice>,
                shutdown_rx: mpsc::Receiver<()>,
            },
            Socket {
                packet_tx: mpsc::Sender<Vec<u8>>,
                packet_rx: mpsc::Receiver<Vec<u8>>,
                shutdown_rx: mpsc::Receiver<()>,
            },
            Bridge {
                packet_tx: mpsc::Sender<Vec<u8>>,
                packet_rx: mpsc::Receiver<Vec<u8>>,
                tap_device: tun::AsyncDevice,
                shutdown_rx: mpsc::Receiver<()>,
            },
            Testing,
        }

        let (setup, fwd_shutdown_tx) = match mode {
            "tun" => {
                let tun_config = TunConfig {
                    name: "svpn0".to_string(),
                    address: gateway_ip,
                    netmask: Ipv4Addr::new(255, 255, 255, 0),
                    mtu: link_mtu,
                };
                let tun_device = tunnel::create_tun(&tun_config)?;
                tunnel::add_route(&config.subnet, &tun_config.name).await?;
                let (reader, writer) = tokio::io::split(tun_device);
                let (tx, rx) = mpsc::channel::<()>(1);
                (ForwardingSetup::Tun { writer, reader, shutdown_rx: rx }, tx)
            }
            "socket" => {
                info!("Starting userspace NAT forwarding (no root required)");
                let (packet_tx, packet_rx) = mpsc::channel::<Vec<u8>>(4096);
                let (tx, rx) = mpsc::channel::<()>(1);
                (ForwardingSetup::Socket { packet_tx, packet_rx, shutdown_rx: rx }, tx)
            }
            "bridge" => {
                info!("Starting L2 bridge forwarding (requires CAP_NET_ADMIN)");
                let phys_iface = match &config.bridge_physical_interface {
                    Some(i) => i.clone(),
                    None => crate::bridge::detect_default_interface().await?,
                };
                let (host_ip, host_prefix) = crate::bridge::get_interface_ip(&phys_iface).await?;

                let bridge_name = "svpn_br0";
                let tap_name = "svpn_tap0";

                // Create TAP + bridge infrastructure
                let tap_device = crate::bridge::create_tap(tap_name, link_mtu)?;
                crate::bridge::create_bridge(bridge_name).await?;
                crate::bridge::set_interface_up(bridge_name).await?;
                crate::bridge::bridge_add_interface(bridge_name, tap_name).await?;
                crate::bridge::set_interface_up(tap_name).await?;
                crate::bridge::bridge_add_interface(bridge_name, &phys_iface).await?;
                crate::bridge::migrate_host_ip_to_bridge(&phys_iface, bridge_name, host_ip, host_prefix).await?;
                crate::bridge::enable_proxy_arp(bridge_name).await?;

                info!("Bridge {} created: TAP={}, physical={}, IP={}/{}", bridge_name, tap_name, phys_iface, host_ip, host_prefix);

                let (packet_tx, packet_rx) = mpsc::channel::<Vec<u8>>(4096);
                let (tx, rx) = mpsc::channel::<()>(1);
                (ForwardingSetup::Bridge { packet_tx, packet_rx, tap_device, shutdown_rx: rx }, tx)
            }
            _ => {
                info!("Forwarding disabled (testing/monitoring mode)");
                let (tx, _rx) = mpsc::channel::<()>(1);
                (ForwardingSetup::Testing, tx)
            }
        };

        // Compute effective MTU from overhead
        let overhead = TunnelOverhead::default_overhead();
        let mtu_config = MtuConfig::new(overhead.effective_tun_mtu(1500).max(link_mtu));

        // Build client registry from config
        let registry = ClientRegistry::from_entries(
            config.clients.clone().unwrap_or_default()
        )?;
        info!("Client registry loaded with {} entries", registry.len());

        let state = Arc::new(ServerState {
            config: config.clone(),
            ip_pool: Mutex::new(ip_pool),
            clients: RwLock::new(HashMap::new()),
            stats: RwLock::new(ServerStatistics::default()),
            rate_limiters: Mutex::new(HashMap::new()),
            mtu_config,
            started_at: std::time::Instant::now(),
            client_registry: RwLock::new(registry),
            forwarding_engine: Mutex::new(ForwardingEngine::Testing),
            tun_routes: RwLock::new(HashMap::new()),
            tun_shutdown: fwd_shutdown_tx,
        });

        // Spawn the forwarding background task and set the engine
        match setup {
            ForwardingSetup::Tun { writer, reader, shutdown_rx } => {
                *state.forwarding_engine.lock().await = ForwardingEngine::Tun(writer);
                let tun_state = state.clone();
                tokio::spawn(async move {
                    if let Err(e) = run_tun_reader(tun_state, reader, shutdown_rx).await {
                        error!("TUN reader error: {}", e);
                    }
                });
            }
            ForwardingSetup::Socket { packet_tx, packet_rx, shutdown_rx } => {
                *state.forwarding_engine.lock().await = ForwardingEngine::Socket(packet_tx);
                let proxy_protocol = config.socket_forward_proxy_protocol.unwrap_or(false);
                let nat_engine = crate::userspace_nat::NatEngine::new(
                    gateway_ip,
                    link_mtu as usize,
                    state.clone(),
                    proxy_protocol,
                    config.destination_policy.clone(),
                );
                tokio::spawn(async move {
                    if let Err(e) = nat_engine.run(packet_rx, shutdown_rx).await {
                        error!("NAT engine error: {}", e);
                    }
                });
            }
            ForwardingSetup::Bridge { packet_tx, packet_rx, tap_device, shutdown_rx } => {
                *state.forwarding_engine.lock().await = ForwardingEngine::Bridge(packet_tx);
                let bridge_engine = crate::bridge::BridgeEngine::new(state.clone());
                tokio::spawn(async move {
                    if let Err(e) = bridge_engine.run(tap_device, packet_rx, shutdown_rx).await {
                        error!("Bridge engine error: {}", e);
                    }
                });
            }
            ForwardingSetup::Testing => {}
        }

        self.state = Some(state.clone());

        let transport_mode = config.transport_mode.as_deref().unwrap_or("all");
        let listen_addr = config.listen_addr.clone();

        // Determine if WG should be included
        let include_wg = config.wg_private_key.is_some()
            && matches!(transport_mode, "all" | "wireguard");

        // Collect shutdown senders for all listeners
        let mut listener_shutdown_txs: Vec<mpsc::Sender<()>> = Vec::new();

        // Spawn transport listeners based on mode
        let spawn_ws = matches!(transport_mode, "all" | "both" | "websocket");
        let spawn_quic = matches!(transport_mode, "all" | "both" | "quic");

        if spawn_ws {
            let (tx, mut rx) = mpsc::channel::<()>(1);
            listener_shutdown_txs.push(tx);
            let ws_state = state.clone();
            let ws_addr = listen_addr.clone();
            tokio::spawn(async move {
                if let Err(e) = run_ws_listener(ws_state, ws_addr, &mut rx).await {
                    error!("WebSocket listener error: {}", e);
                }
            });
        }

        if spawn_quic {
            let quic_addr = config.quic_listen_addr.clone().unwrap_or_else(|| listen_addr.clone());
            let idle_timeout = config.quic_idle_timeout_secs.unwrap_or(30);
            let (tx, mut rx) = mpsc::channel::<()>(1);
            listener_shutdown_txs.push(tx);
            let quic_state = state.clone();
            tokio::spawn(async move {
                if let Err(e) = run_quic_listener(quic_state, quic_addr, idle_timeout, &mut rx).await {
                    error!("QUIC listener error: {}", e);
                }
            });
        }

        if include_wg {
            let wg_config = crate::wireguard::WgListenerConfig {
                private_key: config.wg_private_key.clone().unwrap(),
                listen_port: config.wg_listen_port.unwrap_or(51820),
                peers: config.wg_peers.clone().unwrap_or_default(),
            };
            let (tx, rx) = mpsc::channel::<()>(1);
            listener_shutdown_txs.push(tx);
            let (cmd_tx, cmd_rx) = mpsc::channel::<crate::wireguard::WgCommand>(32);
            self.wg_command_tx = Some(cmd_tx);
            let wg_state = state.clone();
            tokio::spawn(async move {
                if let Err(e) = crate::wireguard::run_wg_listener(wg_state, wg_config, rx, cmd_rx).await {
                    error!("WireGuard listener error: {}", e);
                }
            });
        }

        // Replace self.shutdown_tx with a combined sender that fans out to all listeners
        if listener_shutdown_txs.len() > 1 {
            let (combined_tx, mut combined_rx) = mpsc::channel::<()>(1);
            // Take the original shutdown_tx (from line above)
            let _ = self.shutdown_tx.take();
            self.shutdown_tx = Some(combined_tx);
            tokio::spawn(async move {
                combined_rx.recv().await;
                for tx in listener_shutdown_txs {
                    let _ = tx.send(()).await;
                }
            });
        } else if let Some(single_tx) = listener_shutdown_txs.into_iter().next() {
            self.shutdown_tx = Some(single_tx);
        }

        info!("VPN server started (transport: {})", transport_mode);

        // Register pre-loaded clients (from config.clients) as WG peers.
        // The WG listener only starts with config.wg_peers; clients loaded into the
        // registry need to be dynamically added so WG handshakes work.
        if self.wg_command_tx.is_some() {
            let registry = state.client_registry.read().await;
            for entry in registry.list() {
                if let (Some(ref wg_key), Some(ref ip_str)) = (&entry.wg_public_key, &entry.assigned_ip) {
                    let peer_config = crate::wireguard::WgPeerConfig {
                        public_key: wg_key.clone(),
                        preshared_key: None,
                        allowed_ips: vec![format!("{}/32", ip_str)],
                        endpoint: None,
                        persistent_keepalive: Some(25),
                    };
                    if let Err(e) = self.add_wg_peer(peer_config).await {
                        warn!("Failed to register pre-loaded WG peer for {}: {}", entry.client_id, e);
                    }
                }
            }
        }

        Ok(())
    }

    pub async fn stop(&mut self) -> Result<()> {
        if let Some(ref state) = self.state {
            let mode = state.config.forwarding_mode.as_deref().unwrap_or("testing");

            match mode {
                "tun" => {
                    let _ = state.tun_shutdown.send(()).await;
                    *state.forwarding_engine.lock().await = ForwardingEngine::Testing;
                    if let Err(e) = tunnel::remove_route(&state.config.subnet, "svpn0").await {
                        warn!("Failed to remove TUN route: {}", e);
                    }
                }
                "socket" => {
                    let _ = state.tun_shutdown.send(()).await;
                    *state.forwarding_engine.lock().await = ForwardingEngine::Testing;
                }
                _ => {}
            }

            // Clean up NAT rules
            if state.config.enable_nat.unwrap_or(false) {
                if let Ok(iface) = crate::network::get_default_interface() {
                    if let Err(e) = crate::network::remove_nat(&state.config.subnet, &iface).await {
                        warn!("Failed to remove NAT rules: {}", e);
                    }
                }
            }
        }

        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(()).await;
        }
        self.wg_command_tx = None;
        self.state = None;
        info!("VPN server stopped");
        Ok(())
    }

    pub fn get_status(&self) -> serde_json::Value {
        if let Some(ref state) = self.state {
            serde_json::json!({
                "state": "connected",
                "connectedSince": format!("{:?}", state.started_at.elapsed()),
            })
        } else {
            serde_json::json!({ "state": "disconnected" })
        }
    }

    pub async fn get_statistics(&self) -> ServerStatistics {
        if let Some(ref state) = self.state {
            let mut stats = state.stats.read().await.clone();
            stats.uptime_seconds = state.started_at.elapsed().as_secs();
            let clients = state.clients.read().await;
            stats.active_clients = clients.len() as u64;
            // Compute per-transport active counts
            stats.active_clients_websocket = 0;
            stats.active_clients_quic = 0;
            stats.active_clients_wireguard = 0;
            for info in clients.values() {
                match info.transport_type.as_str() {
                    "websocket" => stats.active_clients_websocket += 1,
                    "quic" => stats.active_clients_quic += 1,
                    "wireguard" => stats.active_clients_wireguard += 1,
                    _ => {}
                }
            }
            drop(clients);
            stats
        } else {
            ServerStatistics::default()
        }
    }

    pub async fn list_clients(&self) -> Vec<ClientInfo> {
        if let Some(ref state) = self.state {
            state.clients.read().await.values().cloned().collect()
        } else {
            Vec::new()
        }
    }

    pub async fn disconnect_client(&self, client_id: &str) -> Result<()> {
        if let Some(ref state) = self.state {
            let mut clients = state.clients.write().await;
            if let Some(client) = clients.remove(client_id) {
                let ip: Ipv4Addr = client.assigned_ip.parse()?;
                state.ip_pool.lock().await.release(&ip);
                state.rate_limiters.lock().await.remove(client_id);
                info!("Client {} disconnected", client_id);
            }
        }
        Ok(())
    }

    /// Set a rate limit for a specific client.
    pub async fn set_client_rate_limit(
        &self,
        client_id: &str,
        rate_bytes_per_sec: u64,
        burst_bytes: u64,
    ) -> Result<()> {
        if let Some(ref state) = self.state {
            let mut limiters = state.rate_limiters.lock().await;
            if let Some(limiter) = limiters.get_mut(client_id) {
                limiter.update_limits(rate_bytes_per_sec, burst_bytes);
            } else {
                limiters.insert(
                    client_id.to_string(),
                    TokenBucket::new(rate_bytes_per_sec, burst_bytes),
                );
            }
            // Update client info
            let mut clients = state.clients.write().await;
            if let Some(info) = clients.get_mut(client_id) {
                info.rate_limit_bytes_per_sec = Some(rate_bytes_per_sec);
                info.burst_bytes = Some(burst_bytes);
            }
        }
        Ok(())
    }

    /// Remove rate limit for a specific client (unlimited).
    pub async fn remove_client_rate_limit(&self, client_id: &str) -> Result<()> {
        if let Some(ref state) = self.state {
            state.rate_limiters.lock().await.remove(client_id);
            let mut clients = state.clients.write().await;
            if let Some(info) = clients.get_mut(client_id) {
                info.rate_limit_bytes_per_sec = None;
                info.burst_bytes = None;
            }
        }
        Ok(())
    }

    // ── WireGuard Peer Management ────────────────────────────────────────

    /// Add a WireGuard peer dynamically (delegates to the WG event loop).
    pub async fn add_wg_peer(&self, config: crate::wireguard::WgPeerConfig) -> Result<()> {
        let tx = self.wg_command_tx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("WireGuard listener not running"))?;
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
        tx.send(crate::wireguard::WgCommand::AddPeer(config, resp_tx))
            .await
            .map_err(|_| anyhow::anyhow!("WG event loop closed"))?;
        resp_rx.await.map_err(|_| anyhow::anyhow!("No response from WG loop"))?
    }

    /// Remove a WireGuard peer dynamically (delegates to the WG event loop).
    pub async fn remove_wg_peer(&self, public_key: &str) -> Result<()> {
        let tx = self.wg_command_tx.as_ref()
            .ok_or_else(|| anyhow::anyhow!("WireGuard listener not running"))?;
        let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
        tx.send(crate::wireguard::WgCommand::RemovePeer(public_key.to_string(), resp_tx))
            .await
            .map_err(|_| anyhow::anyhow!("WG event loop closed"))?;
        resp_rx.await.map_err(|_| anyhow::anyhow!("No response from WG loop"))?
    }

    /// List WireGuard peers from the unified client list.
    pub async fn list_wg_peers(&self) -> Vec<crate::wireguard::WgPeerInfo> {
        if let Some(ref state) = self.state {
            state.clients.read().await.values()
                .filter(|c| c.transport_type == "wireguard")
                .map(|c| crate::wireguard::WgPeerInfo {
                    public_key: c.authenticated_key.clone(),
                    allowed_ips: vec![format!("{}/32", c.assigned_ip)],
                    endpoint: c.remote_addr.clone(),
                    persistent_keepalive: None,
                    stats: crate::wireguard::WgPeerStats {
                        bytes_sent: c.bytes_sent,
                        bytes_received: c.bytes_received,
                        packets_sent: 0,
                        packets_received: 0,
                        last_handshake_time: None,
                    },
                })
                .collect()
        } else {
            Vec::new()
        }
    }

    // ── Client Registry (Hub) Methods ───────────────────────────────────

    /// Create a new client entry. Generates keypairs and assigns an IP.
    /// Returns a JSON value with the full config bundle including secrets.
    pub async fn create_client(&self, partial: serde_json::Value) -> Result<serde_json::Value> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;

        let client_id = partial.get("clientId")
            .and_then(|v| v.as_str())
            .ok_or_else(|| anyhow::anyhow!("clientId is required"))?
            .to_string();

        // Generate Noise IK keypair for the client
        let (noise_pub, noise_priv) = crypto::generate_keypair()?;

        // Generate WireGuard keypair for the client
        let (wg_pub, wg_priv) = crate::wireguard::generate_wg_keypair();

        // Allocate a VPN IP
        let assigned_ip = state.ip_pool.lock().await.allocate(&client_id)?;

        // Build entry from partial + generated values
        let entry = ClientEntry {
            client_id: client_id.clone(),
            public_key: noise_pub.clone(),
            wg_public_key: Some(wg_pub.clone()),
            security: serde_json::from_value(
                partial.get("security").cloned().unwrap_or(serde_json::Value::Null)
            ).ok(),
            priority: partial.get("priority").and_then(|v| v.as_u64()).map(|v| v as u32),
            enabled: partial.get("enabled").and_then(|v| v.as_bool()).or(Some(true)),
            server_defined_client_tags: partial.get("serverDefinedClientTags").and_then(|v| {
                v.as_array().map(|a| a.iter().filter_map(|s| s.as_str().map(String::from)).collect())
            }).or_else(|| {
                // Legacy: accept "tags" as serverDefinedClientTags
                partial.get("tags").and_then(|v| {
                    v.as_array().map(|a| a.iter().filter_map(|s| s.as_str().map(String::from)).collect())
                })
            }),
            client_defined_client_tags: None, // Only set by connecting client
            tags: None, // Legacy field — not used for new entries
            description: partial.get("description").and_then(|v| v.as_str()).map(String::from),
            expires_at: partial.get("expiresAt").and_then(|v| v.as_str()).map(String::from),
            assigned_ip: Some(assigned_ip.to_string()),
        };

        // Add to registry
        state.client_registry.write().await.add(entry.clone())?;

        // Register WG peer with the running WG listener (if active)
        if self.wg_command_tx.is_some() {
            let wg_peer_config = crate::wireguard::WgPeerConfig {
                public_key: wg_pub.clone(),
                preshared_key: None,
                allowed_ips: vec![format!("{}/32", assigned_ip)],
                endpoint: None,
                persistent_keepalive: Some(25),
            };
            if let Err(e) = self.add_wg_peer(wg_peer_config).await {
                warn!("Failed to register WG peer for client {}: {}", client_id, e);
            }
        }

        // Build SmartVPN client config
        let smartvpn_server_url = format!("wss://{}",
            state.config.server_endpoint.as_deref()
                .unwrap_or(&state.config.listen_addr)
                .replace("0.0.0.0", "localhost"));
        let smartvpn_config = serde_json::json!({
            "serverUrl": smartvpn_server_url,
            "serverPublicKey": state.config.public_key,
            "clientPrivateKey": noise_priv,
            "clientPublicKey": noise_pub,
            "dns": state.config.dns,
            "mtu": state.config.mtu,
            "keepaliveIntervalSecs": state.config.keepalive_interval_secs,
        });

        // Build WireGuard config string
        let wg_server_pubkey = match &state.config.wg_private_key {
            Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
            None => state.config.public_key.clone(),
        };
        let wg_endpoint = state.config.server_endpoint.as_deref()
            .unwrap_or(&state.config.listen_addr);
        let wg_allowed_ips = state.config.client_allowed_ips.as_ref()
            .map(|ips| ips.join(", "))
            .unwrap_or_else(|| "0.0.0.0/0".to_string());
        let wg_config = format!(
            "[Interface]\nPrivateKey = {}\nAddress = {}/24\n{}\n[Peer]\nPublicKey = {}\nAllowedIPs = {}\nEndpoint = {}\nPersistentKeepalive = 25\n",
            wg_priv,
            assigned_ip,
            state.config.dns.as_ref()
                .map(|d| format!("DNS = {}", d.join(", ")))
                .unwrap_or_default(),
            wg_server_pubkey,
            wg_allowed_ips,
            wg_endpoint,
        );

        let entry_json = serde_json::to_value(&entry)?;

        Ok(serde_json::json!({
            "entry": entry_json,
            "smartvpnConfig": smartvpn_config,
            "wireguardConfig": wg_config,
            "secrets": {
                "noisePrivateKey": noise_priv,
                "wgPrivateKey": wg_priv,
            }
        }))
    }

    /// Remove a registered client from the registry (and disconnect if connected).
    pub async fn remove_registered_client(&self, client_id: &str) -> Result<()> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        let entry = state.client_registry.write().await.remove(client_id)?;
        // Remove WG peer from running listener
        if self.wg_command_tx.is_some() {
            if let Some(ref wg_key) = entry.wg_public_key {
                if let Err(e) = self.remove_wg_peer(wg_key).await {
                    debug!("Failed to remove WG peer for client {}: {}", client_id, e);
                }
            }
        }
        // Release the IP if assigned
        if let Some(ref ip_str) = entry.assigned_ip {
            if let Ok(ip) = ip_str.parse::<Ipv4Addr>() {
                state.ip_pool.lock().await.release(&ip);
            }
        }
        // Disconnect if currently connected
        let _ = self.disconnect_client(client_id).await;
        Ok(())
    }

    /// Get a registered client by ID.
    pub async fn get_registered_client(&self, client_id: &str) -> Result<serde_json::Value> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        let registry = state.client_registry.read().await;
        let entry = registry.get_by_id(client_id)
            .ok_or_else(|| anyhow::anyhow!("Client '{}' not found", client_id))?;
        Ok(serde_json::to_value(entry)?)
    }

    /// List all registered clients.
    pub async fn list_registered_clients(&self) -> Vec<ClientEntry> {
        if let Some(ref state) = self.state {
            state.client_registry.read().await.list().into_iter().cloned().collect()
        } else {
            Vec::new()
        }
    }

    /// Update a registered client's fields.
    pub async fn update_registered_client(&self, client_id: &str, update: serde_json::Value) -> Result<()> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        state.client_registry.write().await.update(client_id, |entry| {
            if let Some(security) = update.get("security") {
                entry.security = serde_json::from_value(security.clone()).ok();
            }
            if let Some(priority) = update.get("priority").and_then(|v| v.as_u64()) {
                entry.priority = Some(priority as u32);
            }
            if let Some(enabled) = update.get("enabled").and_then(|v| v.as_bool()) {
                entry.enabled = Some(enabled);
            }
            if let Some(tags) = update.get("serverDefinedClientTags").and_then(|v| v.as_array()) {
                entry.server_defined_client_tags = Some(tags.iter().filter_map(|s| s.as_str().map(String::from)).collect());
            } else if let Some(tags) = update.get("tags").and_then(|v| v.as_array()) {
                // Legacy: accept "tags" as serverDefinedClientTags
                entry.server_defined_client_tags = Some(tags.iter().filter_map(|s| s.as_str().map(String::from)).collect());
            }
            if let Some(desc) = update.get("description").and_then(|v| v.as_str()) {
                entry.description = Some(desc.to_string());
            }
            if let Some(expires) = update.get("expiresAt").and_then(|v| v.as_str()) {
                entry.expires_at = Some(expires.to_string());
            }
        })?;
        Ok(())
    }

    /// Enable a registered client.
    pub async fn enable_client(&self, client_id: &str) -> Result<()> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        state.client_registry.write().await.update(client_id, |entry| {
            entry.enabled = Some(true);
        })
    }

    /// Disable a registered client (also disconnects if connected).
    pub async fn disable_client(&self, client_id: &str) -> Result<()> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        state.client_registry.write().await.update(client_id, |entry| {
            entry.enabled = Some(false);
        })?;
        // Disconnect if currently connected
        let _ = self.disconnect_client(client_id).await;
        Ok(())
    }

    /// Rotate a client's keys. Returns a new config bundle with fresh keypairs.
    pub async fn rotate_client_key(&self, client_id: &str) -> Result<serde_json::Value> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;

        // Capture old WG key before rotation (needed to remove from WG listener)
        let old_wg_pub = {
            let registry = state.client_registry.read().await;
            let entry = registry.get_by_id(client_id)
                .ok_or_else(|| anyhow::anyhow!("Client '{}' not found", client_id))?;
            entry.wg_public_key.clone()
        };

        let (noise_pub, noise_priv) = crypto::generate_keypair()?;
        let (wg_pub, wg_priv) = crate::wireguard::generate_wg_keypair();

        state.client_registry.write().await.rotate_key(
            client_id,
            noise_pub.clone(),
            Some(wg_pub.clone()),
        )?;

        // Disconnect existing connection (old key is no longer valid)
        let _ = self.disconnect_client(client_id).await;

        // Get updated entry for the config bundle
        let entry_json = self.get_registered_client(client_id).await?;
        let assigned_ip = entry_json.get("assignedIp")
            .and_then(|v| v.as_str())
            .unwrap_or("0.0.0.0");

        // Update WG listener: remove old peer, add new peer
        if self.wg_command_tx.is_some() {
            if let Some(ref old_key) = old_wg_pub {
                if let Err(e) = self.remove_wg_peer(old_key).await {
                    debug!("Failed to remove old WG peer during rotation: {}", e);
                }
            }
            let wg_peer_config = crate::wireguard::WgPeerConfig {
                public_key: wg_pub.clone(),
                preshared_key: None,
                allowed_ips: vec![format!("{}/32", assigned_ip)],
                endpoint: None,
                persistent_keepalive: Some(25),
            };
            if let Err(e) = self.add_wg_peer(wg_peer_config).await {
                warn!("Failed to register new WG peer during rotation: {}", e);
            }
        }

        let smartvpn_server_url = format!("wss://{}",
            state.config.server_endpoint.as_deref()
                .unwrap_or(&state.config.listen_addr)
                .replace("0.0.0.0", "localhost"));
        let smartvpn_config = serde_json::json!({
            "serverUrl": smartvpn_server_url,
            "serverPublicKey": state.config.public_key,
            "clientPrivateKey": noise_priv,
            "clientPublicKey": noise_pub,
            "dns": state.config.dns,
            "mtu": state.config.mtu,
            "keepaliveIntervalSecs": state.config.keepalive_interval_secs,
        });

        let wg_server_pubkey = match &state.config.wg_private_key {
            Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
            None => state.config.public_key.clone(),
        };
        let wg_endpoint = state.config.server_endpoint.as_deref()
            .unwrap_or(&state.config.listen_addr);
        let wg_allowed_ips = state.config.client_allowed_ips.as_ref()
            .map(|ips| ips.join(", "))
            .unwrap_or_else(|| "0.0.0.0/0".to_string());
        let wg_config = format!(
            "[Interface]\nPrivateKey = {}\nAddress = {}/24\n{}\n[Peer]\nPublicKey = {}\nAllowedIPs = {}\nEndpoint = {}\nPersistentKeepalive = 25\n",
            wg_priv, assigned_ip,
            state.config.dns.as_ref()
                .map(|d| format!("DNS = {}", d.join(", ")))
                .unwrap_or_default(),
            wg_server_pubkey,
            wg_allowed_ips,
            wg_endpoint,
        );

        Ok(serde_json::json!({
            "entry": entry_json,
            "smartvpnConfig": smartvpn_config,
            "wireguardConfig": wg_config,
            "secrets": {
                "noisePrivateKey": noise_priv,
                "wgPrivateKey": wg_priv,
            }
        }))
    }

    /// Export a client config (without secrets) in the specified format.
    pub async fn export_client_config(&self, client_id: &str, format: &str) -> Result<serde_json::Value> {
        let state = self.state.as_ref()
            .ok_or_else(|| anyhow::anyhow!("Server not running"))?;
        let registry = state.client_registry.read().await;
        let entry = registry.get_by_id(client_id)
            .ok_or_else(|| anyhow::anyhow!("Client '{}' not found", client_id))?;

        match format {
            "smartvpn" => {
                let smartvpn_server_url = format!("wss://{}",
                    state.config.server_endpoint.as_deref()
                        .unwrap_or(&state.config.listen_addr)
                        .replace("0.0.0.0", "localhost"));
                Ok(serde_json::json!({
                    "config": {
                        "serverUrl": smartvpn_server_url,
                        "serverPublicKey": state.config.public_key,
                        "clientPublicKey": entry.public_key,
                        "dns": state.config.dns,
                        "mtu": state.config.mtu,
                        "keepaliveIntervalSecs": state.config.keepalive_interval_secs,
                    }
                }))
            }
            "wireguard" => {
                let wg_server_pubkey = match &state.config.wg_private_key {
                    Some(wg_priv_key) => crate::wireguard::wg_public_key_from_private(wg_priv_key)?,
                    None => state.config.public_key.clone(),
                };
                let assigned_ip = entry.assigned_ip.as_deref().unwrap_or("0.0.0.0");
                let wg_endpoint = state.config.server_endpoint.as_deref()
                    .unwrap_or(&state.config.listen_addr);
                let wg_allowed_ips = state.config.client_allowed_ips.as_ref()
                    .map(|ips| ips.join(", "))
                    .unwrap_or_else(|| "0.0.0.0/0".to_string());
                let config = format!(
                    "[Interface]\nAddress = {}/24\n{}\n[Peer]\nPublicKey = {}\nAllowedIPs = {}\nEndpoint = {}\nPersistentKeepalive = 25\n",
                    assigned_ip,
                    state.config.dns.as_ref()
                        .map(|d| format!("DNS = {}", d.join(", ")))
                        .unwrap_or_default(),
                    wg_server_pubkey,
                    wg_allowed_ips,
                    wg_endpoint,
                );
                Ok(serde_json::json!({ "config": config }))
            }
            _ => anyhow::bail!("Unknown format: {}", format),
        }
    }
}

/// WebSocket listener — accepts TCP connections, optionally parses PROXY protocol v2,
/// upgrades to WS, then hands off to `handle_client_connection`.
async fn run_ws_listener(
    state: Arc<ServerState>,
    listen_addr: String,
    shutdown_rx: &mut mpsc::Receiver<()>,
) -> Result<()> {
    let listener = TcpListener::bind(&listen_addr).await?;
    info!("WebSocket server listening on {}", listen_addr);

    loop {
        tokio::select! {
            accept = listener.accept() => {
                match accept {
                    Ok((mut tcp_stream, tcp_addr)) => {
                        info!("New connection from {}", tcp_addr);
                        let state = state.clone();
                        tokio::spawn(async move {
                            // Phase 0: Parse PROXY protocol v2 header if enabled
                            let remote_addr = if state.config.proxy_protocol.unwrap_or(false) {
                                match crate::proxy_protocol::read_proxy_header(&mut tcp_stream).await {
                                    Ok(header) if header.is_local => {
                                        info!("PP v2 LOCAL probe from {}", tcp_addr);
                                        return; // Health check — close gracefully
                                    }
                                    Ok(header) => {
                                        info!("PP v2: real client {} (via {})", header.src_addr, tcp_addr);
                                        Some(header.src_addr)
                                    }
                                    Err(e) => {
                                        warn!("PP v2 parse failed from {}: {}", tcp_addr, e);
                                        return; // Drop connection
                                    }
                                }
                            } else {
                                Some(tcp_addr) // Direct connection — use TCP SocketAddr
                            };

                            // Phase 1: Server-level connection IP block list (pre-handshake)
                            if let (Some(ref block_list), Some(ref addr)) = (&state.config.connection_ip_block_list, &remote_addr) {
                                if !block_list.is_empty() {
                                    if let std::net::IpAddr::V4(v4) = addr.ip() {
                                        if acl::is_connection_blocked(v4, block_list) {
                                            warn!("Connection blocked by server IP block list: {}", addr);
                                            return;
                                        }
                                    }
                                }
                            }

                            // Phase 2: WebSocket upgrade + VPN handshake
                            match transport::accept_connection(tcp_stream).await {
                                Ok(ws) => {
                                    let (sink, stream) = transport_trait::split_ws(ws);
                                    if let Err(e) = handle_client_connection(
                                        state,
                                        Box::new(sink),
                                        Box::new(stream),
                                        remote_addr,
                                        "websocket",
                                    ).await {
                                        warn!("Client connection error: {}", e);
                                    }
                                }
                                Err(e) => {
                                    warn!("WebSocket upgrade failed: {}", e);
                                }
                            }
                        });
                    }
                    Err(e) => {
                        error!("Accept error: {}", e);
                    }
                }
            }
            _ = shutdown_rx.recv() => {
                info!("Shutdown signal received");
                break;
            }
        }
    }

    Ok(())
}

/// QUIC listener — accepts QUIC connections and hands off to the transport-agnostic
/// `handle_client_connection`.
async fn run_quic_listener(
    state: Arc<ServerState>,
    listen_addr: String,
    idle_timeout_secs: u64,
    shutdown_rx: &mut mpsc::Receiver<()>,
) -> Result<()> {
    // Generate or use configured TLS certificate for QUIC
    let (cert_chain, private_key) = if let (Some(ref cert_pem), Some(ref key_pem)) =
        (&state.config.tls_cert, &state.config.tls_key)
    {
        // Parse PEM certificates
        let certs: Vec<rustls_pki_types::CertificateDer<'static>> =
            rustls_pemfile::certs(&mut cert_pem.as_bytes())
                .collect::<Result<Vec<_>, _>>()?;
        let key = rustls_pemfile::private_key(&mut key_pem.as_bytes())?
            .ok_or_else(|| anyhow::anyhow!("No private key found in PEM"))?;
        (certs, key)
    } else {
        // Generate self-signed certificate
        let (certs, key) = quic_transport::generate_self_signed_cert()?;
        info!("QUIC using self-signed certificate (hash: {})", quic_transport::cert_hash(&certs[0]));
        (certs, key)
    };

    let endpoint = quic_transport::create_quic_server(quic_transport::QuicServerConfig {
        listen_addr,
        cert_chain,
        private_key,
        idle_timeout_secs,
    })?;

    loop {
        tokio::select! {
            incoming = endpoint.accept() => {
                match incoming {
                    Some(incoming) => {
                        let state = state.clone();
                        tokio::spawn(async move {
                            match incoming.await {
                                Ok(conn) => {
                                    let remote = conn.remote_address();
                                    info!("New QUIC connection from {}", remote);
                                    match quic_transport::accept_quic_connection(conn).await {
                                        Ok((sink, stream)) => {
                                            if let Err(e) = handle_client_connection(
                                                state,
                                                Box::new(sink),
                                                Box::new(stream),
                                                Some(remote),
                                                "quic",
                                            ).await {
                                                warn!("QUIC client error: {}", e);
                                            }
                                        }
                                        Err(e) => {
                                            warn!("QUIC stream accept failed: {}", e);
                                        }
                                    }
                                }
                                Err(e) => {
                                    warn!("QUIC handshake failed: {}", e);
                                }
                            }
                        });
                    }
                    None => {
                        info!("QUIC endpoint closed");
                        break;
                    }
                }
            }
            _ = shutdown_rx.recv() => {
                info!("QUIC shutdown signal received");
                endpoint.close(0u32.into(), b"shutdown");
                break;
            }
        }
    }

    Ok(())
}

/// TUN reader task: reads IP packets from the TUN device and dispatches them
/// to the correct client via the routing table.
async fn run_tun_reader(
    state: Arc<ServerState>,
    mut tun_reader: tokio::io::ReadHalf<tun::AsyncDevice>,
    mut shutdown_rx: mpsc::Receiver<()>,
) -> Result<()> {
    use tokio::io::AsyncReadExt;

    let mut buf = vec![0u8; 65536];

    loop {
        tokio::select! {
            result = tun_reader.read(&mut buf) => {
                let n = match result {
                    Ok(0) => {
                        info!("TUN reader: device closed");
                        break;
                    }
                    Ok(n) => n,
                    Err(e) => {
                        error!("TUN reader error: {}", e);
                        break;
                    }
                };

                // Extract destination IP from the raw IP packet
                let dst_ip = match tunnel::extract_dst_ip(&buf[..n]) {
                    Some(std::net::IpAddr::V4(v4)) => v4,
                    _ => continue, // IPv6 or malformed — skip
                };

                // Look up client by destination IP
                let routes = state.tun_routes.read().await;
                if let Some(sender) = routes.get(&dst_ip) {
                    if sender.try_send(buf[..n].to_vec()).is_err() {
                        // Channel full or closed — drop packet (correct for IP best-effort)
                    }
                }
            }
            _ = shutdown_rx.recv() => {
                info!("TUN reader shutting down");
                break;
            }
        }
    }

    Ok(())
}

/// Transport-agnostic client handler. Performs the Noise IK handshake, authenticates
/// the client against the registry, and runs the main packet forwarding loop.
async fn handle_client_connection(
    state: Arc<ServerState>,
    mut sink: Box<dyn TransportSink>,
    mut stream: Box<dyn TransportStream>,
    remote_addr: Option<std::net::SocketAddr>,
    transport_type: &str,
) -> Result<()> {
    let server_private_key = base64::Engine::decode(
        &base64::engine::general_purpose::STANDARD,
        &state.config.private_key,
    )?;

    // Noise IK handshake (server side = responder)
    let mut responder = crypto::create_responder(&server_private_key)?;
    let mut buf = vec![0u8; 65535];

    // Receive handshake init (-> e, es, s, ss)
    let init_msg = match stream.recv_reliable().await? {
        Some(data) => data,
        None => anyhow::bail!("Connection closed before handshake"),
    };

    let mut frame_buf = BytesMut::from(&init_msg[..]);
    let frame = <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf)?
        .ok_or_else(|| anyhow::anyhow!("Incomplete handshake frame"))?;

    if frame.packet_type != PacketType::HandshakeInit {
        anyhow::bail!("Expected HandshakeInit, got {:?}", frame.packet_type);
    }

    responder.read_message(&frame.payload, &mut buf)?;

    // Extract client's static public key BEFORE entering transport mode
    let client_pub_key_bytes = responder
        .get_remote_static()
        .ok_or_else(|| anyhow::anyhow!("IK handshake: no client static key received"))?
        .to_vec();
    let client_pub_key_b64 = base64::Engine::encode(
        &base64::engine::general_purpose::STANDARD,
        &client_pub_key_bytes,
    );

    // Verify client against registry
    let (registered_client_id, client_security) = {
        let registry = state.client_registry.read().await;
        if !registry.is_authorized(&client_pub_key_b64) {
            warn!("Rejecting unauthorized client with key {}", &client_pub_key_b64[..8]);
            // Send handshake response but then disconnect
            let len = responder.write_message(&[], &mut buf)?;
            let response_frame = Frame {
                packet_type: PacketType::HandshakeResp,
                payload: buf[..len].to_vec(),
            };
            let mut frame_bytes = BytesMut::new();
            <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, response_frame, &mut frame_bytes)?;
            sink.send_reliable(frame_bytes.to_vec()).await?;

            // Send disconnect frame
            let disconnect_frame = Frame {
                packet_type: PacketType::Disconnect,
                payload: Vec::new(),
            };
            let mut frame_bytes = BytesMut::new();
            <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, disconnect_frame, &mut frame_bytes)?;
            let _ = sink.send_reliable(frame_bytes.to_vec()).await;
            anyhow::bail!("Client not authorized");
        }
        let entry = registry.get_by_key(&client_pub_key_b64).unwrap();
        (entry.client_id.clone(), entry.security.clone())
    };

    // Complete handshake (<- e, ee, se)
    let len = responder.write_message(&[], &mut buf)?;
    let response_payload = buf[..len].to_vec();

    let response_frame = Frame {
        packet_type: PacketType::HandshakeResp,
        payload: response_payload,
    };
    let mut frame_bytes = BytesMut::new();
    <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, response_frame, &mut frame_bytes)?;
    sink.send_reliable(frame_bytes.to_vec()).await?;

    let mut noise_transport = responder.into_transport_mode()?;

    // Connection-level ACL: check real client IP against per-client ipAllowList/ipBlockList
    if let (Some(ref sec), Some(ref addr)) = (&client_security, &remote_addr) {
        if let std::net::IpAddr::V4(v4) = addr.ip() {
            if !acl::is_source_allowed(
                v4,
                sec.ip_allow_list.as_deref(),
                sec.ip_block_list.as_deref(),
            ) {
                warn!("Connection-level ACL denied client {} from IP {}", registered_client_id, addr);
                let disconnect_frame = Frame { packet_type: PacketType::Disconnect, payload: Vec::new() };
                let mut frame_bytes = BytesMut::new();
                <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, disconnect_frame, &mut frame_bytes)?;
                let _ = sink.send_reliable(frame_bytes.to_vec()).await;
                anyhow::bail!("Connection denied: source IP {} not allowed for client {}", addr, registered_client_id);
            }
        }
    }

    // Use the registered client ID as the connection ID
    let client_id = registered_client_id.clone();

    // Allocate IP
    let assigned_ip = state.ip_pool.lock().await.allocate(&client_id)?;

    // Create return-packet channel for forwarding engine -> client
    let (tun_return_tx, mut tun_return_rx) = mpsc::channel::<Vec<u8>>(256);
    let fwd_mode = state.config.forwarding_mode.as_deref().unwrap_or("testing");
    let forwarding_active = fwd_mode == "tun" || fwd_mode == "socket";
    if forwarding_active {
        state.tun_routes.write().await.insert(assigned_ip, tun_return_tx);
    }

    // Determine rate limits: per-client security overrides server defaults
    let (rate_limit, burst) = if let Some(ref sec) = client_security {
        if let Some(ref rl) = sec.rate_limit {
            (Some(rl.bytes_per_sec), Some(rl.burst_bytes))
        } else {
            (state.config.default_rate_limit_bytes_per_sec, state.config.default_burst_bytes)
        }
    } else {
        (state.config.default_rate_limit_bytes_per_sec, state.config.default_burst_bytes)
    };

    // Register connected client
    let client_info = ClientInfo {
        client_id: client_id.clone(),
        assigned_ip: assigned_ip.to_string(),
        connected_since: timestamp_now(),
        bytes_sent: 0,
        bytes_received: 0,
        packets_dropped: 0,
        bytes_dropped: 0,
        last_keepalive_at: None,
        keepalives_received: 0,
        rate_limit_bytes_per_sec: rate_limit,
        burst_bytes: burst,
        authenticated_key: client_pub_key_b64.clone(),
        registered_client_id: registered_client_id.clone(),
        remote_addr: remote_addr.map(|a| a.to_string()),
        transport_type: transport_type.to_string(),
    };
    state.clients.write().await.insert(client_id.clone(), client_info);

    // Set up rate limiter
    if let (Some(rate), Some(burst)) = (rate_limit, burst) {
        state
            .rate_limiters
            .lock()
            .await
            .insert(client_id.clone(), TokenBucket::new(rate, burst));
    }

    {
        let mut stats = state.stats.write().await;
        stats.total_connections += 1;
        match transport_type {
            "websocket" => stats.total_connections_websocket += 1,
            "quic" => stats.total_connections_quic += 1,
            _ => {}
        }
    }

    // Send assigned IP info (encrypted), include effective MTU
    let ip_info = serde_json::json!({
        "assignedIp": assigned_ip.to_string(),
        "gateway": state.ip_pool.lock().await.gateway_addr().to_string(),
        "mtu": state.config.mtu.unwrap_or(1420),
        "effectiveMtu": state.mtu_config.effective_mtu,
    });
    let ip_info_bytes = serde_json::to_vec(&ip_info)?;
    let len = noise_transport.write_message(&ip_info_bytes, &mut buf)?;
    let encrypted_info = Frame {
        packet_type: PacketType::IpPacket,
        payload: buf[..len].to_vec(),
    };
    let mut frame_bytes = BytesMut::new();
    <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, encrypted_info, &mut frame_bytes)?;
    sink.send_reliable(frame_bytes.to_vec()).await?;

    info!("Client {} ({}) connected with IP {} from {}",
        registered_client_id, &client_pub_key_b64[..8], assigned_ip,
        remote_addr.map(|a| a.to_string()).unwrap_or_else(|| "unknown".to_string()));

    // Main packet loop with dead-peer detection
    let mut last_activity = tokio::time::Instant::now();

    loop {
        tokio::select! {
            msg = stream.recv_reliable() => {
                match msg {
                    Ok(Some(data)) => {
                        last_activity = tokio::time::Instant::now();
                        let mut frame_buf = BytesMut::from(&data[..]);
                        match <FrameCodec as tokio_util::codec::Decoder>::decode(&mut FrameCodec, &mut frame_buf) {
                            Ok(Some(frame)) => match frame.packet_type {
                                PacketType::IpPacket => {
                                    match noise_transport.read_message(&frame.payload, &mut buf) {
                                        Ok(len) => {
                                            // ACL check on decrypted packet
                                            if let Some(ref sec) = client_security {
                                                if len >= 20 {
                                                    // Extract src/dst from IPv4 header
                                                    let src = Ipv4Addr::new(buf[12], buf[13], buf[14], buf[15]);
                                                    let dst = Ipv4Addr::new(buf[16], buf[17], buf[18], buf[19]);
                                                    let acl_result = acl::check_acl(sec, src, dst);
                                                    if acl_result != acl::AclResult::Allow {
                                                        let mut clients = state.clients.write().await;
                                                        if let Some(info) = clients.get_mut(&client_id) {
                                                            info.packets_dropped += 1;
                                                            info.bytes_dropped += len as u64;
                                                        }
                                                        continue;
                                                    }
                                                }
                                            }

                                            // Rate limiting check
                                            let allowed = {
                                                let mut limiters = state.rate_limiters.lock().await;
                                                if let Some(limiter) = limiters.get_mut(&client_id) {
                                                    limiter.try_consume(len)
                                                } else {
                                                    true
                                                }
                                            };

                                            if !allowed {
                                                let mut clients = state.clients.write().await;
                                                if let Some(info) = clients.get_mut(&client_id) {
                                                    info.packets_dropped += 1;
                                                    info.bytes_dropped += len as u64;
                                                }
                                                continue;
                                            }

                                            let mut stats = state.stats.write().await;
                                            stats.bytes_received += len as u64;
                                            stats.packets_received += 1;

                                            // Update per-client stats
                                            drop(stats);
                                            let mut clients = state.clients.write().await;
                                            if let Some(info) = clients.get_mut(&client_id) {
                                                info.bytes_received += len as u64;
                                            }
                                            drop(clients);

                                            // Forward decrypted packet via the active engine
                                            {
                                                let mut engine = state.forwarding_engine.lock().await;
                                                match &mut *engine {
                                                    ForwardingEngine::Tun(writer) => {
                                                        use tokio::io::AsyncWriteExt;
                                                        if let Err(e) = writer.write_all(&buf[..len]).await {
                                                            warn!("TUN write error for client {}: {}", client_id, e);
                                                        }
                                                    }
                                                    ForwardingEngine::Socket(sender) => {
                                                        let _ = sender.try_send(buf[..len].to_vec());
                                                    }
                                                    ForwardingEngine::Bridge(sender) => {
                                                        let _ = sender.try_send(buf[..len].to_vec());
                                                    }
                                                    ForwardingEngine::Testing => {}
                                                }
                                            }
                                        }
                                        Err(e) => {
                                            warn!("Decrypt error from {}: {}", client_id, e);
                                            break;
                                        }
                                    }
                                }
                                PacketType::Keepalive => {
                                    // Echo the keepalive payload back in the ACK
                                    let ack_frame = Frame {
                                        packet_type: PacketType::KeepaliveAck,
                                        payload: frame.payload.clone(),
                                    };
                                    let mut frame_bytes = BytesMut::new();
                                    <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(&mut FrameCodec, ack_frame, &mut frame_bytes)?;
                                    sink.send_reliable(frame_bytes.to_vec()).await?;

                                    let mut stats = state.stats.write().await;
                                    stats.keepalives_received += 1;
                                    stats.keepalives_sent += 1;

                                    // Update per-client keepalive tracking
                                    drop(stats);
                                    let mut clients = state.clients.write().await;
                                    if let Some(info) = clients.get_mut(&client_id) {
                                        info.last_keepalive_at = Some(timestamp_now());
                                        info.keepalives_received += 1;
                                    }
                                }
                                PacketType::Disconnect => {
                                    info!("Client {} sent disconnect", client_id);
                                    break;
                                }
                                _ => {
                                    warn!("Unexpected packet type from {}: {:?}", client_id, frame.packet_type);
                                }
                            },
                            Ok(None) => {
                                warn!("Incomplete frame from {}", client_id);
                            }
                            Err(e) => {
                                warn!("Frame decode error from {}: {}", client_id, e);
                                break;
                            }
                        }
                    }
                    Ok(None) => {
                        info!("Client {} connection closed", client_id);
                        break;
                    }
                    Err(e) => {
                        warn!("Transport error from {}: {}", client_id, e);
                        break;
                    }
                }
            }
            // Return packets from TUN device destined for this client
            Some(packet) = tun_return_rx.recv() => {
                let pkt_len = packet.len();
                match noise_transport.write_message(&packet, &mut buf) {
                    Ok(len) => {
                        let frame = Frame {
                            packet_type: PacketType::IpPacket,
                            payload: buf[..len].to_vec(),
                        };
                        let mut frame_bytes = BytesMut::new();
                        <FrameCodec as tokio_util::codec::Encoder<Frame>>::encode(
                            &mut FrameCodec, frame, &mut frame_bytes
                        )?;
                        sink.send_reliable(frame_bytes.to_vec()).await?;

                        // Update stats
                        let mut stats = state.stats.write().await;
                        stats.bytes_sent += pkt_len as u64;
                        stats.packets_sent += 1;
                        drop(stats);
                        let mut clients = state.clients.write().await;
                        if let Some(info) = clients.get_mut(&client_id) {
                            info.bytes_sent += pkt_len as u64;
                        }
                    }
                    Err(e) => {
                        warn!("Noise encrypt error for return packet to {}: {}", client_id, e);
                        break;
                    }
                }
            }
            _ = tokio::time::sleep_until(last_activity + DEAD_PEER_TIMEOUT) => {
                warn!("Client {} dead-peer timeout ({}s inactivity)", client_id, DEAD_PEER_TIMEOUT.as_secs());
                break;
            }
        }
    }

    // Cleanup
    if forwarding_active {
        state.tun_routes.write().await.remove(&assigned_ip);
    }
    state.clients.write().await.remove(&client_id);
    state.ip_pool.lock().await.release(&assigned_ip);
    state.rate_limiters.lock().await.remove(&client_id);
    info!("Client {} disconnected, released IP {}", client_id, assigned_ip);

    Ok(())
}

fn timestamp_now() -> String {
    use std::time::SystemTime;
    let duration = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap_or_default();
    format!("{}", duration.as_secs())
}