shoes 0.3.1

//! Direct TCP Stack Manager for smoltcp integration.
//!
//! This module manages the smoltcp TCP/IP stack in a dedicated OS thread,
//! using `select()` on the TUN fd for event-driven I/O instead of polling.

use std::{
    collections::HashMap,
    io, mem,
    net::SocketAddr,
    ops::{Deref, DerefMut},
    os::fd::{AsRawFd, OwnedFd, RawFd},
    panic::{self, AssertUnwindSafe},
    sync::{
        Arc, LazyLock, Mutex,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle, Thread},
    time::Duration,
};

use bytes::BytesMut;

use log::{debug, error, info, trace, warn};
use smoltcp::{
    iface::{Config as InterfaceConfig, Interface, SocketHandle, SocketSet},
    phy::{Device, DeviceCapabilities, Medium, RxToken, TxToken, wait as phy_wait},
    socket::tcp::{
        CongestionControl, Socket as TcpSocket, SocketBuffer as TcpSocketBuffer, State as TcpState,
    },
    time::{Duration as SmolDuration, Instant as SmolInstant},
    wire::{
        HardwareAddress, IpAddress, IpCidr, IpProtocol, Ipv4Address, Ipv4Packet, Ipv6Address,
        Ipv6Packet, TcpPacket,
    },
};
use tokio::sync::mpsc::{self, Receiver, Sender};

use super::TunServerConfig;
use super::tcp_conn::{TcpConnection, TcpConnectionControl, TcpSocketState};

pub type PacketBuffer = Vec<u8>;
pub const PACKET_QUEUE_CAPACITY: usize = 64;

/// Maximum number of buffers cached globally.
/// Each buffer has capacity ~65536, so 64 * 65536 = 4MB max.
const BUFFER_POOL_MAX_SIZE: usize = 64;

static BUFFER_POOL: LazyLock<Mutex<Vec<BytesMut>>> = LazyLock::new(|| Mutex::new(Vec::new()));

/// Pooled buffer that returns to pool on drop instead of deallocating.
pub struct PooledBuffer {
    buffer: BytesMut,
}

impl Drop for PooledBuffer {
    fn drop(&mut self) {
        if let Ok(mut pool) = BUFFER_POOL.lock()
            && pool.len() < BUFFER_POOL_MAX_SIZE
        {
            let empty = BytesMut::new();
            let mut buffer = mem::replace(&mut self.buffer, empty);
            buffer.clear();
            pool.push(buffer);
        }
    }
}

impl PooledBuffer {
    /// Get a buffer from the pool or create a new one.
    pub fn with_capacity(cap: usize) -> Self {
        if let Ok(mut pool) = BUFFER_POOL.lock()
            && let Some(mut buffer) = pool.pop()
        {
            buffer.reserve(cap);
            return Self { buffer };
        }
        Self {
            buffer: BytesMut::with_capacity(cap),
        }
    }
}

impl Deref for PooledBuffer {
    type Target = BytesMut;

    fn deref(&self) -> &Self::Target {
        &self.buffer
    }
}

impl DerefMut for PooledBuffer {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.buffer
    }
}

/// Tracks socket info including addresses for proper cleanup.
struct SocketInfo {
    control: Arc<TcpConnectionControl>,
    src_addr: SocketAddr,
    dst_addr: SocketAddr,
}

impl Drop for SocketInfo {
    fn drop(&mut self) {
        self.control.set_closed();
    }
}

/// Information about a new TCP connection from the stack.
pub struct NewTcpConnection {
    pub connection: TcpConnection,
    pub remote_addr: SocketAddr,
}

/// Shared state for communication between main thread and stack thread.
struct SharedState {
    /// Channel for UDP responses to write to TUN
    udp_response_rx: Option<Receiver<PacketBuffer>>,
    /// Channel for notifying tokio about new TCP connections
    new_conn_tx: Option<mpsc::UnboundedSender<NewTcpConnection>>,
}

/// Direct TCP Stack Manager.
///
/// Manages the smoltcp interface with direct fd access for efficient I/O.
pub struct TcpStackDirect {
    /// Handle to the stack thread
    thread_handle: Option<JoinHandle<()>>,
    /// Thread handle for waking the stack thread
    stack_thread: Thread,
    /// Flag to signal thread shutdown
    running: Arc<AtomicBool>,
    /// Receiver for UDP packets (filtered from TUN by the stack thread)
    udp_rx: Option<Receiver<PacketBuffer>>,
    /// Shared state with the stack thread
    shared_state: Arc<Mutex<SharedState>>,
    /// TUN file descriptor (owned, will be closed on drop)
    _tun_fd: OwnedFd,
}

impl Drop for TcpStackDirect {
    fn drop(&mut self) {
        // Signal thread to stop
        self.running.store(false, Ordering::Relaxed);
        self.stack_thread.unpark();

        // Wait for thread to finish
        if let Some(handle) = self.thread_handle.take() {
            let _ = handle.join();
        }
    }
}

impl TcpStackDirect {
    /// Create a new direct TCP stack.
    ///
    /// # Arguments
    /// * `fd` - Raw file descriptor for the TUN device
    /// * `mtu` - Maximum transmission unit
    ///
    /// This spawns a dedicated OS thread for running the smoltcp interface.
    /// The thread uses `select()` on the fd for efficient event-driven I/O.
    #[cfg(test)]
    pub fn new(tun_fd: OwnedFd, mtu: usize) -> Self {
        Self::with_config(tun_fd, TunServerConfig::new().mtu(mtu as u16))
    }

