zlayer-overlay 0.14.2

//! smoltcp userspace netstack: TCP bridge over the `WireGuard` tunnel.
//!
//! This module terminates TCP entirely in userspace using [`smoltcp`] `0.12`
//! driven over a `Medium::Ip` device whose rx/tx queues are the cleartext-IP
//! packet stream on the inside of the boringtun tunnel. It is the smoltcp
//! counterpart to the Wintun transport: [`NetstackChannel`] implements
//! [`crate::tunn_loop::IpChannel`] so the same `Tunn` ingress/egress loops
//! feed and drain it.
//!
//! Architecture (mirrors `aramperes/onetun`):
//! - [`NetstackDevice`] is a [`smoltcp::phy::Device`] with two `VecDeque`s:
//!   `rx` holds inbound cleartext IP packets (pushed by the tunnel), `tx`
//!   holds egress IP packets smoltcp produced (drained to the tunnel).
//! - A single **poll task** owns the `Interface`, the `SocketSet`, the device,
//!   and the per-connection tables. Nothing else ever touches those; all
//!   interaction is via bounded channels. This keeps smoltcp's `!Send`-ish
//!   single-threaded invariants trivially satisfied.
//! - [`NetstackChannel`] bridges the tunnel loops to the poll task: inbound
//!   packets are pushed into `from_tunnel` (never back-pressuring the ingress
//!   loop — TUN-ring drop semantics), egress packets are pulled from
//!   `to_tunnel`.
//! - [`Netstack`] is the control handle: [`Netstack::connect`] and
//!   [`Netstack::listen`] send commands to the poll task and await a reply.

// The public names below (`Netstack`, `NetstackChannel`, `NetstackConfig`,
// `NetstackCmd`) are a frozen part of the edge surface — they intentionally
// repeat the module name and cannot be renamed.
#![allow(clippy::module_name_repetitions)]

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

use smoltcp::iface::{Config, Interface, SocketHandle, SocketSet};
use smoltcp::phy::{self, Device, DeviceCapabilities, Medium};
use smoltcp::socket::tcp;
use smoltcp::time::Instant;
use smoltcp::wire::{HardwareAddress, IpCidr, Ipv4Address, Ipv6Address};

use tokio::sync::mpsc::error::{TryRecvError, TrySendError};
use tokio::sync::{mpsc, oneshot, Mutex, Notify};
use tokio::task::JoinHandle;

use crate::edge::{EdgeError, TcpConn, EDGE_EPHEMERAL_PORT_START};
use crate::tunn_loop::IpChannel;
use crate::OverlayError;

/// Depth of the inbound/outbound raw-packet channels between the tunnel loops
/// and the poll task.
const PACKET_CHANNEL_BOUND: usize = 1024;
/// Depth of the control-command channel.
const CMD_CHANNEL_BOUND: usize = 64;
/// Depth of each per-connection byte channel (the socket ring buffers provide
/// the real flow-control window; these just hand chunks between tasks).
const CONN_CHANNEL_BOUND: usize = 128;
/// Maximum bytes moved out of a socket's rx buffer per pump iteration.
const RECV_CHUNK: usize = 16 * 1024;
/// Upper bound on how long the poll task will idle between polls.
const MAX_POLL_INTERVAL: Duration = Duration::from_secs(1);
/// Short retry when egress packets are stuck behind a full `to_tunnel` queue.
const EGRESS_RETRY_INTERVAL: Duration = Duration::from_millis(5);

// ---------------------------------------------------------------------------
// Device
// ---------------------------------------------------------------------------

/// A `Medium::Ip` smoltcp device backed by two in-memory packet queues.
///
/// `rx` is drained by smoltcp during `poll` (inbound, tunnel → stack); `tx`
/// is filled by smoltcp during `poll` (egress, stack → tunnel).
struct NetstackDevice {
    rx: VecDeque<Vec<u8>>,
    tx: VecDeque<Vec<u8>>,
    mtu: usize,
}

impl Device for NetstackDevice {
    type RxToken<'a> = NsRxToken;
    type TxToken<'a> = NsTxToken<'a>;

    fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
        let buffer = self.rx.pop_front()?;
        let rx = NsRxToken { buffer };
        let tx = NsTxToken { tx: &mut self.tx };
        Some((rx, tx))
    }

    fn transmit(&mut self, _timestamp: Instant) -> Option<Self::TxToken<'_>> {
        Some(NsTxToken { tx: &mut self.tx })
    }

    fn capabilities(&self) -> DeviceCapabilities {
        // Default checksum capabilities => smoltcp computes/verifies every
        // checksum in software, which is required because these packets leave
        // the host over WireGuard rather than a NIC that could offload them.
        // `DeviceCapabilities` is `#[non_exhaustive]`, so start from default
        // and set the two fields we care about.
        let mut caps = DeviceCapabilities::default();
        caps.medium = Medium::Ip;
        caps.max_transmission_unit = self.mtu;
        caps
    }
}

/// Receive token: owns one popped inbound packet.
struct NsRxToken {
    buffer: Vec<u8>,
}

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

/// Transmit token: borrows the device's egress queue and appends the packet
/// smoltcp constructs.
struct NsTxToken<'a> {
    tx: &'a mut VecDeque<Vec<u8>>,
}

