smartvpn/rust/src/userspace_nat.rs

use std::collections::{HashMap, VecDeque};
use std::net::{Ipv4Addr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;

use anyhow::Result;
use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet};
use smoltcp::phy::{self, Device, DeviceCapabilities, Medium};
use smoltcp::socket::{tcp, udp};
use smoltcp::wire::{HardwareAddress, IpAddress, IpCidr, IpEndpoint};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpStream, UdpSocket};
use tokio::sync::mpsc;
use tracing::{debug, info, warn};

use crate::server::ServerState;
use crate::tunnel;

// ============================================================================
// Virtual IP device for smoltcp
// ============================================================================

pub struct VirtualIpDevice {
    rx_queue: VecDeque<Vec<u8>>,
    tx_queue: VecDeque<Vec<u8>>,
    mtu: usize,
}

impl VirtualIpDevice {
    pub fn new(mtu: usize) -> Self {
        Self {
            rx_queue: VecDeque::new(),
            tx_queue: VecDeque::new(),
            mtu,
        }
    }

    pub fn inject_packet(&mut self, packet: Vec<u8>) {
        self.rx_queue.push_back(packet);
    }

    pub fn drain_tx(&mut self) -> impl Iterator<Item = Vec<u8>> + '_ {
        self.tx_queue.drain(..)
    }
}

pub struct VirtualRxToken {
    buffer: Vec<u8>,
}

impl phy::RxToken for VirtualRxToken {
    fn consume<R, F>(self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(&self.buffer)
    }
}

pub struct VirtualTxToken<'a> {
    queue: &'a mut VecDeque<Vec<u8>>,
}

impl<'a> phy::TxToken for VirtualTxToken<'a> {
    fn consume<R, F>(self, len: usize, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        let mut buffer = vec![0u8; len];
        let result = f(&mut buffer);
        self.queue.push_back(buffer);
        result
    }
}

impl Device for VirtualIpDevice {
    type RxToken<'a> = VirtualRxToken;
    type TxToken<'a> = VirtualTxToken<'a>;

    fn receive(
        &mut self,
        _timestamp: smoltcp::time::Instant,
    ) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
        self.rx_queue.pop_front().map(|buffer| {
            let rx = VirtualRxToken { buffer };
            let tx = VirtualTxToken {
                queue: &mut self.tx_queue,
            };
            (rx, tx)
        })
    }

    fn transmit(&mut self, _timestamp: smoltcp::time::Instant) -> Option<Self::TxToken<'_>> {
        Some(VirtualTxToken {
            queue: &mut self.tx_queue,
        })
    }

    fn capabilities(&self) -> DeviceCapabilities {
        let mut caps = DeviceCapabilities::default();
        caps.medium = Medium::Ip;
        caps.max_transmission_unit = self.mtu;
        caps.max_burst_size = Some(1);
        caps
    }
}

// ============================================================================
// Session tracking
// ============================================================================

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct SessionKey {
    src_ip: Ipv4Addr,
    src_port: u16,
    dst_ip: Ipv4Addr,
    dst_port: u16,
    protocol: u8,
}

struct TcpSession {
    smoltcp_handle: SocketHandle,
    bridge_data_tx: mpsc::Sender<Vec<u8>>,
    #[allow(dead_code)]
    client_ip: Ipv4Addr,
}

struct UdpSession {
    smoltcp_handle: SocketHandle,
    bridge_data_tx: mpsc::Sender<Vec<u8>>,
    #[allow(dead_code)]
    client_ip: Ipv4Addr,
    last_activity: tokio::time::Instant,
}

enum BridgeMessage {
    TcpData { key: SessionKey, data: Vec<u8> },
    TcpClosed { key: SessionKey },
    UdpData { key: SessionKey, data: Vec<u8> },
}

// ============================================================================
// IP packet parsing helpers
// ============================================================================

fn parse_ipv4_header(packet: &[u8]) -> Option<(u8, Ipv4Addr, Ipv4Addr, u8)> {
    if packet.len() < 20 {
        return None;
    }
    let version = packet[0] >> 4;
    if version != 4 {
        return None;
    }
    let ihl = (packet[0] & 0x0F) as usize * 4;
    let protocol = packet[9];
    let src = Ipv4Addr::new(packet[12], packet[13], packet[14], packet[15]);
    let dst = Ipv4Addr::new(packet[16], packet[17], packet[18], packet[19]);
    Some((ihl as u8, src, dst, protocol))
}

