fips_core/node/handlers/
rx_loop.rs

1//! RX event loop and packet dispatch.
2
3use crate::control::queries;
4use crate::control::{ControlMessage, ControlSenders, ControlSocket, commands};
5use crate::discovery::is_punch_packet;
6use crate::node::decrypt_worker::{
7    DecryptFailureReport, DecryptFallback, DecryptJobBatcher, DecryptWorkerEvent,
8    DecryptWorkerFallbackReceivers,
9};
10use crate::node::handlers::encrypted::EncryptedFrameFastPath;
11use crate::node::wire::{
12    COMMON_PREFIX_SIZE, CommonPrefix, FMP_VERSION, PHASE_ESTABLISHED, PHASE_MSG1, PHASE_MSG2,
13};
14use crate::node::{
15    AuthenticatedFmpPlaintext, EndpointSendBatchCommand, Node, NodeEndpointCommand, NodeError,
16};
17use crate::transport::PacketRx;
18use crate::transport::ReceivedPacket;
19use crate::upper::tun::TunOutboundRx;
20use std::time::{Duration, Instant};
21use tokio::sync::mpsc::Receiver;
22use tracing::{debug, info, trace, warn};
23
24mod budget;
25mod drain;
26
27#[cfg(test)]
28mod tests;
29
30use budget::*;
31use drain::*;
32
33impl Node {
34    /// Run the receive event loop.
35    ///
36    /// Processes packets from all transports, dispatching based on
37    /// the phase field in the 4-byte common prefix:
38    /// - Phase 0x0: Encrypted frame (session data)
39    /// - Phase 0x1: Handshake message 1 (initiator -> responder)
40    /// - Phase 0x2: Handshake message 2 (responder -> initiator)
41    ///
42    /// Also processes outbound IPv6 packets from the TUN reader for session
43    /// encapsulation and routing through the mesh.
44    ///
45    /// Also processes DNS-resolved identities for identity cache population.
46    ///
47    /// Also runs a periodic tick (1s) to clean up stale handshake connections
48    /// that never received a response. This prevents resource leaks when peers
49    /// are unreachable.
50    ///
51    /// This method takes ownership of the packet_rx channel and runs
52    /// until the channel is closed (typically when stop() is called).
53    pub async fn run_rx_loop(&mut self) -> Result<(), NodeError> {
54        let mut packet_rx = self.packet_rx.take().ok_or(NodeError::NotStarted)?;
55
56        // Take the TUN outbound receiver, or create a dummy channel that never
57        // produces messages (when TUN is disabled). Holding the sender prevents
58        // the channel from closing.
59        let (mut tun_outbound_rx, _tun_guard) = match self.tun_outbound_rx.take() {
60            Some(rx) => (rx, None),
61            None => {
62                let (tx, rx) = tokio::sync::mpsc::channel(1);
63                (rx, Some(tx))
64            }
65        };
66
67        // Take the DNS identity receiver, or create a dummy channel (when DNS
68        // is disabled). Same pattern as TUN outbound.
69        let (mut dns_identity_rx, _dns_guard) = match self.dns_identity_rx.take() {
70            Some(rx) => (rx, None),
71            None => {
72                let (tx, rx) = tokio::sync::mpsc::channel(1);
73                (rx, Some(tx))
74            }
75        };
76
77        // Take the endpoint-data command receiver, or create a dummy channel
78        // when the embedded endpoint API is not in use.
79        let (mut endpoint_priority_command_rx, _endpoint_priority_command_guard) =
80            match self.endpoint_priority_command_rx.take() {
81                Some(rx) => (rx, None),
82                None => {
83                    let (tx, rx) = tokio::sync::mpsc::channel(1);
84                    (rx, Some(tx))
85                }
86            };
87        let (mut endpoint_command_rx, _endpoint_command_guard) =
88            match self.endpoint_command_rx.take() {
89                Some(rx) => (rx, None),
90                None => {
91                    let (tx, rx) = tokio::sync::mpsc::channel(1);
92                    (rx, Some(tx))
93                }
94            };
95        let (mut endpoint_bulk_feedback_rx, _endpoint_bulk_feedback_guard) =
96            match self.endpoint_bulk_feedback_rx.take() {
97                Some(rx) => (rx, None),
98                None => {
99                    let (tx, rx) = tokio::sync::mpsc::channel(1);
100                    (rx, Some(tx))
101                }
102            };
103
104        // Take the decrypt worker fallback receiver if a worker pool
105        // is in use. The worker pushes non-fast-path packets (anything
106        // that's not bulk EndpointData) here for the legacy dispatch.
107        let (mut decrypt_fallback_rx, _decrypt_fallback_guard) =
108            match self.decrypt_fallback_rx.take() {
109                Some(rx) => (rx, None),
110                None => {
111                    let (tx, rx) = crate::node::decrypt_worker::decrypt_worker_fallback_channels();
112                    (rx, Some(tx))
113                }
114            };
115
116        let mut tick =
117            tokio::time::interval(Duration::from_secs(self.config.node.tick_interval_secs));
118        tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
119        let mut maintenance_state = RxLoopMaintenanceState::default();
120
121        // Set up control socket channels. Read-only queries are separated
122        // from mutating commands so operator status reads can get reserved
123        // progress while a command awaits slower discovery/transport work.
124        let (control_query_tx, mut control_query_rx) =
125            tokio::sync::mpsc::channel::<ControlMessage>(32);
126        let (control_command_tx, mut control_command_rx) =
127            tokio::sync::mpsc::channel::<ControlMessage>(32);
128
129        if self.config.node.control.enabled {
130            let config = self.config.node.control.clone();
131            let senders = ControlSenders::new(control_query_tx.clone(), control_command_tx.clone());
132            tokio::spawn(async move {
133                match ControlSocket::bind(&config) {
134                    Ok(socket) => {
135                        socket.accept_loop_split(senders).await;
136                    }
137                    Err(e) => {
138                        warn!(error = %e, "Failed to bind control socket");
139                    }
140                }
141            });
142        }
143        // Drop unused sender to avoid keeping channel open if control is disabled
144        drop(control_query_tx);
145        drop(control_command_tx);
146
147        info!("RX event loop started");
148        // Optional perf profiler (FIPS_PERF=1). No-op otherwise.
149        crate::perf_profile::maybe_spawn_reporter();
150        // Tokio intervals tick immediately on first poll. Consume that startup
151        // tick so the reserved-progress branch below represents a due periodic
152        // maintenance turn, not an eager pre-data maintenance pass.
153        tick.tick().await;
154
155        loop {
156            tokio::select! {
157                biased;
158                // Priority decrypt-worker fallback drains first. The
159                // previous packet-first ordering could hold small ACK,
160                // heartbeat, and failure-report plaintexts behind a hot
161                // raw-packet drain long enough to collapse TCP. Bulk
162                // fallback is intentionally below `packet_rx`: bulk
163                // plaintext must keep making bounded progress, but it
164                // should not stop fresh transport priority packets from
165                // being dequeued. `drain_packet_rx` interleaves fallback
166                // turns every few dozen packets to keep that progress
167                // reserve while avoiding a bulk-fallback convoy.
168                Some(event) = decrypt_fallback_rx.priority.recv() => {
169                    let fallback_drained = self.drain_decrypt_priority_fallback(
170                        &mut decrypt_fallback_rx.priority,
171                        Some(event),
172                        PACKET_DRAIN_BUDGET,
173                    ).await;
174                    let side_drained = self.drain_rx_loop_side_queues(
175                        &mut control_query_rx,
176                        &mut tun_outbound_rx,
177                        &mut endpoint_priority_command_rx,
178                        &mut endpoint_command_rx,
179                        SIDE_QUEUE_INTERLEAVE_BUDGET,
180                    ).await;
181                    if fallback_drained > 0 || side_drained.has_data_drained() {
182                        maintenance_state.record_data_activity(Instant::now());
183                    }
184                }
185                // Timer-driven liveness is a reserved-progress branch. It
186                // performs bounded pre/post data drains and timeboxes slow
187                // discovery/status work, so hot packet or bulk-fallback
188                // queues cannot indefinitely postpone heartbeat, rekey, MMP,
189                // route aging, or path maintenance.
190                _ = tick.tick() => {
191                    let drained = self.drain_rx_loop_data_queues(
192                        &mut packet_rx,
193                        &mut decrypt_fallback_rx,
194                        &mut tun_outbound_rx,
195                        &mut endpoint_priority_command_rx,
196                        &mut endpoint_command_rx,
197                        NON_PACKET_DRAIN_BUDGET,
198                    ).await;
199                    if drained.has_drained() {
200                        maintenance_state.record_data_activity(Instant::now());
201                        debug!(
202                            drained = drained.total(),
203                            drained_packets = drained.packets,
204                            drained_decrypt = drained.decrypt,
205                            drained_tun = drained.tun,
206                            drained_endpoint = drained.endpoint,
207                            "Drained queued packets before rx-loop maintenance"
208                        );
209                    }
210                    let maintenance_plan = maintenance_state.plan_maintenance(
211                        drained,
212                        Instant::now(),
213                        RX_LOOP_RECENT_DATA_ACTIVITY_WINDOW,
214                        RX_LOOP_SLOW_MAINTENANCE_IDLE_TIMEOUT,
215                        RX_LOOP_SLOW_MAINTENANCE_BUSY_TIMEOUT,
216                    );
217
218                    let slow_timed_out = self.run_rx_loop_maintenance_tick(
219                        maintenance_plan,
220                    ).await;
221                    maintenance_state.record_maintenance_result(
222                        maintenance_plan.data_pressure(),
223                        slow_timed_out,
224                    );
225
226                    let post_drained = self.drain_rx_loop_data_queues(
227                        &mut packet_rx,
228                        &mut decrypt_fallback_rx,
229                        &mut tun_outbound_rx,
230                        &mut endpoint_priority_command_rx,
231                        &mut endpoint_command_rx,
232                        PACKET_DRAIN_BUDGET,
233                    ).await;
234                    if post_drained.has_drained() {
235                        maintenance_state.record_data_activity(Instant::now());
236                        debug!(
237                            drained = post_drained.total(),
238                            drained_packets = post_drained.packets,
239                            drained_decrypt = post_drained.decrypt,
240                            drained_tun = post_drained.tun,
241                            drained_endpoint = post_drained.endpoint,
242                            "Drained queued packets after rx-loop maintenance"
243                        );
244                    }
245                }
246                Some(event) = decrypt_fallback_rx.authenticated_bulk.recv(),
247                    if authenticated_bulk_preempts_packet_rx(packet_rx.priority_ready_packets()) =>
248                {
249                    let fallback_drained = self.drain_decrypt_fallback(
250                        &mut decrypt_fallback_rx,
251                        None,
252                        Some(event),
253                        None,
254                        NON_PACKET_DRAIN_BUDGET,
255                    ).await;
256                    let side_drained = self.drain_rx_loop_side_queues(
257                        &mut control_query_rx,
258                        &mut tun_outbound_rx,
259                        &mut endpoint_priority_command_rx,
260                        &mut endpoint_command_rx,
261                        SIDE_QUEUE_INTERLEAVE_BUDGET,
262                    ).await;
263                    if fallback_drained > 0 || side_drained.has_data_drained() {
264                        maintenance_state.record_data_activity(Instant::now());
265                    }
266                }
267                Some(message) = control_query_rx.recv() => {
268                    self.drain_control_queries(
269                        &mut control_query_rx,
270                        Some(message),
271                        NON_PACKET_DRAIN_BUDGET,
272                    ).await;
273                }
274                Some(feedback) = endpoint_bulk_feedback_rx.recv() => {
275                    let drained = self.drain_endpoint_bulk_send_feedback(
276                        &mut endpoint_bulk_feedback_rx,
277                        Some(feedback),
278                        NON_PACKET_DRAIN_BUDGET,
279                    );
280                    if drained > 0 {
281                        maintenance_state.record_data_activity(Instant::now());
282                    }
283                }
284                // Endpoint priority is app-owned latency-sensitive traffic
285                // (ICMP, TCP ACK/SYN, tiny TCP data). On platforms without the
286                // unix encrypt-worker fast path, this branch is the outbound
287                // dataplane path, so give it an explicit turn before hot raw
288                // receive. Bulk endpoint commands intentionally remain below
289                // packet_rx.
290                Some(command) = endpoint_priority_command_rx.recv() => {
291                    let drained = self.drain_endpoint_commands(
292                        &mut endpoint_priority_command_rx,
293                        &mut endpoint_command_rx,
294                        Some(command),
295                        None,
296                        NON_PACKET_DRAIN_BUDGET,
297                    ).await;
298                    if drained > 0 {
299                        maintenance_state.record_data_activity(Instant::now());
300                    }
301                }
302                packet = packet_rx.recv() => {
303                    match packet {
304                        Some(p) => {
305                            let drained = self.drain_packet_rx(
306                                &mut packet_rx,
307                                &mut decrypt_fallback_rx,
308                                Some(RxLoopSideQueues {
309                                    control_query_rx: &mut control_query_rx,
310                                    tun_outbound_rx: &mut tun_outbound_rx,
311                                    endpoint_priority_command_rx: &mut endpoint_priority_command_rx,
312                                    endpoint_command_rx: &mut endpoint_command_rx,
313                                }),
314                                Some(p),
315                                PACKET_DRAIN_BUDGET,
316                            ).await;
317                            if drained > 0 {
318                                maintenance_state.record_data_activity(Instant::now());
319                            }
320                        }
321                        None => break, // channel closed
322                    }
323                }
324                Some(event) = decrypt_fallback_rx.bulk.recv() => {
325                    let fallback_plan = fallback_drain_plan();
326                    let fallback_drained = self.drain_decrypt_fallback(
327                        &mut decrypt_fallback_rx,
328                        None,
329                        None,
330                        Some(event),
331                        fallback_plan.trailing_budget,
332                    ).await;
333                    let side_drained = self.drain_rx_loop_side_queues(
334                        &mut control_query_rx,
335                        &mut tun_outbound_rx,
336                        &mut endpoint_priority_command_rx,
337                        &mut endpoint_command_rx,
338                        SIDE_QUEUE_INTERLEAVE_BUDGET,
339                    ).await;
340                    if fallback_drained > 0 || side_drained.has_data_drained() {
341                        maintenance_state.record_data_activity(Instant::now());
342                    }
343                }
344                Some(ipv6_packet) = tun_outbound_rx.recv() => {
345                    let drained = self.drain_tun_outbound(
346                        &mut tun_outbound_rx,
347                        Some(ipv6_packet),
348                        NON_PACKET_DRAIN_BUDGET,
349                    ).await;
350                    if drained > 0 {
351                        maintenance_state.record_data_activity(Instant::now());
352                    }
353                }
354                Some(identity) = dns_identity_rx.recv() => {
355                    debug!(
356                        node_addr = %identity.node_addr,
357                        "Registering identity from DNS resolution"
358                    );
359                    self.register_identity(identity.node_addr, identity.pubkey);
360                }
361                Some(command) = endpoint_command_rx.recv() => {
362                    let drained = self.drain_endpoint_commands(
363                        &mut endpoint_priority_command_rx,
364                        &mut endpoint_command_rx,
365                        None,
366                        Some(command),
367                        NON_PACKET_DRAIN_BUDGET,
368                    ).await;
369                    if drained > 0 {
370                        maintenance_state.record_data_activity(Instant::now());
371                    }
372                }
373                Some((request, response_tx)) = control_command_rx.recv() => {
374                    let response = commands::dispatch(
375                        self,
376                        &request.command,
377                        request.params.as_ref(),
378                    ).await;
379                    let _ = response_tx.send(response);
380                }
381            }
382        }
383
384        info!("RX event loop stopped (channel closed)");
385        Ok(())
386    }
387
388    async fn drain_rx_loop_data_queues(
389        &mut self,
390        packet_rx: &mut PacketRx,
391        decrypt_fallback_rx: &mut DecryptWorkerFallbackReceivers,
392        tun_outbound_rx: &mut TunOutboundRx,
393        endpoint_priority_command_rx: &mut Receiver<NodeEndpointCommand>,
394        endpoint_command_rx: &mut Receiver<NodeEndpointCommand>,
395        budget: usize,
396    ) -> RxLoopDataDrainStats {
397        let drained_packets = self
398            .drain_packet_rx(packet_rx, decrypt_fallback_rx, None, None, budget)
399            .await;
400        let non_packet_budget = non_packet_drain_budget(budget);
401        let drained_decrypt = if decrypt_fallback_has_ready(decrypt_fallback_rx) {
402            self.drain_decrypt_fallback(decrypt_fallback_rx, None, None, None, non_packet_budget)
403                .await
404        } else {
405            0
406        };
407        let drained_tun = self
408            .drain_tun_outbound(tun_outbound_rx, None, non_packet_budget)
409            .await;
410        let drained_endpoint = self
411            .drain_endpoint_commands(
412                endpoint_priority_command_rx,
413                endpoint_command_rx,
414                None,
415                None,
416                non_packet_budget,
417            )
418            .await;
419        RxLoopDataDrainStats::with_decrypt(
420            drained_packets,
421            drained_decrypt,
422            drained_tun,
423            drained_endpoint,
424        )
425    }
426
427    async fn drain_packet_rx(
428        &mut self,
429        packet_rx: &mut PacketRx,
430        decrypt_fallback_rx: &mut DecryptWorkerFallbackReceivers,
431        mut side_queues: Option<RxLoopSideQueues<'_>>,
432        first_packet: Option<ReceivedPacket>,
433        budget: usize,
434    ) -> usize {
435        // Drain remaining ready inbound packets in a tight loop before
436        // yielding back to select! Every yield is a scheduler hop, and at
437        // line rate transports typically have several packets available per
438        // wake. Caps at a batch boundary so other branches eventually get a
439        // turn even under sustained load.
440        self.begin_endpoint_event_batch();
441        let side_queue_interleave_every = if side_queues.is_some() {
442            SIDE_QUEUE_INTERLEAVE_EVERY
443        } else {
444            0
445        };
446        let fallback_plan = fallback_drain_plan();
447        let mut drain = PacketDrainCursor::new(
448            first_packet,
449            budget,
450            fallback_plan.interleave_every,
451            side_queue_interleave_every,
452        );
453        let mut decrypt_jobs = DecryptJobBatcher::new();
454        while let Some(action) = drain.next(packet_rx) {
455            match action {
456                PacketDrainAction::Packet(packet) => {
457                    let action = self.begin_process_packet(packet);
458                    match action {
459                        PacketProcessAction::DecryptJob { job } => {
460                            if let Some(workers) = self.decrypt_workers.as_ref() {
461                                decrypt_jobs.push(workers, job);
462                            }
463                        }
464                        PacketProcessAction::Done => {}
465                        action => {
466                            self.flush_decrypt_job_batcher(&mut decrypt_jobs);
467                            self.finish_packet_process(action).await;
468                        }
469                    }
470                }
471                PacketDrainAction::InterleaveFallback => {
472                    self.flush_decrypt_job_batcher(&mut decrypt_jobs);
473                    let drained = if decrypt_fallback_has_ready(decrypt_fallback_rx) {
474                        self.drain_decrypt_fallback(
475                            decrypt_fallback_rx,
476                            None,
477                            None,
478                            None,
479                            fallback_plan.interleave_budget,
480                        )
481                        .await
482                    } else {
483                        0
484                    };
485                    if drained == 0 {
486                        drain.refund_empty_interleave_turn();
487                    }
488                }
489                PacketDrainAction::InterleaveSideQueues => {
490                    self.flush_decrypt_job_batcher(&mut decrypt_jobs);
491                    let drained = if let Some(side_queues) = side_queues.as_mut() {
492                        if rx_loop_side_queues_have_ready(side_queues) {
493                            self.drain_rx_loop_side_queues(
494                                side_queues.control_query_rx,
495                                side_queues.tun_outbound_rx,
496                                side_queues.endpoint_priority_command_rx,
497                                side_queues.endpoint_command_rx,
498                                SIDE_QUEUE_INTERLEAVE_BUDGET,
499                            )
500                            .await
501                        } else {
502                            RxLoopDataDrainStats::default()
503                        }
504                    } else {
505                        RxLoopDataDrainStats::default()
506                    };
507                    if !drained.has_drained() {
508                        drain.refund_empty_interleave_turn();
509                    }
510                }
511            }
512        }
513
514        self.flush_decrypt_job_batcher(&mut decrypt_jobs);
515        let drained = drain.drained();
516        if drained > 0 {
517            // One trailing fallback slice so the last bounced packets of the
518            // burst aren't held up by the post-burst send flush. Keep it a
519            // non-packet turn: bulk fallback should not convoy ahead of fresh
520            // transport receive work after every hot packet drain.
521            self.drain_decrypt_fallback(
522                decrypt_fallback_rx,
523                None,
524                None,
525                None,
526                fallback_plan.trailing_budget.min(budget),
527            )
528            .await;
529            self.finish_endpoint_event_batch();
530        } else {
531            self.finish_endpoint_event_batch();
532        }
533        drained
534    }
535
536    async fn drain_rx_loop_side_queues(
537        &mut self,
538        control_query_rx: &mut Receiver<ControlMessage>,
539        tun_outbound_rx: &mut TunOutboundRx,
540        endpoint_priority_command_rx: &mut Receiver<NodeEndpointCommand>,
541        endpoint_command_rx: &mut Receiver<NodeEndpointCommand>,
542        budget: usize,
543    ) -> RxLoopDataDrainStats {
544        let control_budget = budget.min(CONTROL_QUERY_INTERLEAVE_BUDGET);
545        let drained_control = self
546            .drain_control_queries(control_query_rx, None, control_budget)
547            .await;
548        let remaining_budget = budget.saturating_sub(drained_control);
549        let (endpoint_budget, tun_budget) = split_side_queue_budget(remaining_budget);
550        let mut drained_endpoint = self
551            .drain_endpoint_commands(
552                endpoint_priority_command_rx,
553                endpoint_command_rx,
554                None,
555                None,
556                endpoint_budget,
557            )
558            .await;
559        let mut drained_tun = self
560            .drain_tun_outbound(tun_outbound_rx, None, tun_budget)
561            .await;
562
563        let endpoint_remainder = remaining_side_queue_budget(endpoint_budget, drained_endpoint);
564        let tun_remainder = remaining_side_queue_budget(tun_budget, drained_tun);
565        if endpoint_remainder > 0 && !tun_outbound_rx.is_empty() {
566            drained_tun += self
567                .drain_tun_outbound(tun_outbound_rx, None, endpoint_remainder)
568                .await;
569        }
570        if tun_remainder > 0
571            && (!endpoint_priority_command_rx.is_empty() || !endpoint_command_rx.is_empty())
572        {
573            drained_endpoint += self
574                .drain_endpoint_commands(
575                    endpoint_priority_command_rx,
576                    endpoint_command_rx,
577                    None,
578                    None,
579                    tun_remainder,
580                )
581                .await;
582        }
583
584        RxLoopDataDrainStats::with_control(0, drained_tun, drained_endpoint, drained_control)
585    }
586
587    async fn drain_control_queries(
588        &mut self,
589        control_query_rx: &mut Receiver<ControlMessage>,
590        first_message: Option<ControlMessage>,
591        budget: usize,
592    ) -> usize {
593        let mut drain = SingleLaneDrainCursor::new(first_message, budget);
594        while let Some((request, response_tx)) = drain.next(control_query_rx) {
595            let response = queries::dispatch(self, &request.command, request.params.as_ref());
596            let _ = response_tx.send(response);
597        }
598
599        drain.drained()
600    }
601
602    async fn drain_tun_outbound(
603        &mut self,
604        tun_outbound_rx: &mut TunOutboundRx,
605        first_packet: Option<Vec<u8>>,
606        budget: usize,
607    ) -> usize {
608        let mut drain = SingleLaneDrainCursor::new(first_packet, budget);
609        while let Some(packet) = drain.next(tun_outbound_rx) {
610            self.handle_tun_outbound(packet).await;
611        }
612
613        drain.drained()
614    }
615
616    async fn drain_endpoint_commands(
617        &mut self,
618        endpoint_priority_command_rx: &mut Receiver<NodeEndpointCommand>,
619        endpoint_command_rx: &mut Receiver<NodeEndpointCommand>,
620        first_priority_command: Option<NodeEndpointCommand>,
621        first_bulk_command: Option<NodeEndpointCommand>,
622        budget: usize,
623    ) -> usize {
624        let mut drain =
625            PriorityBulkDrainCursor::new(first_priority_command, first_bulk_command, budget);
626        while let Some(command) = drain.next(endpoint_priority_command_rx, endpoint_command_rx) {
627            let drain_cost = command.drain_cost();
628            match command.into_send_batch_oneway() {
629                Ok((batch, _lane)) => {
630                    let mut batch_commands = vec![batch];
631                    self.coalesce_endpoint_send_batch_commands(
632                        &mut drain,
633                        endpoint_priority_command_rx,
634                        endpoint_command_rx,
635                        &mut batch_commands,
636                    );
637                    self.handle_endpoint_send_batch_commands(batch_commands)
638                        .await;
639                }
640                Err(command) => {
641                    self.handle_endpoint_data_command(command).await;
642                }
643            }
644            drain.charge_extra(drain_cost.saturating_sub(1));
645        }
646
647        drain.drained()
648    }
649
650    fn drain_endpoint_bulk_send_feedback(
651        &mut self,
652        endpoint_bulk_feedback_rx: &mut Receiver<crate::node::EndpointBulkSendFeedback>,
653        first_feedback: Option<crate::node::EndpointBulkSendFeedback>,
654        budget: usize,
655    ) -> usize {
656        let mut drain = SingleLaneDrainCursor::new(first_feedback, budget);
657        while let Some(feedback) = drain.next(endpoint_bulk_feedback_rx) {
658            self.apply_endpoint_bulk_send_feedback(feedback);
659        }
660
661        drain.drained()
662    }
663
664    fn coalesce_endpoint_send_batch_commands(
665        &mut self,
666        drain: &mut PriorityBulkDrainCursor<NodeEndpointCommand>,
667        endpoint_priority_command_rx: &mut Receiver<NodeEndpointCommand>,
668        endpoint_command_rx: &mut Receiver<NodeEndpointCommand>,
669        batch_commands: &mut Vec<EndpointSendBatchCommand>,
670    ) {
671        let mut payloads = batch_commands
672            .iter()
673            .fold(0usize, |total, command| total.saturating_add(command.len()));
674        while payloads < ENDPOINT_COMMAND_COALESCE_MAX_PACKETS {
675            let Some(command) =
676                drain.next_bulk_if_no_priority(endpoint_priority_command_rx, endpoint_command_rx)
677            else {
678                break;
679            };
680            let drain_cost = command.drain_cost();
681            match command.into_send_batch_oneway() {
682                Ok((batch, _lane))
683                    if batch_commands.last().is_some_and(|last| {
684                        last.can_coalesce_with(&batch, ENDPOINT_COMMAND_COALESCE_MAX_PACKETS)
685                    }) =>
686                {
687                    payloads = payloads.saturating_add(batch.len());
688                    batch_commands.push(batch);
689                    drain.charge_extra(drain_cost.saturating_sub(1));
690                }
691                Ok((batch, lane)) => {
692                    drain.defer_bulk(NodeEndpointCommand::SendBatchOneway {
693                        command: batch,
694                        lane,
695                    });
696                    break;
697                }
698                Err(command) => {
699                    drain.defer_bulk(command);
700                    break;
701                }
702            }
703        }
704    }
705
706    async fn run_rx_loop_maintenance_tick(&mut self, plan: RxLoopMaintenancePlan) -> bool {
707        self.check_timeouts();
708        let now_ms = Self::now_ms();
709        // Link/session liveness must run before slower retry/discovery work:
710        // under bulk send pressure a late heartbeat or MMP report is
711        // indistinguishable from a dead direct path on the remote peer.
712        self.check_link_heartbeats().await;
713        self.reload_peer_acl();
714        self.resend_pending_handshakes(now_ms).await;
715        self.resend_pending_rekeys(now_ms).await;
716        self.resend_pending_session_handshakes(now_ms).await;
717        self.resend_pending_session_msg3(now_ms).await;
718        self.purge_idle_sessions(now_ms);
719        self.purge_learned_routes(now_ms);
720        self.check_mmp_reports().await;
721        self.check_session_mmp_reports().await;
722        self.check_rekey().await;
723        self.check_session_rekey().await;
724        self.check_pending_lookups(now_ms).await;
725        self.poll_pending_connects().await;
726        self.process_pending_retries(now_ms).await;
727        self.poll_transport_discovery().await;
728        self.activate_connected_udp_sessions().await;
729        self.sample_transport_congestion();
730
731        let Some(slow_timeout) = plan.slow_timeout() else {
732            crate::perf_profile::record_event(
733                crate::perf_profile::Event::RxLoopSlowMaintenanceSkipped,
734            );
735            return false;
736        };
737
738        if tokio::time::timeout(slow_timeout, self.run_rx_loop_slow_maintenance_tick())
739            .await
740            .is_err()
741        {
742            crate::perf_profile::record_event(
743                crate::perf_profile::Event::RxLoopSlowMaintenanceTimeout,
744            );
745            self.mark_rx_loop_maintenance_timeout();
746            warn!(
747                timeout_ms = slow_timeout.as_millis() as u64,
748                data_pressure = plan.data_pressure(),
749                "RX loop slow maintenance timed out; continuing packet processing"
750            );
751            return true;
752        }
753        false
754    }
755
756    async fn run_rx_loop_slow_maintenance_tick(&mut self) {
757        if let Some(delay) = rx_loop_slow_maintenance_fault_delay() {
758            tokio::time::sleep(delay).await;
759        }
760
761        // Discovery and graph/stat maintenance can involve relay work or
762        // larger scans. Keep it bounded after direct-path liveness and session
763        // upkeep so a slow Nostr/LAN tick degrades discovery freshness, not
764        // packet flow.
765        self.poll_nostr_discovery().await;
766        self.poll_lan_discovery().await;
767        self.poll_local_instance_discovery().await;
768        self.check_tree_state().await;
769        self.check_bloom_state().await;
770        self.compute_mesh_size();
771        self.record_stats_history();
772    }
773
774    /// Hand a decrypt-worker fallback to the canonical post-FMP-decrypt
775    /// processor as one authenticated receive envelope. The envelope keeps the
776    /// worker-captured source peer, FMP flags, packet facts, and plaintext slice
777    /// together so peer bookkeeping and link dispatch cannot drift apart.
778    async fn process_decrypt_worker_event(&mut self, event: DecryptWorkerEvent) {
779        event.record_queue_wait();
780        match event {
781            DecryptWorkerEvent::Plaintext(fallback) => {
782                self.process_decrypt_fallback(fallback).await;
783            }
784            DecryptWorkerEvent::PlaintextBatch(fallbacks) => {
785                for fallback in fallbacks {
786                    self.process_decrypt_fallback(fallback).await;
787                }
788            }
789            DecryptWorkerEvent::AuthenticatedFmpReceive(receive) => {
790                self.process_authenticated_fmp_receive_from_worker(receive);
791            }
792            DecryptWorkerEvent::AuthenticatedSession(session) => {
793                self.process_authenticated_session_from_worker(session)
794                    .await;
795            }
796            DecryptWorkerEvent::AuthenticatedSessionBatch(sessions) => {
797                self.process_authenticated_session_batch_from_worker(sessions)
798                    .await;
799            }
800            DecryptWorkerEvent::DirectSessionCommit(commit) => {
801                self.process_direct_session_commit_from_worker(commit).await;
802            }
803            DecryptWorkerEvent::DirectSessionCommitBatch(commits) => {
804                self.process_direct_session_commit_batch_from_worker(commits)
805                    .await;
806            }
807            DecryptWorkerEvent::DirectSessionData(direct) => {
808                self.process_direct_session_data_from_worker(direct).await;
809            }
810            DecryptWorkerEvent::DirectSessionDataBatch(directs) => {
811                self.process_direct_session_data_batch_from_worker(directs)
812                    .await;
813            }
814            DecryptWorkerEvent::FspDecryptFailure(report) => {
815                self.process_fsp_decrypt_failure_from_worker(report).await;
816            }
817            DecryptWorkerEvent::DecryptFailure(report) => {
818                self.process_decrypt_failure_report(report).await;
819            }
820        }
821    }
822
823    async fn process_decrypt_fallback(&mut self, fallback: DecryptFallback) {
824        let plaintext = &fallback.packet_data[fallback.fmp_plaintext_offset
825            ..fallback.fmp_plaintext_offset + fallback.fmp_plaintext_len];
826        self.process_authentic_fmp_plaintext(AuthenticatedFmpPlaintext::new(
827            fallback.source_peer,
828            fallback.transport_id,
829            &fallback.remote_addr,
830            fallback.timestamp_ms,
831            fallback.packet_len,
832            fallback.fmp_counter,
833            fallback.fmp_flags,
834            plaintext,
835        ))
836        .await;
837    }
838
839    async fn process_decrypt_failure_report(&mut self, report: DecryptFailureReport) {
840        debug!(
841            peer = %self.peer_display_name(report.source_peer.node_addr()),
842            counter = report.fmp_counter,
843            replay_highest = report.fmp_replay_highest,
844            "Worker FMP AEAD decryption failed"
845        );
846        self.handle_decrypt_failure_report(&report).await;
847    }
848
849    /// Drain only the priority decrypt-worker fallback lane.
850    ///
851    /// This is the top-level reserved-progress arm: priority plaintext and
852    /// decrypt failures get first service, but bulk fallback stays behind
853    /// `packet_rx` unless it is explicitly interleaved inside a packet drain
854    /// or selected by its own lower-priority branch.
855    async fn drain_decrypt_priority_fallback(
856        &mut self,
857        priority_rx: &mut Receiver<DecryptWorkerEvent>,
858        first_event: Option<DecryptWorkerEvent>,
859        budget: usize,
860    ) -> usize {
861        self.begin_endpoint_event_batch();
862        let mut drain = SingleLaneDrainCursor::new(first_event, budget);
863        while let Some(event) = drain.next(priority_rx) {
864            self.process_decrypt_worker_event(event).await;
865        }
866        let drained = drain.drained();
867        self.finish_endpoint_event_batch();
868        drained
869    }
870
871    /// Drain up to `budget` queued fallbacks without yielding back to
872    /// `select!`. Returns the number processed. Called both from the
873    /// bulk-fallback select arm (after the selected head item) and interleaved
874    /// inside the packet_rx drain loop so bounced FMP plaintexts can't
875    /// accumulate behind a hot inbound packet turn.
876    async fn drain_decrypt_fallback(
877        &mut self,
878        rx: &mut DecryptWorkerFallbackReceivers,
879        first_priority_event: Option<DecryptWorkerEvent>,
880        first_authenticated_bulk_event: Option<DecryptWorkerEvent>,
881        first_bulk_event: Option<DecryptWorkerEvent>,
882        budget: usize,
883    ) -> usize {
884        self.begin_endpoint_event_batch();
885        let mut drain = DecryptReturnDrainCursor::new(
886            first_priority_event,
887            first_authenticated_bulk_event,
888            first_bulk_event,
889            budget,
890        );
891        while let Some(event) =
892            drain.next(&mut rx.priority, &mut rx.authenticated_bulk, &mut rx.bulk)
893        {
894            rx.release_dequeued_event(&event);
895            let extra = event.packet_count().saturating_sub(1);
896            self.process_decrypt_worker_event(event).await;
897            drain.charge_extra(extra);
898        }
899        let drained = drain.drained();
900        self.finish_endpoint_event_batch();
901        drained
902    }
903
904    /// Process a single received packet.
905    ///
906    /// Dispatches based on the phase field in the 4-byte common prefix.
907    #[cfg(test)]
908    pub(in crate::node) async fn process_packet(&mut self, packet: ReceivedPacket) {
909        let action = self.begin_process_packet(packet);
910        self.finish_packet_process(action).await;
911    }
912
913    fn begin_process_packet(&mut self, packet: ReceivedPacket) -> PacketProcessAction {
914        let timer = crate::perf_profile::Timer::start(crate::perf_profile::Stage::ProcessPacket);
915        let priority_sized = packet.is_priority_sized();
916        let priority_count = u64::from(priority_sized);
917        let bulk_count = u64::from(!priority_sized);
918        crate::perf_profile::record_since_split_count(
919            crate::perf_profile::Stage::TransportQueueWait,
920            crate::perf_profile::Stage::TransportPriorityQueueWait,
921            crate::perf_profile::Stage::TransportBulkQueueWait,
922            packet.trace_enqueued_at,
923            1,
924            priority_count,
925            bulk_count,
926        );
927        crate::perf_profile::record_since_split_count(
928            crate::perf_profile::Stage::TransportRxLoopWait,
929            crate::perf_profile::Stage::TransportPriorityRxLoopWait,
930            crate::perf_profile::Stage::TransportBulkRxLoopWait,
931            packet.trace_rx_loop_owned_at,
932            1,
933            priority_count,
934            bulk_count,
935        );
936        if is_punch_packet(&packet.data) {
937            trace!(
938                transport_id = %packet.transport_id,
939                remote_addr = %packet.remote_addr,
940                bytes = packet.data.len(),
941                "Dropping stray punch probe/ack in FMP rx loop"
942            );
943            return PacketProcessAction::Done;
944        }
945        if packet.data.len() < COMMON_PREFIX_SIZE {
946            return PacketProcessAction::Done; // Drop packets too short for common prefix
947        }
948
949        let prefix = match CommonPrefix::parse(&packet.data) {
950            Some(p) => p,
951            None => return PacketProcessAction::Done, // Malformed prefix
952        };
953        if matches!(prefix.phase, PHASE_MSG1 | PHASE_MSG2) {
954            debug!(
955                transport_id = %packet.transport_id,
956                remote_addr = %packet.remote_addr,
957                bytes = packet.data.len(),
958                phase = prefix.phase,
959                version = prefix.version,
960                "FMP handshake packet dispatch"
961            );
962        } else {
963            trace!(
964                transport_id = %packet.transport_id,
965                remote_addr = %packet.remote_addr,
966                bytes = packet.data.len(),
967                phase = prefix.phase,
968                version = prefix.version,
969                "FMP packet dispatch"
970            );
971        }
972
973        if prefix.version != FMP_VERSION {
974            debug!(
975                version = prefix.version,
976                transport_id = %packet.transport_id,
977                "Unknown FMP version, dropping"
978            );
979
980            // If the packet arrived on an adopted Nostr-NAT bootstrap
981            // transport, the originating peer is necessarily on a
982            // different FMP-protocol version than us — the discovery
983            // sweep would otherwise re-traverse them every cycle even
984            // though no msg1/msg2 exchange can ever succeed. Bump the
985            // discovery-layer cooldown to the long protocol-mismatch
986            // window and emit a single WARN per fresh observation.
987            let looks_like_fmp_phase =
988                matches!(prefix.phase, PHASE_ESTABLISHED | PHASE_MSG1 | PHASE_MSG2);
989            if looks_like_fmp_phase
990                && self.bootstrap_transports.contains(&packet.transport_id)
991                && let Some(npub) = self.bootstrap_transports.peer_npub(&packet.transport_id)
992                && let Some(handle) = self.nostr_discovery_handle()
993            {
994                let now_ms = Self::now_ms();
995                let cooldown_secs = handle.protocol_mismatch_cooldown_secs();
996                if handle.record_protocol_mismatch(npub, now_ms) {
997                    warn!(
998                        peer_npub = %npub,
999                        transport_id = %packet.transport_id,
1000                        peer_version = prefix.version,
1001                        our_version = FMP_VERSION,
1002                        cooldown_secs,
1003                        "Nostr-discovered peer speaks a different FMP version; suppressing retraversal"
1004                    );
1005                }
1006            }
1007            return PacketProcessAction::Done;
1008        }
1009
1010        match prefix.phase {
1011            PHASE_ESTABLISHED => match self.try_prepare_encrypted_frame_for_worker(packet) {
1012                EncryptedFrameFastPath::Dispatch(job) => PacketProcessAction::DecryptJob { job },
1013                EncryptedFrameFastPath::Dropped => PacketProcessAction::Done,
1014                EncryptedFrameFastPath::Slow(packet) => {
1015                    PacketProcessAction::EncryptedSlow { packet, timer }
1016                }
1017            },
1018            PHASE_MSG1 => PacketProcessAction::Msg1 { packet, timer },
1019            PHASE_MSG2 => PacketProcessAction::Msg2 { packet, timer },
1020            _ => {
1021                debug!(
1022                    phase = prefix.phase,
1023                    transport_id = %packet.transport_id,
1024                    "Unknown FMP phase, dropping"
1025                );
1026                PacketProcessAction::Done
1027            }
1028        }
1029    }
1030
1031    async fn finish_packet_process(&mut self, action: PacketProcessAction) {
1032        match action {
1033            PacketProcessAction::Done => {}
1034            PacketProcessAction::DecryptJob { job } => {
1035                if let Some(workers) = self.decrypt_workers.as_ref() {
1036                    workers.dispatch_job(job);
1037                }
1038            }
1039            PacketProcessAction::EncryptedSlow {
1040                packet,
1041                timer: _timer,
1042            } => {
1043                self.handle_encrypted_frame_slow(packet).await;
1044            }
1045            PacketProcessAction::Msg1 {
1046                packet,
1047                timer: _timer,
1048            } => {
1049                self.handle_msg1(packet).await;
1050            }
1051            PacketProcessAction::Msg2 {
1052                packet,
1053                timer: _timer,
1054            } => {
1055                self.handle_msg2(packet).await;
1056            }
1057        }
1058    }
1059
1060    fn flush_decrypt_job_batcher(&self, batcher: &mut DecryptJobBatcher) {
1061        if let Some(workers) = self.decrypt_workers.as_ref() {
1062            batcher.flush(workers);
1063        }
1064    }
1065}
fips_core/node/handlers/rx_loop.rs

fips_core/node/handlers/
rx_loop.rs