impl phy::TxToken for NsTxToken<'_> {
    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.tx.push_back(buffer);
        result
    }
}

// ---------------------------------------------------------------------------
// IpChannel bridge
// ---------------------------------------------------------------------------

/// The tunnel-facing side of the netstack: an [`IpChannel`] the `Tunn` loops
/// drive.
///
/// `send_packet` injects a decapsulated inbound IP packet into the stack's rx
/// queue (dropping on a full queue, never back-pressuring the ingress loop).
/// `recv_packet` yields the next egress IP packet the stack produced.
pub struct NetstackChannel {
    from_tunnel: mpsc::Sender<Vec<u8>>,
    to_tunnel: Mutex<mpsc::Receiver<Vec<u8>>>,
    wake: Arc<Notify>,
    mtu: u32,
}

#[async_trait::async_trait]
impl IpChannel for NetstackChannel {
    async fn recv_packet(&self, buf: &mut [u8]) -> Result<usize, OverlayError> {
        let packet = {
            let mut rx = self.to_tunnel.lock().await;
            rx.recv().await
        };
        match packet {
            Some(pkt) => {
                let n = pkt.len().min(buf.len());
                buf[..n].copy_from_slice(&pkt[..n]);
                Ok(n)
            }
            // The poll task dropped its egress sender: end the egress loop
            // cleanly (same contract as `WindowsTun::recv` erroring).
            None => Err(OverlayError::NetworkConfig(
                "netstack egress channel closed".into(),
            )),
        }
    }

    async fn send_packet(&self, pkt: &[u8]) -> Result<(), OverlayError> {
        match self.from_tunnel.try_send(pkt.to_vec()) {
            Ok(()) => {
                self.wake.notify_one();
                Ok(())
            }
            // Full queue: drop like a TUN ring overflow. The ingress loop must
            // never block on us.
            Err(TrySendError::Full(_)) => {
                tracing::trace!(
                    len = pkt.len(),
                    "netstack ingress queue full; dropping packet"
                );
                Ok(())
            }
            Err(TrySendError::Closed(_)) => Err(OverlayError::NetworkConfig(
                "netstack ingress channel closed".into(),
            )),
        }
    }

    fn mtu(&self) -> u32 {
        self.mtu
    }
}

// ---------------------------------------------------------------------------
// Control handle + config
// ---------------------------------------------------------------------------

/// Configuration for a [`Netstack`] instance.
#[derive(Clone, Debug)]
pub struct NetstackConfig {
    /// The overlay IP assigned to this edge client.
    pub local_ip: IpAddr,
    /// Prefix length of the overlay subnet (for on-link routing).
    pub prefix_len: u8,
    /// Advisory device MTU reported to smoltcp and the tunnel loops.
    pub mtu: u32,
    /// Per-direction ring-buffer size for each TCP socket.
    pub tcp_buffer_bytes: usize,
    /// Idle timeout applied to every socket.
    pub tcp_timeout: Duration,
}

/// A command sent to the poll task.
enum NetstackCmd {
    Connect {
        dst: SocketAddr,
        reply: oneshot::Sender<Result<TcpConn, EdgeError>>,
    },
    Listen {
        port: u16,
        reply: oneshot::Sender<Result<mpsc::Receiver<TcpConn>, EdgeError>>,
    },
}

/// Control handle to a running netstack poll task.
pub struct Netstack {
    cmd_tx: mpsc::Sender<NetstackCmd>,
    wake: Arc<Notify>,
}

impl Netstack {
    /// Build a netstack and spawn its poll task.
    ///
    /// Returns the control handle, the [`NetstackChannel`] to hand to the
    /// `Tunn` loops, and the poll-task join handle (abort it to tear the
    /// netstack down).
    #[must_use]
    pub fn spawn(cfg: NetstackConfig) -> (Self, NetstackChannel, JoinHandle<()>) {
        let (from_tunnel_tx, from_tunnel_rx) = mpsc::channel(PACKET_CHANNEL_BOUND);
        let (to_tunnel_tx, to_tunnel_rx) = mpsc::channel(PACKET_CHANNEL_BOUND);
        let (cmd_tx, cmd_rx) = mpsc::channel(CMD_CHANNEL_BOUND);
        let wake = Arc::new(Notify::new());
        let mtu = cfg.mtu;

        let channel = NetstackChannel {
            from_tunnel: from_tunnel_tx,
            to_tunnel: Mutex::new(to_tunnel_rx),
            wake: Arc::clone(&wake),
            mtu,
        };

        let task = PollTask::new(cfg, from_tunnel_rx, to_tunnel_tx, cmd_rx, Arc::clone(&wake));
        let handle = tokio::spawn(task.run());

        (Self { cmd_tx, wake }, channel, handle)
    }