    pub fn with_config(tun_fd: OwnedFd, config: TunServerConfig) -> Self {
        let fd = tun_fd.as_raw_fd();
        let (udp_tx, udp_rx) = mpsc::channel(PACKET_QUEUE_CAPACITY);

        let running = Arc::new(AtomicBool::new(true));
        let shared_state = Arc::new(Mutex::new(SharedState {
            udp_response_rx: None,
            new_conn_tx: None,
        }));

        let thread_handle = {
            let running = running.clone();
            let shared_state = shared_state.clone();

            thread::Builder::new()
                .name("shoes-smoltcp-direct".to_owned())
                .spawn(move || {
                    let result = panic::catch_unwind(AssertUnwindSafe(|| {
                        run_direct_stack_thread(fd, config, udp_tx, running.clone(), shared_state);
                    }));

                    match result {
                        Ok(()) => {
                            info!("smoltcp direct stack thread exited normally");
                        }
                        Err(panic_info) => {
                            let msg = if let Some(s) = panic_info.downcast_ref::<&str>() {
                                s.to_string()
                            } else if let Some(s) = panic_info.downcast_ref::<String>() {
                                s.clone()
                            } else {
                                "Unknown panic".to_string()
                            };
                            error!("smoltcp direct stack thread PANICKED: {}", msg);
                        }
                    }

                    running.store(false, Ordering::Relaxed);
                })
                .expect("failed to spawn smoltcp direct thread")
        };

        let stack_thread = thread_handle.thread().clone();

        Self {
            thread_handle: Some(thread_handle),
            stack_thread,
            running,
            udp_rx: Some(udp_rx),
            shared_state,
            _tun_fd: tun_fd,
        }
    }

    /// Take the receiver for UDP packets (filtered from TUN by the stack).
    pub fn take_udp_rx(&mut self) -> Option<Receiver<PacketBuffer>> {
        self.udp_rx.take()
    }

    /// Set the channel for UDP responses to write back to TUN.
    pub fn set_udp_response_tx(&mut self, rx: Receiver<PacketBuffer>) {
        if let Ok(mut state) = self.shared_state.lock() {
            state.udp_response_rx = Some(rx);
        }
        self.stack_thread.unpark();
    }

    /// Set the channel for notifying about new TCP connections.
    pub fn set_new_conn_tx(&mut self, tx: mpsc::UnboundedSender<NewTcpConnection>) {
        if let Ok(mut state) = self.shared_state.lock() {
            state.new_conn_tx = Some(tx);
        }
        self.stack_thread.unpark();
    }

    /// Check if the stack thread is still running.
    pub fn is_running(&self) -> bool {
        self.running.load(Ordering::Relaxed)
    }
}

/// Direct TUN device that reads/writes directly to fd.
struct DirectDevice {
    fd: RawFd,
    mtu: usize,
    pending_rx: Option<PooledBuffer>,
    packet_information: bool,
}

impl DirectDevice {
    fn new(fd: RawFd, mtu: usize, packet_information: bool) -> Self {
        Self {
            fd,
            mtu,
            pending_rx: None,
            packet_information,
        }
    }

    /// Try to read a packet (non-blocking) using pooled buffer.
    /// Returns:
    /// - Ok(Some(packet)) if a packet was read
    /// - Ok(None) if no packet was available (WouldBlock)
    /// - Err(e) if a fatal error occurred (including EOF)
    fn try_recv(&mut self) -> io::Result<Option<PooledBuffer>> {
        if let Some(pkt) = self.pending_rx.take() {
            return Ok(Some(pkt));
        }

        // Get a buffer from the pool
        let mut buffer = PooledBuffer::with_capacity(self.mtu + 4);
        buffer.resize(self.mtu + 4, 0);

        match read_nonblocking(self.fd, &mut buffer) {
            Ok(n) if n > 0 => {
                buffer.truncate(n);
                if self.packet_information {
                    if n < 5 || buffer[..4] != packet_header(buffer[4])? {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "invalid TUN packet information",
                        ));
                    }
                    let _ = buffer.split_to(4);
                }
                Ok(Some(buffer))
            }
            Ok(_) => {
                // n == 0 means EOF
                Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "TUN device closed (EOF)",
                ))
            }
            Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                // Buffer is returned to pool when dropped
                Ok(None)
            }
            Err(e) => {
                // Fatal error
                Err(e)
            }
        }
    }

    /// Store a packet for later processing by smoltcp.
    fn store_packet(&mut self, pkt: PooledBuffer) {
        self.pending_rx = Some(pkt);
    }

    /// Write a packet to TUN.
    fn write_packet(&self, data: &[u8]) -> io::Result<()> {
        write_packet(self.fd, data, self.packet_information)
    }
}

impl Device for DirectDevice {
    type RxToken<'a> = DirectRxToken;
    type TxToken<'a> = DirectTxToken;

    fn receive(
        &mut self,
        _timestamp: SmolInstant,
    ) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
        if let Some(buffer) = self.pending_rx.take() {
            let rx = DirectRxToken { buffer };
            let tx = DirectTxToken {
                fd: self.fd,
                packet_information: self.packet_information,
            };
            Some((rx, tx))
        } else {
            None
        }
    }

    fn transmit(&mut self, _timestamp: SmolInstant) -> Option<Self::TxToken<'_>> {
        Some(DirectTxToken {
            fd: self.fd,
            packet_information: self.packet_information,
        })
    }

    fn capabilities(&self) -> DeviceCapabilities {
        let mut caps = DeviceCapabilities::default();
        caps.medium = Medium::Ip;
        caps.max_transmission_unit = self.mtu;
        caps.checksum.ipv4 = smoltcp::phy::Checksum::Tx;
        caps.checksum.tcp = smoltcp::phy::Checksum::Tx;
        caps.checksum.udp = smoltcp::phy::Checksum::Tx;
        caps.checksum.icmpv4 = smoltcp::phy::Checksum::Tx;
        caps.checksum.icmpv6 = smoltcp::phy::Checksum::Tx;
        caps
    }
}

struct DirectRxToken {
    buffer: PooledBuffer,
}

impl RxToken for DirectRxToken {
    fn consume<R, F>(self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(&self.buffer)
        // buffer is returned to pool when dropped
    }
}

struct DirectTxToken {
    fd: RawFd,
    packet_information: bool,
}