fn parse_tcp_ports(packet: &[u8], ihl: usize) -> Option<(u16, u16, u8)> {
    if packet.len() < ihl + 14 {
        return None;
    }
    let src_port = u16::from_be_bytes([packet[ihl], packet[ihl + 1]]);
    let dst_port = u16::from_be_bytes([packet[ihl + 2], packet[ihl + 3]]);
    let flags = packet[ihl + 13];
    Some((src_port, dst_port, flags))
}

fn parse_udp_ports(packet: &[u8], ihl: usize) -> Option<(u16, u16)> {
    if packet.len() < ihl + 4 {
        return None;
    }
    let src_port = u16::from_be_bytes([packet[ihl], packet[ihl + 1]]);
    let dst_port = u16::from_be_bytes([packet[ihl + 2], packet[ihl + 3]]);
    Some((src_port, dst_port))
}

// ============================================================================
// NAT Engine
// ============================================================================

pub struct NatEngine {
    device: VirtualIpDevice,
    iface: Interface,
    sockets: SocketSet<'static>,
    tcp_sessions: HashMap<SessionKey, TcpSession>,
    udp_sessions: HashMap<SessionKey, UdpSession>,
    state: Arc<ServerState>,
    bridge_rx: mpsc::Receiver<BridgeMessage>,
    bridge_tx: mpsc::Sender<BridgeMessage>,
    start_time: std::time::Instant,
    /// When true, outbound TCP connections prepend PROXY protocol v2 headers
    /// with the VPN client's tunnel IP as source address.
    proxy_protocol: bool,
}

impl NatEngine {
    pub fn new(gateway_ip: Ipv4Addr, mtu: usize, state: Arc<ServerState>, proxy_protocol: bool) -> Self {
        let mut device = VirtualIpDevice::new(mtu);
        let config = Config::new(HardwareAddress::Ip);
        let now = smoltcp::time::Instant::from_millis(0);
        let mut iface = Interface::new(config, &mut device, now);

        // Accept packets to ANY destination IP (essential for NAT)
        iface.set_any_ip(true);

        // Assign the gateway IP as the interface address
        iface.update_ip_addrs(|addrs| {
            addrs
                .push(IpCidr::new(IpAddress::Ipv4(gateway_ip.into()), 24))
                .unwrap();
        });

        // Add a default route so smoltcp knows where to send packets
        iface.routes_mut().add_default_ipv4_route(gateway_ip.into()).unwrap();

        let sockets = SocketSet::new(Vec::with_capacity(256));
        let (bridge_tx, bridge_rx) = mpsc::channel(4096);

        Self {
            device,
            iface,
            sockets,
            tcp_sessions: HashMap::new(),
            udp_sessions: HashMap::new(),
            state,
            bridge_rx,
            bridge_tx,
            start_time: std::time::Instant::now(),
            proxy_protocol,
        }
    }

    fn smoltcp_now(&self) -> smoltcp::time::Instant {
        smoltcp::time::Instant::from_millis(self.start_time.elapsed().as_millis() as i64)
    }

    /// Inject a raw IP packet from a VPN client and handle new session creation.
    fn inject_packet(&mut self, packet: Vec<u8>) {
        let Some((ihl, src_ip, dst_ip, protocol)) = parse_ipv4_header(&packet) else {
            return;
        };
        let ihl = ihl as usize;

        match protocol {
            6 => {
                // TCP
                let Some((src_port, dst_port, flags)) = parse_tcp_ports(&packet, ihl) else {
                    return;
                };
                let key = SessionKey {
                    src_ip,
                    src_port,
                    dst_ip,
                    dst_port,
                    protocol: 6,
                };

                // SYN without ACK = new connection
                let is_syn = (flags & 0x02) != 0 && (flags & 0x10) == 0;
                if is_syn && !self.tcp_sessions.contains_key(&key) {
                    self.create_tcp_session(&key);
                }
            }
            17 => {
                // UDP
                let Some((src_port, dst_port)) = parse_udp_ports(&packet, ihl) else {
                    return;
                };
                let key = SessionKey {
                    src_ip,
                    src_port,
                    dst_ip,
                    dst_port,
                    protocol: 17,
                };

                if !self.udp_sessions.contains_key(&key) {
                    self.create_udp_session(&key);
                }

                // Update last_activity for existing sessions
                if let Some(session) = self.udp_sessions.get_mut(&key) {
                    session.last_activity = tokio::time::Instant::now();
                }
            }
            _ => {
                // ICMP and other protocols — not forwarded in socket mode
                return;
            }
        }

        self.device.inject_packet(packet);
    }