    /// Open an outbound TCP connection to `dst` across the overlay.
    ///
    /// The returned [`TcpConn`] is delivered as soon as the SYN is queued; the
    /// handshake completes asynchronously and a failure surfaces later as the
    /// connection closing.
    ///
    /// # Errors
    /// Returns [`EdgeError::Netstack`] if the poll task has stopped or the
    /// socket could not enter the connecting state.
    pub async fn connect(&self, dst: SocketAddr) -> Result<TcpConn, EdgeError> {
        let (reply, rx) = oneshot::channel();
        self.cmd_tx
            .send(NetstackCmd::Connect { dst, reply })
            .await
            .map_err(|_| EdgeError::Netstack("netstack task stopped".into()))?;
        self.wake.notify_one();
        match rx.await {
            Ok(result) => result,
            Err(_) => Err(EdgeError::Netstack("netstack dropped connect reply".into())),
        }
    }

    /// Listen for inbound TCP connections on `port`.
    ///
    /// Each accepted connection is delivered as a [`TcpConn`] on the returned
    /// channel.
    ///
    /// # Errors
    /// Returns [`EdgeError::Netstack`] if the poll task has stopped or the
    /// socket could not enter the listening state.
    pub async fn listen(&self, port: u16) -> Result<mpsc::Receiver<TcpConn>, EdgeError> {
        let (reply, rx) = oneshot::channel();
        self.cmd_tx
            .send(NetstackCmd::Listen { port, reply })
            .await
            .map_err(|_| EdgeError::Netstack("netstack task stopped".into()))?;
        self.wake.notify_one();
        match rx.await {
            Ok(result) => result,
            Err(_) => Err(EdgeError::Netstack("netstack dropped listen reply".into())),
        }
    }
}

// ---------------------------------------------------------------------------
// Poll task
// ---------------------------------------------------------------------------

/// Per-connection state held by the poll task.
struct Conn {
    /// Overlay → host bytes. `None` once the peer has closed and the rx buffer
    /// has drained (dropping the sender surfaces close to the host).
    to_host: Option<mpsc::Sender<Vec<u8>>>,
    /// Host → overlay bytes.
    from_host: mpsc::Receiver<Vec<u8>>,
    /// Bytes accepted from the host but not yet fully enqueued into the socket.
    pending: Option<Vec<u8>>,
    /// Whether we have already closed the socket's send half.
    send_closed: bool,
}

/// A listening socket + the channel accepted connections are delivered on.
struct Listener {
    port: u16,
    accept_tx: mpsc::Sender<TcpConn>,
}

/// Owns the smoltcp interface, sockets, device, and per-connection tables.
struct PollTask {
    iface: Interface,
    device: NetstackDevice,
    sockets: SocketSet<'static>,
    conns: HashMap<SocketHandle, Conn>,
    listeners: HashMap<SocketHandle, Listener>,
    from_tunnel: mpsc::Receiver<Vec<u8>>,
    to_tunnel: mpsc::Sender<Vec<u8>>,
    cmd_rx: mpsc::Receiver<NetstackCmd>,
    wake: Arc<Notify>,
    cfg: NetstackConfig,
    next_port: u16,
    start: std::time::Instant,
    ingress_open: bool,
    cmd_open: bool,
}

impl PollTask {
    fn new(
        cfg: NetstackConfig,
        from_tunnel: mpsc::Receiver<Vec<u8>>,
        to_tunnel: mpsc::Sender<Vec<u8>>,
        cmd_rx: mpsc::Receiver<NetstackCmd>,
        wake: Arc<Notify>,
    ) -> Self {
        let mtu = usize::try_from(cfg.mtu).unwrap_or(usize::MAX);
        let mut device = NetstackDevice {
            rx: VecDeque::new(),
            tx: VecDeque::new(),
            mtu,
        };
        let start = std::time::Instant::now();
        let iface = build_interface(&cfg, &mut device, start);

        Self {
            iface,
            device,
            sockets: SocketSet::new(Vec::new()),
            conns: HashMap::new(),
            listeners: HashMap::new(),
            from_tunnel,
            to_tunnel,
            cmd_rx,
            wake,
            next_port: EDGE_EPHEMERAL_PORT_START,
            start,
            cfg,
            ingress_open: true,
            cmd_open: true,
        }
    }

    /// Main loop: do a work pass, then wait for the next stimulus.
    async fn run(mut self) {
        loop {
            self.poll_once();
            let sleep = self.compute_sleep();
            tokio::select! {
                biased;
                () = self.wake.notified() => {}
                packet = self.from_tunnel.recv(), if self.ingress_open => match packet {
                    Some(pkt) => self.device.rx.push_back(pkt),
                    None => self.ingress_open = false,
                },
                cmd = self.cmd_rx.recv(), if self.cmd_open => match cmd {
                    Some(cmd) => self.handle_command(cmd),
                    None => self.cmd_open = false,
                },
                () = tokio::time::sleep(sleep) => {}
            }
        }
    }

    /// One work pass: ingest packets, poll smoltcp, service sockets, emit.
    fn poll_once(&mut self) {
        self.drain_ingress();
        let now = self.now();
        let _ = self.iface.poll(now, &mut self.device, &mut self.sockets);
        self.accept_listeners();
        self.pump_connections();
        self.drain_egress();
    }

    /// Move any queued inbound packets from the tunnel channel into the device.
    fn drain_ingress(&mut self) {
        loop {
            match self.from_tunnel.try_recv() {
                Ok(pkt) => self.device.rx.push_back(pkt),
                Err(TryRecvError::Empty) => break,
                Err(TryRecvError::Disconnected) => {
                    self.ingress_open = false;
                    break;
                }
            }
        }
    }

