proxychains-masq 0.1.5

//! TUN tunnel loop — intercepts all TCP flows and relays them through the proxy chain.
//!
//! # Architecture
//!
//! A dedicated OS thread runs the smoltcp poll loop (because `TunTapInterface` is not
//! `Send`). For every new TCP connection that appears on the TUN device, the poll loop
//! spawns an async *relay task* via a Tokio [`Handle`].
//!
//! Data flows between the poll loop and relay tasks through `tokio::sync::mpsc` channels:
//!
//! ```text
//! TUN fd  ──packets──►  IpDevice  ──►  smoltcp::Interface::poll()
//!                                             │
//!                          ┌──────────────────┘
//!                          │  per-flow mpsc channels
//!                          ▼
//!                     relay_flow()  ──►  ChainEngine::connect()  ──►  proxy
//! ```

use std::{
    collections::HashMap,
    fs::File,
    io::{Read, Write},
    net::{IpAddr, Ipv4Addr},
    os::unix::io::AsRawFd,
    sync::Arc,
    thread,
};

use anyhow::{Context, Result};
use rand::{rngs::OsRng, RngCore};
use smoltcp::{
    iface::{Config, Interface, SocketHandle, SocketSet},
    phy::{wait as phy_wait, Device, DeviceCapabilities, Medium, RxToken, TxToken},
    socket::tcp::{self as tcp, Socket as TcpSocket},
    time::{Duration as SmolDuration, Instant as SmolInstant},
    wire::{
        HardwareAddress, IpAddress, IpCidr, IpEndpoint, IpListenEndpoint, IpProtocol, Ipv4Address,
        Ipv4Packet, TcpPacket, UdpPacket,
    },
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    runtime::Handle,
    sync::mpsc,
};
use tracing::{debug, warn};

use crate::{
    chain::ChainEngine,
    dns::DnsMap,
    proxy::{BoxStream, Target},
};

pub use crate::chain::ChainConfig;

// ─── Device ──────────────────────────────────────────────────────────────────

/// A minimal smoltcp [`Device`] that wraps a raw TUN file descriptor.
///
/// Packets to feed into smoltcp are placed into `rx` before calling
/// `Interface::poll`.  Packets smoltcp wants to transmit accumulate in `tx`
/// and are drained by the poll loop after each call to `Interface::poll`.
struct IpDevice {
    rx: Option<Vec<u8>>,
    tx: Vec<Vec<u8>>,
    mtu: usize,
}

impl IpDevice {
    fn new(mtu: usize) -> Self {
        Self {
            rx: None,
            tx: Vec::new(),
            mtu,
        }
    }
}

struct OwnedRxToken(Vec<u8>);
impl RxToken for OwnedRxToken {
    fn consume<R, F: FnOnce(&mut [u8]) -> R>(mut self, f: F) -> R {
        f(&mut self.0)
    }
}

struct OwnedTxToken<'a>(&'a mut Vec<Vec<u8>>);
impl<'a> TxToken for OwnedTxToken<'a> {
    fn consume<R, F: FnOnce(&mut [u8]) -> R>(self, len: usize, f: F) -> R {
        let mut buf = vec![0u8; len];
        let r = f(&mut buf);
        self.0.push(buf);
        r
    }
}

impl Device for IpDevice {
    type RxToken<'a> = OwnedRxToken;
    type TxToken<'a> = OwnedTxToken<'a>;

    fn receive(&mut self, _: SmolInstant) -> Option<(OwnedRxToken, OwnedTxToken<'_>)> {
        self.rx
            .take()
            .map(|p| (OwnedRxToken(p), OwnedTxToken(&mut self.tx)))
    }

    fn transmit(&mut self, _: SmolInstant) -> Option<OwnedTxToken<'_>> {
        Some(OwnedTxToken(&mut self.tx))
    }

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

// ─── Configuration ────────────────────────────────────────────────────────────

/// Configuration for [`ProxyChainTunnel`].
#[derive(Debug, Clone)]
pub struct TunnelConfig {
    /// Proxy chain setup.
    pub chain: ChainConfig,
    /// Bidirectional map used to resolve fake IPs back to hostnames.
    pub dns_map: DnsMap,
    /// IP address to assign to the smoltcp interface (must match the namespace TUN IP).
    pub tun_ip: Ipv4Addr,
    /// Subnet prefix length for `tun_ip`.
    pub prefix_len: u8,
    /// If set, UDP packets to this IP on port 53 are intercepted and handled as
    /// DNS A queries: the response assigns a fake IP from the `dns_map`.
    pub dns_ip: Option<Ipv4Addr>,
}

// ─── Per-flow state ───────────────────────────────────────────────────────────

