microsandbox_network/

stack.rs

//! smoltcp interface setup, frame classification, and poll loop.
//!
//! This module contains the core networking event loop that runs on a
//! dedicated OS thread. It bridges guest ethernet frames (via
//! [`SmoltcpDevice`]) to smoltcp's TCP/IP stack and services connections
//! through tokio proxy tasks.

use std::collections::HashSet;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::Ordering;

#[cfg(windows)]
use msb_krun_utils::event::{EventSet, EventSource, WaitContext, WaitEvent};
use smoltcp::iface::{Config, Interface, SocketSet};
use smoltcp::time::Instant;

use smoltcp::wire::{
    EthernetAddress, EthernetFrame, EthernetProtocol, HardwareAddress, Icmpv4Packet, Icmpv4Repr,
    Icmpv6Packet, Icmpv6Repr, IpAddress, IpCidr, IpProtocol, Ipv4Packet, Ipv4Repr, Ipv6Packet,
    Ipv6Repr, TcpPacket, UdpPacket,
};

use crate::config::{DnsConfig, PublishedPort};
use crate::conn::ConnectionTracker;
use crate::device::SmoltcpDevice;
use crate::dns::common::ports::DnsPortType;
use crate::dns::{
    interceptor::DnsInterceptor,
    proxies::{dot::DotProxy, tcp::DnsTcpProxy},
};
use crate::icmp_relay::IcmpRelay;
use crate::policy::{EgressEvaluation, HostnameSource, NetworkPolicy, Protocol};
use crate::proxy;
use crate::publisher::PortPublisher;
use crate::secrets::handle::SecretsHandle;
use crate::shared::SharedState;
use crate::tls::{proxy as tls_proxy, state::TlsState};
use crate::udp_fragments::{
    Ipv4UdpFragmentReassembler, Ipv6UdpFragmentReassembler, ReassembledUdpDatagram,
    is_ipv4_udp_fragment, is_ipv6_fragment, is_ipv6_udp_fragment,
};
use crate::udp_relay::UdpRelay;

//--------------------------------------------------------------------------------------------------
// Constants
//--------------------------------------------------------------------------------------------------

#[cfg(windows)]
const TX_WAKE_TOKEN: u64 = 1;

#[cfg(windows)]
const PROXY_WAKE_TOKEN: u64 = 2;

//--------------------------------------------------------------------------------------------------
// Types
//--------------------------------------------------------------------------------------------------

/// Result of classifying a guest ethernet frame before smoltcp processes it.
///
/// Pre-inspection allows the poll loop to:
/// - Create TCP sockets before smoltcp sees a SYN (preventing auto-RST).
/// - Handle non-DNS UDP outside smoltcp (smoltcp lacks wildcard port binding).
/// - Route DNS queries to the interception handler.
pub enum FrameAction {
    /// TCP SYN to a new destination — create a smoltcp socket before
    /// letting smoltcp process the frame.
    TcpSyn { src: SocketAddr, dst: SocketAddr },

    /// Non-DNS UDP datagram — handle entirely outside smoltcp via the UDP
    /// relay.
    UdpRelay { src: SocketAddr, dst: SocketAddr },

    /// DNS query (UDP to port 53) — let smoltcp's bound UDP socket handle it.
    Dns,

    /// IPv4 UDP fragment — reassemble before the UDP relay sees it.
    Ipv4UdpFragment,

    /// IPv6 UDP fragment — reassemble before the UDP relay sees it.
    Ipv6UdpFragment,

    /// IPv6 fragment for a protocol this relay cannot safely classify.
    Ipv6UnsupportedFragment,

    /// Everything else (ARP, NDP, ICMP, TCP data/ACK/FIN, etc.) — let
    /// smoltcp process normally.
    Passthrough,
}

/// Local ICMP echo reply plus the destination policy input it answers.
struct GatewayIcmpReply {
    dst: IpAddr,
    protocol: Protocol,
    frame: Vec<u8>,
}

/// Resolved network parameters for the poll loop. Created by
/// `SmoltcpNetwork::new()` from `NetworkConfig` + sandbox slot.
pub struct PollLoopConfig {
    /// Gateway MAC address (smoltcp's identity on the virtual LAN).
    pub gateway_mac: [u8; 6],
    /// Guest MAC address.
    pub guest_mac: [u8; 6],
    /// Gateway addresses owned by the smoltcp virtual stack. Each family
    /// is `Some` when that family is active for this sandbox (host has a
    /// route, or the user supplied an explicit address).
    pub gateway: GatewayIps,
    /// Guest IPv4 address. `None` when IPv4 is inactive for this sandbox.
    pub guest_ipv4: Option<Ipv4Addr>,
    /// Guest IPv6 address. `None` when IPv6 is inactive for this sandbox.
    pub guest_ipv6: Option<Ipv6Addr>,
    /// IP-level MTU (e.g. 1500).
    pub mtu: usize,
}

/// Per-sandbox gateway addresses owned by the smoltcp virtual stack.
///
/// Each family is `Some` when active for this sandbox and `None` otherwise.
/// `resolve_host_dst` rewrites gateway-bound connections to loopback at dial time.
#[derive(Debug, Clone, Copy)]
pub struct GatewayIps {
    /// Gateway IPv4.
    pub ipv4: Option<Ipv4Addr>,
    /// Gateway IPv6.
    pub ipv6: Option<Ipv6Addr>,
}

//--------------------------------------------------------------------------------------------------
// Functions
//--------------------------------------------------------------------------------------------------

/// Classify a raw ethernet frame for pre-inspection.
///
/// Uses smoltcp's wire module for zero-copy parsing. Returns
/// [`FrameAction::Passthrough`] for any frame that cannot be parsed or
/// doesn't match a special case.
pub fn classify_frame(frame: &[u8]) -> FrameAction {
    let Ok(eth) = EthernetFrame::new_checked(frame) else {
        return FrameAction::Passthrough;
    };

    match eth.ethertype() {
        EthernetProtocol::Ipv4 => classify_ipv4(eth.payload()),
        EthernetProtocol::Ipv6 => classify_ipv6(eth.payload()),
        _ => FrameAction::Passthrough, // ARP, etc.
    }
}

/// Create and configure the smoltcp [`Interface`].
///
/// The interface is configured as the **gateway**: it owns the gateway IP
/// addresses and responds to ARP/NDP for them. `any_ip` mode is enabled so
/// smoltcp accepts traffic destined for arbitrary remote IPs (not just the
/// gateway), combined with default routes.
pub fn create_interface(device: &mut SmoltcpDevice, config: &PollLoopConfig) -> Interface {
    let hw_addr = HardwareAddress::Ethernet(EthernetAddress(config.gateway_mac));
    let iface_config = Config::new(hw_addr);
    let mut iface = Interface::new(iface_config, device, smoltcp_now());

    // Configure gateway IP addresses for the active families.
    iface.update_ip_addrs(|addrs| {
        if let Some(ipv4) = config.gateway.ipv4 {
            addrs
                .push(IpCidr::new(IpAddress::Ipv4(ipv4), 30)) // 30 subnet: gateway + guest.
                .expect("failed to add gateway IPv4 address");
        }
        if let Some(ipv6) = config.gateway.ipv6 {
            addrs
                .push(IpCidr::new(IpAddress::Ipv6(ipv6), 64))
                .expect("failed to add gateway IPv6 address");
        }
    });

    // Default routes so smoltcp accepts traffic for all destinations.
    if let Some(ipv4) = config.gateway.ipv4 {
        iface
            .routes_mut()
            .add_default_ipv4_route(ipv4)
            .expect("failed to add default IPv4 route");
    }
    if let Some(ipv6) = config.gateway.ipv6 {
        iface
            .routes_mut()
            .add_default_ipv6_route(ipv6)
            .expect("failed to add default IPv6 route");
    }

    // Accept traffic destined for any IP, not just gateway addresses.
    iface.set_any_ip(true);

    iface
}

