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