    /// Move egress packets from the device into the tunnel channel. On a full
    /// tunnel queue, re-queue and retry on the next (shortened) poll.
    fn drain_egress(&mut self) {
        while let Some(pkt) = self.device.tx.pop_front() {
            match self.to_tunnel.try_send(pkt) {
                Ok(()) => {}
                Err(TrySendError::Full(pkt)) => {
                    self.device.tx.push_front(pkt);
                    break;
                }
                Err(TrySendError::Closed(_)) => {
                    self.device.tx.clear();
                    break;
                }
            }
        }
    }

    /// Convert listening sockets that have received a connection into live
    /// connections and re-arm a fresh listener on the same port.
    fn accept_listeners(&mut self) {
        let ready: Vec<SocketHandle> = self
            .listeners
            .keys()
            .copied()
            .filter(|h| !self.sockets.get::<tcp::Socket>(*h).is_listening())
            .collect();

        for handle in ready {
            let Some(listener) = self.listeners.remove(&handle) else {
                continue;
            };
            let (conn, tcp_conn) = new_conn_pair();
            if listener.accept_tx.try_send(tcp_conn).is_ok() {
                self.conns.insert(handle, conn);
                if let Err(e) = self.arm_listener(listener.port, listener.accept_tx.clone()) {
                    tracing::warn!(error = %e, port = listener.port, "failed to re-arm listener");
                }
            } else {
                // Acceptor gone or saturated: drop the half-open socket.
                self.sockets.remove(handle);
            }
        }
    }

    /// Pump every tracked connection in both directions, then reap the closed.
    fn pump_connections(&mut self) {
        let handles: Vec<SocketHandle> = self.conns.keys().copied().collect();
        let mut closed: Vec<SocketHandle> = Vec::new();
        for handle in handles {
            let socket = self.sockets.get_mut::<tcp::Socket>(handle);
            let Some(conn) = self.conns.get_mut(&handle) else {
                continue;
            };
            pump_recv(socket, conn);
            pump_send(socket, conn);
            if socket.state() == tcp::State::Closed {
                closed.push(handle);
            }
        }
        for handle in closed {
            self.conns.remove(&handle);
            self.sockets.remove(handle);
        }
    }

    fn handle_command(&mut self, cmd: NetstackCmd) {
        match cmd {
            NetstackCmd::Connect { dst, reply } => {
                let _ = reply.send(self.do_connect(dst));
            }
            NetstackCmd::Listen { port, reply } => {
                let _ = reply.send(self.do_listen(port));
            }
        }
    }

    fn do_connect(&mut self, dst: SocketAddr) -> Result<TcpConn, EdgeError> {
        let mut socket = new_tcp_socket(&self.cfg);
        let local_port = self.alloc_ephemeral_port();
        socket
            .connect(self.iface.context(), dst, (self.cfg.local_ip, local_port))
            .map_err(|e| EdgeError::Netstack(format!("connect {dst}: {e:?}")))?;
        let handle = self.sockets.add(socket);
        let (conn, tcp_conn) = new_conn_pair();
        self.conns.insert(handle, conn);
        Ok(tcp_conn)
    }

    fn do_listen(&mut self, port: u16) -> Result<mpsc::Receiver<TcpConn>, EdgeError> {
        let (accept_tx, accept_rx) = mpsc::channel(CONN_CHANNEL_BOUND);
        self.arm_listener(port, accept_tx)?;
        Ok(accept_rx)
    }

    fn arm_listener(
        &mut self,
        port: u16,
        accept_tx: mpsc::Sender<TcpConn>,
    ) -> Result<(), EdgeError> {
        let mut socket = new_tcp_socket(&self.cfg);
        socket
            .listen(port)
            .map_err(|e| EdgeError::Netstack(format!("listen {port}: {e:?}")))?;
        let handle = self.sockets.add(socket);
        self.listeners.insert(handle, Listener { port, accept_tx });
        Ok(())
    }

    /// Allocate the next ephemeral source port, wrapping within the dynamic
    /// range. Collisions with a still-open port are statistically negligible
    /// across a 16k range and are not tracked.
    fn alloc_ephemeral_port(&mut self) -> u16 {
        let port = self.next_port;
        self.next_port = next_ephemeral(self.next_port);
        port
    }

    fn now(&self) -> Instant {
        smoltcp_now(self.start)
    }

    /// How long to idle before the next forced poll: smoltcp's own request,
    /// capped at [`MAX_POLL_INTERVAL`], shortened while egress is back-logged.
    fn compute_sleep(&mut self) -> Duration {
        let now = self.now();
        let delay = self
            .iface
            .poll_delay(now, &self.sockets)
            .map_or(MAX_POLL_INTERVAL, |d| {
                Duration::from(d).min(MAX_POLL_INTERVAL)
            });
        if self.device.tx.is_empty() {
            delay
        } else {
            delay.min(EGRESS_RETRY_INTERVAL)
        }
    }
}

// ---------------------------------------------------------------------------
// Free helpers (kept small so the poll-task methods stay under clippy limits)
// ---------------------------------------------------------------------------