/// Main smoltcp poll loop. Runs on a dedicated OS thread.
///
/// Processes guest frames with pre-inspection, drives smoltcp's TCP/IP stack,
/// and sleeps via `poll(2)` between events.
///
/// # Phases per iteration
///
/// 1. **Drain guest frames** — pop from `tx_ring`, classify, pre-inspect.
/// 2. **smoltcp egress + maintenance** — transmit queued packets, run timers.
/// 3. **Service connections** — relay data between smoltcp sockets and proxy
///    tasks (added by later tasks).
/// 4. **Sleep** — wait on `tx_wake` + `proxy_wake` with smoltcp's requested
///    timeout.
///
/// # Arguments
///
/// * `shared` - Stack-wide shared state: `tx_ring` / `rx_ring` for the virtio-net boundary
///   and the wake eventfds.
/// * `config` - Resolved per-sandbox parameters (gateway / guest MAC + IPv4 + IPv6, MTU).
/// * `network_policy` - User-provided egress policy. Evaluated against the sandbox's
///   gateway IPs (stored on [`SharedState`]) so `DestinationGroup::Host` rules match.
/// * `platform_policy` - Optional host-owned policy floor. Traffic must pass both policies.
/// * `dns_config` - DNS interception settings (block lists, upstreams, timeout).
/// * `tls_state` - Optional TLS MITM state; drives interception of intercepted ports and DoT
///   when present.
/// * `published_ports` - Host → guest port publishes; the publisher accepts inbound
///   connections on the host-bind address and forwards into the guest.
/// * `max_connections` - Optional cap on concurrent guest connections tracked by
///   [`ConnectionTracker`]; `None` uses the default.
/// * `tokio_handle` - Runtime handle used for proxy tasks, DNS forwarding, port publishing,
///   and ICMP relays.
#[allow(clippy::too_many_arguments)]
pub fn smoltcp_poll_loop(
    shared: Arc<SharedState>,
    config: PollLoopConfig,
    network_policy: NetworkPolicy,
    platform_policy: Option<NetworkPolicy>,
    dns_config: DnsConfig,
    tls_state: Option<Arc<TlsState>>,
    published_ports: Vec<PublishedPort>,
    max_connections: Option<usize>,
    tokio_handle: tokio::runtime::Handle,
    secrets: SecretsHandle,
) {
    let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
    let mut iface = create_interface(&mut device, &config);
    let mut sockets = SocketSet::new(vec![]);
    let mut conn_tracker = ConnectionTracker::new(max_connections);

    // The DNS forwarder needs to know which IPs count as "the gateway"
    // (so it routes guest queries to those addresses through the
    // configured upstream) and a policy evaluator (so guest-chosen
    // `@target` resolvers are gated by egress rules just like any
    // other outbound).
    let gateway_ips: Arc<HashSet<IpAddr>> = Arc::new(
        config
            .gateway
            .ipv4
            .map(IpAddr::V4)
            .into_iter()
            .chain(config.gateway.ipv6.map(IpAddr::V6))
            .collect(),
    );
    // Gateway IPs must be on SharedState before any egress evaluation runs,
    // so `DestinationGroup::Host` rules can resolve to the right address.
    shared.set_gateway_ips(config.gateway.ipv4, config.gateway.ipv6);
    let network_policy = Arc::new(network_policy);
    let platform_policy = platform_policy.map(Arc::new);

    let (mut dns_interceptor, dns_forwarder_handle) = DnsInterceptor::new(
        &mut sockets,
        dns_config,
        shared.clone(),
        &tokio_handle,
        gateway_ips,
        network_policy.clone(),
        platform_policy.clone(),
        config.gateway,
        config.gateway_mac,
        config.guest_mac,
    );
    let mut port_publisher = PortPublisher::new(
        &published_ports,
        config.guest_ipv4,
        config.guest_ipv6,
        config.gateway.ipv4,
        config.gateway.ipv6,
        config.gateway_mac,
        config.guest_mac,
        network_policy.clone(),
        shared.clone(),
        &tokio_handle,
    );
    let mut udp_relay = UdpRelay::new(
        shared.clone(),
        config.gateway_mac,
        config.guest_mac,
        config.mtu,
        tokio_handle.clone(),
    );
    let mut udp_fragments = Ipv4UdpFragmentReassembler::new();
    let mut ipv6_udp_fragments = Ipv6UdpFragmentReassembler::new();
    let icmp_relay = IcmpRelay::new(
        shared.clone(),
        config.gateway_mac,
        config.guest_mac,
        tokio_handle.clone(),
    );

    // Rate-limit cleanup operations: run at most once per second.
    let mut last_cleanup = std::time::Instant::now();

    // Wake sources for sleeping.
    #[cfg(unix)]
    let mut poll_fds = [
        libc::pollfd {
            fd: shared.tx_wake.as_raw_fd(),
            events: libc::POLLIN,
            revents: 0,
        },
        libc::pollfd {
            fd: shared.proxy_wake.as_raw_fd(),
            events: libc::POLLIN,
            revents: 0,
        },
    ];
    #[cfg(windows)]
    let wait_context = match windows_stack_wait_context(&shared) {
        Ok(context) => context,
        Err(err) => {
            tracing::error!(error = %err, "network poll loop: failed to create wait context");
            return;
        }
    };

    loop {
        let now = smoltcp_now();

        // ── Phase 1: Drain all guest frames with pre-inspection ──────────
        while let Some(frame) = device.stage_next_frame() {
            if handle_gateway_icmp_echo(
                frame,
                &config,
                &shared,
                &network_policy,
                platform_policy.as_deref(),
            ) {
                device.drop_staged_frame();
                continue;
            }

            if icmp_relay.relay_outbound_if_echo(
                frame,
                &config,
                &network_policy,
                platform_policy.as_deref(),
            ) {
                device.drop_staged_frame();
                continue;
            }

            match classify_frame(frame) {
                FrameAction::TcpSyn { src, dst } => {
                    let allow = match DnsPortType::from_tcp(dst.port()) {
                        // Plain DNS: the interceptor enforces policy at
                        // the application layer (block list + rebind
                        // protection); bypass the network egress check.
                        DnsPortType::Dns => true,
                        // DoT: intercept only when TLS MITM is
                        // configured. Without it, the block list can't
                        // apply (traffic is encrypted end-to-end), so
                        // we refuse to force a fall-back to plain
                        // TCP/53. When TLS MITM is configured, bypass
                        // egress policy the same way plain DNS does —
                        // policy for the upstream resolver is applied
                        // per query by the forwarder.
                        DnsPortType::EncryptedDns => {
                            if tls_state.is_some() {
                                true
                            } else {
                                tracing::debug!(%dst, "DoT port refused (TLS interception not configured); stub should fall back to TCP/53");
                                false
                            }
                        }
                        // Alternative DNS protocol we can't proxy:
                        // refuse outright — no socket means smoltcp
                        // emits RST, which the guest's stub treats as
                        // "upstream unavailable" and falls back to
                        // plain TCP/53.
                        DnsPortType::AlternativeDns => {
                            tracing::debug!(%dst, "alternative-DNS TCP port refused; stub should fall back to TCP/53");
                            false
                        }
                        // Other: regular outbound — defer Domain rules to first-flight;
                        // accept unless an IP-layer rule denies.
                        DnsPortType::Other => {
                            let platform_allows = platform_policy.as_deref().is_none_or(|policy| {
                                policy
                                    .evaluate_egress(dst, Protocol::Tcp, &shared)
                                    .is_allow()
                            });
                            platform_allows
                                && matches!(
                                    network_policy.evaluate_egress_with_source(
                                        dst,
                                        Protocol::Tcp,
                                        &shared,
                                        HostnameSource::Deferred,
                                    ),
                                    EgressEvaluation::Allow | EgressEvaluation::DeferUntilHostname
                                )
                        }
                    };
                    if allow && !conn_tracker.has_socket_for(&src, &dst) {
                        conn_tracker.create_tcp_socket(src, dst, &mut sockets);
                    }
                    // Let smoltcp process — matching socket completes
                    // handshake, no socket means auto-RST.
                    iface.poll_ingress_single(now, &mut device, &mut sockets);
                }

                FrameAction::UdpRelay { src, dst } => {
                    relay_udp_frame(
                        frame,
                        src,
                        dst,
                        &config,
                        &network_policy,
                        platform_policy.as_deref(),
                        &shared,
                        &mut port_publisher,
                        tls_state.as_deref(),
                        &mut udp_relay,
                    );
                    device.drop_staged_frame();
                }

                FrameAction::Ipv4UdpFragment => {
                    if let Some(datagram) = udp_fragments.push(frame) {
                        handle_reassembled_udp_datagram(
                            datagram,
                            &mut device,
                            &mut iface,
                            now,
                            &mut sockets,
                            &config,
                            &network_policy,
                            platform_policy.as_deref(),
                            &shared,
                            &mut port_publisher,
                            tls_state.as_deref(),
                            &mut udp_relay,
                        );
                    } else {
                        device.drop_staged_frame();
                    }
                }

                FrameAction::Ipv6UdpFragment => {
                    if let Some(datagram) = ipv6_udp_fragments.push(frame) {
                        handle_reassembled_udp_datagram(
                            datagram,
                            &mut device,
                            &mut iface,
                            now,
                            &mut sockets,
                            &config,
                            &network_policy,
                            platform_policy.as_deref(),
                            &shared,
                            &mut port_publisher,
                            tls_state.as_deref(),
                            &mut udp_relay,
                        );
                    } else {
                        device.drop_staged_frame();
                    }
                }

                FrameAction::Ipv6UnsupportedFragment => {
                    // Fragmented UDP is only forwarded after reassembly and policy evaluation.
                    // Other fragmented IPv6 traffic is dropped rather than passed through with
                    // an unknown transport tuple.
                    device.drop_staged_frame();
                }

                FrameAction::Dns | FrameAction::Passthrough => {
                    // ARP, ICMP, DNS (port 53), TCP data — smoltcp handles.
                    iface.poll_ingress_single(now, &mut device, &mut sockets);
                }
            }
        }

        // ── Phase 2: Ingress egress + maintenance ─────────────────────────
        // Flush frames generated by Phase 1 ingress (ACKs, SYN-ACKs, etc.)
        // before relaying data so smoltcp has up-to-date state.
        loop {
            let result = iface.poll_egress(now, &mut device, &mut sockets);
            if matches!(result, smoltcp::iface::PollResult::None) {
                break;
            }
        }
        iface.poll_maintenance(now);

        // Coalesced wake: if Phase 1/2 emitted any frames, wake the
        // NetWorker once instead of per-frame.
        if device.frames_emitted.swap(false, Ordering::Relaxed) {
            shared.rx_wake.wake();
        }

        // ── Phase 3: Service connections + relay data ────────────────────
        // Relay proxy data INTO smoltcp sockets first, then a single egress
        // pass flushes everything. This eliminates the former "Phase 2b"
        // double-egress pattern.
        conn_tracker.relay_data(&mut sockets);
        dns_interceptor.process(&mut sockets);

        // Accept queued inbound connections from published port listeners.
        port_publisher.accept_inbound(&mut iface, &mut sockets, &shared, &tokio_handle);
        port_publisher.relay_data(&mut sockets);

        // Detect newly-established connections and spawn proxy tasks.
        let new_conns = conn_tracker.take_new_connections(&mut sockets);
        for conn in new_conns {
            if let Some(ref tls_state) = tls_state
                && tls_state
                    .config
                    .intercepted_ports
                    .contains(&conn.dst.port())
            {
                // TLS-intercepted port — spawn TLS MITM proxy.
                let connect_dst = resolve_host_dst(conn.dst, config.gateway);
                tls_proxy::spawn_tls_proxy(
                    &tokio_handle,
                    conn.dst,
                    connect_dst,
                    conn.from_smoltcp,
                    conn.to_smoltcp,
                    shared.clone(),
                    tls_state.clone(),
                    network_policy.clone(),
                    conn.proxy_connect,
                );
                continue;
            }
            if conn.dst.port() == 53 {
                // DNS proxies have no guest-visible
                // "upstream-unreachable" failure mode — even an
                // upstream DNS failure yields SERVFAIL responses
                // rather than a silently-closed connection. Mark the
                // connection as connected so normal task exit
                // produces FIN, not RST.
                conn.proxy_connect.mark_connected();

                // DNS over TCP: route through the same forwarder the UDP
                // path uses. The forwarder applies the domain block list
                // and rebind protection to every query and routes
                // upstream based on `conn.dst.ip()` — the configured
                // upstream for queries to the gateway, direct forward
                // to the chosen `@target` (subject to egress policy)
                // otherwise. No gateway→loopback rewrite here: the
                // forwarder dials the configured upstream, not the
                // gateway.
                DnsTcpProxy::spawn(
                    &tokio_handle,
                    conn.dst,
                    conn.from_smoltcp,
                    conn.to_smoltcp,
                    dns_forwarder_handle.clone(),
                    shared.clone(),
                );
                continue;
            }
            if conn.dst.port() == 853
                && let Some(ref tls_state) = tls_state
            {
                // Same "always upstream-connected" reasoning as plain DNS over TCP.
                conn.proxy_connect.mark_connected();

                // DNS over TLS: terminate TLS at the gateway with a
                // per-domain cert, hand the inner DNS frames to the
                // same forwarder plain DNS uses. Policy for the
                // chosen `@target` resolver is applied per-query by
                // the forwarder (block list + rebind + egress).
                DotProxy::spawn(
                    &tokio_handle,
                    conn.dst,
                    conn.from_smoltcp,
                    conn.to_smoltcp,
                    dns_forwarder_handle.clone(),
                    tls_state.clone(),
                    shared.clone(),
                );
                continue;
            }
            // Plain TCP proxy.
            let connect_dst = resolve_host_dst(conn.dst, config.gateway);
            proxy::spawn_tcp_proxy(
                &tokio_handle,
                conn.dst,
                connect_dst,
                conn.from_smoltcp,
                conn.to_smoltcp,
                shared.clone(),
                network_policy.clone(),
                // Load the current snapshot per connection so live secret
                // updates apply to traffic the guest starts afterwards.
                secrets.load(),
                tls_state.clone(),
                conn.proxy_connect,
            );
        }

        // Rate-limited cleanup: TIME_WAIT is 60s, session timeout is 60s,
        // so checking once per second is more than sufficient.
        if last_cleanup.elapsed() >= std::time::Duration::from_secs(1) {
            conn_tracker.cleanup_closed(&mut sockets);
            port_publisher.cleanup_closed(&mut sockets);
            udp_relay.cleanup_expired();
            udp_fragments.cleanup_expired();
            ipv6_udp_fragments.cleanup_expired();
            shared.cleanup_resolved_hostnames();
            last_cleanup = std::time::Instant::now();
        }

        // ── Phase 4: Flush relay data + sleep ────────────────────────────
        // Single egress pass flushes all data written by Phase 3.
        loop {
            let result = iface.poll_egress(now, &mut device, &mut sockets);
            if matches!(result, smoltcp::iface::PollResult::None) {
                break;
            }
        }

        // Coalesced wake: if Phase 3/4 emitted any frames, wake once.
        if device.frames_emitted.swap(false, Ordering::Relaxed) {
            shared.rx_wake.wake();
        }

        let timeout_ms = iface
            .poll_delay(now, &sockets)
            .map(|d| d.total_millis().min(i32::MAX as u64) as i32)
            .unwrap_or(100); // 100ms fallback when no timers pending.

        #[cfg(unix)]
        sleep_until_stack_wake(&shared, timeout_ms, &mut poll_fds);
        #[cfg(windows)]
        sleep_until_stack_wake_windows(&shared, timeout_ms, &wait_context);
    }
}