impl TxToken for DirectTxToken {
    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);

        if let Err(e) = write_packet(self.fd, &buffer, self.packet_information) {
            warn!("Failed to write to TUN: {}", e);
        }

        result
    }
}

const MAX_PACKET_BATCH: usize = 64; // Process more packets per poll iteration

/// Run the direct smoltcp stack thread.
fn run_direct_stack_thread(
    fd: RawFd,
    config: TunServerConfig,
    udp_tx: Sender<PacketBuffer>,
    running: Arc<AtomicBool>,
    shared_state: Arc<Mutex<SharedState>>,
) {
    info!("smoltcp direct stack thread initializing...");

    // Sets fd to non-blocking mode once at startup for performance.
    if let Err(e) = set_nonblocking(fd) {
        error!("Failed to set TUN fd to non-blocking: {}", e);
        return;
    }

    let limits = &config.resource_limits;
    let mut device = DirectDevice::new(fd, config.mtu as usize, config.packet_information);

    let mut iface_config = InterfaceConfig::new(HardwareAddress::Ip);
    iface_config.random_seed = rand::random();

    let mut iface = Interface::new(iface_config, &mut device, SmolInstant::now());

    iface.update_ip_addrs(|addrs| {
        if let Err(e) = addrs.push(IpCidr::new(IpAddress::v4(0, 0, 0, 1), 0)) {
            warn!("Failed to add IPv4 address: {:?}", e);
        }
        if let Err(e) = addrs.push(IpCidr::new(IpAddress::v6(0, 0, 0, 0, 0, 0, 0, 1), 0)) {
            warn!("Failed to add IPv6 address: {:?}", e);
        }
    });

    if let Err(e) = iface
        .routes_mut()
        .add_default_ipv4_route(Ipv4Address::new(0, 0, 0, 1))
    {
        warn!("Failed to add IPv4 route: {:?}", e);
    }
    if let Err(e) = iface
        .routes_mut()
        .add_default_ipv6_route(Ipv6Address::new(0, 0, 0, 0, 0, 0, 0, 1))
    {
        warn!("Failed to add IPv6 route: {:?}", e);
    }

    iface.set_any_ip(true);

    let mut socket_set = SocketSet::new(vec![]);
    let mut sockets: HashMap<SocketHandle, SocketInfo> = HashMap::new();
    let mut active_connections: std::collections::HashSet<(SocketAddr, SocketAddr)> =
        std::collections::HashSet::new();

    let mut poll_count: u64 = 0;
    let mut last_log_time = std::time::Instant::now();

    let stack_thread = thread::current();

    let mut phy_wait_error_count: u32 = 0;
    const MAX_PHY_WAIT_ERRORS: u32 = 10;

    info!("smoltcp direct stack thread started, entering main loop");

    while running.load(Ordering::Relaxed) {
        // Checks for UDP responses to write to TUN.
        if let Ok(mut state) = shared_state.try_lock()
            && let Some(ref mut udp_rx) = state.udp_response_rx
        {
            for _ in 0..MAX_PACKET_BATCH {
                let Ok(pkt) = udp_rx.try_recv() else { break };
                if let Err(e) = device.write_packet(&pkt) {
                    warn!("Failed to write UDP response to TUN: {}", e);
                }
            }
        }

        // Reads packets from TUN and filters by protocol (batch processing).
        let mut tcp_packets: Vec<PooledBuffer> = Vec::new();
        let mut packets_read = 0;

        while packets_read < MAX_PACKET_BATCH {
            let pkt = match device.try_recv() {
                Ok(Some(p)) => p,
                Ok(None) => break,
                Err(e) => {
                    // Critical error reading from TUN (EOF or EIO)
                    error!("TUN device read failed: {}. Stack thread stopping.", e);
                    running.store(false, Ordering::Relaxed);
                    break;
                }
            };
            packets_read += 1;

            if should_filter_packet(&pkt) {
                trace!("Filtered packet, len={}", pkt.len());
                continue;
            }

            if let Some(protocol) = get_ip_protocol(&pkt) {
                trace!(
                    "Received packet: protocol={:?}, len={}",
                    protocol,
                    pkt.len()
                );
                match protocol {
                    IpProtocol::Tcp if config.tcp_enabled => {
                        match extract_tcp_info(&pkt) {
                            Some((src_addr, dst_addr, is_syn)) => {
                                trace!("TCP packet: {} -> {}, SYN={}", src_addr, dst_addr, is_syn);
                                if is_syn && !active_connections.contains(&(src_addr, dst_addr)) {
                                    // Check connection limit
                                    if let Some(limit) = limits.tcp_connection_limit()
                                        && sockets.len() >= limit
                                    {
                                        warn!(
                                            "Connection limit reached ({}), dropping SYN from {}",
                                            limit, src_addr
                                        );
                                        continue;
                                    }

                                    info!("New TCP SYN: {} -> {}", src_addr, dst_addr);

                                    if let Some((new_conn, control)) = create_tcp_connection(
                                        src_addr,
                                        dst_addr,
                                        &mut socket_set,
                                        &stack_thread,
                                        limits.tcp_buffer_size,
                                    ) {
                                        sockets.insert(
                                            new_conn.handle,
                                            SocketInfo {
                                                control,
                                                src_addr,
                                                dst_addr,
                                            },
                                        );
                                        active_connections.insert((src_addr, dst_addr));

                                        if let Ok(state) = shared_state.try_lock()
                                            && let Some(ref tx) = state.new_conn_tx
                                        {
                                            let _ = tx.send(new_conn.new_tcp_conn);
                                        }
                                    }
                                }
                            }
                            None => {
                                warn!("Failed to parse TCP packet, len={}", pkt.len());
                            }
                        }

                        tcp_packets.push(pkt);
                    }
                    IpProtocol::Icmp | IpProtocol::Icmpv6 if config.icmp_enabled => {
                        // ICMP goes to smoltcp immediately
                        tcp_packets.push(pkt);
                    }
                    IpProtocol::Udp if config.udp_enabled => {
                        // UDP goes to tokio - convert to Vec since it leaves our pool
                        let _ = udp_tx.try_send(pkt.to_vec());
                    }
                    _ => {
                        trace!("ignoring packet with protocol {:?}", protocol);
                    }
                }
            }
        }

        if packets_read > 0 {
            phy_wait_error_count = 0;
        }

        // Skip remaining work if a fatal read error was detected above.
        if !running.load(Ordering::Relaxed) {
            break;
        }

        // Processes batched TCP/ICMP packets through smoltcp.
        let has_tcp_packet = !tcp_packets.is_empty();
        for pkt in tcp_packets {
            device.store_packet(pkt);
            let now = SmolInstant::now();
            iface.poll(now, &mut device, &mut socket_set);
        }

        let now = SmolInstant::now();
        iface.poll(now, &mut device, &mut socket_set);

        let mut sockets_to_remove = Vec::new();

        for (handle, socket_info) in sockets.iter() {
            let handle = *handle;
            let control = &socket_info.control;
            let socket = socket_set.get_mut::<TcpSocket>(handle);

            if control.is_abandoned() {
                socket.abort();
            }

            // Remove socket only when smoltcp reports Closed state
            if socket.state() == TcpState::Closed {
                sockets_to_remove.push(handle);
                control.set_closed();
                trace!("socket {:?} closed", handle);
                continue;
            }

            // Handle SHUT_WR: Close -> Closing transition
            // Must check send_queue() to ensure smoltcp has transmitted all data
            if control.send_state() == TcpSocketState::Close
                && socket.send_queue() == 0
                && control.send_buffer_empty()
            {
                trace!(
                    "socket {:?}: closing write half, state={:?}",
                    handle,
                    socket.state()
                );
                socket.close();
                control.set_send_state(TcpSocketState::Closing);
            }

            // Receive data from smoltcp into our buffer
            let mut wake_receiver = false;
            while socket.can_recv() && !control.recv_buffer_full() {
                match socket.recv(|data| {
                    let n = control.enqueue_recv_data(data);
                    (n, n)
                }) {
                    Ok(n) if n > 0 => {
                        wake_receiver = true;
                    }
                    Ok(_) => break,
                    Err(e) => {
                        error!(
                            "socket {:?} recv error: {:?}, state={:?}",
                            handle,
                            e,
                            socket.state()
                        );
                        socket.abort();
                        if control.recv_state() == TcpSocketState::Normal {
                            control.set_recv_state(TcpSocketState::Closed);
                        }
                        wake_receiver = true;
                        break;
                    }
                }
            }

            // Detect recv half close using negative state matching.
            // If socket can't receive and is not in an active receiving state, mark recv closed.
            if control.recv_state() == TcpSocketState::Normal
                && !socket.may_recv()
                && !matches!(
                    socket.state(),
                    TcpState::Listen
                        | TcpState::SynReceived
                        | TcpState::Established
                        | TcpState::FinWait1
                        | TcpState::FinWait2
                )
            {
                trace!(
                    "socket {:?}: recv half closed, state={:?}",
                    handle,
                    socket.state()
                );
                control.set_recv_state(TcpSocketState::Closed);
                wake_receiver = true;
            }

            if wake_receiver {
                control.wake_receiver();
            }

            // Send data from our buffer to smoltcp
            let mut wake_sender = false;
            while socket.can_send() && !control.send_buffer_empty() {
                match socket.send(|buf| {
                    let n = control.dequeue_send_data(buf);
                    (n, n)
                }) {
                    Ok(n) if n > 0 => {
                        wake_sender = true;
                    }
                    Ok(_) => break,
                    Err(e) => {
                        error!(
                            "socket {:?} send error: {:?}, state={:?}",
                            handle,
                            e,
                            socket.state()
                        );
                        socket.abort();
                        if control.send_state() == TcpSocketState::Normal {
                            control.set_send_state(TcpSocketState::Closed);
                        }
                        wake_sender = true;
                        break;
                    }
                }
            }

            if wake_sender {
                control.wake_sender();
            }
        }

        // Dispatch queued data and abort resets before removing closed sockets.
        let after_transfer = SmolInstant::now();
        iface.poll(after_transfer, &mut device, &mut socket_set);

        for handle in sockets_to_remove {
            if let Some(socket_info) = sockets.remove(&handle) {
                active_connections.remove(&(socket_info.src_addr, socket_info.dst_addr));
                trace!(
                    "Cleaned up connection: {} -> {}",
                    socket_info.src_addr, socket_info.dst_addr
                );
            }
            socket_set.remove(handle);
        }

        poll_count += 1;
        if last_log_time.elapsed() >= Duration::from_secs(30) {
            debug!(
                "smoltcp direct stack: polls={}, active_sockets={}",
                poll_count,
                sockets.len()
            );
            last_log_time = std::time::Instant::now();
        }

        // Wait for data using select() - this is the key for event-driven I/O
        if !has_tcp_packet && device.pending_rx.is_none() {
            // Cap poll_delay at 10ms to balance CPU usage vs throughput
            let delay = iface.poll_delay(after_transfer, &socket_set);
            let wait_duration = Some(delay.map_or(SmolDuration::from_millis(10), |d| {
                let millis = d.total_millis().min(10);
                SmolDuration::from_millis(millis)
            }));

            // phy_wait calls select() on the TUN fd to sleep until data
            // arrives. If the fd becomes invalid (e.g. device removed),
            // select() returns EBADF immediately with no sleep, creating a
            // hot spin loop. The try_recv path usually catches this first,
            // but this counter acts as a backstop: after 10 consecutive
            // non-EINTR errors with no successful reads in between, treat
            // the fd as dead.
            if let Err(e) = phy_wait(fd, wait_duration)
                && e.kind() != io::ErrorKind::Interrupted
            {
                phy_wait_error_count += 1;
                if phy_wait_error_count >= MAX_PHY_WAIT_ERRORS {
                    error!(
                        "select() failed {} consecutive times (last: {}). Stack thread stopping.",
                        phy_wait_error_count, e
                    );
                    running.store(false, Ordering::Relaxed);
                } else {
                    warn!("select() error ({}): {}", phy_wait_error_count, e);
                }
            }
        }
    }

    info!("smoltcp direct stack thread stopped");
}