/// Build the smoltcp interface for a `Medium::Ip` device: assign the overlay
/// address and install unspecified default routes (no neighbor discovery on
/// `Medium::Ip`, so the gateway address is never resolved — packets leave via
/// `WireGuard` regardless).
fn build_interface(
    cfg: &NetstackConfig,
    device: &mut NetstackDevice,
    start: std::time::Instant,
) -> Interface {
    let mut config = Config::new(HardwareAddress::Ip);
    config.random_seed = rand::random();
    let mut iface = Interface::new(config, device, smoltcp_now(start));
    iface.update_ip_addrs(|addrs| {
        let _ = addrs.push(IpCidr::new(cfg.local_ip.into(), cfg.prefix_len));
    });
    let _ = iface
        .routes_mut()
        .add_default_ipv4_route(Ipv4Address::UNSPECIFIED);
    let _ = iface
        .routes_mut()
        .add_default_ipv6_route(Ipv6Address::UNSPECIFIED);
    iface
}

/// A monotonic smoltcp [`Instant`] derived from a fixed start point.
fn smoltcp_now(start: std::time::Instant) -> Instant {
    let micros = start.elapsed().as_micros();
    Instant::from_micros(i64::try_from(micros).unwrap_or(i64::MAX))
}

/// Build a fresh TCP socket with the configured buffer sizes and timeout.
fn new_tcp_socket(cfg: &NetstackConfig) -> tcp::Socket<'static> {
    let rx = tcp::SocketBuffer::new(vec![0u8; cfg.tcp_buffer_bytes]);
    let tx = tcp::SocketBuffer::new(vec![0u8; cfg.tcp_buffer_bytes]);
    let mut socket = tcp::Socket::new(rx, tx);
    socket.set_timeout(Some(cfg.tcp_timeout.into()));
    socket
}

/// Advance an ephemeral port, wrapping to [`EDGE_EPHEMERAL_PORT_START`] after
/// the top of the dynamic range.
fn next_ephemeral(current: u16) -> u16 {
    if current == u16::MAX {
        EDGE_EPHEMERAL_PORT_START
    } else {
        current + 1
    }
}

/// Create the poll-task-side [`Conn`] and the host-side [`TcpConn`] wired to
/// each other.
fn new_conn_pair() -> (Conn, TcpConn) {
    let (overlay_to_host, host_rx) = mpsc::channel(CONN_CHANNEL_BOUND);
    let (host_tx, host_to_overlay) = mpsc::channel(CONN_CHANNEL_BOUND);
    let conn = Conn {
        to_host: Some(overlay_to_host),
        from_host: host_to_overlay,
        pending: None,
        send_closed: false,
    };
    let tcp_conn = TcpConn {
        tx: host_tx,
        rx: host_rx,
    };
    (conn, tcp_conn)
}

/// Drain a socket's receive buffer into the host channel, honoring the
/// channel's capacity as flow control, and signal peer-close when the receive
/// half is finished.
fn pump_recv(socket: &mut tcp::Socket, conn: &mut Conn) {
    let mut host_gone = false;
    if let Some(tx) = conn.to_host.as_ref() {
        while socket.can_recv() {
            match tx.try_reserve() {
                Ok(permit) => {
                    let mut buf = vec![0u8; RECV_CHUNK];
                    match socket.recv_slice(&mut buf) {
                        Ok(n) if n > 0 => {
                            buf.truncate(n);
                            permit.send(buf);
                        }
                        _ => break,
                    }
                }
                // Host not reading fast enough: leave data in the socket so
                // smoltcp shrinks the advertised window (real back-pressure).
                Err(TrySendError::Full(())) => break,
                // Host dropped the receiver: reset the connection.
                Err(TrySendError::Closed(())) => {
                    host_gone = true;
                    break;
                }
            }
        }
    }
    if host_gone {
        socket.abort();
        conn.to_host = None;
    }
    // `may_recv()` is also false during the opening handshake (SYN-SENT /
    // SYN-RECEIVED / LISTEN), so guard against tearing the host channel down
    // before the connection establishes. Past those states, `!may_recv()`
    // means the remote closed and the rx buffer has drained — drop the sender
    // so the host observes end-of-stream.
    let opening = matches!(
        socket.state(),
        tcp::State::Listen | tcp::State::SynSent | tcp::State::SynReceived
    );
    if !opening && !socket.may_recv() {
        conn.to_host = None;
    }
}