    fn create_tcp_session(&mut self, key: &SessionKey) {
        // Create smoltcp TCP socket
        let tcp_rx_buf = tcp::SocketBuffer::new(vec![0u8; 65535]);
        let tcp_tx_buf = tcp::SocketBuffer::new(vec![0u8; 65535]);
        let mut socket = tcp::Socket::new(tcp_rx_buf, tcp_tx_buf);

        // Listen on the destination address so smoltcp accepts the SYN
        let endpoint = IpEndpoint::new(
            IpAddress::Ipv4(key.dst_ip.into()),
            key.dst_port,
        );
        if socket.listen(endpoint).is_err() {
            warn!("NAT: failed to listen on {:?}", endpoint);
            return;
        }

        let handle = self.sockets.add(socket);

        // Channel for sending data from NAT engine to bridge task
        let (data_tx, data_rx) = mpsc::channel::<Vec<u8>>(256);

        let session = TcpSession {
            smoltcp_handle: handle,
            bridge_data_tx: data_tx,
            client_ip: key.src_ip,
        };
        self.tcp_sessions.insert(key.clone(), session);

        // Spawn bridge task that connects to the real destination
        let bridge_tx = self.bridge_tx.clone();
        let key_clone = key.clone();
        let proxy_protocol = self.proxy_protocol;
        tokio::spawn(async move {
            tcp_bridge_task(key_clone, data_rx, bridge_tx, proxy_protocol).await;
        });

        debug!(
            "NAT: new TCP session {}:{} -> {}:{}",
            key.src_ip, key.src_port, key.dst_ip, key.dst_port
        );
    }

    fn create_udp_session(&mut self, key: &SessionKey) {
        // Create smoltcp UDP socket
        let udp_rx_buf = udp::PacketBuffer::new(
            vec![udp::PacketMetadata::EMPTY; 32],
            vec![0u8; 65535],
        );
        let udp_tx_buf = udp::PacketBuffer::new(
            vec![udp::PacketMetadata::EMPTY; 32],
            vec![0u8; 65535],
        );
        let mut socket = udp::Socket::new(udp_rx_buf, udp_tx_buf);

        let endpoint = IpEndpoint::new(
            IpAddress::Ipv4(key.dst_ip.into()),
            key.dst_port,
        );
        if socket.bind(endpoint).is_err() {
            warn!("NAT: failed to bind UDP on {:?}", endpoint);
            return;
        }

        let handle = self.sockets.add(socket);

        let (data_tx, data_rx) = mpsc::channel::<Vec<u8>>(256);

        let session = UdpSession {
            smoltcp_handle: handle,
            bridge_data_tx: data_tx,
            client_ip: key.src_ip,
            last_activity: tokio::time::Instant::now(),
        };
        self.udp_sessions.insert(key.clone(), session);

        let bridge_tx = self.bridge_tx.clone();
        let key_clone = key.clone();
        tokio::spawn(async move {
            udp_bridge_task(key_clone, data_rx, bridge_tx).await;
        });

        debug!(
            "NAT: new UDP session {}:{} -> {}:{}",
            key.src_ip, key.src_port, key.dst_ip, key.dst_port
        );
    }

