Skip to main content

microsandbox_network/engine/netstack/
poll.rs

1//! smoltcp interface setup, frame classification, and poll loop.
2//!
3//! This module contains the core networking event loop that runs on a
4//! dedicated OS thread. It bridges guest ethernet frames (via
5//! [`SmoltcpDevice`]) to smoltcp's TCP/IP stack and services connections
6//! through tokio proxy tasks.
7
8use std::collections::HashSet;
9use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
10use std::num::NonZeroUsize;
11use std::sync::Arc;
12use std::sync::atomic::Ordering;
13
14#[cfg(windows)]
15use msb_krun_utils::event::{EventSet, EventSource, WaitContext, WaitEvent};
16use smoltcp::iface::{Config, Interface, SocketSet};
17use smoltcp::time::Instant;
18
19use smoltcp::wire::{
20    EthernetAddress, EthernetFrame, EthernetProtocol, HardwareAddress, Icmpv4Packet, Icmpv4Repr,
21    Icmpv6Packet, Icmpv6Repr, IpAddress, IpCidr, IpProtocol, Ipv4Packet, Ipv4Repr, Ipv6Packet,
22    Ipv6Repr, TcpPacket, UdpPacket,
23};
24
25use crate::config::{ConnectionLimit, DnsConfig, PublishedPort, TcpAcceptQueueSize};
26use crate::engine::dns::common::ports::DnsPortType;
27use crate::engine::dns::{
28    interceptor::DnsInterceptor,
29    proxies::{dot::DotProxy, tcp::DnsTcpProxy},
30};
31use crate::engine::tls::{proxy::TlsProxy, state::TlsState};
32use crate::icmp::relay::IcmpRelay;
33use crate::policy::{EgressEvaluation, HostnameSource, NetworkPolicy, Protocol};
34use crate::ports::PortPublisher;
35use crate::proxy::ResolvedOutboundProxy;
36use crate::secrets::handle::SecretsHandle;
37use crate::tcp::{
38    connection::{TcpAcceptMode, TcpConnectionTracker},
39    deny as tcp_deny,
40    pending::PendingConnect,
41    proxy::TcpProxy,
42    upstream::UpstreamTcpTarget,
43};
44use crate::udp::fragments::{
45    Ipv4UdpFragmentReassembler, Ipv6UdpFragmentReassembler, ReassembledUdpDatagram,
46    is_ipv4_udp_fragment, is_ipv6_fragment, is_ipv6_udp_fragment,
47};
48use crate::udp::relay::UdpRelay;
49
50use super::{device::SmoltcpDevice, shared::SharedState};
51
52//--------------------------------------------------------------------------------------------------
53// Constants
54//--------------------------------------------------------------------------------------------------
55
56const CLEANUP_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);
57
58#[cfg(windows)]
59const TX_WAKE_TOKEN: u64 = 1;
60
61#[cfg(windows)]
62const PROXY_WAKE_TOKEN: u64 = 2;
63
64//--------------------------------------------------------------------------------------------------
65// Types
66//--------------------------------------------------------------------------------------------------
67
68/// Result of classifying a guest ethernet frame before smoltcp processes it.
69///
70/// Pre-inspection allows the poll loop to:
71/// - Create TCP sockets before smoltcp sees a SYN (preventing auto-RST).
72/// - Handle non-DNS UDP outside smoltcp (smoltcp lacks wildcard port binding).
73/// - Route DNS queries to the interception handler.
74pub enum FrameAction {
75    /// TCP SYN to a new destination — create a smoltcp socket before
76    /// letting smoltcp process the frame.
77    TcpSyn { src: SocketAddr, dst: SocketAddr },
78
79    /// Non-DNS UDP datagram — handle entirely outside smoltcp via the UDP
80    /// relay.
81    UdpRelay { src: SocketAddr, dst: SocketAddr },
82
83    /// DNS query (UDP to port 53) — let smoltcp's bound UDP socket handle it.
84    Dns,
85
86    /// IPv4 UDP fragment — reassemble before the UDP relay sees it.
87    Ipv4UdpFragment,
88
89    /// IPv6 UDP fragment — reassemble before the UDP relay sees it.
90    Ipv6UdpFragment,
91
92    /// IPv6 fragment for a protocol this relay cannot safely classify.
93    Ipv6UnsupportedFragment,
94
95    /// Everything else (ARP, NDP, ICMP, TCP data/ACK/FIN, etc.) — let
96    /// smoltcp process normally.
97    Passthrough,
98}
99
100/// Local ICMP echo reply plus the destination policy input it answers.
101struct GatewayIcmpReply {
102    dst: IpAddr,
103    protocol: Protocol,
104    frame: Vec<u8>,
105}
106
107/// Resolved network parameters for the poll loop. Created by
108/// `SmoltcpNetwork::new()` from a resolved network configuration and sandbox slot.
109pub struct PollLoopConfig {
110    /// Gateway MAC address (smoltcp's identity on the virtual LAN).
111    pub gateway_mac: [u8; 6],
112    /// Guest MAC address.
113    pub guest_mac: [u8; 6],
114    /// Gateway addresses owned by the smoltcp virtual stack. Each family
115    /// is `Some` when that family is active for this sandbox (host has a
116    /// route, or the user supplied an explicit address).
117    pub gateway: GatewayIps,
118    /// Guest IPv4 address. `None` when IPv4 is inactive for this sandbox.
119    pub guest_ipv4: Option<Ipv4Addr>,
120    /// Guest IPv6 address. `None` when IPv6 is inactive for this sandbox.
121    pub guest_ipv6: Option<Ipv6Addr>,
122    /// IP-level MTU (e.g. 1500).
123    pub mtu: usize,
124}
125
126/// Per-sandbox gateway addresses owned by the smoltcp virtual stack.
127///
128/// Each family is `Some` when active for this sandbox and `None` otherwise.
129/// `resolve_host_dst` rewrites gateway-bound connections to loopback at dial time.
130#[derive(Debug, Clone, Copy)]
131pub struct GatewayIps {
132    /// Gateway IPv4.
133    pub ipv4: Option<Ipv4Addr>,
134    /// Gateway IPv6.
135    pub ipv6: Option<Ipv6Addr>,
136}
137
138//--------------------------------------------------------------------------------------------------
139// Functions
140//--------------------------------------------------------------------------------------------------
141
142/// Classify a raw ethernet frame for pre-inspection.
143///
144/// Uses smoltcp's wire module for zero-copy parsing. Returns
145/// [`FrameAction::Passthrough`] for any frame that cannot be parsed or
146/// doesn't match a special case.
147pub fn classify_frame(frame: &[u8]) -> FrameAction {
148    let Ok(eth) = EthernetFrame::new_checked(frame) else {
149        return FrameAction::Passthrough;
150    };
151
152    match eth.ethertype() {
153        EthernetProtocol::Ipv4 => classify_ipv4(eth.payload()),
154        EthernetProtocol::Ipv6 => classify_ipv6(eth.payload()),
155        _ => FrameAction::Passthrough, // ARP, etc.
156    }
157}
158
159/// Create and configure the smoltcp [`Interface`].
160///
161/// The interface is configured as the **gateway**: it owns the gateway IP
162/// addresses and responds to ARP/NDP for them. `any_ip` mode is enabled so
163/// smoltcp accepts traffic destined for arbitrary remote IPs (not just the
164/// gateway), combined with default routes.
165pub fn create_interface(device: &mut SmoltcpDevice, config: &PollLoopConfig) -> Interface {
166    let hw_addr = HardwareAddress::Ethernet(EthernetAddress(config.gateway_mac));
167    let iface_config = Config::new(hw_addr);
168    let mut iface = Interface::new(iface_config, device, smoltcp_now());
169
170    // Configure gateway IP addresses for the active families.
171    iface.update_ip_addrs(|addrs| {
172        if let Some(ipv4) = config.gateway.ipv4 {
173            addrs
174                .push(IpCidr::new(IpAddress::Ipv4(ipv4), 30)) // 30 subnet: gateway + guest.
175                .expect("failed to add gateway IPv4 address");
176        }
177        if let Some(ipv6) = config.gateway.ipv6 {
178            addrs
179                .push(IpCidr::new(IpAddress::Ipv6(ipv6), 64))
180                .expect("failed to add gateway IPv6 address");
181        }
182    });
183
184    // Default routes so smoltcp accepts traffic for all destinations.
185    if let Some(ipv4) = config.gateway.ipv4 {
186        iface
187            .routes_mut()
188            .add_default_ipv4_route(ipv4)
189            .expect("failed to add default IPv4 route");
190    }
191    if let Some(ipv6) = config.gateway.ipv6 {
192        iface
193            .routes_mut()
194            .add_default_ipv6_route(ipv6)
195            .expect("failed to add default IPv6 route");
196    }
197
198    // Accept traffic destined for any IP, not just gateway addresses.
199    iface.set_any_ip(true);
200
201    iface
202}
203
204/// Main smoltcp poll loop. Runs on a dedicated OS thread.
205///
206/// Processes guest frames with pre-inspection, drives smoltcp's TCP/IP stack,
207/// and sleeps via `poll(2)` between events.
208///
209/// # Phases per iteration
210///
211/// 1. **Drain guest frames** — pop from `tx_ring`, classify, pre-inspect.
212/// 2. **smoltcp egress + maintenance** — transmit queued packets, run timers.
213/// 3. **Service connections** — relay data between smoltcp sockets and proxy
214///    tasks (added by later tasks).
215/// 4. **Sleep** — wait on `tx_wake` + `proxy_wake` with smoltcp's requested
216///    timeout.
217///
218/// # Arguments
219///
220/// * `shared` - Stack-wide shared state: `tx_ring` / `rx_ring` for the virtio-net boundary
221///   and the wake eventfds.
222/// * `config` - Resolved per-sandbox parameters (gateway / guest MAC + IPv4 + IPv6, MTU).
223/// * `network_policy` - User-provided egress policy. Evaluated against the sandbox's
224///   gateway IPs (stored on [`SharedState`]) so `DestinationGroup::Host` rules match.
225/// * `platform_policy` - Optional host-owned policy floor. Traffic must pass both policies.
226/// * `dns_config` - DNS interception settings (block lists, upstreams, timeout).
227/// * `tls_state` - Optional TLS MITM state; drives interception of intercepted ports and DoT
228///   when present.
229/// * `published_ports` - Host → guest port publishes; the publisher accepts inbound
230///   connections on the host-bind address and forwards into the guest.
231/// * `max_tcp_connections` - Optional cap on concurrent guest connections tracked by
232///   [`TcpConnectionTracker`]; `None` uses the default.
233/// * `max_udp_connections` - Optional cap on concurrent UDP relay sessions.
234/// * `tcp_accept_queue_size` - Accept-queue depth requested for each published TCP port listener.
235/// * `tokio_handle` - Runtime handle used for proxy tasks, DNS forwarding, port publishing,
236///   and ICMP relays.
237#[allow(clippy::too_many_arguments)]
238pub fn smoltcp_poll_loop(
239    shared: Arc<SharedState>,
240    config: PollLoopConfig,
241    network_policy: NetworkPolicy,
242    platform_policy: Option<NetworkPolicy>,
243    dns_config: DnsConfig,
244    tls_state: Option<Arc<TlsState>>,
245    published_ports: Vec<PublishedPort>,
246    strict: bool,
247    max_tcp_connections: Option<NonZeroUsize>,
248    max_udp_connections: Option<ConnectionLimit>,
249    tcp_accept_queue_size: TcpAcceptQueueSize,
250    max_inbound_tcp_connections: Option<ConnectionLimit>,
251    tokio_handle: tokio::runtime::Handle,
252    secrets: SecretsHandle,
253    outbound_proxy: Option<Arc<ResolvedOutboundProxy>>,
254) {
255    let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
256    let mut iface = create_interface(&mut device, &config);
257    let mut sockets = SocketSet::new(vec![]);
258    let mut conn_tracker = TcpConnectionTracker::new(max_tcp_connections);
259
260    // The DNS forwarder needs to know which IPs count as "the gateway"
261    // (so it routes guest queries to those addresses through the
262    // configured upstream) and a policy evaluator (so guest-chosen
263    // `@target` resolvers are gated by egress rules just like any
264    // other outbound).
265    let gateway_ips: Arc<HashSet<IpAddr>> = Arc::new(
266        config
267            .gateway
268            .ipv4
269            .map(IpAddr::V4)
270            .into_iter()
271            .chain(config.gateway.ipv6.map(IpAddr::V6))
272            .collect(),
273    );
274    // Gateway IPs must be on SharedState before any egress evaluation runs,
275    // so `DestinationGroup::Host` rules can resolve to the right address.
276    shared.set_gateway_ips(config.gateway.ipv4, config.gateway.ipv6);
277    let network_policy = Arc::new(network_policy);
278    let platform_policy = platform_policy.map(Arc::new);
279
280    let (mut dns_interceptor, dns_forwarder_handle) = DnsInterceptor::new(
281        &mut sockets,
282        dns_config,
283        shared.clone(),
284        &tokio_handle,
285        gateway_ips,
286        network_policy.clone(),
287        platform_policy.clone(),
288        config.gateway,
289        config.gateway_mac,
290        config.guest_mac,
291    );
292    let mut port_publisher = PortPublisher::new(
293        &published_ports,
294        tcp_accept_queue_size,
295        max_inbound_tcp_connections,
296        config.guest_ipv4,
297        config.guest_ipv6,
298        config.gateway.ipv4,
299        config.gateway.ipv6,
300        config.gateway_mac,
301        config.guest_mac,
302        network_policy.clone(),
303        shared.clone(),
304        &tokio_handle,
305    );
306    let mut udp_relay = UdpRelay::new(
307        shared.clone(),
308        config.gateway_mac,
309        config.guest_mac,
310        config.mtu,
311        tokio_handle.clone(),
312        outbound_proxy.clone(),
313    );
314    udp_relay.set_max_sessions(max_udp_connections.and_then(ConnectionLimit::cap));
315    udp_relay.attach_dns_forwarder(dns_forwarder_handle.clone());
316    let mut udp_fragments = Ipv4UdpFragmentReassembler::new();
317    let mut ipv6_udp_fragments = Ipv6UdpFragmentReassembler::new();
318    let icmp_relay = IcmpRelay::new(
319        shared.clone(),
320        config.gateway_mac,
321        config.guest_mac,
322        tokio_handle.clone(),
323    );
324
325    // Rate-limit cleanup operations: run at most once per second.
326    let mut last_cleanup = std::time::Instant::now();
327
328    // Wake sources for sleeping.
329    #[cfg(unix)]
330    let mut poll_fds = [
331        libc::pollfd {
332            fd: shared.tx_wake.as_raw_fd(),
333            events: libc::POLLIN,
334            revents: 0,
335        },
336        libc::pollfd {
337            fd: shared.proxy_wake.as_raw_fd(),
338            events: libc::POLLIN,
339            revents: 0,
340        },
341    ];
342    #[cfg(windows)]
343    let wait_context = match windows_stack_wait_context(&shared) {
344        Ok(context) => context,
345        Err(err) => {
346            tracing::error!(error = %err, "network poll loop: failed to create wait context");
347            return;
348        }
349    };
350
351    loop {
352        let now = smoltcp_now();
353
354        // ── Phase 1: Drain all guest frames with pre-inspection ──────────
355        while let Some(frame) = device.stage_next_frame() {
356            if handle_gateway_icmp_echo(
357                frame,
358                &config,
359                &shared,
360                &network_policy,
361                platform_policy.as_deref(),
362            ) {
363                device.drop_staged_frame();
364                continue;
365            }
366
367            if icmp_relay.relay_outbound_if_echo(
368                frame,
369                &config,
370                &network_policy,
371                platform_policy.as_deref(),
372            ) {
373                device.drop_staged_frame();
374                continue;
375            }
376
377            match classify_frame(frame) {
378                FrameAction::TcpSyn { src, dst } => {
379                    let accept_mode = match DnsPortType::from_tcp(dst.port()) {
380                        // Plain DNS: the interceptor enforces policy at
381                        // the application layer (block list + rebind
382                        // protection); bypass the network egress check.
383                        DnsPortType::Dns => TcpAcceptMode::AcceptImmediately,
384                        // DoT: intercept only when TLS MITM is
385                        // configured. Without it, the block list can't
386                        // apply (traffic is encrypted end-to-end), so
387                        // we refuse to force a fall-back to plain
388                        // TCP/53. When TLS MITM is configured, bypass
389                        // egress policy the same way plain DNS does —
390                        // policy for the upstream resolver is applied
391                        // per query by the forwarder.
392                        DnsPortType::EncryptedDns => {
393                            if tls_state.is_some() {
394                                TcpAcceptMode::AcceptImmediately
395                            } else {
396                                tracing::debug!(%dst, "DoT port refused (TLS interception not configured); stub should fall back to TCP/53");
397                                TcpAcceptMode::Reject
398                            }
399                        }
400                        // Alternative DNS protocol we can't proxy:
401                        // refuse outright — no socket means smoltcp
402                        // emits RST, which the guest's stub treats as
403                        // "upstream unavailable" and falls back to
404                        // plain TCP/53.
405                        DnsPortType::AlternativeDns => {
406                            tracing::debug!(%dst, "alternative-DNS TCP port refused; stub should fall back to TCP/53");
407                            TcpAcceptMode::Reject
408                        }
409                        // Other: regular outbound — defer Domain rules to first-flight;
410                        // accept unless an IP-layer rule denies.
411                        DnsPortType::Other => {
412                            let platform_allows = platform_policy.as_deref().is_none_or(|policy| {
413                                policy
414                                    .evaluate_egress(dst, Protocol::Tcp, &shared)
415                                    .is_allow()
416                            });
417                            let evaluation = network_policy.evaluate_egress_with_source(
418                                dst,
419                                Protocol::Tcp,
420                                &shared,
421                                HostnameSource::Deferred,
422                            );
423
424                            match (platform_allows, evaluation) {
425                                (false, _) => TcpAcceptMode::Reject,
426                                (true, EgressEvaluation::Allow) => TcpAcceptMode::WaitForUpstream,
427                                // Hostname rules need guest bytes before they
428                                // can authorize a host connection.
429                                (true, EgressEvaluation::DeferUntilHostname) => {
430                                    TcpAcceptMode::AcceptImmediately
431                                }
432                                // Only default denies may receive an HTTP
433                                // response; explicit denies remain resets.
434                                (true, EgressEvaluation::Deny)
435                                    if shared.http_deny_response_enabled()
436                                        && tcp_deny::answers_denied_http(
437                                            dst.port(),
438                                            tls_state.as_deref(),
439                                        )
440                                        && network_policy.egress_denied_by_default(
441                                            dst,
442                                            Protocol::Tcp,
443                                            &shared,
444                                            HostnameSource::Deferred,
445                                        ) =>
446                                {
447                                    TcpAcceptMode::RespondWithDenial
448                                }
449                                (true, EgressEvaluation::Deny) => TcpAcceptMode::Reject,
450                            }
451                        }
452                    };
453                    // A reset can return a passive socket to Listen. Retire it
454                    // before another SYN can reuse its previous host connection,
455                    // including when both packets arrive in the same batch.
456                    conn_tracker.cleanup_closed(&mut sockets);
457                    if !conn_tracker.has_socket_for(&src, &dst) {
458                        conn_tracker.create_tcp_socket(src, dst, &mut sockets, accept_mode);
459                    }
460
461                    // smoltcp validates the packet before any host-side dial.
462                    // A paused socket accepts SYN but withholds SYN-ACK.
463                    iface.poll_ingress_single(now, &mut device, &mut sockets);
464                }
465
466                FrameAction::UdpRelay { src, dst } => {
467                    relay_udp_frame(
468                        frame,
469                        src,
470                        dst,
471                        &config,
472                        &network_policy,
473                        platform_policy.as_deref(),
474                        &shared,
475                        &mut port_publisher,
476                        tls_state.as_deref(),
477                        &mut udp_relay,
478                    );
479                    device.drop_staged_frame();
480                }
481
482                FrameAction::Ipv4UdpFragment => {
483                    if let Some(datagram) = udp_fragments.push(frame) {
484                        handle_reassembled_udp_datagram(
485                            datagram,
486                            &mut device,
487                            &mut iface,
488                            now,
489                            &mut sockets,
490                            &config,
491                            &network_policy,
492                            platform_policy.as_deref(),
493                            &shared,
494                            &mut port_publisher,
495                            tls_state.as_deref(),
496                            &mut udp_relay,
497                        );
498                    } else {
499                        device.drop_staged_frame();
500                    }
501                }
502
503                FrameAction::Ipv6UdpFragment => {
504                    if let Some(datagram) = ipv6_udp_fragments.push(frame) {
505                        handle_reassembled_udp_datagram(
506                            datagram,
507                            &mut device,
508                            &mut iface,
509                            now,
510                            &mut sockets,
511                            &config,
512                            &network_policy,
513                            platform_policy.as_deref(),
514                            &shared,
515                            &mut port_publisher,
516                            tls_state.as_deref(),
517                            &mut udp_relay,
518                        );
519                    } else {
520                        device.drop_staged_frame();
521                    }
522                }
523
524                FrameAction::Ipv6UnsupportedFragment => {
525                    // Fragmented UDP is only forwarded after reassembly and policy evaluation.
526                    // Other fragmented IPv6 traffic is dropped rather than passed through with
527                    // an unknown transport tuple.
528                    device.drop_staged_frame();
529                }
530
531                FrameAction::Dns | FrameAction::Passthrough => {
532                    // ARP, ICMP, DNS (port 53), TCP data — smoltcp handles.
533                    iface.poll_ingress_single(now, &mut device, &mut sockets);
534                }
535            }
536        }
537
538        for (handle, dst) in conn_tracker.poll_connects(&mut sockets) {
539            let target = resolve_tcp_host_target(dst, config.gateway);
540            let proxy = ResolvedOutboundProxy::select_for_destination(
541                &outbound_proxy,
542                dst,
543                target.primary(),
544            );
545            let pending = PendingConnect::start(target, proxy, shared.clone(), &tokio_handle);
546            conn_tracker.track_connect(handle, pending);
547        }
548
549        // ── Phase 2: Ingress egress + maintenance ─────────────────────────
550        // Flush frames generated by Phase 1 ingress (ACKs, SYN-ACKs, etc.)
551        // before relaying data so smoltcp has up-to-date state.
552        loop {
553            let result = iface.poll_egress(now, &mut device, &mut sockets);
554            if matches!(result, smoltcp::iface::PollResult::None) {
555                break;
556            }
557        }
558        iface.poll_maintenance(now);
559
560        // Coalesced wake: if Phase 1/2 emitted any frames, wake the
561        // NetWorker once instead of per-frame.
562        if device.frames_emitted.swap(false, Ordering::Relaxed) {
563            shared.rx_wake.wake();
564        }
565
566        // ── Phase 3: Service connections + relay data ────────────────────
567        // Relay proxy data INTO smoltcp sockets first, then a single egress
568        // pass flushes everything. This eliminates the former "Phase 2b"
569        // double-egress pattern.
570        conn_tracker.relay_data(&mut sockets);
571        dns_interceptor.process(&mut sockets);
572
573        // Accept queued inbound connections from published port listeners.
574        port_publisher.accept_inbound(&mut iface, &mut sockets, &shared, &tokio_handle);
575        port_publisher.relay_data(&mut sockets);
576
577        // Detect newly-established connections and spawn proxy tasks.
578        let new_conns = conn_tracker.take_new_connections(&mut sockets);
579        for conn in new_conns {
580            if conn.policy_denied {
581                // Accepted only to answer 403; never dials upstream.
582                tcp_deny::spawn_deny_responder(
583                    &tokio_handle,
584                    conn.dst,
585                    conn.from_smoltcp,
586                    conn.to_smoltcp,
587                    shared.clone(),
588                    tls_state.clone(),
589                    conn.proxy_connect,
590                );
591                continue;
592            }
593            if let Some(ref tls_state) = tls_state
594                && tls_state
595                    .config
596                    .intercepted_ports
597                    .contains(&conn.dst.port())
598            {
599                // TLS-intercepted port — spawn TLS MITM proxy.
600                let connect_target = resolve_tcp_host_target(conn.dst, config.gateway);
601                let connection_outbound_proxy = ResolvedOutboundProxy::select_for_destination(
602                    &outbound_proxy,
603                    conn.dst,
604                    connect_target.primary(),
605                );
606                let proxy = TlsProxy::new(
607                    conn.upstream,
608                    conn.dst,
609                    connect_target,
610                    conn.from_smoltcp,
611                    conn.to_smoltcp,
612                    shared.clone(),
613                    tls_state.clone(),
614                    network_policy.clone(),
615                    strict,
616                    conn.proxy_connect,
617                    connection_outbound_proxy,
618                );
619                tokio_handle.spawn(proxy.run());
620                continue;
621            }
622            if conn.dst.port() == 53 {
623                // DNS proxies have no guest-visible
624                // "upstream-unreachable" failure mode — even an
625                // upstream DNS failure yields SERVFAIL responses
626                // rather than a silently-closed connection. Mark the
627                // connection as connected so normal task exit
628                // produces FIN, not RST.
629                conn.proxy_connect.mark_connected();
630
631                // DNS over TCP: route through the same forwarder the UDP
632                // path uses. The forwarder applies the domain block list
633                // and rebind protection to every query and routes
634                // upstream based on `conn.dst.ip()` — the configured
635                // upstream for queries to the gateway, direct forward
636                // to the chosen `@target` (subject to egress policy)
637                // otherwise. No gateway→loopback rewrite here: the
638                // forwarder dials the configured upstream, not the
639                // gateway.
640                let proxy = DnsTcpProxy::new(
641                    conn.dst,
642                    conn.from_smoltcp,
643                    conn.to_smoltcp,
644                    dns_forwarder_handle.clone(),
645                    shared.clone(),
646                );
647                tokio_handle.spawn(proxy.run());
648                continue;
649            }
650            if conn.dst.port() == 853
651                && let Some(ref tls_state) = tls_state
652            {
653                // Same "always upstream-connected" reasoning as plain DNS over TCP.
654                conn.proxy_connect.mark_connected();
655
656                // DNS over TLS: terminate TLS at the gateway with a
657                // per-domain cert, hand the inner DNS frames to the
658                // same forwarder plain DNS uses. Policy for the
659                // chosen `@target` resolver is applied per-query by
660                // the forwarder (block list + rebind + egress).
661                let proxy = DotProxy::new(
662                    conn.dst,
663                    conn.from_smoltcp,
664                    conn.to_smoltcp,
665                    dns_forwarder_handle.clone(),
666                    tls_state.clone(),
667                    shared.clone(),
668                );
669                tokio_handle.spawn(proxy.run());
670                continue;
671            }
672            // Plain TCP proxy.
673            let connect_target = resolve_tcp_host_target(conn.dst, config.gateway);
674            let connection_outbound_proxy = ResolvedOutboundProxy::select_for_destination(
675                &outbound_proxy,
676                conn.dst,
677                connect_target.primary(),
678            );
679            let proxy = TcpProxy::new(
680                conn.upstream,
681                conn.dst,
682                connect_target,
683                conn.from_smoltcp,
684                conn.to_smoltcp,
685                shared.clone(),
686                network_policy.clone(),
687                // Load the current snapshot per connection so live secret
688                // updates apply to traffic the guest starts afterwards.
689                secrets.load(),
690                tls_state.clone(),
691                strict,
692                conn.proxy_connect,
693                connection_outbound_proxy,
694            );
695            tokio_handle.spawn(proxy.run());
696        }
697
698        // Periodic cleanup is the idle fallback. TCP creation also reclaims
699        // completed flows when the table is full, before rejecting a new SYN.
700        if last_cleanup.elapsed() >= CLEANUP_INTERVAL {
701            conn_tracker.cleanup_closed(&mut sockets);
702            conn_tracker.trace_stats(&sockets);
703            port_publisher.cleanup_closed(&mut sockets);
704            udp_relay.cleanup_expired();
705            udp_fragments.cleanup_expired();
706            ipv6_udp_fragments.cleanup_expired();
707            shared.cleanup_resolved_hostnames();
708            last_cleanup = std::time::Instant::now();
709        }
710
711        // ── Phase 4: Flush relay data + sleep ────────────────────────────
712        // Single egress pass flushes all data written by Phase 3.
713        loop {
714            let result = iface.poll_egress(now, &mut device, &mut sockets);
715            if matches!(result, smoltcp::iface::PollResult::None) {
716                break;
717            }
718        }
719
720        // Coalesced wake: if Phase 3/4 emitted any frames, wake once.
721        if device.frames_emitted.swap(false, Ordering::Relaxed) {
722            shared.rx_wake.wake();
723        }
724
725        let stack_delay = iface
726            .poll_delay(now, &sockets)
727            .map(|delay| std::time::Duration::from_millis(delay.total_millis()));
728
729        // Packet/proxy events wake us immediately. Otherwise sleep to the next
730        // actual deadline, including cleanup even when the stack timer is later.
731        let timeout_ms = poll_timeout_ms(
732            stack_delay,
733            conn_tracker.deferred_close_delay(),
734            port_publisher.deferred_close_delay(),
735            CLEANUP_INTERVAL.saturating_sub(last_cleanup.elapsed()),
736        );
737
738        #[cfg(unix)]
739        sleep_until_stack_wake(&shared, timeout_ms, &mut poll_fds);
740        #[cfg(windows)]
741        sleep_until_stack_wake_windows(&shared, timeout_ms, &wait_context);
742    }
743}
744
745//--------------------------------------------------------------------------------------------------
746// Functions: Helpers
747//--------------------------------------------------------------------------------------------------
748
749/// Bound sleep by every maintenance deadline without polling an idle stack.
750fn poll_timeout_ms(
751    stack_delay: Option<std::time::Duration>,
752    connection_close_delay: Option<std::time::Duration>,
753    port_close_delay: Option<std::time::Duration>,
754    cleanup_delay: std::time::Duration,
755) -> i32 {
756    let delay = [stack_delay, connection_close_delay, port_close_delay]
757        .into_iter()
758        .flatten()
759        .fold(cleanup_delay, std::cmp::min);
760    // Round up fractional milliseconds to avoid spinning just before cleanup.
761    // A deadline that is already due still requests an immediate poll.
762    delay.as_nanos().div_ceil(1_000_000).min(i32::MAX as u128) as i32
763}
764
765#[cfg(unix)]
766fn sleep_until_stack_wake(shared: &SharedState, timeout_ms: i32, poll_fds: &mut [libc::pollfd; 2]) {
767    // SAFETY: poll_fds is a valid array of pollfd structs with valid fds.
768    unsafe {
769        libc::poll(
770            poll_fds.as_mut_ptr(),
771            poll_fds.len() as libc::nfds_t,
772            timeout_ms,
773        );
774    }
775
776    if poll_fds[0].revents & libc::POLLIN != 0 {
777        shared.tx_wake.drain();
778    }
779    if poll_fds[1].revents & libc::POLLIN != 0 {
780        shared.proxy_wake.drain();
781    }
782}
783
784#[cfg(windows)]
785fn windows_stack_wait_context(shared: &SharedState) -> std::io::Result<WaitContext> {
786    let mut context = WaitContext::new();
787    context.add(
788        EventSource::waitable_handle(shared.tx_wake.as_raw_handle(), TX_WAKE_TOKEN),
789        EventSet::IN,
790    )?;
791    context.add(
792        EventSource::waitable_handle(shared.proxy_wake.as_raw_handle(), PROXY_WAKE_TOKEN),
793        EventSet::IN,
794    )?;
795    Ok(context)
796}
797
798#[cfg(windows)]
799fn sleep_until_stack_wake_windows(
800    shared: &SharedState,
801    timeout_ms: i32,
802    wait_context: &WaitContext,
803) {
804    let mut events = [WaitEvent::default(); 2];
805    let count = match wait_context.wait(timeout_ms, &mut events) {
806        Ok(count) => count,
807        Err(err) => {
808            tracing::warn!(error = %err, "network poll loop: wait failed");
809            return;
810        }
811    };
812
813    for event in events.iter().take(count) {
814        match event.token() {
815            TX_WAKE_TOKEN => shared.tx_wake.drain(),
816            PROXY_WAKE_TOKEN => shared.proxy_wake.drain(),
817            token => tracing::warn!(token, "network poll loop: unknown wake token"),
818        }
819    }
820}
821
822/// Apply the common non-DNS UDP dispatch path to a complete guest datagram.
823#[allow(clippy::too_many_arguments)]
824fn relay_udp_frame(
825    frame: &[u8],
826    src: SocketAddr,
827    dst: SocketAddr,
828    config: &PollLoopConfig,
829    network_policy: &NetworkPolicy,
830    platform_policy: Option<&NetworkPolicy>,
831    shared: &Arc<SharedState>,
832    port_publisher: &mut PortPublisher,
833    tls_state: Option<&TlsState>,
834    udp_relay: &mut UdpRelay,
835) {
836    if port_publisher.relay_udp_outbound(frame, src, dst) {
837        return;
838    }
839
840    // QUIC blocking: drop UDP to intercepted ports when TLS interception is active.
841    if let Some(tls) = tls_state
842        && tls.config.intercepted_ports.contains(&dst.port())
843        && tls.config.block_quic_on_intercept
844    {
845        return;
846    }
847
848    match DnsPortType::from_udp(dst.port()) {
849        // Dns: unreachable here — classify_transport routes UDP/53 to
850        // FrameAction::Dns, not UdpRelay. Defensive drop covers regressions.
851        DnsPortType::Dns | DnsPortType::EncryptedDns => return,
852        // Alternative DNS protocols on well-known UDP ports are dropped —
853        // forces fall-back to UDP/53.
854        DnsPortType::AlternativeDns => {
855            tracing::debug!(%dst, "alternative-DNS UDP port dropped; stub should fall back to UDP/53");
856            return;
857        }
858        DnsPortType::Other => {}
859    }
860
861    // Policy is applied after reassembly, when the UDP destination port is known.
862    if platform_policy
863        .is_some_and(|policy| policy.evaluate_egress(dst, Protocol::Udp, shared).is_deny())
864        || network_policy
865            .evaluate_egress(dst, Protocol::Udp, shared)
866            .is_deny()
867    {
868        return;
869    }
870
871    // Resolve the host-side destination for the dial. `dst` stays unchanged so
872    // reply frames are stamped with the IP the guest expects.
873    let host_dst = resolve_host_dst(dst, config.gateway);
874    udp_relay.relay_outbound(frame, src, dst, host_dst);
875}
876
877/// Dispatch a complete datagram produced by fragment reassembly.
878#[allow(clippy::too_many_arguments)]
879fn handle_reassembled_udp_datagram(
880    datagram: ReassembledUdpDatagram,
881    device: &mut SmoltcpDevice,
882    iface: &mut Interface,
883    now: Instant,
884    sockets: &mut SocketSet<'_>,
885    config: &PollLoopConfig,
886    network_policy: &NetworkPolicy,
887    platform_policy: Option<&NetworkPolicy>,
888    shared: &Arc<SharedState>,
889    port_publisher: &mut PortPublisher,
890    tls_state: Option<&TlsState>,
891    udp_relay: &mut UdpRelay,
892) {
893    if DnsPortType::from_udp(datagram.dst.port()) == DnsPortType::Dns {
894        device.replace_staged_frame(datagram.frame);
895        iface.poll_ingress_single(now, device, sockets);
896        return;
897    }
898
899    relay_udp_frame(
900        &datagram.frame,
901        datagram.src,
902        datagram.dst,
903        config,
904        network_policy,
905        platform_policy,
906        shared,
907        port_publisher,
908        tls_state,
909        udp_relay,
910    );
911    device.drop_staged_frame();
912}
913
914/// Resolve host-side TCP destination candidates for a guest connection.
915///
916/// A guest connection to either gateway family first dials the matching host
917/// loopback address, then may fall back to the other loopback family. Regular
918/// outbound destinations have no fallback.
919fn resolve_tcp_host_target(dst: SocketAddr, gateway: GatewayIps) -> UpstreamTcpTarget {
920    let port = dst.port();
921    match dst.ip() {
922        IpAddr::V4(v4) if gateway.ipv4 == Some(v4) => UpstreamTcpTarget::with_fallback(
923            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port),
924            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), port),
925        ),
926        IpAddr::V6(v6) if gateway.ipv6 == Some(v6) => UpstreamTcpTarget::with_fallback(
927            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), port),
928            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port),
929        ),
930        _ => UpstreamTcpTarget::direct(dst),
931    }
932}
933
934/// Map a guest-wire UDP destination to its host-socket equivalent.
935///
936/// Gateway IPs rewrite to loopback (`127.0.0.1` / `::1`); everything else
937/// passes through.
938///
939/// # Arguments
940///
941/// * `dst` - Destination from the guest's packet.
942/// * `gateway` - Per-sandbox gateway IPs that trigger the loopback rewrite.
943pub(crate) fn resolve_host_dst(dst: SocketAddr, gateway: GatewayIps) -> SocketAddr {
944    match dst.ip() {
945        IpAddr::V4(v4) if gateway.ipv4 == Some(v4) => {
946            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), dst.port())
947        }
948        IpAddr::V6(v6) if gateway.ipv6 == Some(v6) => {
949            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), dst.port())
950        }
951        _ => dst,
952    }
953}
954
955/// Get the current time as a smoltcp [`Instant`] using a monotonic clock.
956///
957/// Uses `std::time::Instant` (monotonic) instead of `SystemTime` (wall
958/// clock) to avoid issues with NTP clock step corrections that could
959/// cause smoltcp timers to misbehave.
960fn smoltcp_now() -> Instant {
961    static EPOCH: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
962    let epoch = EPOCH.get_or_init(std::time::Instant::now);
963    let elapsed = epoch.elapsed();
964    Instant::from_millis(elapsed.as_millis() as i64)
965}
966
967/// Reply locally to ICMP echo requests aimed at the sandbox gateway.
968///
969/// `any_ip` is required so smoltcp accepts guest traffic for arbitrary remote
970/// destinations, but that would make smoltcp's automatic ICMP echo replies
971/// spoof remote hosts. Handle only the real gateway IPs here and leave all
972/// other ICMP traffic untouched.
973fn handle_gateway_icmp_echo(
974    frame: &[u8],
975    config: &PollLoopConfig,
976    shared: &SharedState,
977    network_policy: &NetworkPolicy,
978    platform_policy: Option<&NetworkPolicy>,
979) -> bool {
980    let Ok(eth) = EthernetFrame::new_checked(frame) else {
981        return false;
982    };
983
984    let reply = match eth.ethertype() {
985        EthernetProtocol::Ipv4 => gateway_icmpv4_echo_reply(&eth, config),
986        EthernetProtocol::Ipv6 => gateway_icmpv6_echo_reply(&eth, config),
987        _ => None,
988    };
989    let Some(reply) = reply else {
990        return false;
991    };
992
993    // This path bypasses the normal proxy evaluators, so it must independently
994    // satisfy both the tenant policy and the optional host-owned policy floor.
995    let tenant_denied = network_policy
996        .evaluate_egress_ip(reply.dst, reply.protocol, shared)
997        .is_deny();
998    let platform_denied = platform_policy.is_some_and(|policy| {
999        policy
1000            .evaluate_egress_ip(reply.dst, reply.protocol, shared)
1001            .is_deny()
1002    });
1003    if tenant_denied || platform_denied {
1004        tracing::debug!(
1005            dst = %reply.dst,
1006            tenant_denied,
1007            platform_denied,
1008            "gateway ICMP echo denied by policy",
1009        );
1010        return true;
1011    }
1012
1013    shared.push_rx_frame_and_wake(reply.frame);
1014
1015    true
1016}
1017
1018/// Build an IPv4 ICMP echo reply when the guest pings the gateway IPv4.
1019fn gateway_icmpv4_echo_reply(
1020    eth: &EthernetFrame<&[u8]>,
1021    config: &PollLoopConfig,
1022) -> Option<GatewayIcmpReply> {
1023    let gateway_ipv4 = config.gateway.ipv4?;
1024    let ipv4 = Ipv4Packet::new_checked(eth.payload()).ok()?;
1025    if ipv4.dst_addr() != gateway_ipv4 || ipv4.next_header() != IpProtocol::Icmp {
1026        return None;
1027    }
1028
1029    let icmp = Icmpv4Packet::new_checked(ipv4.payload()).ok()?;
1030    let Icmpv4Repr::EchoRequest {
1031        ident,
1032        seq_no,
1033        data,
1034    } = Icmpv4Repr::parse(&icmp, &smoltcp::phy::ChecksumCapabilities::default()).ok()?
1035    else {
1036        return None;
1037    };
1038
1039    let ipv4_repr = Ipv4Repr {
1040        src_addr: gateway_ipv4,
1041        dst_addr: ipv4.src_addr(),
1042        next_header: IpProtocol::Icmp,
1043        payload_len: 8 + data.len(),
1044        hop_limit: 64,
1045    };
1046    let icmp_repr = Icmpv4Repr::EchoReply {
1047        ident,
1048        seq_no,
1049        data,
1050    };
1051    let mut reply = vec![0u8; 14 + ipv4_repr.buffer_len() + icmp_repr.buffer_len()];
1052
1053    let mut reply_eth = EthernetFrame::new_unchecked(&mut reply);
1054    reply_eth.set_src_addr(EthernetAddress(config.gateway_mac));
1055    reply_eth.set_dst_addr(eth.src_addr());
1056    reply_eth.set_ethertype(EthernetProtocol::Ipv4);
1057
1058    ipv4_repr.emit(
1059        &mut Ipv4Packet::new_unchecked(&mut reply[14..34]),
1060        &smoltcp::phy::ChecksumCapabilities::default(),
1061    );
1062    icmp_repr.emit(
1063        &mut Icmpv4Packet::new_unchecked(&mut reply[34..]),
1064        &smoltcp::phy::ChecksumCapabilities::default(),
1065    );
1066
1067    Some(GatewayIcmpReply {
1068        dst: IpAddr::V4(gateway_ipv4),
1069        protocol: Protocol::Icmpv4,
1070        frame: reply,
1071    })
1072}
1073
1074/// Build an IPv6 ICMP echo reply when the guest pings the gateway IPv6.
1075fn gateway_icmpv6_echo_reply(
1076    eth: &EthernetFrame<&[u8]>,
1077    config: &PollLoopConfig,
1078) -> Option<GatewayIcmpReply> {
1079    let gateway_ipv6 = config.gateway.ipv6?;
1080    let ipv6 = Ipv6Packet::new_checked(eth.payload()).ok()?;
1081    if ipv6.dst_addr() != gateway_ipv6 || ipv6.next_header() != IpProtocol::Icmpv6 {
1082        return None;
1083    }
1084
1085    let icmp = Icmpv6Packet::new_checked(ipv6.payload()).ok()?;
1086    let Icmpv6Repr::EchoRequest {
1087        ident,
1088        seq_no,
1089        data,
1090    } = Icmpv6Repr::parse(
1091        &ipv6.src_addr(),
1092        &ipv6.dst_addr(),
1093        &icmp,
1094        &smoltcp::phy::ChecksumCapabilities::default(),
1095    )
1096    .ok()?
1097    else {
1098        return None;
1099    };
1100
1101    let ipv6_repr = Ipv6Repr {
1102        src_addr: gateway_ipv6,
1103        dst_addr: ipv6.src_addr(),
1104        next_header: IpProtocol::Icmpv6,
1105        payload_len: icmp_repr_buffer_len_v6(data),
1106        hop_limit: 64,
1107    };
1108    let icmp_repr = Icmpv6Repr::EchoReply {
1109        ident,
1110        seq_no,
1111        data,
1112    };
1113    let ipv6_hdr_len = 40;
1114    let mut reply = vec![0u8; 14 + ipv6_hdr_len + icmp_repr.buffer_len()];
1115
1116    let mut reply_eth = EthernetFrame::new_unchecked(&mut reply);
1117    reply_eth.set_src_addr(EthernetAddress(config.gateway_mac));
1118    reply_eth.set_dst_addr(eth.src_addr());
1119    reply_eth.set_ethertype(EthernetProtocol::Ipv6);
1120
1121    ipv6_repr.emit(&mut Ipv6Packet::new_unchecked(&mut reply[14..54]));
1122    icmp_repr.emit(
1123        &gateway_ipv6,
1124        &ipv6.src_addr(),
1125        &mut Icmpv6Packet::new_unchecked(&mut reply[54..]),
1126        &smoltcp::phy::ChecksumCapabilities::default(),
1127    );
1128
1129    Some(GatewayIcmpReply {
1130        dst: IpAddr::V6(gateway_ipv6),
1131        protocol: Protocol::Icmpv6,
1132        frame: reply,
1133    })
1134}
1135
1136fn icmp_repr_buffer_len_v6(data: &[u8]) -> usize {
1137    Icmpv6Repr::EchoReply {
1138        ident: 0,
1139        seq_no: 0,
1140        data,
1141    }
1142    .buffer_len()
1143}
1144
1145/// Classify an IPv4 packet payload (after stripping the Ethernet header).
1146fn classify_ipv4(payload: &[u8]) -> FrameAction {
1147    let Ok(ipv4) = Ipv4Packet::new_checked(payload) else {
1148        return FrameAction::Passthrough;
1149    };
1150    if is_ipv4_udp_fragment(&ipv4) {
1151        return FrameAction::Ipv4UdpFragment;
1152    }
1153    classify_transport(
1154        ipv4.next_header(),
1155        ipv4.src_addr().into(),
1156        ipv4.dst_addr().into(),
1157        ipv4.payload(),
1158    )
1159}
1160
1161/// Classify an IPv6 packet payload (after stripping the Ethernet header).
1162fn classify_ipv6(payload: &[u8]) -> FrameAction {
1163    let Ok(ipv6) = Ipv6Packet::new_checked(payload) else {
1164        return FrameAction::Passthrough;
1165    };
1166    if is_ipv6_udp_fragment(&ipv6) {
1167        return FrameAction::Ipv6UdpFragment;
1168    }
1169    if is_ipv6_fragment(&ipv6) {
1170        return FrameAction::Ipv6UnsupportedFragment;
1171    }
1172    classify_transport(
1173        ipv6.next_header(),
1174        ipv6.src_addr().into(),
1175        ipv6.dst_addr().into(),
1176        ipv6.payload(),
1177    )
1178}
1179
1180/// Classify the transport-layer protocol (shared by IPv4 and IPv6).
1181fn classify_transport(
1182    protocol: IpProtocol,
1183    src_ip: std::net::IpAddr,
1184    dst_ip: std::net::IpAddr,
1185    transport_payload: &[u8],
1186) -> FrameAction {
1187    match protocol {
1188        IpProtocol::Tcp => {
1189            let Ok(tcp) = TcpPacket::new_checked(transport_payload) else {
1190                return FrameAction::Passthrough;
1191            };
1192            if tcp.syn() && !tcp.ack() {
1193                FrameAction::TcpSyn {
1194                    src: SocketAddr::new(src_ip, tcp.src_port()),
1195                    dst: SocketAddr::new(dst_ip, tcp.dst_port()),
1196                }
1197            } else {
1198                FrameAction::Passthrough
1199            }
1200        }
1201        IpProtocol::Udp => {
1202            let Ok(udp) = UdpPacket::new_checked(transport_payload) else {
1203                return FrameAction::Passthrough;
1204            };
1205            // The plain-DNS port (UDP/53) lives in dns::common::ports so
1206            // the alternative-DNS refusal logic and this dispatcher
1207            // share one source of truth for "which UDP ports are DNS".
1208            if DnsPortType::from_udp(udp.dst_port()) == DnsPortType::Dns {
1209                FrameAction::Dns
1210            } else {
1211                FrameAction::UdpRelay {
1212                    src: SocketAddr::new(src_ip, udp.src_port()),
1213                    dst: SocketAddr::new(dst_ip, udp.dst_port()),
1214                }
1215            }
1216        }
1217        _ => FrameAction::Passthrough, // ICMP, etc.
1218    }
1219}
1220
1221//--------------------------------------------------------------------------------------------------
1222// Tests
1223//--------------------------------------------------------------------------------------------------
1224
1225#[cfg(test)]
1226mod tests {
1227    use super::*;
1228    use std::sync::Arc;
1229
1230    use smoltcp::phy::ChecksumCapabilities;
1231    use smoltcp::wire::{
1232        ArpOperation, ArpPacket, ArpRepr, EthernetRepr, Icmpv4Packet, Icmpv4Repr, Ipv4Repr,
1233    };
1234
1235    use super::super::{device::SmoltcpDevice, shared::SharedState};
1236    use crate::tcp::connection::NewConnection;
1237
1238    /// Build a minimal Ethernet + IPv4 + TCP SYN frame.
1239    fn build_tcp_syn_frame(
1240        src_ip: [u8; 4],
1241        dst_ip: [u8; 4],
1242        src_port: u16,
1243        dst_port: u16,
1244    ) -> Vec<u8> {
1245        let mut frame = vec![0u8; 14 + 20 + 20]; // eth + ipv4 + tcp
1246
1247        // Ethernet header.
1248        frame[12] = 0x08; // EtherType: IPv4
1249        frame[13] = 0x00;
1250
1251        // IPv4 header.
1252        let ip = &mut frame[14..34];
1253        ip[0] = 0x45; // Version + IHL
1254        let total_len = 40u16; // 20 (IP) + 20 (TCP)
1255        ip[2..4].copy_from_slice(&total_len.to_be_bytes());
1256        ip[6] = 0x40; // Don't Fragment
1257        ip[8] = 64; // TTL
1258        ip[9] = 6; // Protocol: TCP
1259        ip[12..16].copy_from_slice(&src_ip);
1260        ip[16..20].copy_from_slice(&dst_ip);
1261
1262        // TCP header.
1263        let tcp = &mut frame[34..54];
1264        tcp[0..2].copy_from_slice(&src_port.to_be_bytes());
1265        tcp[2..4].copy_from_slice(&dst_port.to_be_bytes());
1266        tcp[12] = 0x50; // Data offset: 5 words
1267        tcp[13] = 0x02; // SYN flag
1268
1269        frame
1270    }
1271
1272    /// Build a minimal Ethernet + IPv4 + UDP frame.
1273    fn build_udp_frame(src_ip: [u8; 4], dst_ip: [u8; 4], src_port: u16, dst_port: u16) -> Vec<u8> {
1274        let mut frame = vec![0u8; 14 + 20 + 8]; // eth + ipv4 + udp
1275
1276        // Ethernet header.
1277        frame[12] = 0x08;
1278        frame[13] = 0x00;
1279
1280        // IPv4 header.
1281        let ip = &mut frame[14..34];
1282        ip[0] = 0x45;
1283        let total_len = 28u16; // 20 (IP) + 8 (UDP)
1284        ip[2..4].copy_from_slice(&total_len.to_be_bytes());
1285        ip[8] = 64;
1286        ip[9] = 17; // Protocol: UDP
1287        ip[12..16].copy_from_slice(&src_ip);
1288        ip[16..20].copy_from_slice(&dst_ip);
1289
1290        // UDP header.
1291        let udp = &mut frame[34..42];
1292        udp[0..2].copy_from_slice(&src_port.to_be_bytes());
1293        udp[2..4].copy_from_slice(&dst_port.to_be_bytes());
1294        let udp_len = 8u16;
1295        udp[4..6].copy_from_slice(&udp_len.to_be_bytes());
1296
1297        frame
1298    }
1299
1300    /// Build a minimal Ethernet + IPv4 + ICMP echo request frame.
1301    fn build_icmpv4_echo_frame(
1302        src_mac: [u8; 6],
1303        dst_mac: [u8; 6],
1304        src_ip: [u8; 4],
1305        dst_ip: [u8; 4],
1306        ident: u16,
1307        seq_no: u16,
1308        data: &[u8],
1309    ) -> Vec<u8> {
1310        let ipv4_repr = Ipv4Repr {
1311            src_addr: Ipv4Addr::from(src_ip),
1312            dst_addr: Ipv4Addr::from(dst_ip),
1313            next_header: IpProtocol::Icmp,
1314            payload_len: 8 + data.len(),
1315            hop_limit: 64,
1316        };
1317        let icmp_repr = Icmpv4Repr::EchoRequest {
1318            ident,
1319            seq_no,
1320            data,
1321        };
1322        let frame_len = 14 + ipv4_repr.buffer_len() + icmp_repr.buffer_len();
1323        let mut frame = vec![0u8; frame_len];
1324
1325        let mut eth_frame = EthernetFrame::new_unchecked(&mut frame);
1326        EthernetRepr {
1327            src_addr: EthernetAddress(src_mac),
1328            dst_addr: EthernetAddress(dst_mac),
1329            ethertype: EthernetProtocol::Ipv4,
1330        }
1331        .emit(&mut eth_frame);
1332
1333        ipv4_repr.emit(
1334            &mut Ipv4Packet::new_unchecked(&mut frame[14..34]),
1335            &ChecksumCapabilities::default(),
1336        );
1337        icmp_repr.emit(
1338            &mut Icmpv4Packet::new_unchecked(&mut frame[34..]),
1339            &ChecksumCapabilities::default(),
1340        );
1341
1342        frame
1343    }
1344
1345    /// Build a minimal Ethernet + ARP request frame.
1346    fn build_arp_request_frame(src_mac: [u8; 6], src_ip: [u8; 4], target_ip: [u8; 4]) -> Vec<u8> {
1347        let mut frame = vec![0u8; 14 + 28];
1348
1349        let mut eth_frame = EthernetFrame::new_unchecked(&mut frame);
1350        EthernetRepr {
1351            src_addr: EthernetAddress(src_mac),
1352            dst_addr: EthernetAddress([0xff; 6]),
1353            ethertype: EthernetProtocol::Arp,
1354        }
1355        .emit(&mut eth_frame);
1356
1357        ArpRepr::EthernetIpv4 {
1358            operation: ArpOperation::Request,
1359            source_hardware_addr: EthernetAddress(src_mac),
1360            source_protocol_addr: Ipv4Addr::from(src_ip),
1361            target_hardware_addr: EthernetAddress([0x00; 6]),
1362            target_protocol_addr: Ipv4Addr::from(target_ip),
1363        }
1364        .emit(&mut ArpPacket::new_unchecked(&mut frame[14..]));
1365
1366        frame
1367    }
1368
1369    #[test]
1370    fn poll_sleep_honors_cleanup_and_earlier_transport_deadlines() {
1371        use std::time::Duration;
1372
1373        let second = Duration::from_secs(1);
1374        assert_eq!(poll_timeout_ms(None, None, None, second), 1000);
1375        assert_eq!(poll_timeout_ms(Some(second * 60), None, None, second), 1000);
1376        for index in 0..3 {
1377            let mut timers = [None; 3];
1378            timers[index] = Some(Duration::from_millis(25));
1379            assert_eq!(poll_timeout_ms(timers[0], timers[1], timers[2], second), 25);
1380            assert_eq!(
1381                poll_timeout_ms(timers[0], timers[1], timers[2], Duration::ZERO),
1382                0
1383            );
1384        }
1385        assert_eq!(poll_timeout_ms(Some(Duration::ZERO), None, None, second), 0);
1386        assert_eq!(
1387            poll_timeout_ms(None, None, None, Duration::from_nanos(1)),
1388            1
1389        );
1390        assert_eq!(poll_timeout_ms(None, None, None, Duration::MAX), i32::MAX);
1391    }
1392
1393    #[test]
1394    fn classify_tcp_syn() {
1395        let frame = build_tcp_syn_frame([10, 0, 0, 2], [93, 184, 216, 34], 54321, 443);
1396        match classify_frame(&frame) {
1397            FrameAction::TcpSyn { src, dst } => {
1398                assert_eq!(
1399                    src,
1400                    SocketAddr::new(Ipv4Addr::new(10, 0, 0, 2).into(), 54321)
1401                );
1402                assert_eq!(
1403                    dst,
1404                    SocketAddr::new(Ipv4Addr::new(93, 184, 216, 34).into(), 443)
1405                );
1406            }
1407            _ => panic!("expected TcpSyn"),
1408        }
1409    }
1410
1411    #[test]
1412    fn classify_tcp_ack_is_passthrough() {
1413        let mut frame = build_tcp_syn_frame([10, 0, 0, 2], [93, 184, 216, 34], 54321, 443);
1414        // Change flags to ACK only (not SYN).
1415        frame[34 + 13] = 0x10; // ACK flag
1416        assert!(matches!(classify_frame(&frame), FrameAction::Passthrough));
1417    }
1418
1419    #[test]
1420    fn classify_udp_dns() {
1421        let frame = build_udp_frame([10, 0, 0, 2], [10, 0, 0, 1], 12345, 53);
1422        assert!(matches!(classify_frame(&frame), FrameAction::Dns));
1423    }
1424
1425    #[test]
1426    fn classify_udp_non_dns() {
1427        let frame = build_udp_frame([10, 0, 0, 2], [8, 8, 8, 8], 12345, 443);
1428        match classify_frame(&frame) {
1429            FrameAction::UdpRelay { src, dst } => {
1430                assert_eq!(src.port(), 12345);
1431                assert_eq!(dst.port(), 443);
1432            }
1433            _ => panic!("expected UdpRelay"),
1434        }
1435    }
1436
1437    #[test]
1438    fn classify_ipv4_udp_fragment() {
1439        let mut frame = build_udp_frame([10, 0, 0, 2], [8, 8, 8, 8], 12345, 443);
1440        frame[14 + 6] = 0x20; // More Fragments flag.
1441        assert!(matches!(
1442            classify_frame(&frame),
1443            FrameAction::Ipv4UdpFragment
1444        ));
1445    }
1446
1447    #[test]
1448    fn classify_ipv6_udp_fragment() {
1449        let mut frame = vec![0u8; 14 + 40 + 8];
1450
1451        frame[12] = 0x86;
1452        frame[13] = 0xdd;
1453
1454        let ip = &mut frame[14..54];
1455        ip[0] = 0x60;
1456        ip[4..6].copy_from_slice(&8u16.to_be_bytes());
1457        ip[6] = u8::from(IpProtocol::Ipv6Frag);
1458        ip[7] = 64;
1459        ip[8..24].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
1460        ip[24..40].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
1461
1462        let fragment = &mut frame[54..62];
1463        fragment[0] = u8::from(IpProtocol::Udp);
1464        fragment[3] = 1; // More Fragments flag.
1465
1466        assert!(matches!(
1467            classify_frame(&frame),
1468            FrameAction::Ipv6UdpFragment
1469        ));
1470    }
1471
1472    #[test]
1473    fn classify_ipv6_non_udp_fragment_is_unsupported() {
1474        let mut frame = vec![0u8; 14 + 40 + 8];
1475
1476        frame[12] = 0x86;
1477        frame[13] = 0xdd;
1478
1479        let ip = &mut frame[14..54];
1480        ip[0] = 0x60;
1481        ip[4..6].copy_from_slice(&8u16.to_be_bytes());
1482        ip[6] = u8::from(IpProtocol::Ipv6Frag);
1483        ip[7] = 64;
1484        ip[8..24].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
1485        ip[24..40].copy_from_slice(&Ipv6Addr::LOCALHOST.octets());
1486
1487        let fragment = &mut frame[54..62];
1488        fragment[0] = u8::from(IpProtocol::Tcp);
1489        fragment[3] = 1; // More Fragments flag.
1490
1491        assert!(matches!(
1492            classify_frame(&frame),
1493            FrameAction::Ipv6UnsupportedFragment
1494        ));
1495    }
1496
1497    #[test]
1498    fn classify_arp_is_passthrough() {
1499        let mut frame = vec![0u8; 42]; // ARP frame
1500        frame[12] = 0x08;
1501        frame[13] = 0x06; // EtherType: ARP
1502        assert!(matches!(classify_frame(&frame), FrameAction::Passthrough));
1503    }
1504
1505    #[test]
1506    fn classify_garbage_is_passthrough() {
1507        assert!(matches!(classify_frame(&[]), FrameAction::Passthrough));
1508        assert!(matches!(classify_frame(&[0; 5]), FrameAction::Passthrough));
1509    }
1510
1511    #[test]
1512    fn gateway_replies_to_icmp_echo_requests() {
1513        fn drive_one_frame(
1514            device: &mut SmoltcpDevice,
1515            iface: &mut Interface,
1516            sockets: &mut SocketSet<'_>,
1517            shared: &Arc<SharedState>,
1518            poll_config: &PollLoopConfig,
1519            now: Instant,
1520        ) {
1521            let frame = device.stage_next_frame().expect("expected staged frame");
1522            if handle_gateway_icmp_echo(
1523                frame,
1524                poll_config,
1525                shared,
1526                &NetworkPolicy::allow_all(),
1527                None,
1528            ) {
1529                device.drop_staged_frame();
1530                return;
1531            }
1532            let _ = iface.poll_ingress_single(now, device, sockets);
1533            let _ = iface.poll_egress(now, device, sockets);
1534        }
1535
1536        let shared = Arc::new(SharedState::new(4));
1537        let poll_config = PollLoopConfig {
1538            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
1539            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
1540            gateway: GatewayIps {
1541                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
1542                ipv6: Some(Ipv6Addr::LOCALHOST),
1543            },
1544            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
1545            guest_ipv6: None,
1546            mtu: 1500,
1547        };
1548        let guest_ipv4 = poll_config.guest_ipv4.unwrap();
1549        let gateway_ipv4 = poll_config.gateway.ipv4.unwrap();
1550        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
1551        let mut iface = create_interface(&mut device, &poll_config);
1552        let mut sockets = SocketSet::new(vec![]);
1553        let now = smoltcp_now();
1554
1555        // Mirror the real guest flow: resolve the gateway MAC before sending
1556        // the ICMP echo request.
1557        shared
1558            .tx_ring
1559            .push(build_arp_request_frame(
1560                poll_config.guest_mac,
1561                guest_ipv4.octets(),
1562                gateway_ipv4.octets(),
1563            ))
1564            .unwrap();
1565        shared
1566            .tx_ring
1567            .push(build_icmpv4_echo_frame(
1568                poll_config.guest_mac,
1569                poll_config.gateway_mac,
1570                guest_ipv4.octets(),
1571                gateway_ipv4.octets(),
1572                0x1234,
1573                0xABCD,
1574                b"ping",
1575            ))
1576            .unwrap();
1577
1578        drive_one_frame(
1579            &mut device,
1580            &mut iface,
1581            &mut sockets,
1582            &shared,
1583            &poll_config,
1584            now,
1585        );
1586        let _ = shared.rx_ring.pop().expect("expected ARP reply");
1587
1588        drive_one_frame(
1589            &mut device,
1590            &mut iface,
1591            &mut sockets,
1592            &shared,
1593            &poll_config,
1594            now,
1595        );
1596
1597        let reply = shared.rx_ring.pop().expect("expected ICMP echo reply");
1598        let eth = EthernetFrame::new_checked(&reply).expect("valid ethernet frame");
1599        assert_eq!(eth.src_addr(), EthernetAddress(poll_config.gateway_mac));
1600        assert_eq!(eth.dst_addr(), EthernetAddress(poll_config.guest_mac));
1601        assert_eq!(eth.ethertype(), EthernetProtocol::Ipv4);
1602
1603        let ipv4 = Ipv4Packet::new_checked(eth.payload()).expect("valid IPv4 packet");
1604        assert_eq!(ipv4.src_addr(), gateway_ipv4);
1605        assert_eq!(ipv4.dst_addr(), guest_ipv4);
1606        assert_eq!(ipv4.next_header(), IpProtocol::Icmp);
1607
1608        let icmp = Icmpv4Packet::new_checked(ipv4.payload()).expect("valid ICMP packet");
1609        let icmp_repr = Icmpv4Repr::parse(&icmp, &ChecksumCapabilities::default())
1610            .expect("valid ICMP echo reply");
1611        assert_eq!(
1612            icmp_repr,
1613            Icmpv4Repr::EchoReply {
1614                ident: 0x1234,
1615                seq_no: 0xABCD,
1616                data: b"ping",
1617            }
1618        );
1619    }
1620
1621    #[test]
1622    fn gateway_icmp_echo_respects_deny_policy() {
1623        let shared = SharedState::new(4);
1624        let poll_config = PollLoopConfig {
1625            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
1626            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
1627            gateway: GatewayIps {
1628                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
1629                ipv6: None,
1630            },
1631            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
1632            guest_ipv6: None,
1633            mtu: 1500,
1634        };
1635        let policy = NetworkPolicy::builder().default_deny().build().unwrap();
1636        let frame = build_icmpv4_echo_frame(
1637            poll_config.guest_mac,
1638            poll_config.gateway_mac,
1639            poll_config.guest_ipv4.unwrap().octets(),
1640            poll_config.gateway.ipv4.unwrap().octets(),
1641            0x1234,
1642            0xABCD,
1643            b"ping",
1644        );
1645
1646        assert!(handle_gateway_icmp_echo(
1647            &frame,
1648            &poll_config,
1649            &shared,
1650            &policy,
1651            None,
1652        ));
1653        assert!(
1654            shared.rx_ring.pop().is_none(),
1655            "denied gateway ICMP should not queue a reply"
1656        );
1657    }
1658
1659    #[test]
1660    fn platform_public_floor_consumes_gateway_echo_without_replying() {
1661        let shared = SharedState::new(4);
1662        let gateway = Ipv4Addr::new(100, 96, 0, 1);
1663        let guest = Ipv4Addr::new(100, 96, 0, 2);
1664        shared.set_gateway_ips(Some(gateway), None);
1665        let config = PollLoopConfig {
1666            gateway_mac: [0x02, 0, 0, 0, 0, 1],
1667            guest_mac: [0x02, 0, 0, 0, 0, 2],
1668            gateway: GatewayIps {
1669                ipv4: Some(gateway),
1670                ipv6: None,
1671            },
1672            guest_ipv4: Some(guest),
1673            guest_ipv6: None,
1674            mtu: 1500,
1675        };
1676        let frame = build_icmpv4_echo_frame(
1677            config.guest_mac,
1678            config.gateway_mac,
1679            guest.octets(),
1680            gateway.octets(),
1681            1,
1682            1,
1683            b"ping",
1684        );
1685        let platform = NetworkPolicy::from_profiles([crate::policy::NetworkProfile::Public]);
1686
1687        assert!(handle_gateway_icmp_echo(
1688            &frame,
1689            &config,
1690            &shared,
1691            &NetworkPolicy::allow_all(),
1692            Some(&platform),
1693        ));
1694        assert!(shared.rx_ring.pop().is_none());
1695    }
1696
1697    fn test_gateway() -> GatewayIps {
1698        GatewayIps {
1699            ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
1700            ipv6: Some("fd42:6d73:62::1".parse().unwrap()),
1701        }
1702    }
1703
1704    #[test]
1705    fn resolve_tcp_host_target_ipv4_can_fall_back_to_ipv6() {
1706        let gw = test_gateway();
1707        let dst = SocketAddr::new(IpAddr::V4(gw.ipv4.unwrap()), 8080);
1708
1709        assert_eq!(
1710            resolve_tcp_host_target(dst, gw),
1711            UpstreamTcpTarget::with_fallback(
1712                SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080),
1713                SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 8080),
1714            )
1715        );
1716    }
1717
1718    #[test]
1719    fn resolve_tcp_host_target_ipv6_can_fall_back_to_ipv4() {
1720        let gw = test_gateway();
1721        let dst = SocketAddr::new(IpAddr::V6(gw.ipv6.unwrap()), 8080);
1722
1723        assert_eq!(
1724            resolve_tcp_host_target(dst, gw),
1725            UpstreamTcpTarget::with_fallback(
1726                SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 8080),
1727                SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080),
1728            )
1729        );
1730    }
1731
1732    #[test]
1733    fn resolve_tcp_host_target_external_has_no_fallback() {
1734        let dst = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 443);
1735
1736        assert_eq!(
1737            resolve_tcp_host_target(dst, test_gateway()),
1738            UpstreamTcpTarget::direct(dst)
1739        );
1740    }
1741
1742    #[test]
1743    fn resolve_host_dst_matches_ipv4() {
1744        let gw = test_gateway();
1745        let dst = SocketAddr::new(IpAddr::V4(gw.ipv4.unwrap()), 8080);
1746        assert_eq!(
1747            resolve_host_dst(dst, gw),
1748            SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080)
1749        );
1750    }
1751
1752    #[test]
1753    fn resolve_host_dst_matches_ipv6() {
1754        let gw = test_gateway();
1755        let dst = SocketAddr::new(IpAddr::V6(gw.ipv6.unwrap()), 8080);
1756        assert_eq!(
1757            resolve_host_dst(dst, gw),
1758            SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 8080)
1759        );
1760    }
1761
1762    #[test]
1763    fn resolve_host_dst_passes_through_when_family_absent() {
1764        let gw = GatewayIps {
1765            ipv4: None,
1766            ipv6: Some("fd42:6d73:62::1".parse().unwrap()),
1767        };
1768        // IPv4 dst with no IPv4 gateway must not be rewritten to loopback.
1769        let dst = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(100, 96, 0, 1)), 8080);
1770        assert_eq!(resolve_host_dst(dst, gw), dst);
1771    }
1772
1773    #[test]
1774    fn resolve_host_dst_passes_through_non_gateway() {
1775        let gw = test_gateway();
1776        let dst = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 443);
1777        assert_eq!(resolve_host_dst(dst, gw), dst);
1778    }
1779
1780    #[test]
1781    fn outbound_proxy_is_skipped_for_host_destination() {
1782        let gw = test_gateway();
1783        let guest_dst = SocketAddr::new(IpAddr::V4(gw.ipv4.unwrap()), 8080);
1784        let connect_target = resolve_tcp_host_target(guest_dst, gw);
1785        let proxy = Some(Arc::new(ResolvedOutboundProxy::Socks5 {
1786            address: "192.0.2.1:1080".parse().unwrap(),
1787            credentials: None,
1788        }));
1789
1790        assert!(
1791            ResolvedOutboundProxy::select_for_destination(
1792                &proxy,
1793                guest_dst,
1794                connect_target.primary(),
1795            )
1796            .is_none()
1797        );
1798    }
1799
1800    #[test]
1801    fn outbound_proxy_is_preserved_for_external_destination() {
1802        let gw = test_gateway();
1803        let guest_dst = "198.51.100.10:443".parse().unwrap();
1804        let connect_target = resolve_tcp_host_target(guest_dst, gw);
1805        let proxy = Some(Arc::new(ResolvedOutboundProxy::Socks5 {
1806            address: "192.0.2.1:1080".parse().unwrap(),
1807            credentials: None,
1808        }));
1809
1810        assert!(
1811            ResolvedOutboundProxy::select_for_destination(
1812                &proxy,
1813                guest_dst,
1814                connect_target.primary(),
1815            )
1816            .is_some()
1817        );
1818    }
1819
1820    #[test]
1821    fn external_icmp_echo_requests_are_not_answered_locally() {
1822        fn drive_one_frame(
1823            device: &mut SmoltcpDevice,
1824            iface: &mut Interface,
1825            sockets: &mut SocketSet<'_>,
1826            shared: &Arc<SharedState>,
1827            poll_config: &PollLoopConfig,
1828            now: Instant,
1829        ) {
1830            let frame = device.stage_next_frame().expect("expected staged frame");
1831            if handle_gateway_icmp_echo(
1832                frame,
1833                poll_config,
1834                shared,
1835                &NetworkPolicy::allow_all(),
1836                None,
1837            ) {
1838                device.drop_staged_frame();
1839                return;
1840            }
1841            let _ = iface.poll_ingress_single(now, device, sockets);
1842            let _ = iface.poll_egress(now, device, sockets);
1843        }
1844
1845        let shared = Arc::new(SharedState::new(4));
1846        let poll_config = PollLoopConfig {
1847            gateway_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x01],
1848            guest_mac: [0x02, 0x00, 0x00, 0x00, 0x00, 0x02],
1849            gateway: GatewayIps {
1850                ipv4: Some(Ipv4Addr::new(100, 96, 0, 1)),
1851                ipv6: Some(Ipv6Addr::LOCALHOST),
1852            },
1853            guest_ipv4: Some(Ipv4Addr::new(100, 96, 0, 2)),
1854            guest_ipv6: None,
1855            mtu: 1500,
1856        };
1857        let guest_ipv4 = poll_config.guest_ipv4.unwrap();
1858        let gateway_ipv4 = poll_config.gateway.ipv4.unwrap();
1859        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
1860        let mut iface = create_interface(&mut device, &poll_config);
1861        let mut sockets = SocketSet::new(vec![]);
1862        let now = smoltcp_now();
1863
1864        shared
1865            .tx_ring
1866            .push(build_arp_request_frame(
1867                poll_config.guest_mac,
1868                guest_ipv4.octets(),
1869                gateway_ipv4.octets(),
1870            ))
1871            .unwrap();
1872        shared
1873            .tx_ring
1874            .push(build_icmpv4_echo_frame(
1875                poll_config.guest_mac,
1876                poll_config.gateway_mac,
1877                guest_ipv4.octets(),
1878                [142, 251, 216, 46],
1879                0x1234,
1880                0xABCD,
1881                b"ping",
1882            ))
1883            .unwrap();
1884
1885        drive_one_frame(
1886            &mut device,
1887            &mut iface,
1888            &mut sockets,
1889            &shared,
1890            &poll_config,
1891            now,
1892        );
1893        let _ = shared.rx_ring.pop().expect("expected ARP reply");
1894
1895        drive_one_frame(
1896            &mut device,
1897            &mut iface,
1898            &mut sockets,
1899            &shared,
1900            &poll_config,
1901            now,
1902        );
1903        assert!(
1904            shared.rx_ring.pop().is_none(),
1905            "external ICMP should not be answered locally"
1906        );
1907    }
1908
1909    // ─────────────────────────────────────────────────────────────────────
1910    // Guest-initiated TCP teardown.
1911    //
1912    // These tests drive the real `TcpConnectionTracker` + smoltcp interface
1913    // through a full TCP handshake and then a guest-initiated teardown,
1914    // asserting the observable the proxy task sees on its channel:
1915    //   - guest FIN => half-close propagated (channel EOF), server → guest
1916    //     stays open, and the slot is reclaimed once the proxy task exits
1917    //   - guest RST => immediate clean teardown
1918    // Regression tests for the CLOSE_WAIT orphan leak: without close
1919    // propagation, a guest FIN left the socket in CLOSE_WAIT forever, the
1920    // proxy task blocked on `from_smoltcp.recv()`, the upstream socket
1921    // open, and the connection-table slot consumed until the table filled.
1922    // ─────────────────────────────────────────────────────────────────────
1923
1924    use smoltcp::socket::tcp;
1925    use smoltcp::wire::{TcpControl, TcpPacket, TcpRepr, TcpSeqNumber};
1926
1927    const GUEST_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02];
1928    const GATEWAY_MAC: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x01];
1929    const GUEST_IP: [u8; 4] = [100, 96, 0, 2];
1930    const GATEWAY_IP: [u8; 4] = [100, 96, 0, 1];
1931    // Off-subnet external destination reached via the default route.
1932    const SERVER_IP: [u8; 4] = [93, 184, 216, 34];
1933
1934    fn leak_poll_config() -> PollLoopConfig {
1935        PollLoopConfig {
1936            gateway_mac: GATEWAY_MAC,
1937            guest_mac: GUEST_MAC,
1938            gateway: GatewayIps {
1939                ipv4: Some(Ipv4Addr::from(GATEWAY_IP)),
1940                ipv6: None,
1941            },
1942            guest_ipv4: Some(Ipv4Addr::from(GUEST_IP)),
1943            guest_ipv6: None,
1944            mtu: 1500,
1945        }
1946    }
1947
1948    /// Build an Ethernet+IPv4+TCP frame from the guest with correct checksums.
1949    #[allow(clippy::too_many_arguments)]
1950    fn build_tcp_frame(
1951        src_port: u16,
1952        dst_port: u16,
1953        control: TcpControl,
1954        seq: i32,
1955        ack: Option<i32>,
1956        payload: &[u8],
1957    ) -> Vec<u8> {
1958        let src_ip = Ipv4Addr::from(GUEST_IP);
1959        let dst_ip = Ipv4Addr::from(SERVER_IP);
1960
1961        let tcp_repr = TcpRepr {
1962            src_port,
1963            dst_port,
1964            control,
1965            seq_number: TcpSeqNumber(seq),
1966            ack_number: ack.map(TcpSeqNumber),
1967            window_len: 65535,
1968            window_scale: None,
1969            max_seg_size: None,
1970            sack_permitted: false,
1971            sack_ranges: [None, None, None],
1972            timestamp: None,
1973            payload,
1974        };
1975        let ipv4_repr = Ipv4Repr {
1976            src_addr: src_ip,
1977            dst_addr: dst_ip,
1978            next_header: IpProtocol::Tcp,
1979            payload_len: tcp_repr.buffer_len(),
1980            hop_limit: 64,
1981        };
1982
1983        let frame_len = 14 + ipv4_repr.buffer_len() + tcp_repr.buffer_len();
1984        let mut frame = vec![0u8; frame_len];
1985
1986        let mut eth = EthernetFrame::new_unchecked(&mut frame);
1987        EthernetRepr {
1988            src_addr: EthernetAddress(GUEST_MAC),
1989            dst_addr: EthernetAddress(GATEWAY_MAC),
1990            ethertype: EthernetProtocol::Ipv4,
1991        }
1992        .emit(&mut eth);
1993
1994        let ip_end = 14 + ipv4_repr.buffer_len();
1995        ipv4_repr.emit(
1996            &mut Ipv4Packet::new_unchecked(&mut frame[14..ip_end]),
1997            &ChecksumCapabilities::default(),
1998        );
1999        tcp_repr.emit(
2000            &mut TcpPacket::new_unchecked(&mut frame[ip_end..]),
2001            &IpAddress::Ipv4(src_ip),
2002            &IpAddress::Ipv4(dst_ip),
2003            &ChecksumCapabilities::default(),
2004        );
2005
2006        frame
2007    }
2008
2009    /// An abandoned handshake must not donate its host stream to a retry.
2010    #[test]
2011    fn reset_before_guest_ack_requires_a_fresh_host_connection() {
2012        const CHILD_ENV: &str = "MSB_TEST_RESET_PRECONNECTED_HANDSHAKE";
2013        if std::env::var_os(CHILD_ENV).is_none() {
2014            let output = std::process::Command::new(std::env::current_exe().unwrap())
2015                .args([
2016                    "--exact",
2017                    "engine::netstack::poll::tests::reset_before_guest_ack_requires_a_fresh_host_connection",
2018                    "--nocapture",
2019                ])
2020                .env(CHILD_ENV, "1")
2021                .output()
2022                .unwrap();
2023            assert!(
2024                output.status.success(),
2025                "{}\n{}",
2026                String::from_utf8_lossy(&output.stdout),
2027                String::from_utf8_lossy(&output.stderr)
2028            );
2029            return;
2030        }
2031        let runtime = tokio::runtime::Builder::new_current_thread()
2032            .enable_all()
2033            .build()
2034            .unwrap();
2035        runtime.block_on(async {
2036            // Cover both reuse of the old four-tuple and a new ephemeral port.
2037            for retry_port in [40000, 40001] {
2038                tokio::time::timeout(
2039                    std::time::Duration::from_secs(5),
2040                    check_reset_before_guest_ack(retry_port),
2041                )
2042                .await
2043                .expect("reset/retry regression timed out");
2044            }
2045        });
2046    }
2047
2048    async fn check_reset_before_guest_ack(retry_port: u16) {
2049        use std::time::Duration;
2050        use tokio::io::AsyncReadExt;
2051        use tokio::net::TcpListener;
2052
2053        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
2054        let port = listener.local_addr().unwrap().port();
2055        let shared = Arc::new(SharedState::new(256));
2056        let loop_shared = shared.clone();
2057        let handle = tokio::runtime::Handle::current();
2058        std::thread::spawn(move || {
2059            smoltcp_poll_loop(
2060                loop_shared,
2061                leak_poll_config(),
2062                NetworkPolicy::allow_all(),
2063                None,
2064                DnsConfig::default(),
2065                None,
2066                vec![],
2067                false,
2068                None,
2069                None,
2070                TcpAcceptQueueSize::DEFAULT,
2071                None,
2072                handle,
2073                SecretsHandle::new(Default::default()),
2074                None,
2075            );
2076        });
2077        let send = |guest_port, control, seq, ack, payload: &[u8]| {
2078            let mut frame = build_tcp_frame(guest_port, port, control, seq, ack, payload);
2079            let guest = IpAddress::Ipv4(Ipv4Addr::from(GUEST_IP));
2080            let gateway = IpAddress::Ipv4(Ipv4Addr::from(GATEWAY_IP));
2081            let mut ip = Ipv4Packet::new_unchecked(&mut frame[14..34]);
2082            ip.set_dst_addr(Ipv4Addr::from(GATEWAY_IP));
2083            ip.fill_checksum();
2084            TcpPacket::new_unchecked(&mut frame[34..]).fill_checksum(&guest, &gateway);
2085            shared.tx_ring.push(frame).unwrap();
2086        };
2087        shared
2088            .tx_ring
2089            .push(build_arp_request_frame(GUEST_MAC, GUEST_IP, GATEWAY_IP))
2090            .unwrap();
2091        send(40000, TcpControl::Syn, 1000, None, &[]);
2092        shared.tx_wake.wake();
2093        let (mut abandoned, _) = listener.accept().await.unwrap();
2094        let first_isn = loop {
2095            if let Some((seq, ack, syn, _, rst)) = last_tcp_reply(&shared) {
2096                assert!(!rst);
2097                if syn {
2098                    assert_eq!(ack, 1001);
2099                    break seq;
2100                }
2101            }
2102            tokio::time::sleep(Duration::from_millis(1)).await;
2103        };
2104        // Queue both packets before waking the loop: a fix must work even
2105        // when reset and retry are processed in a single ingress batch.
2106        send(
2107            40000,
2108            TcpControl::Rst,
2109            1001,
2110            Some(first_isn.wrapping_add(1)),
2111            &[],
2112        );
2113        send(retry_port, TcpControl::Syn, 2000, None, &[]);
2114        shared.tx_wake.wake();
2115        let retry_isn = loop {
2116            if let Some((seq, ack, syn, _, rst)) = last_tcp_reply(&shared) {
2117                assert!(!rst, "retry was unexpectedly reset");
2118                if syn && ack == 2001 {
2119                    break seq;
2120                }
2121            }
2122            tokio::time::sleep(Duration::from_millis(1)).await;
2123        };
2124        send(
2125            retry_port,
2126            TcpControl::None,
2127            2001,
2128            Some(retry_isn.wrapping_add(1)),
2129            b"fresh",
2130        );
2131        shared.tx_wake.wake();
2132        let (mut fresh, _) = tokio::time::timeout(Duration::from_secs(1), listener.accept())
2133            .await
2134            .expect("retry reused the abandoned host stream instead of opening a fresh connection")
2135            .unwrap();
2136        let mut request = [0; 5];
2137        fresh.read_exact(&mut request).await.unwrap();
2138        assert_eq!(&request, b"fresh");
2139        assert_eq!(
2140            abandoned.read(&mut request).await.unwrap(),
2141            0,
2142            "abandoned host socket must close without receiving retry data"
2143        );
2144    }
2145
2146    /// Run the real OS-thread loop in a subprocess because it has no shutdown API.
2147    #[test]
2148    fn stalled_drain_wakes_poll_loop_without_network_events() {
2149        const CHILD_ENV: &str = "MSB_TEST_STALLED_DRAIN_WAKE_CHILD";
2150        if std::env::var_os(CHILD_ENV).is_none() {
2151            let output = std::process::Command::new(std::env::current_exe().unwrap())
2152                .args([
2153                    "--exact",
2154                    "engine::netstack::poll::tests::stalled_drain_wakes_poll_loop_without_network_events",
2155                    "--nocapture",
2156                ])
2157                .env(CHILD_ENV, "1")
2158                .output()
2159                .unwrap();
2160            assert!(
2161                output.status.success(),
2162                "poll-loop child failed:\n{}\n{}",
2163                String::from_utf8_lossy(&output.stdout),
2164                String::from_utf8_lossy(&output.stderr),
2165            );
2166            return;
2167        }
2168
2169        let runtime = tokio::runtime::Builder::new_current_thread()
2170            .enable_all()
2171            .build()
2172            .unwrap();
2173        runtime.block_on(async {
2174            tokio::time::timeout(std::time::Duration::from_secs(65), async {
2175                check_stalled_drain_wakeup().await;
2176            })
2177            .await
2178            .expect("poll-loop regression timed out");
2179        });
2180    }
2181
2182    async fn check_stalled_drain_wakeup() {
2183        use std::time::Duration;
2184
2185        use tokio::io::{AsyncReadExt, AsyncWriteExt};
2186        use tokio::net::TcpListener;
2187
2188        let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
2189        let port = listener.local_addr().unwrap().port();
2190        let (sent_tx, sent_rx) = tokio::sync::oneshot::channel();
2191        let host = tokio::spawn(async move {
2192            let (mut stream, _) = listener.accept().await.unwrap();
2193            let mut request = [0; 6];
2194            stream.read_exact(&mut request).await.unwrap();
2195            assert_eq!(&request, b"hello\n");
2196            stream.write_all(&vec![b'x'; 131072]).await.unwrap();
2197            sent_tx.send(()).unwrap();
2198
2199            // Let zero-window retransmission backoff grow before host EOF starts
2200            // the drain timer. The next TCP timer alone would wake too late.
2201            tokio::time::sleep(Duration::from_secs(10)).await;
2202        });
2203
2204        let shared = Arc::new(SharedState::new(256));
2205        let loop_shared = shared.clone();
2206        let handle = tokio::runtime::Handle::current();
2207        std::thread::spawn(move || {
2208            smoltcp_poll_loop(
2209                loop_shared,
2210                leak_poll_config(),
2211                NetworkPolicy::allow_all(),
2212                None,
2213                DnsConfig::default(),
2214                None,
2215                vec![],
2216                false,
2217                None,
2218                None,
2219                TcpAcceptQueueSize::DEFAULT,
2220                None,
2221                handle,
2222                SecretsHandle::new(Default::default()),
2223                None,
2224            );
2225        });
2226
2227        let send = |control, seq, ack, payload: &[u8]| {
2228            let mut frame = build_tcp_frame(54321, port, control, seq, ack, payload);
2229            let guest = IpAddress::Ipv4(Ipv4Addr::from(GUEST_IP));
2230            let gateway = IpAddress::Ipv4(Ipv4Addr::from(GATEWAY_IP));
2231            let mut ip = Ipv4Packet::new_unchecked(&mut frame[14..34]);
2232            ip.set_dst_addr(Ipv4Addr::from(GATEWAY_IP));
2233            ip.fill_checksum();
2234            let mut tcp = TcpPacket::new_unchecked(&mut frame[34..]);
2235            tcp.set_window_len(0);
2236            tcp.fill_checksum(&guest, &gateway);
2237            shared.tx_ring.push(frame).unwrap();
2238            shared.tx_wake.wake();
2239        };
2240        shared
2241            .tx_ring
2242            .push(build_arp_request_frame(GUEST_MAC, GUEST_IP, GATEWAY_IP))
2243            .unwrap();
2244        send(TcpControl::Syn, 1000, None, &[]);
2245        let server_isn = tokio::time::timeout(Duration::from_secs(5), async {
2246            loop {
2247                if let Some((seq, _, syn, _, rst)) = last_tcp_reply(&shared) {
2248                    assert!(!rst, "unexpected reset during handshake");
2249                    if syn {
2250                        break seq;
2251                    }
2252                }
2253                tokio::time::sleep(Duration::from_millis(1)).await;
2254            }
2255        })
2256        .await
2257        .expect("guest handshake timed out");
2258        send(
2259            TcpControl::None,
2260            1001,
2261            Some(server_isn.wrapping_add(1)),
2262            b"hello\n",
2263        );
2264        tokio::time::timeout(Duration::from_secs(5), sent_rx)
2265            .await
2266            .expect("host did not receive request")
2267            .unwrap();
2268        let started = std::time::Instant::now();
2269
2270        // From here on, never send guest packets or wake the network thread.
2271        // Reading its output ring does not signal either of its wake sources.
2272        tokio::time::timeout(Duration::from_secs(55), async {
2273            loop {
2274                if let Some((_, _, _, fin, rst)) = last_tcp_reply(&shared) {
2275                    assert!(!fin, "queued response must not be discarded with FIN");
2276                    if rst {
2277                        break;
2278                    }
2279                }
2280                tokio::time::sleep(Duration::from_millis(5)).await;
2281            }
2282        })
2283        .await
2284        .expect("drain deadline did not wake the idle poll loop");
2285        assert!(
2286            started.elapsed() >= Duration::from_secs(30),
2287            "premature reset"
2288        );
2289        host.await.unwrap();
2290    }
2291
2292    /// Push one guest frame, run a single ingress pass, then drain egress.
2293    fn ingress(
2294        frame: Vec<u8>,
2295        device: &mut SmoltcpDevice,
2296        iface: &mut Interface,
2297        sockets: &mut SocketSet<'_>,
2298        shared: &Arc<SharedState>,
2299        now: Instant,
2300    ) {
2301        shared.tx_ring.push(frame).unwrap();
2302        device.stage_next_frame().expect("frame should stage");
2303        iface.poll_ingress_single(now, device, sockets);
2304        loop {
2305            let r = iface.poll_egress(now, device, sockets);
2306            if matches!(r, smoltcp::iface::PollResult::None) {
2307                break;
2308            }
2309        }
2310    }
2311
2312    /// Pop the newest smoltcp→guest reply and return its (seq, ack, is_syn,
2313    /// is_fin, is_rst). Drains all queued replies, returning the last TCP one.
2314    fn last_tcp_reply(shared: &Arc<SharedState>) -> Option<(i32, i32, bool, bool, bool)> {
2315        let mut out = None;
2316        while let Some(frame) = shared.rx_ring.pop() {
2317            if frame.len() < 34 {
2318                continue;
2319            }
2320            // eth(14) + ipv4(20) then TCP.
2321            if frame[23] != 6 {
2322                continue; // not TCP (e.g. ARP reply passes as non-IPv4)
2323            }
2324            let tcp = match TcpPacket::new_checked(&frame[34..]) {
2325                Ok(p) => p,
2326                Err(_) => continue,
2327            };
2328            out = Some((
2329                tcp.seq_number().0,
2330                tcp.ack_number().0,
2331                tcp.syn(),
2332                tcp.fin(),
2333                tcp.rst(),
2334            ));
2335        }
2336        out
2337    }
2338
2339    /// Find the single tracked TCP socket's state.
2340    fn only_tcp_state(sockets: &SocketSet<'_>) -> Option<tcp::State> {
2341        for (_h, sock) in sockets.iter() {
2342            if let smoltcp::socket::Socket::Tcp(s) = sock {
2343                return Some(s.state());
2344            }
2345        }
2346        None
2347    }
2348
2349    /// Drive guest→server handshake to ESTABLISHED, returning (server_isn,
2350    /// guest_seq_after_handshake).
2351    fn handshake(
2352        tracker: &mut TcpConnectionTracker,
2353        device: &mut SmoltcpDevice,
2354        iface: &mut Interface,
2355        sockets: &mut SocketSet<'_>,
2356        shared: &Arc<SharedState>,
2357        now: Instant,
2358        guest_port: u16,
2359    ) -> (i32, i32) {
2360        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), guest_port);
2361        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);
2362
2363        // Pre-populate neighbor cache so smoltcp can address replies to the
2364        // guest without stalling on ARP.
2365        ingress(
2366            build_arp_request_frame(GUEST_MAC, GUEST_IP, GATEWAY_IP),
2367            device,
2368            iface,
2369            sockets,
2370            shared,
2371            now,
2372        );
2373        let _ = shared.rx_ring.pop(); // ARP reply
2374
2375        let guest_isn = 1000i32;
2376
2377        // 1) Guest SYN — tracker creates the listening socket first (as the
2378        //    real poll loop does), then smoltcp completes the handshake.
2379        assert!(
2380            tracker.create_tcp_socket(src, dst, sockets, TcpAcceptMode::AcceptImmediately),
2381            "socket creation should succeed under the limit"
2382        );
2383        ingress(
2384            build_tcp_frame(guest_port, 443, TcpControl::Syn, guest_isn, None, &[]),
2385            device,
2386            iface,
2387            sockets,
2388            shared,
2389            now,
2390        );
2391        let (server_isn, ack, is_syn, _, _) =
2392            last_tcp_reply(shared).expect("expected SYN-ACK from smoltcp");
2393        assert!(is_syn, "expected SYN flag on handshake reply");
2394        assert_eq!(ack, guest_isn + 1, "SYN-ACK should ack guest ISN+1");
2395
2396        // 2) Guest ACK — completes the handshake.
2397        ingress(
2398            build_tcp_frame(
2399                guest_port,
2400                443,
2401                TcpControl::None,
2402                guest_isn + 1,
2403                Some(server_isn + 1),
2404                &[],
2405            ),
2406            device,
2407            iface,
2408            sockets,
2409            shared,
2410            now,
2411        );
2412        assert_eq!(
2413            only_tcp_state(sockets),
2414            Some(tcp::State::Established),
2415            "socket should be ESTABLISHED after handshake",
2416        );
2417
2418        (server_isn, guest_isn + 1)
2419    }
2420
2421    /// Complete a guest→server handshake and hand the connection to a proxy.
2422    fn establish(
2423        tracker: &mut TcpConnectionTracker,
2424        device: &mut SmoltcpDevice,
2425        iface: &mut Interface,
2426        sockets: &mut SocketSet<'_>,
2427        shared: &Arc<SharedState>,
2428        now: Instant,
2429        guest_port: u16,
2430    ) -> (i32, i32, Vec<NewConnection>) {
2431        let (server_isn, guest_seq) =
2432            handshake(tracker, device, iface, sockets, shared, now, guest_port);
2433
2434        // Poll loop detects the established connection and spawns a proxy.
2435        let new_conns = tracker.take_new_connections(sockets);
2436        assert_eq!(
2437            new_conns.len(),
2438            1,
2439            "one new connection should be handed off"
2440        );
2441
2442        (server_isn, guest_seq, new_conns)
2443    }
2444
2445    #[test]
2446    fn guest_fin_propagates_half_close_without_killing_the_connection() {
2447        let shared = Arc::new(SharedState::new(64));
2448        let poll_config = leak_poll_config();
2449        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
2450        let mut iface = create_interface(&mut device, &poll_config);
2451        let mut sockets = SocketSet::new(vec![]);
2452        let mut tracker = TcpConnectionTracker::new(None);
2453        let now = smoltcp_now();
2454
2455        let (server_isn, guest_seq, mut new_conns) = establish(
2456            &mut tracker,
2457            &mut device,
2458            &mut iface,
2459            &mut sockets,
2460            &shared,
2461            now,
2462            54321,
2463        );
2464        // The proxy side of the channel: a real proxy task blocks on
2465        // `from_smoltcp.recv()` while the connection is idle.
2466        let conn = new_conns.remove(0);
2467        let mut from_smoltcp = conn.from_smoltcp;
2468        let to_smoltcp = conn.to_smoltcp;
2469
2470        // Guest half-closes (shutdown(SHUT_WR), or its process exits while
2471        // holding an idle keep-alive: no unread data => FIN, not RST).
2472        ingress(
2473            build_tcp_frame(
2474                54321,
2475                443,
2476                TcpControl::Fin,
2477                guest_seq,
2478                Some(server_isn + 1),
2479                &[],
2480            ),
2481            &mut device,
2482            &mut iface,
2483            &mut sockets,
2484            &shared,
2485            now,
2486        );
2487        assert_eq!(
2488            only_tcp_state(&sockets),
2489            Some(tcp::State::CloseWait),
2490            "guest FIN should move the smoltcp socket to CLOSE_WAIT",
2491        );
2492
2493        // One relay pass propagates the half-close: the proxy task's
2494        // receiver disconnects, so its `recv()` returns `None` and it can
2495        // shut down the guest → server direction upstream.
2496        tracker.relay_data(&mut sockets);
2497        assert!(
2498            matches!(
2499                from_smoltcp.try_recv(),
2500                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected)
2501            ),
2502            "guest FIN must propagate EOF to the proxy task",
2503        );
2504
2505        // The connection must NOT be torn down: a half-closed guest can
2506        // still receive. Pending server → guest data is still delivered.
2507        assert!(
2508            tracker.has_socket_for(
2509                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54321),
2510                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
2511            ),
2512            "half-closed connection must stay tracked while the proxy runs",
2513        );
2514        let payload = b"pending server response";
2515        to_smoltcp
2516            .try_send(bytes::Bytes::from_static(payload))
2517            .expect("server → guest channel should accept data");
2518        while shared.rx_ring.pop().is_some() {} // drain handshake/ACK frames
2519        tracker.relay_data(&mut sockets);
2520        loop {
2521            let r = iface.poll_egress(now, &mut device, &mut sockets);
2522            if matches!(r, smoltcp::iface::PollResult::None) {
2523                break;
2524            }
2525        }
2526        let mut delivered = false;
2527        while let Some(frame) = shared.rx_ring.pop() {
2528            if frame.windows(payload.len()).any(|w| w == payload) {
2529                delivered = true;
2530            }
2531        }
2532        assert!(
2533            delivered,
2534            "server data must still reach a half-closed guest",
2535        );
2536        assert_eq!(
2537            only_tcp_state(&sockets),
2538            Some(tcp::State::CloseWait),
2539            "socket must stay open (CLOSE_WAIT) while the proxy is alive",
2540        );
2541    }
2542
2543    #[test]
2544    fn guest_fin_connection_is_reaped_after_proxy_exit() {
2545        let shared = Arc::new(SharedState::new(64));
2546        let poll_config = leak_poll_config();
2547        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
2548        let mut iface = create_interface(&mut device, &poll_config);
2549        let mut sockets = SocketSet::new(vec![]);
2550        let mut tracker = TcpConnectionTracker::new(None);
2551        let now = smoltcp_now();
2552
2553        let (server_isn, guest_seq, mut new_conns) = establish(
2554            &mut tracker,
2555            &mut device,
2556            &mut iface,
2557            &mut sockets,
2558            &shared,
2559            now,
2560            54321,
2561        );
2562        let conn = new_conns.remove(0);
2563        let from_smoltcp = conn.from_smoltcp;
2564        let to_smoltcp = conn.to_smoltcp;
2565
2566        // Guest FIN, then a relay pass to propagate the half-close.
2567        ingress(
2568            build_tcp_frame(
2569                54321,
2570                443,
2571                TcpControl::Fin,
2572                guest_seq,
2573                Some(server_isn + 1),
2574                &[],
2575            ),
2576            &mut device,
2577            &mut iface,
2578            &mut sockets,
2579            &shared,
2580            now,
2581        );
2582        tracker.relay_data(&mut sockets);
2583
2584        // The proxy task exits (upstream closed after seeing our FIN) and
2585        // drops its channel ends.
2586        drop(from_smoltcp);
2587        drop(to_smoltcp);
2588
2589        // The tracker detects the proxy exit and closes the socket: the
2590        // guest gets our FIN (CLOSE_WAIT → LAST_ACK)...
2591        while shared.rx_ring.pop().is_some() {} // drain handshake/ACK frames
2592        tracker.relay_data(&mut sockets);
2593        loop {
2594            let r = iface.poll_egress(now, &mut device, &mut sockets);
2595            if matches!(r, smoltcp::iface::PollResult::None) {
2596                break;
2597            }
2598        }
2599        let (fin_seq, _, _, is_fin, _) =
2600            last_tcp_reply(&shared).expect("expected FIN toward the guest");
2601        assert!(is_fin, "proxy exit after guest FIN must FIN the guest side");
2602
2603        // ...and the guest's final ACK completes the close.
2604        ingress(
2605            build_tcp_frame(
2606                54321,
2607                443,
2608                TcpControl::None,
2609                guest_seq + 1,
2610                Some(fin_seq + 1),
2611                &[],
2612            ),
2613            &mut device,
2614            &mut iface,
2615            &mut sockets,
2616            &shared,
2617            now,
2618        );
2619        tracker.relay_data(&mut sockets);
2620        tracker.cleanup_closed(&mut sockets);
2621
2622        // The socket is reaped and the table slot is free again — no
2623        // CLOSE_WAIT orphan pinning a slot until the table fills.
2624        assert!(
2625            !tracker.has_socket_for(
2626                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54321),
2627                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
2628            ),
2629            "connection must be evicted after FIN + proxy exit",
2630        );
2631        assert_eq!(
2632            only_tcp_state(&sockets),
2633            None,
2634            "socket must be removed from the socket set",
2635        );
2636    }
2637
2638    #[test]
2639    fn guest_rst_is_cleaned_up() {
2640        let shared = Arc::new(SharedState::new(64));
2641        let poll_config = leak_poll_config();
2642        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
2643        let mut iface = create_interface(&mut device, &poll_config);
2644        let mut sockets = SocketSet::new(vec![]);
2645        let mut tracker = TcpConnectionTracker::new(None);
2646        let now = smoltcp_now();
2647
2648        let (server_isn, guest_seq, mut new_conns) = establish(
2649            &mut tracker,
2650            &mut device,
2651            &mut iface,
2652            &mut sockets,
2653            &shared,
2654            now,
2655            54322,
2656        );
2657        let conn = new_conns.remove(0);
2658        let mut from_smoltcp = conn.from_smoltcp;
2659        let _to_smoltcp = conn.to_smoltcp;
2660
2661        // Guest aborts the connection (RST) instead of closing it cleanly.
2662        ingress(
2663            build_tcp_frame(
2664                54322,
2665                443,
2666                TcpControl::Rst,
2667                guest_seq,
2668                Some(server_isn + 1),
2669                &[],
2670            ),
2671            &mut device,
2672            &mut iface,
2673            &mut sockets,
2674            &shared,
2675            now,
2676        );
2677
2678        // Maintenance reaps the Closed socket and drops the proxy channel.
2679        for _ in 0..8 {
2680            tracker.relay_data(&mut sockets);
2681            tracker.cleanup_closed(&mut sockets);
2682            let _ = iface.poll_egress(now, &mut device, &mut sockets);
2683        }
2684
2685        assert!(
2686            !tracker.has_socket_for(
2687                &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 54322),
2688                &SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443),
2689            ),
2690            "RST connection should be evicted from the tracker",
2691        );
2692        // The proxy task's receiver IS disconnected => `recv()` returns None
2693        // => the proxy breaks its relay loop => the upstream socket is closed.
2694        assert!(
2695            matches!(
2696                from_smoltcp.try_recv(),
2697                Err(tokio::sync::mpsc::error::TryRecvError::Disconnected)
2698            ),
2699            "RST teardown must close the proxy channel (clean, no orphan)",
2700        );
2701    }
2702
2703    #[test]
2704    fn closed_slot_is_reused_under_pressure_but_pending_reset_is_sent_first() {
2705        let shared = Arc::new(SharedState::new(64));
2706        let config = leak_poll_config();
2707        let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
2708        let mut iface = create_interface(&mut device, &config);
2709        let mut sockets = SocketSet::new(vec![]);
2710        let mut tracker = TcpConnectionTracker::new(NonZeroUsize::new(1));
2711        let now = smoltcp_now();
2712        handshake(
2713            &mut tracker,
2714            &mut device,
2715            &mut iface,
2716            &mut sockets,
2717            &shared,
2718            now,
2719            40000,
2720        );
2721        let handle = sockets.iter().next().unwrap().0;
2722        sockets.get_mut::<tcp::Socket>(handle).abort();
2723        tracker.cleanup_closed(&mut sockets);
2724        assert_eq!(
2725            only_tcp_state(&sockets),
2726            Some(tcp::State::Closed),
2727            "pending RST must not be lost during cleanup"
2728        );
2729        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40001);
2730        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);
2731        assert!(
2732            !tracker.create_tcp_socket(src, dst, &mut sockets, TcpAcceptMode::AcceptImmediately),
2733            "the pending reset still owns its socket budget"
2734        );
2735        loop {
2736            if matches!(
2737                iface.poll_egress(now, &mut device, &mut sockets),
2738                smoltcp::iface::PollResult::None
2739            ) {
2740                break;
2741            }
2742        }
2743        let (_, _, _, _, rst) = last_tcp_reply(&shared).expect("RST must reach guest");
2744        assert!(rst);
2745        // No periodic cleanup call: allocating the next SYN reclaims the slot.
2746        assert!(tracker.create_tcp_socket(
2747            src,
2748            dst,
2749            &mut sockets,
2750            TcpAcceptMode::AcceptImmediately
2751        ));
2752        assert!(!tracker.has_socket_for(
2753            &SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40000),
2754            &dst
2755        ));
2756    }
2757
2758    #[test]
2759    fn reset_during_handshake_does_not_leave_an_idle_listener_at_capacity() {
2760        let shared = Arc::new(SharedState::new(64));
2761        let config = leak_poll_config();
2762        let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
2763        let mut iface = create_interface(&mut device, &config);
2764        let mut sockets = SocketSet::new(vec![]);
2765        let mut tracker = TcpConnectionTracker::new(NonZeroUsize::new(1));
2766        let now = smoltcp_now();
2767        ingress(
2768            build_arp_request_frame(GUEST_MAC, GUEST_IP, GATEWAY_IP),
2769            &mut device,
2770            &mut iface,
2771            &mut sockets,
2772            &shared,
2773            now,
2774        );
2775        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40000);
2776        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);
2777        assert!(tracker.create_tcp_socket(
2778            src,
2779            dst,
2780            &mut sockets,
2781            TcpAcceptMode::AcceptImmediately
2782        ));
2783        ingress(
2784            build_tcp_frame(40000, 443, TcpControl::Syn, 1000, None, &[]),
2785            &mut device,
2786            &mut iface,
2787            &mut sockets,
2788            &shared,
2789            now,
2790        );
2791        let (server_seq, _, syn, _, _) = last_tcp_reply(&shared).unwrap();
2792        assert!(syn);
2793        assert_eq!(only_tcp_state(&sockets), Some(tcp::State::SynReceived));
2794        ingress(
2795            build_tcp_frame(40000, 443, TcpControl::Rst, 1001, Some(server_seq + 1), &[]),
2796            &mut device,
2797            &mut iface,
2798            &mut sockets,
2799            &shared,
2800            now,
2801        );
2802        assert_eq!(only_tcp_state(&sockets), Some(tcp::State::Listen));
2803        let next_src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40001);
2804        assert!(tracker.create_tcp_socket(
2805            next_src,
2806            dst,
2807            &mut sockets,
2808            TcpAcceptMode::AcceptImmediately
2809        ));
2810        assert!(!tracker.has_socket_for(&src, &dst));
2811    }
2812
2813    #[test]
2814    fn connection_pressure_preserves_time_wait_protection() {
2815        let shared = Arc::new(SharedState::new(64));
2816        let config = leak_poll_config();
2817        let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
2818        let mut iface = create_interface(&mut device, &config);
2819        let mut sockets = SocketSet::new(vec![]);
2820        let mut tracker = TcpConnectionTracker::new(NonZeroUsize::new(1));
2821        let now = smoltcp_now();
2822        let (_, guest_seq) = handshake(
2823            &mut tracker,
2824            &mut device,
2825            &mut iface,
2826            &mut sockets,
2827            &shared,
2828            now,
2829            40000,
2830        );
2831        let handle = sockets.iter().next().unwrap().0;
2832        sockets.get_mut::<tcp::Socket>(handle).close();
2833        loop {
2834            if matches!(
2835                iface.poll_egress(now, &mut device, &mut sockets),
2836                smoltcp::iface::PollResult::None
2837            ) {
2838                break;
2839            }
2840        }
2841        let (fin_seq, _, _, fin, _) = last_tcp_reply(&shared).unwrap();
2842        assert!(fin);
2843        ingress(
2844            build_tcp_frame(
2845                40000,
2846                443,
2847                TcpControl::None,
2848                guest_seq,
2849                Some(fin_seq + 1),
2850                &[],
2851            ),
2852            &mut device,
2853            &mut iface,
2854            &mut sockets,
2855            &shared,
2856            now,
2857        );
2858        ingress(
2859            build_tcp_frame(
2860                40000,
2861                443,
2862                TcpControl::Fin,
2863                guest_seq,
2864                Some(fin_seq + 1),
2865                &[],
2866            ),
2867            &mut device,
2868            &mut iface,
2869            &mut sockets,
2870            &shared,
2871            now,
2872        );
2873        assert_eq!(only_tcp_state(&sockets), Some(tcp::State::TimeWait));
2874        tracker.cleanup_closed(&mut sockets);
2875        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40001);
2876        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);
2877        assert!(!tracker.create_tcp_socket(
2878            src,
2879            dst,
2880            &mut sockets,
2881            TcpAcceptMode::AcceptImmediately
2882        ));
2883        assert_eq!(only_tcp_state(&sockets), Some(tcp::State::TimeWait));
2884    }
2885
2886    #[test]
2887    fn full_connection_table_refuses_new_sockets() {
2888        // Once the table is full, new guest connections are refused. Uses a
2889        // small max to avoid 256 full handshakes; the gating logic is
2890        // identical to the 256 default.
2891        let mut tracker = TcpConnectionTracker::new(NonZeroUsize::new(4));
2892        let mut sockets = SocketSet::new(vec![]);
2893        let shared = Arc::new(SharedState::new(64));
2894        let config = leak_poll_config();
2895        let mut device = SmoltcpDevice::new(shared.clone(), config.mtu);
2896        let mut iface = create_interface(&mut device, &config);
2897        let now = smoltcp_now();
2898        let dst = SocketAddr::new(Ipv4Addr::from(SERVER_IP).into(), 443);
2899
2900        for port in 40000u16..40004 {
2901            handshake(
2902                &mut tracker,
2903                &mut device,
2904                &mut iface,
2905                &mut sockets,
2906                &shared,
2907                now,
2908                port,
2909            );
2910        }
2911        // Table full (4 slots held). The 5th guest SYN gets no socket — which
2912        // in the poll loop means smoltcp emits RST / no reply => the guest
2913        // sees egress as unreachable.
2914        let src = SocketAddr::new(Ipv4Addr::from(GUEST_IP).into(), 40004);
2915        assert!(
2916            !tracker.create_tcp_socket(src, dst, &mut sockets, TcpAcceptMode::AcceptImmediately),
2917            "creation at the limit must be refused",
2918        );
2919    }
2920
2921    #[test]
2922    fn guest_fin_before_proxy_spawn_is_handed_off() {
2923        let shared = Arc::new(SharedState::new(64));
2924        let poll_config = leak_poll_config();
2925        let mut device = SmoltcpDevice::new(shared.clone(), poll_config.mtu);
2926        let mut iface = create_interface(&mut device, &poll_config);
2927        let mut sockets = SocketSet::new(vec![]);
2928        let mut tracker = TcpConnectionTracker::new(None);
2929        let now = smoltcp_now();
2930        let guest_port = 54323;
2931        let (server_isn, guest_seq) = handshake(
2932            &mut tracker,
2933            &mut device,
2934            &mut iface,
2935            &mut sockets,
2936            &shared,
2937            now,
2938            guest_port,
2939        );
2940
2941        // The production poll loop drains every queued guest frame before
2942        // taking new connections, so the FIN can be processed before the
2943        // proxy task is spawned.
2944        ingress(
2945            build_tcp_frame(
2946                guest_port,
2947                443,
2948                TcpControl::Fin,
2949                guest_seq,
2950                Some(server_isn + 1),
2951                &[],
2952            ),
2953            &mut device,
2954            &mut iface,
2955            &mut sockets,
2956            &shared,
2957            now,
2958        );
2959        assert_eq!(
2960            only_tcp_state(&sockets),
2961            Some(tcp::State::CloseWait),
2962            "guest FIN should arrive before the proxy handoff",
2963        );
2964
2965        let new_conns = tracker.take_new_connections(&mut sockets);
2966        assert_eq!(
2967            new_conns.len(),
2968            1,
2969            "a connection that reached CLOSE_WAIT still needs a proxy task",
2970        );
2971    }
2972}