/// Result of creating a TCP connection.
struct CreateConnectionResult {
    handle: SocketHandle,
    new_tcp_conn: NewTcpConnection,
}

/// Create a new TCP connection in the smoltcp stack.
fn create_tcp_connection(
    src_addr: SocketAddr,
    dst_addr: SocketAddr,
    socket_set: &mut SocketSet<'static>,
    stack_thread: &Thread,
    buffer_size: usize,
) -> Option<(CreateConnectionResult, Arc<TcpConnectionControl>)> {
    let mut socket = TcpSocket::new(
        TcpSocketBuffer::new(vec![0u8; buffer_size]),
        TcpSocketBuffer::new(vec![0u8; buffer_size]),
    );

    // Matched to netstack-smoltcp settings for optimal performance
    socket.set_congestion_control(CongestionControl::Cubic);
    socket.set_keep_alive(Some(SmolDuration::from_secs(28)));
    // 7200s matches Linux default (tcp_keepalive_time) and shadowsocks-rust
    socket.set_timeout(Some(SmolDuration::from_secs(7200)));
    socket.set_nagle_enabled(false);
    socket.set_ack_delay(None);

    if let Err(e) = socket.listen(dst_addr) {
        warn!("Failed to listen on socket for {}: {:?}", dst_addr, e);
        return None;
    }

    debug!("Creating TCP connection: {} -> {}", src_addr, dst_addr);

    let control = Arc::new(TcpConnectionControl::new(buffer_size, buffer_size));

    let handle = socket_set.add(socket);
    let connection = TcpConnection::new(control.clone(), stack_thread.clone());

    Some((
        CreateConnectionResult {
            handle,
            new_tcp_conn: NewTcpConnection {
                connection,
                remote_addr: dst_addr,
            },
        },
        control,
    ))
}