    /// Poll smoltcp, bridge data between smoltcp sockets and bridge tasks,
    /// and dispatch outgoing packets to VPN clients.
    async fn process(&mut self) {
        let now = self.smoltcp_now();
        self.iface
            .poll(now, &mut self.device, &mut self.sockets);

        // Bridge: read data from smoltcp TCP sockets → send to bridge tasks
        let mut closed_tcp: Vec<SessionKey> = Vec::new();
        for (key, session) in &self.tcp_sessions {
            let socket = self.sockets.get_mut::<tcp::Socket>(session.smoltcp_handle);
            if socket.can_recv() {
                let _ = socket.recv(|data| {
                    let _ = session.bridge_data_tx.try_send(data.to_vec());
                    (data.len(), ())
                });
            }
            // Detect closed connections
            if !socket.is_open() && !socket.is_listening() {
                closed_tcp.push(key.clone());
            }
        }

        // Clean up closed TCP sessions
        for key in closed_tcp {
            if let Some(session) = self.tcp_sessions.remove(&key) {
                self.sockets.remove(session.smoltcp_handle);
                debug!("NAT: TCP session closed {}:{} -> {}:{}", key.src_ip, key.src_port, key.dst_ip, key.dst_port);
            }
        }

        // Bridge: read data from smoltcp UDP sockets → send to bridge tasks
        for (_key, session) in &self.udp_sessions {
            let socket = self.sockets.get_mut::<udp::Socket>(session.smoltcp_handle);
            while let Ok((data, _meta)) = socket.recv() {
                let _ = session.bridge_data_tx.try_send(data.to_vec());
            }
        }

        // Dispatch outgoing packets from smoltcp to VPN clients
        let routes = self.state.tun_routes.read().await;
        for packet in self.device.drain_tx() {
            if let Some(std::net::IpAddr::V4(dst_ip)) = tunnel::extract_dst_ip(&packet) {
                if let Some(sender) = routes.get(&dst_ip) {
                    let _ = sender.try_send(packet);
                }
            }
        }
    }

    fn handle_bridge_message(&mut self, msg: BridgeMessage) {
        match msg {
            BridgeMessage::TcpData { key, data } => {
                if let Some(session) = self.tcp_sessions.get(&key) {
                    let socket =
                        self.sockets.get_mut::<tcp::Socket>(session.smoltcp_handle);
                    if socket.can_send() {
                        let _ = socket.send_slice(&data);
                    }
                }
            }
            BridgeMessage::TcpClosed { key } => {
                if let Some(session) = self.tcp_sessions.remove(&key) {
                    let socket =
                        self.sockets.get_mut::<tcp::Socket>(session.smoltcp_handle);
                    socket.close();
                    // Don't remove from SocketSet yet — let smoltcp send FIN
                    // It will be cleaned up in process() when is_open() returns false
                    self.tcp_sessions.insert(key, session);
                }
            }
            BridgeMessage::UdpData { key, data } => {
                if let Some(session) = self.udp_sessions.get_mut(&key) {
                    session.last_activity = tokio::time::Instant::now();
                    let socket =
                        self.sockets.get_mut::<udp::Socket>(session.smoltcp_handle);
                    let dst_endpoint = IpEndpoint::new(
                        IpAddress::Ipv4(key.src_ip.into()),
                        key.src_port,
                    );
                    // Send response: from the "server" (dst) back to the "client" (src)
                    let _ = socket.send_slice(&data, dst_endpoint);
                }
            }
        }
    }

    fn cleanup_idle_udp_sessions(&mut self) {
        let timeout = Duration::from_secs(60);
        let now = tokio::time::Instant::now();
        let expired: Vec<SessionKey> = self
            .udp_sessions
            .iter()
            .filter(|(_, s)| now.duration_since(s.last_activity) > timeout)
            .map(|(k, _)| k.clone())
            .collect();

        for key in expired {
            if let Some(session) = self.udp_sessions.remove(&key) {
                self.sockets.remove(session.smoltcp_handle);
                debug!(
                    "NAT: UDP session timed out {}:{} -> {}:{}",
                    key.src_ip, key.src_port, key.dst_ip, key.dst_port
                );
            }
        }
    }

