fips/src/node/dataplane/rx_loop.rs

//! RX event loop and packet dispatch.

use crate::control::{ControlSocket, commands};
use crate::node::{Node, NodeError};
use crate::proto::fmp::wire::{
    COMMON_PREFIX_SIZE, CommonPrefix, FMP_VERSION, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2,
    PHASE_MSG3,
};
use crate::transport::ReceivedPacket;
use std::time::Duration;
use tracing::{debug, info, warn};

/// Inside the packet_rx burst drain, run a fallback drain every
/// N packets so bounced FMP plaintexts can't sit behind a full
/// 256-packet UDP burst. Used on unix; on Windows the decrypt-worker
/// pool isn't spawned so the fallback channel is always empty —
/// hold the constants at module scope anyway so the burst-loop
/// dispatch in `run_rx_loop` doesn't need a `#[cfg]` on every site.
const FALLBACK_INTERLEAVE_EVERY: usize = 32;
/// How many fallback events to drain per interleave step. Bounded so
/// the inner loop can keep making forward progress on packet_rx.
#[cfg(unix)]
const FALLBACK_INTERLEAVE_BUDGET: usize = 32;

impl Node {
    /// Run the receive event loop.
    ///
    /// Processes packets from all transports, dispatching based on
    /// the phase field in the 4-byte common prefix:
    /// - Phase 0x0: Encrypted frame (session data)
    /// - Phase 0x1: Handshake message 1 (initiator -> responder)
    /// - Phase 0x2: Handshake message 2 (responder -> initiator)
    /// - Phase 0x3: Handshake message 3 (initiator -> responder, XX completion)
    ///
    /// Also processes outbound IPv6 packets from the TUN reader for session
    /// encapsulation and routing through the mesh.
    ///
    /// Also processes DNS-resolved identities for identity cache population.
    ///
    /// Also runs a periodic tick (1s) to clean up stale handshake connections
    /// that never received a response. This prevents resource leaks when peers
    /// are unreachable.
    ///
    /// This method takes ownership of the packet_rx channel and runs
    /// until the channel is closed (typically when stop() is called).
    pub async fn run_rx_loop(&mut self) -> Result<(), NodeError> {
        // No shutdown observer → today's infinite loop, byte-identical. All
        // existing callers/tests use this; `pending()` never fires, so the
        // shutdown/deadline arms below stay permanently disabled.
        self.run_rx_loop_with_shutdown(std::future::pending()).await
    }