struct FlowInfo {
    /// Forward smoltcp-received bytes to the relay task.
    to_relay: mpsc::Sender<Vec<u8>>,
    /// Receive bytes from the relay task to write into the smoltcp socket.
    from_relay: mpsc::Receiver<Vec<u8>>,
}

// ─── ProxyChainTunnel ─────────────────────────────────────────────────────────

/// Intercepts all TCP flows arriving on a TUN device and relays them through
/// the configured proxy chain.
pub struct ProxyChainTunnel {
    config: TunnelConfig,
}

impl ProxyChainTunnel {
    /// Create a new tunnel with the given configuration.
    pub fn new(config: TunnelConfig) -> Self {
        Self { config }
    }

    /// Spawn the blocking poll loop on a new OS thread and return immediately.
    ///
    /// `tun_file` must be a `File` wrapping the TUN device fd with `O_NONBLOCK` set.
    /// The returned [`thread::JoinHandle`] can be used to detect if the loop exits unexpectedly.
    pub fn spawn(self, tun_file: File, rt: Handle) -> thread::JoinHandle<Result<()>> {
        thread::spawn(move || poll_loop(self.config, tun_file, rt))
    }

    /// Run the tunnel loop on an async task, blocking the current async context
    /// until the tunnel exits.
    ///
    /// The blocking work runs on a dedicated OS thread (not the Tokio thread pool)
    /// so that the non-`Send` smoltcp state stays on its own thread.
    ///
    /// # Errors
    ///
    /// Propagates any error returned by the poll loop.
    pub async fn run(self, tun_file: File) -> Result<()> {
        let rt = Handle::current();
        let config = self.config;
        // `std::thread::spawn` (rather than `spawn_blocking`) so the non-Send
        // smoltcp state never needs to be moved between threads.
        let join = thread::spawn(move || poll_loop(config, tun_file, rt));
        tokio::task::spawn_blocking(move || join.join())
            .await
            .context("tunnel thread panicked")?
            .map_err(|_| anyhow::anyhow!("tunnel thread panicked"))
            .and_then(|r| r)
    }
}

// ─── Poll loop ────────────────────────────────────────────────────────────────

/// The maximum number of bytes per channel message.
const CHANNEL_CAPACITY: usize = 64;
/// Receive and transmit buffer sizes for each smoltcp TCP socket.
const SOCKET_BUF: usize = 64 * 1024;
/// Maximum time to wait on the TUN fd between smoltcp polls.
const MAX_POLL_WAIT_MS: u64 = 5;