/// Extract IP protocol from a raw IP packet.
fn get_ip_protocol(packet: &[u8]) -> Option<IpProtocol> {
    if packet.is_empty() {
        return None;
    }

    let version = packet[0] >> 4;
    match version {
        4 => Ipv4Packet::new_checked(packet)
            .ok()
            .map(|p| p.next_header()),
        6 => Ipv6Packet::new_checked(packet)
            .ok()
            .map(|p| p.next_header()),
        _ => None,
    }
}

/// Extract TCP connection info from a raw IP packet.
fn extract_tcp_info(packet: &[u8]) -> Option<(SocketAddr, SocketAddr, bool)> {
    if packet.is_empty() {
        return None;
    }

    let version = packet[0] >> 4;
    match version {
        4 => {
            let ip = Ipv4Packet::new_checked(packet).ok()?;
            if ip.next_header() != IpProtocol::Tcp {
                return None;
            }
            let tcp = TcpPacket::new_checked(ip.payload()).ok()?;
            let src_addr = SocketAddr::new(
                std::net::IpAddr::V4(std::net::Ipv4Addr::from(ip.src_addr().octets())),
                tcp.src_port(),
            );
            let dst_addr = SocketAddr::new(
                std::net::IpAddr::V4(std::net::Ipv4Addr::from(ip.dst_addr().octets())),
                tcp.dst_port(),
            );
            let is_syn = tcp.syn() && !tcp.ack();
            Some((src_addr, dst_addr, is_syn))
        }
        6 => {
            let ip = Ipv6Packet::new_checked(packet).ok()?;
            if ip.next_header() != IpProtocol::Tcp {
                return None;
            }
            let tcp = TcpPacket::new_checked(ip.payload()).ok()?;
            let src_addr = SocketAddr::new(
                std::net::IpAddr::V6(std::net::Ipv6Addr::from(ip.src_addr().octets())),
                tcp.src_port(),
            );
            let dst_addr = SocketAddr::new(
                std::net::IpAddr::V6(std::net::Ipv6Addr::from(ip.dst_addr().octets())),
                tcp.dst_port(),
            );
            let is_syn = tcp.syn() && !tcp.ack();
            Some((src_addr, dst_addr, is_syn))
        }
        _ => None,
    }
}

/// Check if an IP packet should be filtered.
fn should_filter_packet(packet: &[u8]) -> bool {
    if packet.is_empty() {
        return true;
    }

    let version = packet[0] >> 4;
    match version {
        4 => {
            if let Ok(ip) = Ipv4Packet::new_checked(packet) {
                let src = ip.src_addr();
                let dst = ip.dst_addr();

                let src_bytes = src.octets();
                let dst_bytes = dst.octets();

                // Filter unspecified source
                if src_bytes == [0, 0, 0, 0] {
                    return true;
                }
                // Filter multicast source
                if src_bytes[0] >= 224 && src_bytes[0] <= 239 {
                    return true;
                }
                // Filter broadcast destination
                if dst_bytes == [255, 255, 255, 255] {
                    return true;
                }
                // Filter multicast destination
                if dst_bytes[0] >= 224 && dst_bytes[0] <= 239 {
                    return true;
                }
                // Filter unspecified destination
                if dst_bytes == [0, 0, 0, 0] {
                    return true;
                }

                false
            } else {
                true
            }
        }
        6 => {
            if let Ok(ip) = Ipv6Packet::new_checked(packet) {
                let src = ip.src_addr();
                let dst = ip.dst_addr();

                let src_bytes = src.octets();
                let dst_bytes = dst.octets();

                // Filter unspecified source
                if src_bytes == [0u8; 16] {
                    return true;
                }
                // Filter multicast destination
                if dst_bytes[0] == 0xff {
                    return true;
                }
                // Filter unspecified destination
                if dst_bytes == [0u8; 16] {
                    return true;
                }

                false
            } else {
                true
            }
        }
        _ => true,
    }
}