    /// The rx event loop, which serves until `shutdown` fires and then drains
    /// **in place** before returning.
    ///
    /// The channel receivers are moved into this frame's locals and live across
    /// both serve and drain, so — unlike a `select!`-cancelled loop — they are
    /// never destructively dropped mid-flight; they are released only on clean
    /// exit, after which teardown does not need them.
    ///
    /// - While serving (`drain_deadline == None`) the loop is behaviorally
    ///   identical to before: the shutdown arm, the deadline arm, and the
    ///   peers-empty early-exit are all guarded off, so the hot per-packet path
    ///   and the `biased` order of the real arms are unchanged.
    /// - When `shutdown` fires, the loop calls [`Node::enter_drain`] once
    ///   (broadcast Disconnect, gate the reconciler off) and arms the bounded
    ///   deadline, then keeps servicing inbound/tick/peer-removal until all
    ///   peers clear or the deadline elapses, then returns. The caller
    ///   ([`Node::finish_shutdown`]) closes the window and tears down.
    pub async fn run_rx_loop_with_shutdown(
        &mut self,
        shutdown: impl std::future::Future<Output = ()>,
    ) -> Result<(), NodeError> {
        tokio::pin!(shutdown);
        // `None` = serving; `Some(deadline)` = draining (bounded window).
        let mut drain_deadline: Option<tokio::time::Instant> = None;
        let mut packet_rx = self.packet_rx.take().ok_or(NodeError::NotStarted)?;

        // Take the TUN outbound receiver, or create a dummy channel that never
        // produces messages (when TUN is disabled). Holding the sender prevents
        // the channel from closing.
        let (mut tun_outbound_rx, _tun_guard) = match self.supervisor.tun_outbound_rx.take() {
            Some(rx) => (rx, None),
            None => {
                let (tx, rx) = tokio::sync::mpsc::channel(1);
                (rx, Some(tx))
            }
        };

        // Take the DNS identity receiver, or create a dummy channel (when DNS
        // is disabled). Same pattern as TUN outbound.
        let (mut dns_identity_rx, _dns_guard) = match self.supervisor.dns_identity_rx.take() {
            Some(rx) => (rx, None),
            None => {
                let (tx, rx) = tokio::sync::mpsc::channel(1);
                (rx, Some(tx))
            }
        };

        // Take the runtime child-liveness receiver, or a dummy channel (when the
        // node was seeded straight into Running without a start()). Holding the
        // dummy sender in the guard keeps the channel open. Same pattern as TUN
        // outbound / DNS identity.
        let (mut child_exit_rx, _child_exit_guard) = match self.child_exit_rx.take() {
            Some(rx) => (rx, None),
            None => {
                let (tx, rx) = tokio::sync::mpsc::channel(1);
                (rx, Some(tx))
            }
        };

        let mut tick =
            tokio::time::interval(Duration::from_secs(self.config().node.tick_interval_secs));

        // Set up control socket channel
        let (control_tx, mut control_rx) =
            tokio::sync::mpsc::channel::<crate::control::ControlMessage>(32);

        if self.config().node.control.enabled {
            let config = self.config().node.control.clone();
            let tx = control_tx.clone();
            let read_handle = self.control_read_handle();
            tokio::spawn(async move {
                match ControlSocket::bind(&config) {
                    Ok(socket) => {
                        socket.accept_loop(tx, read_handle).await;
                    }
                    Err(e) => {
                        warn!(error = %e, "Failed to bind control socket");
                    }
                }
            });
        }
        // Drop unused sender to avoid keeping channel open if control is disabled
        drop(control_tx);

        // Decrypt-worker fallback receiver. The worker pushes each
        // authenticated FMP plaintext here so rx_loop can finish the
        // per-peer side-effects (stats, MMP, ECN, link dispatch).
        // Always declared so the `tokio::select!` arm doesn't need
        // a `cfg` (which the macro doesn't support); on Windows the
        // channel just never sees events.
        let (mut decrypt_fallback_rx, _decrypt_fallback_guard) = {
            #[cfg(unix)]
            {
                match self.decrypt_fallback_rx.take() {
                    Some(rx) => (rx, None),
                    None => {
                        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
                        (rx, Some(tx))
                    }
                }
            }
            #[cfg(not(unix))]
            {
                // On non-unix nothing ever sends, but the macro arm
                // still needs an existing rx. Keep the sender alive to
                // avoid the channel closing into an Err loop.
                let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>();
                (rx, Some(tx))
            }
        };

        info!("RX event loop started");
        // Optional per-stage perf profiler (FIPS_PERF=1). No-op otherwise.
        crate::perf_profile::maybe_spawn_reporter();

        loop {
            // Bounded drain mode: break as soon as all peers have cleared. In
            // normal mode (`None`) this short-circuits before touching
            // `self.peers`, so the loop is byte-identical.
            if drain_deadline.is_some() && self.peers.is_empty() {
                info!("Drain complete: all peers cleared, ending drain loop");
                break;
            }
            tokio::select! {
                biased;
                // Decrypt-worker fallback drains FIRST. Under sustained
                // inbound bursts the packet_rx drain (up to 256 packets)
                // can starve fallback work for tens of ms — TCP doesn't
                // tolerate that (late ACKs → dup-ACK fast retransmits →
                // cwnd collapse). Promoting fallback gives the kernel's
                // TCP machinery a fair chance to ACK in time.
                Some(event) = decrypt_fallback_rx.recv() => {
                    #[cfg(unix)]
                    {
                        self.process_decrypt_worker_event(event).await;
                        let mut drained = 0;
                        while drained < 255 {
                            match decrypt_fallback_rx.try_recv() {
                                Ok(ev) => {
                                    self.process_decrypt_worker_event(ev).await;
                                    drained += 1;
                                }
                                Err(_) => break,
                            }
                        }
                    }
                    #[cfg(not(unix))]
                    let _ = event;
                }
                packet = packet_rx.recv() => {
                    match packet {
                        Some(p) => self.process_packet(p).await,
                        None => break, // channel closed
                    }
                    // Drain remaining ready inbound packets in a tight loop
                    // before yielding back to select! — every yield is a
                    // futex hop on tokio's multi-thread scheduler, and at
                    // line rate the kernel UDP queue typically has several
                    // datagrams available per wake. Caps at a batch
                    // boundary so other branches (tick, control) eventually
                    // get a turn even under sustained load.
                    //
                    // **Interleave fallback drain** every N packets so
                    // bounced FMP plaintexts (heartbeats, post-FMP-
                    // decrypt forwarding payloads, control frames) don't
                    // sit in the fallback queue for a full 256-packet
                    // burst. Even with the priority-first ordering of
                    // the outer select!, once we're inside this inner
                    // loop only this interleave can free queued
                    // fallbacks. On multihop forwarding paths this is
                    // the difference between back-to-back encrypt-
                    // worker dispatches happening promptly vs piling up
                    // behind the rx burst.
                    let mut drained: usize = 1; // count the packet processed above
                    while drained < 256 {
                        if drained.is_multiple_of(FALLBACK_INTERLEAVE_EVERY) {
                            #[cfg(unix)]
                            {
                                let mut fb_drained = 0;
                                while fb_drained < FALLBACK_INTERLEAVE_BUDGET {
                                    match decrypt_fallback_rx.try_recv() {
                                        Ok(ev) => {
                                            self.process_decrypt_worker_event(ev).await;
                                            fb_drained += 1;
                                        }
                                        Err(_) => break,
                                    }
                                }
                            }
                        }
                        match packet_rx.try_recv() {
                            Ok(p) => {
                                self.process_packet(p).await;
                                drained += 1;
                            }
                            Err(_) => break,
                        }
                    }
                    // Trailing fallback drain so the last bounced
                    // packets of the burst aren't held up by the
                    // next select! iteration.
                    #[cfg(unix)]
                    {
                        let mut fb_drained = 0;
                        while fb_drained < 256 {
                            match decrypt_fallback_rx.try_recv() {
                                Ok(ev) => {
                                    self.process_decrypt_worker_event(ev).await;
                                    fb_drained += 1;
                                }
                                Err(_) => break,
                            }
                        }
                    }
                }
                // Runtime child-liveness. Placed AFTER `packet_rx` so the hot
                // inbound path keeps its `biased` priority. A directly-observable
                // child (TUN threads, DNS/mDNS/Nostr) exited on its own; feed the
                // FSM, which republishes health (Degraded here — a Running node
                // always has ≥1 transport up). `on_child_exited` only ever emits
                // `PublishState`; other variants are ignored defensively.
                maybe_child = child_exit_rx.recv() => {
                    if let Some(child) = maybe_child {
                        let actions = self
                            .supervisor
                            .fsm
                            .step(crate::node::lifecycle::supervisor::Event::ChildExited { child });
                        for action in actions {
                            if let crate::node::lifecycle::supervisor::Action::PublishState(ns) =
                                action
                            {
                                self.supervisor.state = ns;
                            }
                        }
                    }
                }
                Some(ipv6_packet) = tun_outbound_rx.recv() => {
                    self.handle_tun_outbound(ipv6_packet).await;
                    let mut drained = 0;
                    while drained < 256 {
                        match tun_outbound_rx.try_recv() {
                            Ok(p) => {
                                self.handle_tun_outbound(p).await;
                                drained += 1;
                            }
                            Err(_) => break,
                        }
                    }
                }
                Some(identity) = dns_identity_rx.recv() => {
                    debug!(
                        node_addr = %identity.node_addr,
                        "Registering identity from DNS resolution"
                    );
                    self.register_identity(identity.node_addr, identity.pubkey);
                }
                Some((request, response_tx)) = control_rx.recv() => {
                    // Only mutating COMMAND requests (`connect` / `disconnect`)
                    // reach the rx_loop now. Every pure-read `show_*` query is
                    // served off-loop from the read handle in the control accept
                    // task (`snapshot_dispatch`), so it never round-trips here —
                    // the data-plane dispatch path carries no `show_*` arm. A
                    // `show_*` that somehow arrives (none does) falls through to
                    // `commands::dispatch`, which returns "unknown command".
                    let response = commands::dispatch(
                        self,
                        &request.command,
                        request.params.as_ref(),
                    ).await;
                    let _ = response_tx.send(response);
                }
                _ = tick.tick() => {
                    self.check_timeouts();
                    let now_ms = Self::now_ms();
                    self.reload_peer_acl().await;
                    // The host map hot-reloads on the same tick as the ACL. It
                    // is polled separately from `reload_peer_acl` because the
                    // ACL's embedded alias reloader and this snapshot are
                    // distinct resources; the `path_mtu_lookup` cache and the
                    // `nostr_rendezvous` subsystem are deliberately excluded
                    // from `Reloadable` since neither reloads from a backing
                    // file (see `node::reloadable`).
                    self.reload_host_map().await;
                    self.poll_pending_connects().await;
                    self.poll_nostr_rendezvous().await;
                    self.poll_lan_rendezvous().await;
                    self.resend_pending_handshakes(now_ms).await;
                    self.resend_pending_rekeys(now_ms).await;
                    self.resend_pending_fmp_rekey_msg3(now_ms).await;
                    self.resend_pending_session_handshakes(now_ms).await;
                    self.resend_pending_session_msg3(now_ms).await;
                    self.purge_idle_sessions(now_ms);
                    self.process_pending_retries(now_ms).await;
                    self.check_tree_state().await;
                    self.check_bloom_state().await;
                    self.compute_mesh_size();
                    self.record_stats_history();
                    self.check_mmp_reports().await;
                    self.check_session_mmp_reports().await;
                    self.check_link_heartbeats().await;
                    self.check_rekey().await;
                    self.check_session_rekey().await;
                    self.check_pending_lookups(now_ms).await;
                    self.poll_transport_discovery().await;
                    self.sample_transport_congestion();
                    #[cfg(any(target_os = "linux", target_os = "macos"))]
                    self.activate_connected_udp_sessions().await;
                }
                // Shutdown signal → enter the bounded drain in place, ONCE.
                // Gated on `is_none()` so it only fires while serving; after
                // entering drain the arm is disabled (the completed signal is
                // never polled again) and the deadline arm below bounds the
                // window. Placed after the real arms so their `biased` priority
                // is unchanged, and inert while serving with `pending()`.
                _ = &mut shutdown, if drain_deadline.is_none() => {
                    self.enter_drain().await;
                    drain_deadline =
                        Some(tokio::time::Instant::now() + self.config().node.drain_timeout());
                }
                // Bounded drain deadline (drain mode only). Placed LAST so the
                // `biased` priority of the normal arms is unchanged, and gated
                // on `is_some()` so in normal mode the branch is disabled — the
                // future is created but never polled and never fires.
                _ = tokio::time::sleep_until(
                    drain_deadline.unwrap_or_else(tokio::time::Instant::now)
                ), if drain_deadline.is_some() => {
                    info!("Drain deadline elapsed, ending drain loop");
                    break;
                }
            }
        }

        info!("RX event loop stopped (channel closed)");
        Ok(())
    }

