Queue server challenges

Author: divagant-martian

quinn-proto/src/connection/mod.rs

@@ -4488,8 +4488,52 @@ impl Connection {
                             )));
                     }
                 }
-                Frame::ReachOut(_frame) => {
-                    // TODO(@divma): handle
+                Frame::ReachOut(reach_out) => {
+                    let Some(hp_state) = self.iroh_hp.as_mut() else {
+                        return Err(TransportError::PROTOCOL_VIOLATION(
+                            "received REACH_OUT frame when iroh's nat traversal was not negotiated",
+                        ));
+                    };
+
+                    match hp_state.handle_reach_out(reach_out) {
+                        Ok(None) => {
+                            // no action required here
+                        }
+                        Ok(Some((challenge_info, is_new_round))) => {
+                            let iroh_hp::ChallengeNeeded {
+                                token,
+                                ip,
+                                port,
+                                round,
+                            } = challenge_info;
+                            if is_new_round {
+                                // TODO(@divma): this depends on round starting on 1 right now,
+                                // because the round should be greater to the default one, which is
+                                // zero
+                                self.spaces[SpaceId::Data].pending.challenges_round = round;
+                                self.spaces[SpaceId::Data].pending.challenges.clear();
+                            }
+
+                            self.spaces[SpaceId::Data]
+                                .pending
+                                .challenges
+                                .push((token, ip, port));
+                        }
+                        Err(iroh_hp::Error::WrongConnectionSide) => {
+                            return Err(TransportError::PROTOCOL_VIOLATION(
+                                "server sent REACH_OUT frames for nat traversal",
+                            ));
+                        }
+                        Err(iroh_hp::Error::TooManyAddresses) => {
+                            return Err(TransportError::PROTOCOL_VIOLATION(
+                                "client exceeded allowed REACH_OUT frames for this round",
+                            ));
+                        }
+                        Err(error) => {
+                            warn!(%error,"error handling REACH_OUT frame");
+                            // TODO(@divma): check if this is reachable
+                        }
+                    }
                 }
             }
         }
@@ -5391,6 +5435,7 @@ impl Connection {
                     max_remote_addresses,
                     max_local_addresses,
                     self.side(),
+                    self.rng.clone(),
                 ));
                 debug!(
                     %max_remote_addresses, %max_local_addresses,

quinn-proto/src/connection/spaces.rs

@@ -564,8 +564,12 @@ pub struct Retransmits {
     pub(super) remove_address: BTreeSet<RemoveAddress>,
     /// Round and local addresses to advertise in `REACH_OUT` frames
     pub(super) reach_out: Option<(VarInt, Vec<(IpAddr, u16)>)>,
-    /// Round and remote addresses to which `PATH_CHALLENGE`s need to be sent
-    pub(super) challenges: Option<(VarInt, Vec<(IpAddr, u16)>)>,
+    /// Round of the nat traversal challenges that are pending
+    ///
+    /// This is only used for bitwise operations on the retransmission data.
+    pub(super) challenges_round: VarInt,
+    /// Tokens and remote addresses to which `PATH_CHALLENGE`s need to be sent
+    pub(super) challenges: Vec<(u64, IpAddr, u16)>,
 }
 
 impl Retransmits {
@@ -592,7 +596,7 @@ impl Retransmits {
             && self.add_address.is_empty()
             && self.remove_address.is_empty()
             && self.reach_out.is_none()
-            && self.challenges.is_none()
+            && self.challenges.is_empty()
     }
 }
 
@@ -633,15 +637,11 @@ impl ::std::ops::BitOrAssign for Retransmits {
             _ => {}
         }
 
-        // if there are two rounds, prefer the most recent pending challenges set
-        let lhs_round = self.challenges.as_ref().map(|(round, _)| *round);
-        let rhs_round = rhs.challenges.as_ref().map(|(round, _)| *round);
-        match (lhs_round, rhs_round) {
-            (None, Some(_)) => self.challenges = rhs.challenges.clone(),
-            (Some(lhs_round), Some(rhs_round)) if rhs_round > lhs_round => {
-                self.challenges = rhs.challenges.clone()
-            }
-            _ => {}
+        if self.challenges_round < rhs.challenges_round {
+            self.challenges_round = rhs.challenges_round;
+            self.challenges = rhs.challenges.clone();
+        } else if self.challenges_round == rhs.challenges_round {
+            self.challenges.extend_from_slice(&rhs.challenges);
         }
     }
 }