/// Set a file descriptor to non-blocking mode (call once at startup).
fn set_nonblocking(fd: RawFd) -> io::Result<()> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags < 0 {
        return Err(io::Error::last_os_error());
    }
    if (flags & libc::O_NONBLOCK) == 0
        && unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } < 0
    {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

fn packet_header(first_byte: u8) -> io::Result<[u8; 4]> {
    let version = first_byte >> 4;
    #[cfg(any(target_os = "macos", target_os = "ios"))]
    let protocol = match version {
        4 => libc::AF_INET as u32,
        6 => libc::AF_INET6 as u32,
        _ => 0,
    };
    #[cfg(not(any(target_os = "macos", target_os = "ios")))]
    let protocol: u32 = match version {
        4 => 0x0800,
        6 => 0x86dd,
        _ => 0,
    };
    if protocol == 0 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "invalid IP version",
        ));
    }
    Ok(protocol.to_be_bytes())
}

fn write_packet(fd: RawFd, data: &[u8], packet_information: bool) -> io::Result<()> {
    if packet_information {
        let Some(&first_byte) = data.first() else {
            return Ok(());
        };
        let mut framed = Vec::with_capacity(data.len() + 4);
        framed.extend_from_slice(&packet_header(first_byte)?);
        framed.extend_from_slice(data);
        write_all(fd, &framed)
    } else {
        write_all(fd, data)
    }
}

/// Non-blocking read from a file descriptor (fd must already be non-blocking).
/// Returns Err(WouldBlock) when no data is available, so callers can
/// distinguish it from Ok(0) which indicates EOF.
fn read_nonblocking(fd: RawFd, buf: &mut [u8]) -> io::Result<usize> {
    let n = unsafe { libc::read(fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) };
    if n < 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(n as usize)
    }
}