//--------------------------------------------------------------------------------------------------
// Functions: Helpers
//--------------------------------------------------------------------------------------------------

#[cfg(unix)]
fn sleep_until_stack_wake(shared: &SharedState, timeout_ms: i32, poll_fds: &mut [libc::pollfd; 2]) {
    // SAFETY: poll_fds is a valid array of pollfd structs with valid fds.
    unsafe {
        libc::poll(
            poll_fds.as_mut_ptr(),
            poll_fds.len() as libc::nfds_t,
            timeout_ms,
        );
    }

    if poll_fds[0].revents & libc::POLLIN != 0 {
        shared.tx_wake.drain();
    }
    if poll_fds[1].revents & libc::POLLIN != 0 {
        shared.proxy_wake.drain();
    }
}

#[cfg(windows)]
fn windows_stack_wait_context(shared: &SharedState) -> std::io::Result<WaitContext> {
    let mut context = WaitContext::new();
    context.add(
        EventSource::waitable_handle(shared.tx_wake.as_raw_handle(), TX_WAKE_TOKEN),
        EventSet::IN,
    )?;
    context.add(
        EventSource::waitable_handle(shared.proxy_wake.as_raw_handle(), PROXY_WAKE_TOKEN),
        EventSet::IN,
    )?;
    Ok(context)
}

#[cfg(windows)]
fn sleep_until_stack_wake_windows(
    shared: &SharedState,
    timeout_ms: i32,
    wait_context: &WaitContext,
) {
    let mut events = [WaitEvent::default(); 2];
    let count = match wait_context.wait(timeout_ms, &mut events) {
        Ok(count) => count,
        Err(err) => {
            tracing::warn!(error = %err, "network poll loop: wait failed");
            return;
        }
    };

    for event in events.iter().take(count) {
        match event.token() {
            TX_WAKE_TOKEN => shared.tx_wake.drain(),
            PROXY_WAKE_TOKEN => shared.proxy_wake.drain(),
            token => tracing::warn!(token, "network poll loop: unknown wake token"),
        }
    }
}

/// Apply the common non-DNS UDP dispatch path to a complete guest datagram.
#[allow(clippy::too_many_arguments)]
fn relay_udp_frame(
    frame: &[u8],
    src: SocketAddr,
    dst: SocketAddr,
    config: &PollLoopConfig,
    network_policy: &NetworkPolicy,
    platform_policy: Option<&NetworkPolicy>,
    shared: &Arc<SharedState>,
    port_publisher: &mut PortPublisher,
    tls_state: Option<&TlsState>,
    udp_relay: &mut UdpRelay,
) {
    if port_publisher.relay_udp_outbound(frame, src, dst) {
        return;
    }

    // QUIC blocking: drop UDP to intercepted ports when TLS interception is active.
    if let Some(tls) = tls_state
        && tls.config.intercepted_ports.contains(&dst.port())
        && tls.config.block_quic_on_intercept
    {
        return;
    }

    match DnsPortType::from_udp(dst.port()) {
        // Dns: unreachable here — classify_transport routes UDP/53 to
        // FrameAction::Dns, not UdpRelay. Defensive drop covers regressions.
        DnsPortType::Dns | DnsPortType::EncryptedDns => return,
        // Alternative DNS protocols on well-known UDP ports are dropped —
        // forces fall-back to UDP/53.
        DnsPortType::AlternativeDns => {
            tracing::debug!(%dst, "alternative-DNS UDP port dropped; stub should fall back to UDP/53");
            return;
        }
        DnsPortType::Other => {}
    }

    // Policy is applied after reassembly, when the UDP destination port is known.
    if platform_policy
        .is_some_and(|policy| policy.evaluate_egress(dst, Protocol::Udp, shared).is_deny())
        || network_policy
            .evaluate_egress(dst, Protocol::Udp, shared)
            .is_deny()
    {
        return;
    }

    // Resolve the host-side destination for the dial. `dst` stays unchanged so
    // reply frames are stamped with the IP the guest expects.
    let host_dst = resolve_host_dst(dst, config.gateway);
    udp_relay.relay_outbound(frame, src, dst, host_dst);
}

