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::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;

use smoltcp::iface::{Config, Interface, SocketSet};
use smoltcp::time::Instant;
use std::sync::atomic::Ordering;

use smoltcp::wire::{
    EthernetAddress, EthernetFrame, EthernetProtocol, HardwareAddress, IpAddress, IpCidr,
    IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket,
};

use crate::config::{DnsConfig, PublishedPort};
use crate::conn::ConnectionTracker;
use crate::device::SmoltcpDevice;
use crate::dns::interceptor::DnsInterceptor;
use crate::policy::{NetworkPolicy, Protocol};
use crate::proxy;
use crate::publisher::PortPublisher;
use crate::shared::SharedState;
use crate::tls::{proxy as tls_proxy, state::TlsState};
use crate::udp_relay::UdpRelay;

//--------------------------------------------------------------------------------------------------
// 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,

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

/// 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 IPv4 address.
    pub gateway_ipv4: Ipv4Addr,
    /// Guest IPv4 address.
    pub guest_ipv4: Ipv4Addr,
    /// Gateway IPv6 address.
    pub gateway_ipv6: Ipv6Addr,
    /// IP-level MTU (e.g. 1500).
    pub mtu: usize,
}

//--------------------------------------------------------------------------------------------------
// 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.
    iface.update_ip_addrs(|addrs| {
        addrs
            .push(IpCidr::new(
                IpAddress::Ipv4(config.gateway_ipv4),
                // /30 subnet: gateway + guest.
                30,
            ))
            .expect("failed to add gateway IPv4 address");
        addrs
            .push(IpCidr::new(IpAddress::Ipv6(config.gateway_ipv6), 64))
            .expect("failed to add gateway IPv6 address");
    });

    // Default routes so smoltcp accepts traffic for all destinations.
    iface
        .routes_mut()
        .add_default_ipv4_route(config.gateway_ipv4)
        .expect("failed to add default IPv4 route");
    iface
        .routes_mut()
        .add_default_ipv6_route(config.gateway_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** — `poll(2)` on `tx_wake` + `proxy_wake` pipes with smoltcp's
///    requested timeout.
#[allow(clippy::too_many_arguments)]
pub fn smoltcp_poll_loop(
    shared: Arc<SharedState>,
    config: PollLoopConfig,
    network_policy: NetworkPolicy,
    dns_config: DnsConfig,
    tls_state: Option<Arc<TlsState>>,
    published_ports: Vec<PublishedPort>,
    max_connections: Option<usize>,
    tokio_handle: tokio::runtime::Handle,
) {
    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);

    let mut dns_interceptor =
        DnsInterceptor::new(&mut sockets, dns_config, shared.clone(), &tokio_handle);
    let mut port_publisher = PortPublisher::new(&published_ports, config.guest_ipv4, &tokio_handle);
    let mut udp_relay = UdpRelay::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();

    // poll(2) file descriptors for sleeping.
    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,
        },
    ];

    loop {
        let now = smoltcp_now();

        // ── Phase 1: Drain all guest frames with pre-inspection ──────────
        while let Some(frame) = device.stage_next_frame() {
            match classify_frame(frame) {
                FrameAction::TcpSyn { src, dst } => {
                    // Policy check before socket creation.
                    if network_policy
                        .evaluate_egress(dst, Protocol::Tcp)
                        .is_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 } => {
                    // QUIC blocking: drop UDP to intercepted ports when
                    // TLS interception is active.
                    if let Some(ref tls) = tls_state
                        && tls.config.intercepted_ports.contains(&dst.port())
                        && tls.config.block_quic_on_intercept
                    {
                        device.drop_staged_frame();
                        continue;
                    }

                    // Policy check.
                    if network_policy.evaluate_egress(dst, Protocol::Udp).is_deny() {
                        device.drop_staged_frame();
                        continue;
                    }

                    udp_relay.relay_outbound(frame, src, dst);
                    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.
                tls_proxy::spawn_tls_proxy(
                    &tokio_handle,
                    conn.dst,
                    conn.from_smoltcp,
                    conn.to_smoltcp,
                    shared.clone(),
                    tls_state.clone(),
                );
                continue;
            }
            // Plain TCP proxy.
            proxy::spawn_tcp_proxy(
                &tokio_handle,
                conn.dst,
                conn.from_smoltcp,
                conn.to_smoltcp,
                shared.clone(),
            );
        }

        // 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();
            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.

        // 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,
            );
        }

        // Conditional drain: only drain pipes that actually have data.
        if poll_fds[0].revents & libc::POLLIN != 0 {
            shared.tx_wake.drain();
        }
        if poll_fds[1].revents & libc::POLLIN != 0 {
            shared.proxy_wake.drain();
        }
    }
}

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

/// 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)
}

/// 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;
    };
    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;
    };
    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;
            };
            if udp.dst_port() == 53 {
                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::*;

    /// 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
    }

    #[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_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));
    }
}