    /// Main async event loop for the NAT engine.
    pub async fn run(
        mut self,
        mut packet_rx: mpsc::Receiver<Vec<u8>>,
        mut shutdown_rx: mpsc::Receiver<()>,
    ) -> Result<()> {
        info!("Userspace NAT engine started");
        let mut timer = tokio::time::interval(Duration::from_millis(50));
        let mut cleanup_timer = tokio::time::interval(Duration::from_secs(10));

        loop {
            tokio::select! {
                Some(packet) = packet_rx.recv() => {
                    self.inject_packet(packet);
                    self.process().await;
                }
                Some(msg) = self.bridge_rx.recv() => {
                    self.handle_bridge_message(msg);
                    self.process().await;
                }
                _ = timer.tick() => {
                    // Periodic poll for smoltcp maintenance (TCP retransmit, etc.)
                    self.process().await;
                }
                _ = cleanup_timer.tick() => {
                    self.cleanup_idle_udp_sessions();
                }
                _ = shutdown_rx.recv() => {
                    info!("Userspace NAT engine shutting down");
                    break;
                }
            }
        }

        Ok(())
    }
}

// ============================================================================
// Bridge tasks
// ============================================================================

async fn tcp_bridge_task(
    key: SessionKey,
    mut data_rx: mpsc::Receiver<Vec<u8>>,
    bridge_tx: mpsc::Sender<BridgeMessage>,
    proxy_protocol: bool,
) {
    let addr = SocketAddr::new(key.dst_ip.into(), key.dst_port);

    // Connect to real destination with timeout
    let stream = match tokio::time::timeout(Duration::from_secs(30), TcpStream::connect(addr)).await
    {
        Ok(Ok(s)) => s,
        Ok(Err(e)) => {
            debug!("NAT TCP connect to {} failed: {}", addr, e);
            let _ = bridge_tx.send(BridgeMessage::TcpClosed { key }).await;
            return;
        }
        Err(_) => {
            debug!("NAT TCP connect to {} timed out", addr);
            let _ = bridge_tx.send(BridgeMessage::TcpClosed { key }).await;
            return;
        }
    };

    let (mut reader, mut writer) = stream.into_split();

    // Send PROXY protocol v2 header with VPN client's tunnel IP as source
    if proxy_protocol {
        let src = SocketAddr::new(key.src_ip.into(), key.src_port);
        let dst = SocketAddr::new(key.dst_ip.into(), key.dst_port);
        let pp_header = crate::proxy_protocol::build_pp_v2_header(src, dst);
        if let Err(e) = writer.write_all(&pp_header).await {
            debug!("NAT: failed to send PP v2 header to {}: {}", addr, e);
            let _ = bridge_tx.send(BridgeMessage::TcpClosed { key }).await;
            return;
        }
    }

    // Read from real socket → send to NAT engine
    let bridge_tx2 = bridge_tx.clone();
    let key2 = key.clone();
    let read_task = tokio::spawn(async move {
        let mut buf = vec![0u8; 65536];
        loop {
            match reader.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    if bridge_tx2
                        .send(BridgeMessage::TcpData {
                            key: key2.clone(),
                            data: buf[..n].to_vec(),
                        })
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
        let _ = bridge_tx2
            .send(BridgeMessage::TcpClosed { key: key2 })
            .await;
    });

    // Receive from NAT engine → write to real socket
    while let Some(data) = data_rx.recv().await {
        if writer.write_all(&data).await.is_err() {
            break;
        }
    }

    read_task.abort();
}

async fn udp_bridge_task(
    key: SessionKey,
    mut data_rx: mpsc::Receiver<Vec<u8>>,
    bridge_tx: mpsc::Sender<BridgeMessage>,
) {
    let socket = match UdpSocket::bind("0.0.0.0:0").await {
        Ok(s) => s,
        Err(e) => {
            warn!("NAT UDP bind failed: {}", e);
            return;
        }
    };
    let dest = SocketAddr::new(key.dst_ip.into(), key.dst_port);

    let socket = Arc::new(socket);
    let socket2 = socket.clone();
    let bridge_tx2 = bridge_tx.clone();
    let key2 = key.clone();

    // Read responses from real socket
    let read_task = tokio::spawn(async move {
        let mut buf = vec![0u8; 65536];
        loop {
            match socket2.recv_from(&mut buf).await {
                Ok((n, _src)) => {
                    if bridge_tx2
                        .send(BridgeMessage::UdpData {
                            key: key2.clone(),
                            data: buf[..n].to_vec(),
                        })
                        .await
                        .is_err()
                    {
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    });

    // Forward data from NAT engine to real socket
    while let Some(data) = data_rx.recv().await {
        let _ = socket.send_to(&data, dest).await;
    }

    read_task.abort();
}