/// Dispatch a complete datagram produced by fragment reassembly.
#[allow(clippy::too_many_arguments)]
fn handle_reassembled_udp_datagram(
    datagram: ReassembledUdpDatagram,
    device: &mut SmoltcpDevice,
    iface: &mut Interface,
    now: Instant,
    sockets: &mut SocketSet<'_>,
    config: &PollLoopConfig,
    network_policy: &NetworkPolicy,
    platform_policy: Option<&NetworkPolicy>,
    shared: &Arc<SharedState>,
    port_publisher: &mut PortPublisher,
    tls_state: Option<&TlsState>,
    udp_relay: &mut UdpRelay,
) {
    if DnsPortType::from_udp(datagram.dst.port()) == DnsPortType::Dns {
        device.replace_staged_frame(datagram.frame);
        iface.poll_ingress_single(now, device, sockets);
        return;
    }

    relay_udp_frame(
        &datagram.frame,
        datagram.src,
        datagram.dst,
        config,
        network_policy,
        platform_policy,
        shared,
        port_publisher,
        tls_state,
        udp_relay,
    );
    device.drop_staged_frame();
}

/// Map a guest-wire destination to its host-socket equivalent.
///
/// Gateway IPs rewrite to loopback (`127.0.0.1` / `::1`); everything else
/// passes through. Shared by the TCP proxy dispatch and the UDP relay.
///
/// # Arguments
///
/// * `dst` - Destination from the guest's packet.
/// * `gateway` - Per-sandbox gateway IPs that trigger the loopback rewrite.
pub(crate) fn resolve_host_dst(dst: SocketAddr, gateway: GatewayIps) -> SocketAddr {
    match dst.ip() {
        IpAddr::V4(v4) if gateway.ipv4 == Some(v4) => {
            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), dst.port())
        }
        IpAddr::V6(v6) if gateway.ipv6 == Some(v6) => {
            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), dst.port())
        }
        _ => dst,
    }
}

/// Get the current time as a smoltcp [`Instant`] using a monotonic clock.
///
/// Uses `std::time::Instant` (monotonic) instead of `SystemTime` (wall
/// clock) to avoid issues with NTP clock step corrections that could
/// cause smoltcp timers to misbehave.
fn smoltcp_now() -> Instant {
    static EPOCH: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
    let epoch = EPOCH.get_or_init(std::time::Instant::now);
    let elapsed = epoch.elapsed();
    Instant::from_millis(elapsed.as_millis() as i64)
}

/// Reply locally to ICMP echo requests aimed at the sandbox gateway.
///
/// `any_ip` is required so smoltcp accepts guest traffic for arbitrary remote
/// destinations, but that would make smoltcp's automatic ICMP echo replies
/// spoof remote hosts. Handle only the real gateway IPs here and leave all
/// other ICMP traffic untouched.
fn handle_gateway_icmp_echo(
    frame: &[u8],
    config: &PollLoopConfig,
    shared: &SharedState,
    network_policy: &NetworkPolicy,
    platform_policy: Option<&NetworkPolicy>,
) -> bool {
    let Ok(eth) = EthernetFrame::new_checked(frame) else {
        return false;
    };

    let reply = match eth.ethertype() {
        EthernetProtocol::Ipv4 => gateway_icmpv4_echo_reply(&eth, config),
        EthernetProtocol::Ipv6 => gateway_icmpv6_echo_reply(&eth, config),
        _ => None,
    };
    let Some(reply) = reply else {
        return false;
    };

    // This path bypasses the normal proxy evaluators, so it must independently
    // satisfy both the tenant policy and the optional host-owned policy floor.
    let tenant_denied = network_policy
        .evaluate_egress_ip(reply.dst, reply.protocol, shared)
        .is_deny();
    let platform_denied = platform_policy.is_some_and(|policy| {
        policy
            .evaluate_egress_ip(reply.dst, reply.protocol, shared)
            .is_deny()
    });
    if tenant_denied || platform_denied {
        tracing::debug!(
            dst = %reply.dst,
            tenant_denied,
            platform_denied,
            "gateway ICMP echo denied by policy",
        );
        return true;
    }

    shared.push_rx_frame_and_wake(reply.frame);

    true
}

/// Build an IPv4 ICMP echo reply when the guest pings the gateway IPv4.
fn gateway_icmpv4_echo_reply(
    eth: &EthernetFrame<&[u8]>,
    config: &PollLoopConfig,
) -> Option<GatewayIcmpReply> {
    let gateway_ipv4 = config.gateway.ipv4?;
    let ipv4 = Ipv4Packet::new_checked(eth.payload()).ok()?;
    if ipv4.dst_addr() != gateway_ipv4 || ipv4.next_header() != IpProtocol::Icmp {
        return None;
    }

    let icmp = Icmpv4Packet::new_checked(ipv4.payload()).ok()?;
    let Icmpv4Repr::EchoRequest {
        ident,
        seq_no,
        data,
    } = Icmpv4Repr::parse(&icmp, &smoltcp::phy::ChecksumCapabilities::default()).ok()?
    else {
        return None;
    };

    let ipv4_repr = Ipv4Repr {
        src_addr: gateway_ipv4,
        dst_addr: ipv4.src_addr(),
        next_header: IpProtocol::Icmp,
        payload_len: 8 + data.len(),
        hop_limit: 64,
    };
    let icmp_repr = Icmpv4Repr::EchoReply {
        ident,
        seq_no,
        data,
    };
    let mut reply = vec![0u8; 14 + ipv4_repr.buffer_len() + icmp_repr.buffer_len()];

    let mut reply_eth = EthernetFrame::new_unchecked(&mut reply);
    reply_eth.set_src_addr(EthernetAddress(config.gateway_mac));
    reply_eth.set_dst_addr(eth.src_addr());
    reply_eth.set_ethertype(EthernetProtocol::Ipv4);

    ipv4_repr.emit(
        &mut Ipv4Packet::new_unchecked(&mut reply[14..34]),
        &smoltcp::phy::ChecksumCapabilities::default(),
    );
    icmp_repr.emit(
        &mut Icmpv4Packet::new_unchecked(&mut reply[34..]),
        &smoltcp::phy::ChecksumCapabilities::default(),
    );

    Some(GatewayIcmpReply {
        dst: IpAddr::V4(gateway_ipv4),
        protocol: Protocol::Icmpv4,
        frame: reply,
    })
}

/// Build an IPv6 ICMP echo reply when the guest pings the gateway IPv6.
fn gateway_icmpv6_echo_reply(
    eth: &EthernetFrame<&[u8]>,
    config: &PollLoopConfig,
) -> Option<GatewayIcmpReply> {
    let gateway_ipv6 = config.gateway.ipv6?;
    let ipv6 = Ipv6Packet::new_checked(eth.payload()).ok()?;
    if ipv6.dst_addr() != gateway_ipv6 || ipv6.next_header() != IpProtocol::Icmpv6 {
        return None;
    }

    let icmp = Icmpv6Packet::new_checked(ipv6.payload()).ok()?;
    let Icmpv6Repr::EchoRequest {
        ident,
        seq_no,
        data,
    } = Icmpv6Repr::parse(
        &ipv6.src_addr(),
        &ipv6.dst_addr(),
        &icmp,
        &smoltcp::phy::ChecksumCapabilities::default(),
    )
    .ok()?
    else {
        return None;
    };

    let ipv6_repr = Ipv6Repr {
        src_addr: gateway_ipv6,
        dst_addr: ipv6.src_addr(),
        next_header: IpProtocol::Icmpv6,
        payload_len: icmp_repr_buffer_len_v6(data),
        hop_limit: 64,
    };
    let icmp_repr = Icmpv6Repr::EchoReply {
        ident,
        seq_no,
        data,
    };
    let ipv6_hdr_len = 40;
    let mut reply = vec![0u8; 14 + ipv6_hdr_len + icmp_repr.buffer_len()];

    let mut reply_eth = EthernetFrame::new_unchecked(&mut reply);
    reply_eth.set_src_addr(EthernetAddress(config.gateway_mac));
    reply_eth.set_dst_addr(eth.src_addr());
    reply_eth.set_ethertype(EthernetProtocol::Ipv6);

    ipv6_repr.emit(&mut Ipv6Packet::new_unchecked(&mut reply[14..54]));
    icmp_repr.emit(
        &gateway_ipv6,
        &ipv6.src_addr(),
        &mut Icmpv6Packet::new_unchecked(&mut reply[54..]),
        &smoltcp::phy::ChecksumCapabilities::default(),
    );

    Some(GatewayIcmpReply {
        dst: IpAddr::V6(gateway_ipv6),
        protocol: Protocol::Icmpv6,
        frame: reply,
    })
}

fn icmp_repr_buffer_len_v6(data: &[u8]) -> usize {
    Icmpv6Repr::EchoReply {
        ident: 0,
        seq_no: 0,
        data,
    }
    .buffer_len()
}

/// Classify an IPv4 packet payload (after stripping the Ethernet header).
fn classify_ipv4(payload: &[u8]) -> FrameAction {
    let Ok(ipv4) = Ipv4Packet::new_checked(payload) else {
        return FrameAction::Passthrough;
    };
    if is_ipv4_udp_fragment(&ipv4) {
        return FrameAction::Ipv4UdpFragment;
    }
    classify_transport(
        ipv4.next_header(),
        ipv4.src_addr().into(),
        ipv4.dst_addr().into(),
        ipv4.payload(),
    )
}