fn poll_loop(config: TunnelConfig, mut tun_file: File, rt: Handle) -> Result<()> {
    let mtu = 1500usize;

    // Raw fd value used only for smoltcp's `phy_wait` (a thin `select(2)` wrapper
    // that needs the underlying integer).  All actual I/O goes through `tun_file`.
    let tun_raw_fd = tun_file.as_raw_fd();

    // ── smoltcp setup ──────────────────────────────────────────────────────
    let mut device = IpDevice::new(mtu);

    let mut iface_cfg = Config::new(HardwareAddress::Ip);
    // Randomise the seed so TCP ISNs are unpredictable on every run.
    iface_cfg.random_seed = OsRng.next_u64();
    let mut iface = Interface::new(iface_cfg, &mut device, SmolInstant::now());
    // Accept packets for any destination IP, not just the interface's own IP.
    iface.set_any_ip(true);
    iface.update_ip_addrs(|addrs| {
        let cidr = IpCidr::new(IpAddress::Ipv4(config.tun_ip.into()), config.prefix_len);
        let _ = addrs.push(cidr);
    });
    // With any_ip enabled, smoltcp still requires a route to exist for the
    // destination IP before it will accept the packet.  A default route
    // pointing at our own IP satisfies the route lookup (has_ip_addr check).
    iface
        .routes_mut()
        .add_default_ipv4_route(Ipv4Address::from(config.tun_ip))
        .expect("smoltcp route table is full");

    let mut sockets = SocketSet::new(vec![]);
    let engine = Arc::new(ChainEngine::new(config.chain));
    let dns_map = config.dns_map;
    let dns_ip = config.dns_ip;

    // pending:  SYN seen, socket created, waiting for smoltcp to complete the
    //           three-way handshake.  Key = 4-tuple, value = socket handle.
    let mut pending: HashMap<(IpEndpoint, IpEndpoint), SocketHandle> = HashMap::new();

    // active:   Established flows with running relay tasks.
    let mut active: HashMap<SocketHandle, FlowInfo> = HashMap::new();

    let mut read_buf = vec![0u8; mtu + 4];

    loop {
        // ── 1. Wait for data on the TUN fd (or smoltcp poll timeout) ───────
        let delay = iface
            .poll_delay(SmolInstant::now(), &sockets)
            .map(|d| d.min(SmolDuration::from_millis(MAX_POLL_WAIT_MS)));
        phy_wait(tun_raw_fd, delay).ok();

        // ── 2. Read one packet from the TUN fd (non-blocking) ──────────────
        // O_NONBLOCK is set on the fd; when no packet is available, `read`
        // returns `WouldBlock` which the `_` arm below treats as no data.
        let maybe_packet = match tun_file.read(&mut read_buf) {
            Ok(n) if n > 0 => Some(read_buf[..n].to_vec()),
            _ => None,
        };

        if let Some(pkt) = maybe_packet {
            // Intercept DNS queries before smoltcp sees them.
            if let Some(dns_ip) = dns_ip {
                if let Some(resp) = try_handle_dns(&pkt, dns_ip, &dns_map) {
                    let _ = tun_file.write_all(&resp);
                    // Skip smoltcp for this packet — it's been handled.
                    continue;
                }
            }

            // Inspect before handing to smoltcp so we can pre-create a
            // listening socket for brand-new TCP flows.
            if let Some((src, dst)) = extract_tcp_syn(&pkt) {
                if let std::collections::hash_map::Entry::Vacant(e) = pending.entry((src, dst)) {
                    // IpEndpoint is Copy, so we can borrow the key from the entry.
                    let listen_dst = e.key().1;
                    let mut sock = TcpSocket::new(
                        tcp::SocketBuffer::new(vec![0u8; SOCKET_BUF]),
                        tcp::SocketBuffer::new(vec![0u8; SOCKET_BUF]),
                    );
                    // Listen on the exact destination endpoint so that
                    // concurrent connections to different IPs on the same port
                    // each get their own socket.
                    let listen_ep = IpListenEndpoint {
                        addr: Some(listen_dst.addr),
                        port: listen_dst.port,
                    };
                    if sock.listen(listen_ep).is_ok() {
                        let handle = sockets.add(sock);
                        e.insert(handle);
                        debug!("new TCP flow {src} → {listen_dst}");
                    }
                }
            }
            // Feed the packet into the smoltcp device buffer.
            device.rx = Some(pkt);
        }

        // ── 3. Drive smoltcp ───────────────────────────────────────────────
        iface.poll(SmolInstant::now(), &mut device, &mut sockets);

        // ── 4. Write smoltcp output packets back to the TUN fd ─────────────
        for pkt in device.tx.drain(..) {
            let _ = tun_file.write_all(&pkt);
        }

        // ── 5. Promote pending sockets that finished the handshake ─────────
        pending.retain(|(_src, dst), &mut handle| {
            let sock = sockets.get_mut::<TcpSocket>(handle);
            // `may_send()` is true only in ESTABLISHED or CLOSE_WAIT, meaning
            // the three-way handshake has fully completed.  `is_active()` is
            // also true in SYN_RECEIVED, which is too early to relay data.
            if !sock.may_send() {
                return true;
            }

            // Build the proxy target: resolve fake IP to hostname if mapped.
            let target = {
                let ip: IpAddr = match dst.addr {
                    IpAddress::Ipv4(a) => IpAddr::V4(Ipv4Addr::from(a)),
                    IpAddress::Ipv6(a) => IpAddr::V6(a.into()),
                };
                let hostname = if let IpAddr::V4(v4) = ip {
                    dns_map.lookup_hostname(v4)
                } else {
                    None
                };
                match hostname {
                    Some(h) => Target::Host(h, dst.port),
                    None => Target::Ip(ip, dst.port),
                }
            };

            // Create the channel pair for this flow.
            let (to_relay_tx, to_relay_rx) = mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
            let (from_relay_tx, from_relay_rx) = mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);

            let eng = engine.clone();
            rt.spawn(relay_flow(eng, target, to_relay_rx, from_relay_tx));

            active.insert(
                handle,
                FlowInfo {
                    to_relay: to_relay_tx,
                    from_relay: from_relay_rx,
                },
            );
            false // remove from pending
        });

        // ── 6. Service active flows ────────────────────────────────────────
        active.retain(|&handle, info| {
            let sock = sockets.get_mut::<TcpSocket>(handle);

            // smoltcp socket → relay task
            if sock.can_recv() {
                if let Ok(data) = sock.recv(|b| (b.len(), b.to_vec())) {
                    if !data.is_empty() {
                        let _ = info.to_relay.try_send(data);
                    }
                }
            }

            // relay task → smoltcp socket
            while let Ok(data) = info.from_relay.try_recv() {
                let _ = sock.send_slice(&data);
            }

            // Clean up closed sockets.
            if matches!(
                sock.state(),
                tcp::State::Closed | tcp::State::CloseWait | tcp::State::TimeWait
            ) && !sock.can_recv()
            {
                sockets.remove(handle);
                return false;
            }
            true
        });
    }
}

// ─── Relay task ───────────────────────────────────────────────────────────────