/// Push host bytes into a socket's send buffer, stashing any tail that did not
/// fit and closing the send half once the host drops its sender.
fn pump_send(socket: &mut tcp::Socket, conn: &mut Conn) {
    while socket.can_send() {
        if let Some(pending) = conn.pending.take() {
            match socket.send_slice(&pending) {
                Ok(n) if n < pending.len() => {
                    conn.pending = Some(pending[n..].to_vec());
                    break;
                }
                Ok(_) => continue,
                Err(_) => {
                    conn.pending = Some(pending);
                    break;
                }
            }
        }
        if conn.send_closed {
            break;
        }
        match conn.from_host.try_recv() {
            Ok(data) => match socket.send_slice(&data) {
                Ok(n) if n < data.len() => {
                    conn.pending = Some(data[n..].to_vec());
                    break;
                }
                Ok(_) => {}
                Err(_) => {
                    conn.pending = Some(data);
                    break;
                }
            },
            Err(TryRecvError::Empty) => break,
            Err(TryRecvError::Disconnected) => {
                socket.close();
                conn.send_closed = true;
                break;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::Ipv4Addr;

    fn test_cfg(ip: Ipv4Addr) -> NetstackConfig {
        NetstackConfig {
            local_ip: IpAddr::V4(ip),
            prefix_len: 24,
            mtu: 1420,
            tcp_buffer_bytes: 64 * 1024,
            tcp_timeout: Duration::from_secs(10),
        }
    }

    /// Copy every egress packet of `src` into `dst`'s ingress until either
    /// channel closes — cross-connecting two netstacks into a loopback fabric.
    async fn wire(src: Arc<NetstackChannel>, dst: Arc<NetstackChannel>) {
        let mut buf = vec![0u8; 65536];
        loop {
            match src.recv_packet(&mut buf).await {
                Ok(n) if n > 0 => {
                    if dst.send_packet(&buf[..n]).await.is_err() {
                        break;
                    }
                }
                Ok(_) => {}
                Err(_) => break,
            }
        }
    }

    /// Full two-stack round trip: a byte payload travels client → smoltcp →
    /// (loopback fabric) → smoltcp server → echo → back, entirely in-process.
    #[tokio::test]
    async fn netstack_echo_loopback() {
        let client_ip = Ipv4Addr::new(10, 0, 0, 1);
        let server_ip = Ipv4Addr::new(10, 0, 0, 2);
        let port = 4000u16;

        let (client, client_chan, client_task) = Netstack::spawn(test_cfg(client_ip));
        let (server, server_chan, server_task) = Netstack::spawn(test_cfg(server_ip));
        let client_chan = Arc::new(client_chan);
        let server_chan = Arc::new(server_chan);

        let w1 = tokio::spawn(wire(Arc::clone(&client_chan), Arc::clone(&server_chan)));
        let w2 = tokio::spawn(wire(Arc::clone(&server_chan), Arc::clone(&client_chan)));

        // Echo server: reflect every chunk back on the same connection.
        let mut accept_rx = server.listen(port).await.expect("listen");
        let echo = tokio::spawn(async move {
            if let Some(conn) = accept_rx.recv().await {
                let TcpConn { tx, mut rx } = conn;
                while let Some(chunk) = rx.recv().await {
                    if tx.send(chunk).await.is_err() {
                        break;
                    }
                }
            }
        });

        // Client connect + send.
        let dst = SocketAddr::new(IpAddr::V4(server_ip), port);
        let conn = tokio::time::timeout(Duration::from_secs(10), client.connect(dst))
            .await
            .expect("connect timed out")
            .expect("connect failed");
        let TcpConn { tx, mut rx } = conn;
        let payload = b"hello smoltcp userspace round trip".to_vec();
        tx.send(payload.clone()).await.expect("send payload");

        // Collect the echo (TCP may chunk; accumulate until we have it all).
        let mut got = Vec::new();
        while got.len() < payload.len() {
            match tokio::time::timeout(Duration::from_secs(10), rx.recv()).await {
                Ok(Some(chunk)) => got.extend_from_slice(&chunk),
                Ok(None) => break,
                Err(e) => panic!("timed out waiting for echo (have {} bytes): {e}", got.len()),
            }
        }
        assert_eq!(got, payload);

        drop(tx);
        for task in [w1, w2, echo, client_task, server_task] {
            task.abort();
        }
    }

    /// Generate an x25519 keypair as raw 32-byte arrays, deriving the public
    /// key exactly the way [`crate::tunn_loop::build_tunn`] consumes it
    /// (boringtun's re-exported `x25519`). Raw bytes are generated with `rand`
    /// (no dependency on x25519-dalek's own rng feature / `rand_core` version).
    fn wg_keypair() -> ([u8; 32], [u8; 32]) {
        use rand::RngCore;
        let mut priv_bytes = [0u8; 32];
        rand::rng().fill_bytes(&mut priv_bytes);
        let public =
            boringtun::x25519::PublicKey::from(&boringtun::x25519::StaticSecret::from(priv_bytes));
        (priv_bytes, public.to_bytes())
    }

    /// Assemble one [`crate::tunn_loop::PeerState`] for a peer identified by
    /// `peer_pub`, returning it alongside a clone of its shared
    /// `last_handshake_sec` clock so the test can assert a real Noise handshake
    /// landed. Mirrors the field-for-field construction in
    /// `edge::client::Client::build_peers`.
    fn build_peer_state(
        our_priv: &[u8; 32],
        peer_pub: &[u8; 32],
        endpoint: Option<SocketAddr>,
        keepalive: Option<u16>,
        allowed_cidr: &str,
    ) -> (
        crate::tunn_loop::PeerState,
        Arc<std::sync::atomic::AtomicU64>,
    ) {
        use crate::tunn_loop::{build_tunn, PeerState};
        let last_hs = Arc::new(std::sync::atomic::AtomicU64::new(0));
        let state = PeerState {
            tunn: Arc::new(Mutex::new(build_tunn(our_priv, peer_pub, None, keepalive))),
            endpoint: Arc::new(parking_lot::RwLock::new(endpoint)),
            last_handshake_sec: Arc::clone(&last_hs),
            allowed_ips: Arc::new(vec![allowed_cidr
                .parse::<ipnet::IpNet>()
                .expect("valid overlay cidr")]),
            persistent_keepalive: keepalive,
        };
        (state, last_hs)
    }

    /// Headline correctness proof of the encrypted edge data plane: two real
    /// boringtun [`Tunn`](boringtun::noise::Tunn) endpoints complete a live
    /// `WireGuard` Noise handshake over loopback UDP, then carry a TCP round
    /// trip through the full pipeline (UDP ↔ `Tunn` encapsulate/decapsulate ↔
    /// [`NetstackChannel`] ↔ smoltcp ↔ [`TcpConn`]).
    ///
    /// Unlike [`netstack_echo_loopback`] — which cross-wires the two stacks in
    /// **cleartext** via `wire()` and never touches boringtun — this drives the
    /// production [`crate::tunn_loop`] machinery (`build_tunn` / `PeerState` /
    /// `TunnDriver`) and asserts (step 7) that a handshake was recorded, so it
    /// cannot pass through a bypass.
    ///
    /// Fully unprivileged: no TUN, no root — boringtun's Noise handshake is
    /// pure userspace over two `127.0.0.1:0` UDP sockets.
    ///
    /// Determinism / no flakiness: [`Netstack::connect`] returns as soon as the
    /// SYN is queued. The WG handshake is driven by that first egress SYN —
    /// `Tunn::encapsulate` with no session queues the packet and returns a
    /// handshake-init (boringtun 0.7 `noise/mod.rs:264-267`) — and completes in
    /// a few loopback RTTs. Even if the first SYN races session establishment,
    /// smoltcp's SYN retransmit plus the now-live session still complete well
    /// inside the generous 20 s timeouts, so no fixed sleep is needed.
    #[tokio::test]
    async fn tunn_pipeline_encrypted_tcp_roundtrip() {
        use crate::tunn_loop::{PeerState, TunnDriver};
        use dashmap::DashMap;
        use std::sync::atomic::Ordering;
        use tokio::net::UdpSocket;

        // 1. Two loopback UDP sockets: the ciphertext transport for each end.
        let udp_edge = Arc::new(UdpSocket::bind("127.0.0.1:0").await.expect("bind edge udp"));
        let udp_node = Arc::new(UdpSocket::bind("127.0.0.1:0").await.expect("bind node udp"));
        let addr_node = udp_node.local_addr().expect("node udp addr");

        // 2. Two x25519 keypairs (edge = A, node = B).
        let (edge_priv, edge_pub) = wg_keypair();
        let (node_priv, node_pub) = wg_keypair();

        // 3. Node side (B, 10.88.0.2): its smoltcp app netstack + a Tunn peer
        //    keyed by the REMOTE (edge) public key — the ingress loop looks
        //    peers up by decapsulated pubkey. Endpoint starts `None`: the node
        //    learns the edge's source address from the handshake (roaming
        //    edge), exactly as `ingress_loop` writes `endpoint` on decap.
        let node_ip = Ipv4Addr::new(10, 88, 0, 2);
        let (netstack_node, chan_node, task_node) = Netstack::spawn(test_cfg(node_ip));
        let chan_node = Arc::new(chan_node);
        let (node_peer, node_last_hs) =
            build_peer_state(&node_priv, &edge_pub, None, None, "10.88.0.1/32");
        let peers_node: Arc<DashMap<[u8; 32], PeerState>> = Arc::new(DashMap::new());
        peers_node.insert(edge_pub, node_peer);
        let driver_node = TunnDriver::spawn(udp_node, chan_node, peers_node);

        // 4. Edge side (A, 10.88.0.1): a Tunn peer keyed by the node's public
        //    key, with a known endpoint + keepalive (mirroring production
        //    `build_peers`). The edge initiates the handshake because its first
        //    egress SYN has an endpoint to encapsulate toward.
        let edge_ip = Ipv4Addr::new(10, 88, 0, 1);
        let (netstack_edge, chan_edge, task_edge) = Netstack::spawn(test_cfg(edge_ip));
        let chan_edge = Arc::new(chan_edge);
        let (edge_peer, edge_last_hs) = build_peer_state(
            &edge_priv,
            &node_pub,
            Some(addr_node),
            Some(25),
            "10.88.0.2/32",
        );
        let peers_edge: Arc<DashMap<[u8; 32], PeerState>> = Arc::new(DashMap::new());
        peers_edge.insert(node_pub, edge_peer);
        let driver_edge = TunnDriver::spawn(udp_edge, chan_edge, peers_edge);

        // 5. Node app: listen + echo every chunk (mirrors netstack_echo_loopback).
        let port = 4100u16;
        let mut accept_rx = netstack_node.listen(port).await.expect("listen");
        let echo = tokio::spawn(async move {
            if let Some(conn) = accept_rx.recv().await {
                let TcpConn { tx, mut rx } = conn;
                while let Some(chunk) = rx.recv().await {
                    if tx.send(chunk).await.is_err() {
                        break;
                    }
                }
            }
        });

        // 6. Edge app: connect across the encrypted overlay, send, collect echo.
        let dst = SocketAddr::new(IpAddr::V4(node_ip), port);
        let conn = tokio::time::timeout(Duration::from_secs(20), netstack_edge.connect(dst))
            .await
            .expect("connect timed out")
            .expect("connect failed");
        let TcpConn { tx, mut rx } = conn;
        let payload = b"boringtun-encrypted round trip across the edge data plane".to_vec();
        tx.send(payload.clone()).await.expect("send payload");

        let mut got = Vec::new();
        while got.len() < payload.len() {
            match tokio::time::timeout(Duration::from_secs(20), rx.recv()).await {
                Ok(Some(chunk)) => got.extend_from_slice(&chunk),
                Ok(None) => break,
                Err(e) => panic!(
                    "timed out waiting for encrypted echo (have {} of {} bytes): {e}",
                    got.len(),
                    payload.len()
                ),
            }
        }
        assert_eq!(
            got, payload,
            "echoed payload mismatch through the encrypted pipeline"
        );

        // 7. Prove a real Noise handshake happened. `last_handshake_sec` is only
        //    bumped when a data packet is successfully DECAPSULATED (ingress
        //    `WriteToTunnelV4`), so a non-zero value is impossible unless
        //    boringtun established a session and decrypted real traffic — the
        //    assertion that separates this from the cleartext `wire()` harness.
        //    Because the payload round-tripped, both clocks are in fact > 0.
        let edge_hs = edge_last_hs.load(Ordering::Relaxed);
        let node_hs = node_last_hs.load(Ordering::Relaxed);
        assert!(
            edge_hs > 0 || node_hs > 0,
            "no WireGuard handshake was recorded (edge={edge_hs}, node={node_hs}); \
             the payload must have bypassed boringtun"
        );

        // 8. Teardown: stop the Tunn drivers, the echo task, and both netstacks.
        drop(tx);
        driver_edge.abort_all();
        driver_node.abort_all();
        echo.abort();
        task_edge.abort();
        task_node.abort();
    }

    /// `send_packet` drops (does not error or back-pressure) when the ingress
    /// queue is full.
    #[tokio::test]
    async fn send_packet_drops_when_full() {
        let (from_tunnel, mut from_tunnel_rx) = mpsc::channel::<Vec<u8>>(1);
        let (_to_tunnel_tx, to_tunnel_rx) = mpsc::channel::<Vec<u8>>(1);
        let chan = NetstackChannel {
            from_tunnel,
            to_tunnel: Mutex::new(to_tunnel_rx),
            wake: Arc::new(Notify::new()),
            mtu: 1420,
        };

        chan.send_packet(&[1, 2, 3]).await.expect("first send");
        // Queue (depth 1) is now full: the second packet is dropped, not errored.
        chan.send_packet(&[4, 5, 6])
            .await
            .expect("second send drops");

        assert_eq!(from_tunnel_rx.recv().await, Some(vec![1, 2, 3]));
        assert!(
            from_tunnel_rx.try_recv().is_err(),
            "second packet must be dropped"
        );
    }

    /// `recv_packet` returns an error (ending the egress loop) once the poll
    /// task's egress sender is gone.
    #[tokio::test]
    async fn recv_packet_errors_on_closed() {
        let (from_tunnel, _from_tunnel_rx) = mpsc::channel::<Vec<u8>>(1);
        let (to_tunnel_tx, to_tunnel_rx) = mpsc::channel::<Vec<u8>>(1);
        drop(to_tunnel_tx);
        let chan = NetstackChannel {
            from_tunnel,
            to_tunnel: Mutex::new(to_tunnel_rx),
            wake: Arc::new(Notify::new()),
            mtu: 1420,
        };

        let mut buf = [0u8; 64];
        assert!(chan.recv_packet(&mut buf).await.is_err());
    }

    #[test]
    fn ephemeral_port_wraps() {
        assert_eq!(
            next_ephemeral(EDGE_EPHEMERAL_PORT_START),
            EDGE_EPHEMERAL_PORT_START + 1
        );
        assert_eq!(next_ephemeral(u16::MAX), EDGE_EPHEMERAL_PORT_START);
        // A full walk from the top wraps back into the dynamic range.
        let mut p = u16::MAX;
        p = next_ephemeral(p);
        assert_eq!(p, EDGE_EPHEMERAL_PORT_START);
    }

    #[test]
    fn config_plumbing() {
        let cfg = test_cfg(Ipv4Addr::new(10, 0, 0, 5));
        assert_eq!(cfg.local_ip, IpAddr::V4(Ipv4Addr::new(10, 0, 0, 5)));
        assert_eq!(cfg.prefix_len, 24);
        assert_eq!(cfg.mtu, 1420);
        assert_eq!(cfg.tcp_buffer_bytes, 64 * 1024);
        let socket = new_tcp_socket(&cfg);
        assert_eq!(socket.recv_capacity(), 64 * 1024);
        assert_eq!(socket.send_capacity(), 64 * 1024);
    }
}