/// Classify an IPv6 packet payload (after stripping the Ethernet header).
fn classify_ipv6(payload: &[u8]) -> FrameAction {
    let Ok(ipv6) = Ipv6Packet::new_checked(payload) else {
        return FrameAction::Passthrough;
    };
    if is_ipv6_udp_fragment(&ipv6) {
        return FrameAction::Ipv6UdpFragment;
    }
    if is_ipv6_fragment(&ipv6) {
        return FrameAction::Ipv6UnsupportedFragment;
    }
    classify_transport(
        ipv6.next_header(),
        ipv6.src_addr().into(),
        ipv6.dst_addr().into(),
        ipv6.payload(),
    )
}

/// Classify the transport-layer protocol (shared by IPv4 and IPv6).
fn classify_transport(
    protocol: IpProtocol,
    src_ip: std::net::IpAddr,
    dst_ip: std::net::IpAddr,
    transport_payload: &[u8],
) -> FrameAction {
    match protocol {
        IpProtocol::Tcp => {
            let Ok(tcp) = TcpPacket::new_checked(transport_payload) else {
                return FrameAction::Passthrough;
            };
            if tcp.syn() && !tcp.ack() {
                FrameAction::TcpSyn {
                    src: SocketAddr::new(src_ip, tcp.src_port()),
                    dst: SocketAddr::new(dst_ip, tcp.dst_port()),
                }
            } else {
                FrameAction::Passthrough
            }
        }
        IpProtocol::Udp => {
            let Ok(udp) = UdpPacket::new_checked(transport_payload) else {
                return FrameAction::Passthrough;
            };
            // The plain-DNS port (UDP/53) lives in dns::common::ports so
            // the alternative-DNS refusal logic and this dispatcher
            // share one source of truth for "which UDP ports are DNS".
            if DnsPortType::from_udp(udp.dst_port()) == DnsPortType::Dns {
                FrameAction::Dns
            } else {
                FrameAction::UdpRelay {
                    src: SocketAddr::new(src_ip, udp.src_port()),
                    dst: SocketAddr::new(dst_ip, udp.dst_port()),
                }
            }
        }
        _ => FrameAction::Passthrough, // ICMP, etc.
    }
}

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

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

    use smoltcp::phy::ChecksumCapabilities;
    use smoltcp::wire::{
        ArpOperation, ArpPacket, ArpRepr, EthernetRepr, Icmpv4Packet, Icmpv4Repr, Ipv4Repr,
    };

    use crate::device::SmoltcpDevice;
    use crate::shared::SharedState;

    /// Build a minimal Ethernet + IPv4 + TCP SYN frame.
    fn build_tcp_syn_frame(
        src_ip: [u8; 4],
        dst_ip: [u8; 4],
        src_port: u16,
        dst_port: u16,
    ) -> Vec<u8> {
        let mut frame = vec![0u8; 14 + 20 + 20]; // eth + ipv4 + tcp

        // Ethernet header.
        frame[12] = 0x08; // EtherType: IPv4
        frame[13] = 0x00;

        // IPv4 header.
        let ip = &mut frame[14..34];
        ip[0] = 0x45; // Version + IHL
        let total_len = 40u16; // 20 (IP) + 20 (TCP)
        ip[2..4].copy_from_slice(&total_len.to_be_bytes());
        ip[6] = 0x40; // Don't Fragment
        ip[8] = 64; // TTL
        ip[9] = 6; // Protocol: TCP
        ip[12..16].copy_from_slice(&src_ip);
        ip[16..20].copy_from_slice(&dst_ip);

        // TCP header.
        let tcp = &mut frame[34..54];
        tcp[0..2].copy_from_slice(&src_port.to_be_bytes());
        tcp[2..4].copy_from_slice(&dst_port.to_be_bytes());
        tcp[12] = 0x50; // Data offset: 5 words
        tcp[13] = 0x02; // SYN flag

        frame
    }

    /// Build a minimal Ethernet + IPv4 + UDP frame.
    fn build_udp_frame(src_ip: [u8; 4], dst_ip: [u8; 4], src_port: u16, dst_port: u16) -> Vec<u8> {
        let mut frame = vec![0u8; 14 + 20 + 8]; // eth + ipv4 + udp

        // Ethernet header.
        frame[12] = 0x08;
        frame[13] = 0x00;

        // IPv4 header.
        let ip = &mut frame[14..34];
        ip[0] = 0x45;
        let total_len = 28u16; // 20 (IP) + 8 (UDP)
        ip[2..4].copy_from_slice(&total_len.to_be_bytes());
        ip[8] = 64;
        ip[9] = 17; // Protocol: UDP
        ip[12..16].copy_from_slice(&src_ip);
        ip[16..20].copy_from_slice(&dst_ip);

        // UDP header.
        let udp = &mut frame[34..42];
        udp[0..2].copy_from_slice(&src_port.to_be_bytes());
        udp[2..4].copy_from_slice(&dst_port.to_be_bytes());
        let udp_len = 8u16;
        udp[4..6].copy_from_slice(&udp_len.to_be_bytes());

        frame
    }

    /// Build a minimal Ethernet + IPv4 + ICMP echo request frame.
    fn build_icmpv4_echo_frame(
        src_mac: [u8; 6],
        dst_mac: [u8; 6],
        src_ip: [u8; 4],
        dst_ip: [u8; 4],
        ident: u16,
        seq_no: u16,
        data: &[u8],
    ) -> Vec<u8> {
        let ipv4_repr = Ipv4Repr {
            src_addr: Ipv4Addr::from(src_ip),
            dst_addr: Ipv4Addr::from(dst_ip),
            next_header: IpProtocol::Icmp,
            payload_len: 8 + data.len(),
            hop_limit: 64,
        };
        let icmp_repr = Icmpv4Repr::EchoRequest {
            ident,
            seq_no,
            data,
        };
        let frame_len = 14 + ipv4_repr.buffer_len() + icmp_repr.buffer_len();
        let mut frame = vec![0u8; frame_len];

        let mut eth_frame = EthernetFrame::new_unchecked(&mut frame);
        EthernetRepr {
            src_addr: EthernetAddress(src_mac),
            dst_addr: EthernetAddress(dst_mac),
            ethertype: EthernetProtocol::Ipv4,
        }
        .emit(&mut eth_frame);

        ipv4_repr.emit(
            &mut Ipv4Packet::new_unchecked(&mut frame[14..34]),
            &ChecksumCapabilities::default(),
        );
        icmp_repr.emit(
            &mut Icmpv4Packet::new_unchecked(&mut frame[34..]),
            &ChecksumCapabilities::default(),
        );

        frame
    }

    /// Build a minimal Ethernet + ARP request frame.
    fn build_arp_request_frame(src_mac: [u8; 6], src_ip: [u8; 4], target_ip: [u8; 4]) -> Vec<u8> {
        let mut frame = vec![0u8; 14 + 28];

        let mut eth_frame = EthernetFrame::new_unchecked(&mut frame);
        EthernetRepr {
            src_addr: EthernetAddress(src_mac),
            dst_addr: EthernetAddress([0xff; 6]),
            ethertype: EthernetProtocol::Arp,
        }
        .emit(&mut eth_frame);

        ArpRepr::EthernetIpv4 {
            operation: ArpOperation::Request,
            source_hardware_addr: EthernetAddress(src_mac),
            source_protocol_addr: Ipv4Addr::from(src_ip),
            target_hardware_addr: EthernetAddress([0x00; 6]),
            target_protocol_addr: Ipv4Addr::from(target_ip),
        }
        .emit(&mut ArpPacket::new_unchecked(&mut frame[14..]));

        frame
    }

    #[test]
    fn classify_tcp_syn() {
        let frame = build_tcp_syn_frame([10, 0, 0, 2], [93, 184, 216, 34], 54321, 443);
        match classify_frame(&frame) {
            FrameAction::TcpSyn { src, dst } => {
                assert_eq!(
                    src,
                    SocketAddr::new(Ipv4Addr::new(10, 0, 0, 2).into(), 54321)
                );
                assert_eq!(
                    dst,
                    SocketAddr::new(Ipv4Addr::new(93, 184, 216, 34).into(), 443)
                );
            }
            _ => panic!("expected TcpSyn"),
        }
    }

    #[test]
    fn classify_tcp_ack_is_passthrough() {
        let mut frame = build_tcp_syn_frame([10, 0, 0, 2], [93, 184, 216, 34], 54321, 443);
        // Change flags to ACK only (not SYN).
        frame[34 + 13] = 0x10; // ACK flag
        assert!(matches!(classify_frame(&frame), FrameAction::Passthrough));
    }

    #[test]
    fn classify_udp_dns() {
        let frame = build_udp_frame([10, 0, 0, 2], [10, 0, 0, 1], 12345, 53);
        assert!(matches!(classify_frame(&frame), FrameAction::Dns));
    }

    #[test]
    fn classify_udp_non_dns() {
        let frame = build_udp_frame([10, 0, 0, 2], [8, 8, 8, 8], 12345, 443);
        match classify_frame(&frame) {
            FrameAction::UdpRelay { src, dst } => {
                assert_eq!(src.port(), 12345);
                assert_eq!(dst.port(), 443);
            }
            _ => panic!("expected UdpRelay"),
        }
    }

    #[test]
    fn classify_ipv4_udp_fragment() {
        let mut frame = build_udp_frame([10, 0, 0, 2], [8, 8, 8, 8], 12345, 443);
        frame[14 + 6] = 0x20; // More Fragments flag.
        assert!(matches!(
            classify_frame(&frame),
            FrameAction::Ipv4UdpFragment
        ));
    }

    #[test]
    fn classify_ipv6_udp_fragment() {
        let mut frame = vec![0u8; 14 + 40 + 8];

        frame[12] = 0x86;
        frame[13] = 0xdd;

        let ip = &mut frame[14..54];
        ip[0] = 0x60;
        ip[4..6].copy_from_slice(&8u16.to_be_bytes());
        ip[6] = u8::from(IpProtocol::Ipv6Frag);
        ip[7] = 64;
        ip[8..24].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
        ip[24..40].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());

        let fragment = &mut frame[54..62];
        fragment[0] = u8::from(IpProtocol::Udp);
        fragment[3] = 1; // More Fragments flag.

        assert!(matches!(
            classify_frame(&frame),
            FrameAction::Ipv6UdpFragment
        ));
    }

    #[test]
    fn classify_ipv6_non_udp_fragment_is_unsupported() {
        let mut frame = vec![0u8; 14 + 40 + 8];

        frame[12] = 0x86;
        frame[13] = 0xdd;

        let ip = &mut frame[14..54];
        ip[0] = 0x60;
        ip[4..6].copy_from_slice(&8u16.to_be_bytes());
        ip[6] = u8::from(IpProtocol::Ipv6Frag);
        ip[7] = 64;
        ip[8..24].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
        ip[24..40].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());

        let fragment = &mut frame[54..62];
        fragment[0] = u8::from(IpProtocol::Tcp);
        fragment[3] = 1; // More Fragments flag.

        assert!(matches!(
            classify_frame(&frame),
            FrameAction::Ipv6UnsupportedFragment
        ));
    }

    #[test]
    fn classify_arp_is_passthrough() {
        let mut frame = vec![0u8; 42]; // ARP frame
        frame[12] = 0x08;
        frame[13] = 0x06; // EtherType: ARP
        assert!(matches!(classify_frame(&frame), FrameAction::Passthrough));
    }

    #[test]
    fn classify_garbage_is_passthrough() {
        assert!(matches!(classify_frame(&[]), FrameAction::Passthrough));
        assert!(matches!(classify_frame(&[0; 5]), FrameAction::Passthrough));
    }

    #[test]
    fn gateway_replies_to_icmp_echo_requests() {
        fn drive_one_frame(
            device: &mut SmoltcpDevice,
            iface: &mut Interface,
            sockets: &mut SocketSet<'_>,
            shared: &Arc<SharedState>,
            poll_config: &PollLoopConfig,
            now: Instant,
        ) {
            let frame = device.stage_next_frame().expect("expected staged frame");
            if handle_gateway_icmp_echo(
                frame,
                poll_config,
                shared,
                &NetworkPolicy::allow_all(),
                None,
            ) {
                device.drop_staged_frame();
                return;
            }
            let _ = iface.poll_ingress_single(now, device, sockets);
            let _ = iface.poll_egress(now, device, sockets);
        }

        let shared = Arc::new(SharedState::new(4));
        let poll_config = PollLoopConfig {
            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
            gateway: GatewayIps {
                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
                ipv6: Some(Ipv6Addr::LOCALHOST),
            },
            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
            guest_ipv6: None,
            mtu: 1500,
        };
        let guest_ipv4 = poll_config.guest_ipv4.unwrap();
        let gateway_ipv4 = poll_config.gateway.ipv4.unwrap();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let now = smoltcp_now();

        // Mirror the real guest flow: resolve the gateway MAC before sending
        // the ICMP echo request.
        shared
            .tx_ring
            .push(build_arp_request_frame(
                poll_config.guest_mac,
                guest_ipv4.octets(),
                gateway_ipv4.octets(),
            ))
            .unwrap();
        shared
            .tx_ring
            .push(build_icmpv4_echo_frame(
                poll_config.guest_mac,
                poll_config.gateway_mac,
                guest_ipv4.octets(),
                gateway_ipv4.octets(),
                0x1234,
                0xABCD,
                b"ping",
            ))
            .unwrap();

        drive_one_frame(
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            &poll_config,
            now,
        );
        let _ = shared.rx_ring.pop().expect("expected ARP reply");

        drive_one_frame(
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            &poll_config,
            now,
        );

        let reply = shared.rx_ring.pop().expect("expected ICMP echo reply");
        let eth = EthernetFrame::new_checked(&reply).expect("valid ethernet frame");
        assert_eq!(eth.src_addr(), EthernetAddress(poll_config.gateway_mac));
        assert_eq!(eth.dst_addr(), EthernetAddress(poll_config.guest_mac));
        assert_eq!(eth.ethertype(), EthernetProtocol::Ipv4);

        let ipv4 = Ipv4Packet::new_checked(eth.payload()).expect("valid IPv4 packet");
        assert_eq!(ipv4.src_addr(), gateway_ipv4);
        assert_eq!(ipv4.dst_addr(), guest_ipv4);
        assert_eq!(ipv4.next_header(), IpProtocol::Icmp);

        let icmp = Icmpv4Packet::new_checked(ipv4.payload()).expect("valid ICMP packet");
        let icmp_repr = Icmpv4Repr::parse(&icmp, &ChecksumCapabilities::default())
            .expect("valid ICMP echo reply");
        assert_eq!(
            icmp_repr,
            Icmpv4Repr::EchoReply {
                ident: 0x1234,
                seq_no: 0xABCD,
                data: b"ping",
            }
        );
    }

    #[test]
    fn gateway_icmp_echo_respects_deny_policy() {
        let shared = SharedState::new(4);
        let poll_config = PollLoopConfig {
            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
            gateway: GatewayIps {
                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
                ipv6: None,
            },
            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
            guest_ipv6: None,
            mtu: 1500,
        };
        let policy = NetworkPolicy::builder().default_deny().build().unwrap();
        let frame = build_icmpv4_echo_frame(
            poll_config.guest_mac,
            poll_config.gateway_mac,
            poll_config.guest_ipv4.unwrap().octets(),
            poll_config.gateway.ipv4.unwrap().octets(),
            0x1234,
            0xABCD,
            b"ping",
        );

        assert!(handle_gateway_icmp_echo(
            &frame,
            &poll_config,
            &shared,
            &policy,
            None,
        ));
        assert!(
            shared.rx_ring.pop().is_none(),
            "denied gateway ICMP should not queue a reply"
        );
    }

    #[test]
    fn platform_public_floor_consumes_gateway_echo_without_replying() {
        let shared = SharedState::new(4);
        let gateway = Ipv4Addr::new(100, 96, 0, 1);
        let guest = Ipv4Addr::new(100, 96, 0, 2);
        shared.set_gateway_ips(Some(gateway), None);
        let config = PollLoopConfig {
            gateway_mac: [0x02, 0, 0, 0, 0, 1],
            guest_mac: [0x02, 0, 0, 0, 0, 2],
            gateway: GatewayIps {
                ipv4: Some(gateway),
                ipv6: None,
            },
            guest_ipv4: Some(guest),
            guest_ipv6: None,
            mtu: 1500,
        };
        let frame = build_icmpv4_echo_frame(
            config.guest_mac,
            config.gateway_mac,
            guest.octets(),
            gateway.octets(),
            1,
            1,
            b"ping",
        );
        let platform = NetworkPolicy::from_profiles([crate::policy::NetworkProfile::Public]);

        assert!(handle_gateway_icmp_echo(
            &frame,
            &config,
            &shared,
            &NetworkPolicy::allow_all(),
            Some(&platform),
        ));
        assert!(shared.rx_ring.pop().is_none());
    }

    fn test_gateway() -> GatewayIps {
        GatewayIps {
            ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
            ipv6: Some("fd42:6d73:62::1".parse().unwrap()),
        }
    }

    #[test]
    fn resolve_host_dst_matches_ipv4() {
        let gw = test_gateway();
        let dst = SocketAddr::new(IpAddr::V4(gw.ipv4.unwrap()), 8080);
        assert_eq!(
            resolve_host_dst(dst, gw),
            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)
        );
    }

    #[test]
    fn resolve_host_dst_matches_ipv6() {
        let gw = test_gateway();
        let dst = SocketAddr::new(IpAddr::V6(gw.ipv6.unwrap()), 8080);
        assert_eq!(
            resolve_host_dst(dst, gw),
            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 8080)
        );
    }

    #[test]
    fn resolve_host_dst_passes_through_when_family_absent() {
        let gw = GatewayIps {
            ipv4: None,
            ipv6: Some("fd42:6d73:62::1".parse().unwrap()),
        };
        // IPv4 dst with no IPv4 gateway must not be rewritten to loopback.
        let dst = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(100, 96, 0, 1)), 8080);
        assert_eq!(resolve_host_dst(dst, gw), dst);
    }

    #[test]
    fn resolve_host_dst_passes_through_non_gateway() {
        let gw = test_gateway();
        let dst = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 443);
        assert_eq!(resolve_host_dst(dst, gw), dst);
    }

    #[test]
    fn external_icmp_echo_requests_are_not_answered_locally() {
        fn drive_one_frame(
            device: &mut SmoltcpDevice,
            iface: &mut Interface,
            sockets: &mut SocketSet<'_>,
            shared: &Arc<SharedState>,
            poll_config: &PollLoopConfig,
            now: Instant,
        ) {
            let frame = device.stage_next_frame().expect("expected staged frame");
            if handle_gateway_icmp_echo(
                frame,
                poll_config,
                shared,
                &NetworkPolicy::allow_all(),
                None,
            ) {
                device.drop_staged_frame();
                return;
            }
            let _ = iface.poll_ingress_single(now, device, sockets);
            let _ = iface.poll_egress(now, device, sockets);
        }

        let shared = Arc::new(SharedState::new(4));
        let poll_config = PollLoopConfig {
            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
            gateway: GatewayIps {
                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
                ipv6: Some(Ipv6Addr::LOCALHOST),
            },
            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
            guest_ipv6: None,
            mtu: 1500,
        };
        let guest_ipv4 = poll_config.guest_ipv4.unwrap();
        let gateway_ipv4 = poll_config.gateway.ipv4.unwrap();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let now = smoltcp_now();

        shared
            .tx_ring
            .push(build_arp_request_frame(
                poll_config.guest_mac,
                guest_ipv4.octets(),
                gateway_ipv4.octets(),
            ))
            .unwrap();
        shared
            .tx_ring
            .push(build_icmpv4_echo_frame(
                poll_config.guest_mac,
                poll_config.gateway_mac,
                guest_ipv4.octets(),
                [142, 251, 216, 46],
                0x1234,
                0xABCD,
                b"ping",
            ))
            .unwrap();

        drive_one_frame(
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            &poll_config,
            now,
        );
        let _ = shared.rx_ring.pop().expect("expected ARP reply");

        drive_one_frame(
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            &poll_config,
            now,
        );
        assert!(
            shared.rx_ring.pop().is_none(),
            "external ICMP should not be answered locally"
        );
    }

    // ─────────────────────────────────────────────────────────────────────
    // Guest-initiated TCP teardown.
    //
    // These tests drive the real `ConnectionTracker` + smoltcp interface
    // through a full TCP handshake and then a guest-initiated teardown,
    // asserting the observable the proxy task sees on its channel:
    //   - guest FIN => half-close propagated (channel EOF), server → guest
    //     stays open, and the slot is reclaimed once the proxy task exits
    //   - guest RST => immediate clean teardown
    // Regression tests for the CLOSE_WAIT orphan leak: without close
    // propagation, a guest FIN left the socket in CLOSE_WAIT forever, the
    // proxy task blocked on `from_smoltcp.recv()`, the upstream socket
    // open, and the connection-table slot consumed until the table filled.
    // ─────────────────────────────────────────────────────────────────────

    use smoltcp::socket::tcp;
    use smoltcp::wire::{TcpControl, TcpPacket, TcpRepr, TcpSeqNumber};

    const GUEST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02];
    const GATEWAY_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01];
    const GUEST_IP: [u8; 4] = [100, 96, 0, 2];
    const GATEWAY_IP: [u8; 4] = [100, 96, 0, 1];
    // Off-subnet external destination reached via the default route.
    const SERVER_IP: [u8; 4] = [93, 184, 216, 34];

    fn leak_poll_config() -> PollLoopConfig {
        PollLoopConfig {
            gateway_mac: GATEWAY_MAC,
            guest_mac: GUEST_MAC,
            gateway: GatewayIps {
                ipv4: Some(Ipv4Addr::from(GATEWAY_IP)),
                ipv6: None,
            },
            guest_ipv4: Some(Ipv4Addr::from(GUEST_IP)),
            guest_ipv6: None,
            mtu: 1500,
        }
    }

    /// Build an Ethernet+IPv4+TCP frame from the guest with correct checksums.
    #[allow(clippy::too_many_arguments)]
    fn build_tcp_frame(
        src_port: u16,
        dst_port: u16,
        control: TcpControl,
        seq: i32,
        ack: Option<i32>,
        payload: &[u8],
    ) -> Vec<u8> {
        let src_ip = Ipv4Addr::from(GUEST_IP);
        let dst_ip = Ipv4Addr::from(SERVER_IP);

        let tcp_repr = TcpRepr {
            src_port,
            dst_port,
            control,
            seq_number: TcpSeqNumber(seq),
            ack_number: ack.map(TcpSeqNumber),
            window_len: 65535,
            window_scale: None,
            max_seg_size: None,
            sack_permitted: false,
            sack_ranges: [None, None, None],
            timestamp: None,
            payload,
        };
        let ipv4_repr = Ipv4Repr {
            src_addr: src_ip,
            dst_addr: dst_ip,
            next_header: IpProtocol::Tcp,
            payload_len: tcp_repr.buffer_len(),
            hop_limit: 64,
        };

        let frame_len = 14 + ipv4_repr.buffer_len() + tcp_repr.buffer_len();
        let mut frame = vec![0u8; frame_len];

        let mut eth = EthernetFrame::new_unchecked(&mut frame);
        EthernetRepr {
            src_addr: EthernetAddress(GUEST_MAC),
            dst_addr: EthernetAddress(GATEWAY_MAC),
            ethertype: EthernetProtocol::Ipv4,
        }
        .emit(&mut eth);

        let ip_end = 14 + ipv4_repr.buffer_len();
        ipv4_repr.emit(
            &mut Ipv4Packet::new_unchecked(&mut frame[14..ip_end]),
            &ChecksumCapabilities::default(),
        );
        tcp_repr.emit(
            &mut TcpPacket::new_unchecked(&mut frame[ip_end..]),
            &IpAddress::Ipv4(src_ip),
            &IpAddress::Ipv4(dst_ip),
            &ChecksumCapabilities::default(),
        );

        frame
    }

    /// Push one guest frame, run a single ingress pass, then drain egress.
    fn ingress(
        frame: Vec<u8>,
        device: &mut SmoltcpDevice,
        iface: &mut Interface,
        sockets: &mut SocketSet<'_>,
        shared: &Arc<SharedState>,
        now: Instant,
    ) {
        shared.tx_ring.push(frame).unwrap();
        device.stage_next_frame().expect("frame should stage");
        iface.poll_ingress_single(now, device, sockets);
        loop {
            let r = iface.poll_egress(now, device, sockets);
            if matches!(r, smoltcp::iface::PollResult::None) {
                break;
            }
        }
    }

    /// Pop the newest smoltcp→guest reply and return its (seq, ack, is_syn,
    /// is_fin, is_rst). Drains all queued replies, returning the last TCP one.
    fn last_tcp_reply(shared: &Arc<SharedState>) -> Option<(i32, i32, bool, bool, bool)> {
        let mut out = None;
        while let Some(frame) = shared.rx_ring.pop() {
            if frame.len() < 34 {
                continue;
            }
            // eth(14) + ipv4(20) then TCP.
            if frame[23] != 6 {
                continue; // not TCP (e.g. ARP reply passes as non-IPv4)
            }
            let tcp = match TcpPacket::new_checked(&frame[34..]) {
                Ok(p) => p,
                Err(_) => continue,
            };
            out = Some((
                tcp.seq_number().0,
                tcp.ack_number().0,
                tcp.syn(),
                tcp.fin(),
                tcp.rst(),
            ));
        }
        out
    }

    /// Find the single tracked TCP socket's state.
    fn only_tcp_state(sockets: &SocketSet<'_>) -> Option<tcp::State> {
        for (_h, sock) in sockets.iter() {
            if let smoltcp::socket::Socket::Tcp(s) = sock {
                return Some(s.state());
            }
        }
        None
    }

    /// Drive guest→server handshake to ESTABLISHED, returning (server_isn,
    /// guest_seq_after_handshake).
    fn handshake(
        tracker: &mut ConnectionTracker,
        device: &mut SmoltcpDevice,
        iface: &mut Interface,
        sockets: &mut SocketSet<'_>,
        shared: &Arc<SharedState>,
        now: Instant,
        guest_port: u16,
    ) -> (i32, i32) {
        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), guest_port);
        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);

        // Pre-populate neighbor cache so smoltcp can address replies to the
        // guest without stalling on ARP.
        ingress(
            build_arp_request_frame(GUEST_MAC, GUEST_IP, GATEWAY_IP),
            device,
            iface,
            sockets,
            shared,
            now,
        );
        let _ = shared.rx_ring.pop(); // ARP reply

        let guest_isn = 1000i32;

        // 1) Guest SYN — tracker creates the listening socket first (as the
        //    real poll loop does), then smoltcp completes the handshake.
        assert!(
            tracker.create_tcp_socket(src, dst, sockets),
            "socket creation should succeed under the limit"
        );
        ingress(
            build_tcp_frame(guest_port, 443, TcpControl::Syn, guest_isn, None, &[]),
            device,
            iface,
            sockets,
            shared,
            now,
        );
        let (server_isn, ack, is_syn, _, _) =
            last_tcp_reply(shared).expect("expected SYN-ACK from smoltcp");
        assert!(is_syn, "expected SYN flag on handshake reply");
        assert_eq!(ack, guest_isn + 1, "SYN-ACK should ack guest ISN+1");

        // 2) Guest ACK — completes the handshake.
        ingress(
            build_tcp_frame(
                guest_port,
                443,
                TcpControl::None,
                guest_isn + 1,
                Some(server_isn + 1),
                &[],
            ),
            device,
            iface,
            sockets,
            shared,
            now,
        );
        assert_eq!(
            only_tcp_state(sockets),
            Some(tcp::State::Established),
            "socket should be ESTABLISHED after handshake",
        );

        (server_isn, guest_isn + 1)
    }

    /// Complete a guest→server handshake and hand the connection to a proxy.
    fn establish(
        tracker: &mut ConnectionTracker,
        device: &mut SmoltcpDevice,
        iface: &mut Interface,
        sockets: &mut SocketSet<'_>,
        shared: &Arc<SharedState>,
        now: Instant,
        guest_port: u16,
    ) -> (i32, i32, Vec<crate::conn::NewConnection>) {
        let (server_isn, guest_seq) =
            handshake(tracker, device, iface, sockets, shared, now, guest_port);

        // Poll loop detects the established connection and spawns a proxy.
        let new_conns = tracker.take_new_connections(sockets);
        assert_eq!(
            new_conns.len(),
            1,
            "one new connection should be handed off"
        );

        (server_isn, guest_seq, new_conns)
    }

    #[test]
    fn guest_fin_propagates_half_close_without_killing_the_connection() {
        let shared = Arc::new(SharedState::new(64));
        let poll_config = leak_poll_config();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let mut tracker = ConnectionTracker::new(None);
        let now = smoltcp_now();

        let (server_isn, guest_seq, mut new_conns) = establish(
            &mut tracker,
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
            54321,
        );
        // The proxy side of the channel: a real proxy task blocks on
        // `from_smoltcp.recv()` while the connection is idle.
        let conn = new_conns.remove(0);
        let mut from_smoltcp = conn.from_smoltcp;
        let to_smoltcp = conn.to_smoltcp;

        // Guest half-closes (shutdown(SHUT_WR), or its process exits while
        // holding an idle keep-alive: no unread data => FIN, not RST).
        ingress(
            build_tcp_frame(
                54321,
                443,
                TcpControl::Fin,
                guest_seq,
                Some(server_isn + 1),
                &[],
            ),
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
        );
        assert_eq!(
            only_tcp_state(&sockets),
            Some(tcp::State::CloseWait),
            "guest FIN should move the smoltcp socket to CLOSE_WAIT",
        );

        // One relay pass propagates the half-close: the proxy task's
        // receiver disconnects, so its `recv()` returns `None` and it can
        // shut down the guest → server direction upstream.
        tracker.relay_data(&mut sockets);
        assert!(
            matches!(
                from_smoltcp.try_recv(),
                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected)
            ),
            "guest FIN must propagate EOF to the proxy task",
        );

        // The connection must NOT be torn down: a half-closed guest can
        // still receive. Pending server → guest data is still delivered.
        assert!(
            tracker.has_socket_for(
                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54321),
                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
            ),
            "half-closed connection must stay tracked while the proxy runs",
        );
        let payload = b"pending server response";
        to_smoltcp
            .try_send(bytes::Bytes::from_static(payload))
            .expect("server → guest channel should accept data");
        while shared.rx_ring.pop().is_some() {} // drain handshake/ACK frames
        tracker.relay_data(&mut sockets);
        loop {
            let r = iface.poll_egress(now, &mut device, &mut sockets);
            if matches!(r, smoltcp::iface::PollResult::None) {
                break;
            }
        }
        let mut delivered = false;
        while let Some(frame) = shared.rx_ring.pop() {
            if frame.windows(payload.len()).any(|w| w == payload) {
                delivered = true;
            }
        }
        assert!(
            delivered,
            "server data must still reach a half-closed guest",
        );
        assert_eq!(
            only_tcp_state(&sockets),
            Some(tcp::State::CloseWait),
            "socket must stay open (CLOSE_WAIT) while the proxy is alive",
        );
    }

    #[test]
    fn guest_fin_connection_is_reaped_after_proxy_exit() {
        let shared = Arc::new(SharedState::new(64));
        let poll_config = leak_poll_config();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let mut tracker = ConnectionTracker::new(None);
        let now = smoltcp_now();

        let (server_isn, guest_seq, mut new_conns) = establish(
            &mut tracker,
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
            54321,
        );
        let conn = new_conns.remove(0);
        let from_smoltcp = conn.from_smoltcp;
        let to_smoltcp = conn.to_smoltcp;

        // Guest FIN, then a relay pass to propagate the half-close.
        ingress(
            build_tcp_frame(
                54321,
                443,
                TcpControl::Fin,
                guest_seq,
                Some(server_isn + 1),
                &[],
            ),
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
        );
        tracker.relay_data(&mut sockets);

        // The proxy task exits (upstream closed after seeing our FIN) and
        // drops its channel ends.
        drop(from_smoltcp);
        drop(to_smoltcp);

        // The tracker detects the proxy exit and closes the socket: the
        // guest gets our FIN (CLOSE_WAIT → LAST_ACK)...
        while shared.rx_ring.pop().is_some() {} // drain handshake/ACK frames
        tracker.relay_data(&mut sockets);
        loop {
            let r = iface.poll_egress(now, &mut device, &mut sockets);
            if matches!(r, smoltcp::iface::PollResult::None) {
                break;
            }
        }
        let (fin_seq, _, _, is_fin, _) =
            last_tcp_reply(&shared).expect("expected FIN toward the guest");
        assert!(is_fin, "proxy exit after guest FIN must FIN the guest side");

        // ...and the guest's final ACK completes the close.
        ingress(
            build_tcp_frame(
                54321,
                443,
                TcpControl::None,
                guest_seq + 1,
                Some(fin_seq + 1),
                &[],
            ),
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
        );
        tracker.relay_data(&mut sockets);
        tracker.cleanup_closed(&mut sockets);

        // The socket is reaped and the table slot is free again — no
        // CLOSE_WAIT orphan pinning a slot until the table fills.
        assert!(
            !tracker.has_socket_for(
                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54321),
                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
            ),
            "connection must be evicted after FIN + proxy exit",
        );
        assert_eq!(
            only_tcp_state(&sockets),
            None,
            "socket must be removed from the socket set",
        );
    }

    #[test]
    fn guest_rst_is_cleaned_up() {
        let shared = Arc::new(SharedState::new(64));
        let poll_config = leak_poll_config();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let mut tracker = ConnectionTracker::new(None);
        let now = smoltcp_now();

        let (server_isn, guest_seq, mut new_conns) = establish(
            &mut tracker,
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
            54322,
        );
        let conn = new_conns.remove(0);
        let mut from_smoltcp = conn.from_smoltcp;
        let _to_smoltcp = conn.to_smoltcp;

        // Guest aborts the connection (RST) instead of closing it cleanly.
        ingress(
            build_tcp_frame(
                54322,
                443,
                TcpControl::Rst,
                guest_seq,
                Some(server_isn + 1),
                &[],
            ),
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
        );

        // Maintenance reaps the Closed socket and drops the proxy channel.
        for _ in 0..8 {
            tracker.relay_data(&mut sockets);
            tracker.cleanup_closed(&mut sockets);
            let _ = iface.poll_egress(now, &mut device, &mut sockets);
        }

        assert!(
            !tracker.has_socket_for(
                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54322),
                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
            ),
            "RST connection should be evicted from the tracker",
        );
        // The proxy task's receiver IS disconnected => `recv()` returns None
        // => the proxy breaks its relay loop => the upstream socket is closed.
        assert!(
            matches!(
                from_smoltcp.try_recv(),
                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected)
            ),
            "RST teardown must close the proxy channel (clean, no orphan)",
        );
    }

    #[test]
    fn full_connection_table_refuses_new_sockets() {
        // Once the table is full, new guest connections are refused. Uses a
        // small max to avoid 256 full handshakes; the gating logic is
        // identical to the 256 default.
        let mut tracker = ConnectionTracker::new(Some(4));
        let mut sockets = SocketSet::new(vec![]);
        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);

        for port in 40000u16..40004 {
            let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), port);
            assert!(
                tracker.create_tcp_socket(src, dst, &mut sockets),
                "creation under the limit must succeed",
            );
        }
        // Table full (4 slots held). The 5th guest SYN gets no socket — which
        // in the poll loop means smoltcp emits RST / no reply => the guest
        // sees egress as unreachable.
        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40004);
        assert!(
            !tracker.create_tcp_socket(src, dst, &mut sockets),
            "creation at the limit must be refused",
        );
    }

    #[test]
    fn guest_fin_before_proxy_spawn_is_handed_off() {
        let shared = Arc::new(SharedState::new(64));
        let poll_config = leak_poll_config();
        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
        let mut iface = create_interface(&mut device, &poll_config);
        let mut sockets = SocketSet::new(vec![]);
        let mut tracker = ConnectionTracker::new(None);
        let now = smoltcp_now();
        let guest_port = 54323;
        let (server_isn, guest_seq) = handshake(
            &mut tracker,
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
            guest_port,
        );

        // The production poll loop drains every queued guest frame before
        // taking new connections, so the FIN can be processed before the
        // proxy task is spawned.
        ingress(
            build_tcp_frame(
                guest_port,
                443,
                TcpControl::Fin,
                guest_seq,
                Some(server_isn + 1),
                &[],
            ),
            &mut device,
            &mut iface,
            &mut sockets,
            &shared,
            now,
        );
        assert_eq!(
            only_tcp_state(&sockets),
            Some(tcp::State::CloseWait),
            "guest FIN should arrive before the proxy handoff",
        );

        let new_conns = tracker.take_new_connections(&mut sockets);
        assert_eq!(
            new_conns.len(),
            1,
            "a connection that reached CLOSE_WAIT still needs a proxy task",
        );
    }
}