/// Async task: connect through the proxy chain and ferry bytes between the
/// smoltcp socket (via channels) and the real proxy connection.
async fn relay_flow(
    engine: Arc<ChainEngine>,
    target: Target,
    mut smol_rx: mpsc::Receiver<Vec<u8>>,
    smol_tx: mpsc::Sender<Vec<u8>>,
) {
    let proxy: BoxStream = match engine.connect(target.clone()).await {
        Ok(s) => s,
        Err(e) => {
            warn!("relay: failed to connect to {target}: {e:#}");
            return;
        }
    };

    let (mut proxy_r, mut proxy_w) = tokio::io::split(proxy);
    let mut buf = vec![0u8; 8192];

    loop {
        tokio::select! {
            // smoltcp → proxy
            data = smol_rx.recv() => {
                match data {
                    Some(d) => { if proxy_w.write_all(&d).await.is_err() { break; } }
                    None => break,
                }
            }
            // proxy → smoltcp
            n = proxy_r.read(&mut buf) => {
                match n {
                    Ok(0) | Err(_) => break,
                    Ok(n) => {
                        if smol_tx.send(buf[..n].to_vec()).await.is_err() {
                            break;
                        }
                    }
                }
            }
        }
    }
    debug!("relay flow ended for {target}");
}

// ─── Packet parsing ───────────────────────────────────────────────────────────

/// If `packet` is a UDP DNS A-query destined for `dns_ip`, return a crafted
/// IP+UDP response with a fake IP from `dns_map`, else `None`.
fn try_handle_dns(packet: &[u8], dns_ip: Ipv4Addr, dns_map: &DnsMap) -> Option<Vec<u8>> {
    let ipv4 = Ipv4Packet::new_checked(packet).ok()?;
    if ipv4.next_header() != IpProtocol::Udp {
        return None;
    }
    if Ipv4Addr::from(ipv4.dst_addr()) != dns_ip {
        return None;
    }
    let udp = UdpPacket::new_checked(ipv4.payload()).ok()?;
    if udp.dst_port() != 53 {
        return None;
    }
    // Handle the DNS query inline via DnsMap and skip forwarding to smoltcp.
    let dns_resp = dns_map.handle_dns_query(udp.payload())?;

    // Craft response IP+UDP packet (swapped src/dst).
    let udp_len = 8u16 + dns_resp.len() as u16;
    let total_len = 20u16 + udp_len;
    let mut pkt = vec![0u8; total_len as usize];

    // IPv4 header
    pkt[0] = 0x45; // version 4, IHL 5
    pkt[2] = (total_len >> 8) as u8;
    pkt[3] = total_len as u8;
    pkt[8] = 64; // TTL
    pkt[9] = 17; // protocol UDP
    pkt[12..16].copy_from_slice(&ipv4.dst_addr().0); // src = DNS server
    pkt[16..20].copy_from_slice(&ipv4.src_addr().0); // dst = client
    let csum = ip_checksum(&pkt[0..20]);
    pkt[10] = (csum >> 8) as u8;
    pkt[11] = csum as u8;

    // UDP header
    pkt[20] = 0;
    pkt[21] = 53; // src port 53
    pkt[22] = (udp.src_port() >> 8) as u8;
    pkt[23] = udp.src_port() as u8; // dst port = client's ephemeral
    pkt[24] = (udp_len >> 8) as u8;
    pkt[25] = udp_len as u8;
    // UDP checksum = 0 (valid for IPv4)
    pkt[28..].copy_from_slice(&dns_resp);

    Some(pkt)
}

fn ip_checksum(data: &[u8]) -> u16 {
    let mut sum = 0u32;
    for chunk in data.chunks(2) {
        let word = if chunk.len() == 2 {
            u16::from_be_bytes([chunk[0], chunk[1]]) as u32
        } else {
            (chunk[0] as u32) << 8
        };
        sum += word;
    }
    while sum >> 16 != 0 {
        sum = (sum & 0xFFFF) + (sum >> 16);
    }
    !(sum as u16)
}

/// If `packet` is an IPv4 TCP SYN (not SYN-ACK), return `(src, dst)` endpoints.
fn extract_tcp_syn(packet: &[u8]) -> Option<(IpEndpoint, IpEndpoint)> {
    let ipv4 = Ipv4Packet::new_checked(packet).ok()?;
    if ipv4.next_header() != IpProtocol::Tcp {
        return None;
    }
    let tcp = TcpPacket::new_checked(ipv4.payload()).ok()?;
    if !tcp.syn() || tcp.ack() {
        return None;
    }
    let src = IpEndpoint::new(IpAddress::Ipv4(ipv4.src_addr()), tcp.src_port());
    let dst = IpEndpoint::new(IpAddress::Ipv4(ipv4.dst_addr()), tcp.dst_port());
    Some((src, dst))
}