/// Write all data to a file descriptor.
fn write_all(fd: RawFd, buf: &[u8]) -> io::Result<()> {
    let mut written = 0;
    while written < buf.len() {
        let n = unsafe {
            libc::write(
                fd,
                buf[written..].as_ptr() as *const libc::c_void,
                buf.len() - written,
            )
        };
        if n < 0 {
            let err = io::Error::last_os_error();
            if err.raw_os_error() == Some(libc::ENOBUFS) || err.kind() == io::ErrorKind::WouldBlock
            {
                trace!("TUN write {}, packet dropped", err);
                return Ok(());
            }
            return Err(err);
        }
        if n == 0 {
            return Err(io::Error::new(
                io::ErrorKind::WriteZero,
                "TUN write made no progress",
            ));
        }
        written += n as usize;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::os::unix::io::IntoRawFd;
    use std::os::unix::net::UnixStream;

    #[test]
    fn default_admission_preserves_bursts_beyond_old_queue_and_socket_caps() {
        use std::os::unix::net::UnixDatagram;

        let (peer, tun) = UnixDatagram::pair().unwrap();
        peer.set_read_timeout(Some(Duration::from_secs(2))).unwrap();
        let mut stack = TcpStackDirect::new(tun.into(), 1500);
        let (tx, mut rx) = mpsc::unbounded_channel();
        stack.set_new_conn_tx(tx);
        let mut buf = [0; 1500];
        for port in 10000..10300 {
            let builder = etherparse::PacketBuilder::ipv4([10, 0, 0, 2], [1, 1, 1, 1], 64)
                .tcp(port, 443, 101, 4096);
            let mut syn = Vec::new();
            builder.syn().write(&mut syn, b"").unwrap();
            peer.send(&syn).unwrap();
            let acknowledgement = loop {
                let n = peer.recv(&mut buf).unwrap();
                let ip = Ipv4Packet::new_checked(&buf[..n]).unwrap();
                let tcp = TcpPacket::new_checked(ip.payload()).unwrap();
                assert!(!tcp.rst(), "burst connection was reset");
                if tcp.dst_port() == port && tcp.syn() && tcp.ack() {
                    break (tcp.seq_number().0 as u32).wrapping_add(1);
                }
            };
            let mut ack = Vec::new();
            etherparse::PacketBuilder::ipv4([10, 0, 0, 2], [1, 1, 1, 1], 64)
                .tcp(port, 443, 102, 4096)
                .ack(acknowledgement)
                .write(&mut ack, b"")
                .unwrap();
            peer.send(&ack).unwrap();
        }
        assert_eq!(rx.len(), 300);
        let connections: Vec<_> = (0..300).map(|_| rx.try_recv().unwrap()).collect();
        drop(connections);
    }

    #[tokio::test]
    async fn abandoned_connection_sends_reset_before_releasing_socket() {
        use std::os::unix::net::UnixDatagram;
        use tokio::io::AsyncReadExt;

        let (peer, tun) = UnixDatagram::pair().unwrap();
        peer.set_read_timeout(Some(Duration::from_secs(1))).unwrap();
        let mut stack = TcpStackDirect::new(tun.into(), 1500);
        let (tx, mut rx) = mpsc::unbounded_channel();
        stack.set_new_conn_tx(tx);
        let packet = |sequence, acknowledgment: Option<u32>, payload: &[u8]| {
            let builder = etherparse::PacketBuilder::ipv4([10, 0, 0, 2], [1, 1, 1, 1], 64)
                .tcp(10001, 443, sequence, 4096);
            let builder = match acknowledgment {
                Some(ack) => builder.ack(ack),
                None => builder.syn(),
            };
            let mut bytes = Vec::new();
            builder.write(&mut bytes, payload).unwrap();
            bytes
        };
        let mut buf = [0; 1500];
        peer.send(&packet(101, None, b"")).unwrap();
        let n = peer.recv(&mut buf).unwrap();
        let ip = Ipv4Packet::new_checked(&buf[..n]).unwrap();
        let syn_ack = TcpPacket::new_checked(ip.payload()).unwrap();
        assert!(syn_ack.syn() && syn_ack.ack());
        let acknowledgment = (syn_ack.seq_number().0 as u32).wrapping_add(1);
        let mut incoming = rx.recv().await.unwrap();
        peer.send(&packet(102, Some(acknowledgment), b"x")).unwrap();
        tokio::time::timeout(
            Duration::from_secs(1),
            incoming.connection.read_exact(&mut [0]),
        )
        .await
        .unwrap()
        .unwrap();
        drop(incoming);

        loop {
            let n = peer
                .recv(&mut buf)
                .expect("abandoned connection did not send a reset");
            let ip = Ipv4Packet::new_checked(&buf[..n]).unwrap();
            let tcp = TcpPacket::new_checked(ip.payload()).unwrap();
            if tcp.rst() {
                assert!(!tcp.fin());
                break;
            }
        }
        peer.send(&packet(1001, None, b"")).unwrap();
        let n = peer.recv(&mut buf).unwrap();
        let ip = Ipv4Packet::new_checked(&buf[..n]).unwrap();
        let syn_ack = TcpPacket::new_checked(ip.payload()).unwrap();
        assert!(syn_ack.syn() && syn_ack.ack());
        assert_eq!(syn_ack.ack_number().0 as u32, 1002);
    }

    #[test]
    fn packet_information_roundtrips_ipv4_and_ipv6() {
        use std::os::unix::net::UnixDatagram;
        let (peer, tun) = UnixDatagram::pair().unwrap();
        peer.set_read_timeout(Some(Duration::from_secs(1))).unwrap();
        set_nonblocking(tun.as_raw_fd()).unwrap();
        let mut device = DirectDevice::new(tun.as_raw_fd(), 1500, true);
        for version in [0x45, 0x60] {
            let packet = [version, 1, 2, 3, 4, 5];
            let mut framed = packet_header(version).unwrap().to_vec();
            framed.extend_from_slice(&packet);
            peer.send(&framed).unwrap();
            assert_eq!(&device.try_recv().unwrap().unwrap()[..], &packet);
            device.write_packet(&packet).unwrap();
            let mut buf = [0; 64];
            let n = peer.recv(&mut buf).unwrap();
            assert_eq!(&buf[..n], &framed);
        }
    }

    #[test]
    fn disabled_udp_does_not_reach_async_handler() {
        use std::os::unix::net::UnixDatagram;
        let (peer, tun) = UnixDatagram::pair().unwrap();
        let mut stack =
            TcpStackDirect::with_config(tun.into(), TunServerConfig::new().udp_enabled(false));
        let mut rx = stack.take_udp_rx().unwrap();
        let builder =
            etherparse::PacketBuilder::ipv4([10, 0, 0, 2], [1, 1, 1, 1], 64).udp(1000, 53);
        let mut packet = Vec::new();
        builder.write(&mut packet, b"dns").unwrap();
        peer.send(&packet).unwrap();
        thread::sleep(Duration::from_millis(40));
        assert!(rx.try_recv().is_err());
    }

    #[test]
    fn icmp_echo_replies_respect_protocol_switch() {
        use smoltcp::wire::{Icmpv4Message, Icmpv4Packet, Icmpv6Message, Icmpv6Packet};
        use std::os::unix::net::UnixDatagram;

        let ipv4_src = [10, 0, 0, 2];
        let ipv4_dst = [1, 1, 1, 1];
        let ipv6_src = Ipv6Address::new(0xfd00, 0, 0, 0, 0, 0, 0, 2);
        let ipv6_dst = Ipv6Address::new(0xfd00, 0, 0, 0, 0, 0, 0, 3);
        let mut ipv4_request = Vec::new();
        etherparse::PacketBuilder::ipv4(ipv4_src, ipv4_dst, 64)
            .icmpv4_echo_request(17, 23)
            .write(&mut ipv4_request, b"ping")
            .unwrap();
        let mut ipv6_request = Vec::new();
        etherparse::PacketBuilder::ipv6(ipv6_src.octets(), ipv6_dst.octets(), 64)
            .icmpv6_echo_request(17, 23)
            .write(&mut ipv6_request, b"ping")
            .unwrap();

        for enabled in [true, false] {
            let (peer, tun) = UnixDatagram::pair().unwrap();
            peer.set_read_timeout(Some(Duration::from_millis(if enabled {
                1000
            } else {
                200
            })))
            .unwrap();
            let stack = TcpStackDirect::with_config(
                tun.into(),
                TunServerConfig::new().icmp_enabled(enabled),
            );
            for request in [&ipv4_request, &ipv6_request] {
                peer.send(request).unwrap();
                let mut response = [0; 1500];
                let result = peer.recv(&mut response);
                if !enabled {
                    assert!(matches!(
                        result.unwrap_err().kind(),
                        io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut
                    ));
                    assert!(stack.is_running());
                    continue;
                }
                let response = &response[..result.unwrap()];
                if request[0] >> 4 == 4 {
                    let ip = Ipv4Packet::new_checked(response).unwrap();
                    assert_eq!(ip.src_addr().octets(), ipv4_dst);
                    assert_eq!(ip.dst_addr().octets(), ipv4_src);
                    assert!(ip.verify_checksum());
                    let icmp = Icmpv4Packet::new_checked(ip.payload()).unwrap();
                    assert_eq!(icmp.msg_type(), Icmpv4Message::EchoReply);
                    assert_eq!((icmp.echo_ident(), icmp.echo_seq_no()), (17, 23));
                    assert_eq!(icmp.data(), b"ping");
                    assert!(icmp.verify_checksum());
                } else {
                    let ip = Ipv6Packet::new_checked(response).unwrap();
                    assert_eq!(ip.src_addr(), ipv6_dst);
                    assert_eq!(ip.dst_addr(), ipv6_src);
                    let icmp = Icmpv6Packet::new_checked(ip.payload()).unwrap();
                    assert_eq!(icmp.msg_type(), Icmpv6Message::EchoReply);
                    assert_eq!((icmp.echo_ident(), icmp.echo_seq_no()), (17, 23));
                    assert_eq!(icmp.payload(), b"ping");
                    assert!(icmp.verify_checksum(&ip.src_addr(), &ip.dst_addr()));
                }
            }
        }
    }

    #[test]
    fn socket_owner_drop_wakes_readers_and_writers() {
        use futures::task::{ArcWake, waker};
        use std::pin::Pin;
        use std::task::Context;
        use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

        struct WakeFlag(AtomicBool);
        impl ArcWake for WakeFlag {
            fn wake_by_ref(flag: &Arc<Self>) {
                flag.0.store(true, Ordering::Relaxed);
            }
        }
        let flag = Arc::new(WakeFlag(AtomicBool::new(false)));
        let waker = waker(flag.clone());
        let mut cx = Context::from_waker(&waker);
        let control = Arc::new(TcpConnectionControl::new(1, 1));
        let mut conn = TcpConnection::new(control.clone(), thread::current());
        let info = SocketInfo {
            control: control.clone(),
            src_addr: "127.0.0.1:1".parse().unwrap(),
            dst_addr: "127.0.0.1:2".parse().unwrap(),
        };
        let mut buf = [0; 1];
        assert!(
            Pin::new(&mut conn)
                .poll_read(&mut cx, &mut ReadBuf::new(&mut buf))
                .is_pending()
        );
        assert!(Pin::new(&mut conn).poll_write(&mut cx, b"a").is_ready());
        assert!(Pin::new(&mut conn).poll_write(&mut cx, b"b").is_pending());
        drop(info);
        assert!(flag.0.load(Ordering::Relaxed));
        assert!(
            Pin::new(&mut conn)
                .poll_read(&mut cx, &mut ReadBuf::new(&mut buf))
                .is_ready()
        );
        assert!(Pin::new(&mut conn).poll_write(&mut cx, b"b").is_ready());
        drop(conn);
        assert!(control.is_abandoned());
    }

    #[test]
    fn idle_stack_delivers_udp_output_without_inbound_traffic() {
        use std::io::Read;
        let (mut peer, client) = UnixStream::pair().unwrap();
        peer.set_read_timeout(Some(Duration::from_secs(1))).unwrap();
        let mut stack = TcpStackDirect::new(client.into(), 1500);
        let (tx, rx) = mpsc::channel(PACKET_QUEUE_CAPACITY);
        stack.set_udp_response_tx(rx);
        thread::sleep(Duration::from_millis(30));
        tx.try_send(b"reply".to_vec()).unwrap();
        let mut reply = [0; 5];
        peer.read_exact(&mut reply).unwrap();
        assert_eq!(&reply, b"reply");
    }

    #[test]
    fn test_stack_shutdown_on_eof() {
        let (server, client) = UnixStream::pair().expect("Failed to create socket pair");
        let stack = TcpStackDirect::new(client.into(), 1500);

        thread::sleep(Duration::from_millis(100));
        assert!(stack.is_running(), "Stack thread should be running");

        // Closing the writer end triggers EOF on the reader.
        drop(server);

        let start = std::time::Instant::now();
        let timeout = Duration::from_secs(2);

        while stack.is_running() {
            if start.elapsed() > timeout {
                panic!("Stack thread did not exit after EOF on FD");
            }
            thread::sleep(Duration::from_millis(50));
        }
    }

    #[test]
    fn idle_stack_drop_does_not_require_incoming_io() {
        let (peer, client) = UnixStream::pair().unwrap();
        let stack = TcpStackDirect::new(client.into(), 1500);
        thread::sleep(Duration::from_millis(100));
        let (tx, rx) = std::sync::mpsc::channel();
        let dropper = thread::spawn(move || {
            drop(stack);
            let _ = tx.send(());
        });
        let result = rx.recv_timeout(Duration::from_millis(300));
        drop(peer);
        dropper.join().unwrap();
        assert!(result.is_ok());
    }

    #[test]
    fn test_stack_exits_promptly() {
        // Verifies the stack exits within 1 second of EOF, catching
        // regressions that would cause CPU spin on a dead fd.
        let (server, client) = UnixStream::pair().expect("Failed to create socket pair");
        let stack = TcpStackDirect::new(client.into(), 1500);

        thread::sleep(Duration::from_millis(100));
        assert!(stack.is_running());

        drop(server);

        let start = std::time::Instant::now();
        let timeout = Duration::from_secs(1);

        while stack.is_running() {
            if start.elapsed() > timeout {
                panic!("Stack thread took >1s to exit after EOF (possible spin)");
            }
            thread::sleep(Duration::from_millis(10));
        }
    }

    #[test]
    fn test_write_all_eagain() {
        // Fill a non-blocking socket's write buffer, then verify write_all
        // treats EAGAIN the same as ENOBUFS (drops the packet, returns Ok).
        let (reader, writer) = UnixStream::pair().expect("Failed to create socket pair");
        let writer_fd = writer.into_raw_fd();

        set_nonblocking(writer_fd).expect("set_nonblocking");

        // Fill the write buffer until WouldBlock
        let big_buf = vec![0u8; 65536];
        loop {
            let n = unsafe {
                libc::write(
                    writer_fd,
                    big_buf.as_ptr() as *const libc::c_void,
                    big_buf.len(),
                )
            };
            if n < 0 {
                let err = io::Error::last_os_error();
                assert_eq!(
                    err.kind(),
                    io::ErrorKind::WouldBlock,
                    "unexpected error: {}",
                    err
                );
                break;
            }
        }

        // Now write_all should drop the packet gracefully
        let result = write_all(writer_fd, &[1, 2, 3]);
        assert!(
            result.is_ok(),
            "write_all should return Ok on EAGAIN, got {:?}",
            result
        );

        unsafe { libc::close(writer_fd) };
        drop(reader);
    }
}