    /// Process a single received packet.
    ///
    /// Dispatches based on the phase field in the 4-byte common prefix.
    async fn process_packet(&mut self, packet: ReceivedPacket) {
        if packet.data.len() < COMMON_PREFIX_SIZE {
            return; // Drop packets too short for common prefix
        }

        let prefix = match CommonPrefix::parse(&packet.data) {
            Some(p) => p,
            None => return, // Malformed prefix
        };

        if prefix.version != FMP_VERSION {
            debug!(
                version = prefix.version,
                transport_id = %packet.transport_id,
                "Unknown FMP version, dropping"
            );

            // If the packet arrived on an adopted Nostr-NAT bootstrap
            // transport, the originating peer is necessarily on a
            // different FMP-protocol version than us — the discovery
            // sweep would otherwise re-traverse them every cycle even
            // though no msg1/msg2 exchange can ever succeed. Bump the
            // discovery-layer cooldown to the long protocol-mismatch
            // window and emit a single WARN per fresh observation.
            if self
                .supervisor
                .nostr_rendezvous
                .is_bootstrap_transport(&packet.transport_id)
                && let Some(npub) = self
                    .supervisor
                    .nostr_rendezvous
                    .bootstrap_transport_npub(&packet.transport_id)
                    .cloned()
                && let Some(handle) = self.nostr_rendezvous_handle()
            {
                let now_ms = Self::now_ms();
                let cooldown_secs = handle.protocol_mismatch_cooldown_secs();
                if handle.record_protocol_mismatch(&npub, now_ms) {
                    warn!(
                        peer_npub = %npub,
                        transport_id = %packet.transport_id,
                        peer_version = prefix.version,
                        our_version = FMP_VERSION,
                        cooldown_secs,
                        "Nostr-discovered peer speaks a different FMP version; suppressing retraversal"
                    );
                }
            }
            return;
        }

        match prefix.phase {
            PHASE_ESTABLISHED => {
                self.handle_encrypted_frame(packet).await;
            }
            PHASE_MSG1 => {
                self.handle_msg1(packet).await;
            }
            PHASE_MSG2 => {
                self.handle_msg2(packet).await;
            }
            PHASE_MSG3 => {
                self.handle_msg3(packet).await;
            }
            _ => {
                debug!(
                    phase = prefix.phase,
                    transport_id = %packet.transport_id,
                    "Unknown FMP phase, dropping"
                );
            }
        }
    }
}