quinn-proto/src/iroh_hp.rs

@@ -3,11 +3,12 @@ use std::{
     net::{IpAddr, SocketAddr},
 };
 
+use rand::{Rng, rngs::StdRng};
 use rustc_hash::FxHashMap;
 
 use crate::{
     PathId, Side, VarInt,
-    frame::{AddAddress, RemoveAddress},
+    frame::{AddAddress, ReachOut, RemoveAddress},
 };
 
 /// Maximum number of addresses to handle, applied both to local and remote addresses, regardless
@@ -45,6 +46,17 @@ pub(crate) struct NatTraversalRound {
     pub(crate) prev_round_path_ids: Vec<PathId>,
 }
 
+pub(crate) struct ChallengeNeeded {
+    /// Token to send on the challenge
+    pub(crate) token: u64,
+    /// Destination address of the challenge
+    pub(crate) ip: IpAddr,
+    /// Destination port of the challenge
+    pub(crate) port: u16,
+    /// Round to which this challenge belongs to
+    pub(crate) round: VarInt,
+}
+
 // TODO(@divma): unclear to me what these events are useful for\
 #[derive(Debug, Clone)]
 pub enum Event {
@@ -63,6 +75,8 @@ pub(crate) struct State {
     ///
     /// This is set by the remote endpoint.
     max_local_addresses: usize,
+    /// Random generator
+    rng: StdRng,
     /// Candidate addresses the remote server reports as potentially reachable, to use for nat
     /// traversal attempts.
     remote_addresses: FxHashMap<VarInt, (IpAddr, u16)>,
@@ -83,7 +97,7 @@ pub(crate) struct State {
     round_path_ids: Vec<PathId>,
     /// Challenges sent by servers to validate client addresses without attempting to open
     /// multipath paths
-    challenges: FxHashMap<u64, (IpAddr, u16)>,
+    challenges: FxHashMap<(IpAddr, u16), u64>,
 }
 
 /// Nat traversal api exclusive to clients
@@ -147,10 +161,16 @@ impl State {
         }
     }
 
-    pub(crate) fn new(max_remote_addresses: u8, max_local_addresses: u8, side: Side) -> Self {
+    pub(crate) fn new(
+        max_remote_addresses: u8,
+        max_local_addresses: u8,
+        side: Side,
+        rng: StdRng,
+    ) -> Self {
         Self {
             remote_addresses: Default::default(),
             local_addresses: Default::default(),
+            rng,
             next_local_addr_id: Default::default(),
             side,
             round: Default::default(),
@@ -209,6 +229,43 @@ impl State {
         self.round_path_ids = path_ids;
         Ok(())
     }
+
+    /// Handles a received [`ReachOut`]
+    ///
+    /// It returns the token that should be sent in response to this frame as a challenge, and
+    /// whether this starts a new nat traversal round.
+    ///
+    /// If this frame was ignored, it returns `None`.
+    pub(crate) fn handle_reach_out(
+        &mut self,
+        reach_out: ReachOut,
+    ) -> Result<Option<(ChallengeNeeded, bool)>, Error> {
+        let ReachOut { round, ip, port } = reach_out;
+        if self.side.is_client() {
+            return Err(Error::WrongConnectionSide);
+        }
+
+        if round >= self.round {
+            let is_new_round = round > self.round;
+            if is_new_round {
+                self.challenges.clear();
+            }
+            if self.challenges.len() >= self.max_remote_addresses {
+                return Err(Error::TooManyAddresses);
+            }
+            let token = self.rng.random();
+            self.challenges.insert((ip, port), token);
+            let challenge_info = ChallengeNeeded {
+                token,
+                ip,
+                port,
+                round,
+            };
+            return Ok(Some((challenge_info, is_new_round)));
+        }
+
+        Ok(None)
+    }
 }
 
 impl<'a> ClientSide<'a> {

quinn-proto/src/lib.rs

@@ -329,7 +329,7 @@ pub struct Transmit {
 //
 
 /// The maximum number of CIDs we bother to issue per path
-const LOC_CID_COUNT: u64 = 8;
+const LOC_CID_COUNT: u64 = 12;
 const RESET_TOKEN_SIZE: usize = 16;
 const MAX_CID_SIZE: usize = 20;
 const MIN_INITIAL_SIZE: u16 = 1200;