fix(udp): repair symmetric hole-punch socket handoff and retries (#2774)

* fix(hole-punch): repair symmetric UDP candidate and handshake flow Move punched sockets out of the candidate array before publishing them, and keep transaction registration and results in one map. This prevents receiver cleanup from racing with socket reuse or port rebinding and removes the unsafe reference-count assertion. Refill hard-symmetric candidates consumed before transport admission. Refresh easy-symmetric candidates before taking a new STUN baseline so prediction covers their new mappings. Keep probing a selected both-easy socket during session setup so a peer still receiving raw HolePunch packets can switch to its listener even after one-way probe loss. Keep predicted ports within 1..=65535, permit a single remaining port, and cover the full prediction interval for every advertised public IP. Preserve the existing random scan budgets, retry policy and wire format. Add regression coverage for candidate depletion, prediction retries, receiver handoff, continued probing, multiple IPs and port boundaries. Validate all 877 core tests, full workspace Clippy and NAT mapping probes that exchange payloads through the production UDP session code. Related to #2766; its field root cause remains unconfirmed. * fix(hole-punch): retain candidate reuse and prediction budgets Remove the unconditional easy-symmetric array reset so retries reuse the existing 84-socket pool. Keep refilling consumed slots to avoid permanent depletion after failed transport admission. Preserve the predictable scan budget for multiple public IPs instead of multiplying it by address count. Defer prediction freshness and coverage changes until they can be addressed within the existing pressure limits. Verify that two failed easy-symmetric attempts reuse all candidates, that consumed candidates still recover, and that a 50-target prediction budget with two public IPs retains the original 300-datagram total.
2026-10-29 08:03:54 +00:00 · 2026-10-29 12:31:38 +08:00
parent 3ed7fcece6
commit c66076b8f3
3 changed files with 295 additions and 65 deletions
@@ -41,12 +41,11 @@ const DST_PORT_OFFSET: u16 = 20;
 const REMOTE_WAIT_TIME_MS: u64 = 5000;

 pub fn apply_peer_easy_sym_port_offset(base_port: u16, peer_is_incremental: bool) -> u16 {
-    let port = if peer_is_incremental {
-        (base_port as u32).saturating_add(DST_PORT_OFFSET as u32)
+    if peer_is_incremental {
+        base_port.saturating_add(DST_PORT_OFFSET)
    } else {
-        (base_port as u32).saturating_sub(DST_PORT_OFFSET as u32)
-    };
-    port as u16
+        base_port.saturating_sub(DST_PORT_OFFSET).max(1)
+    }
 }

 #[tracing::instrument(skip(runtime, signaling), fields(dst_peer_id), err)]
@@ -213,24 +212,19 @@ where
    }

    async fn prepare_udp_array(&self) -> anyhow::Result<Arc<UdpSocketArray<R>>> {
-        let rlocked = self.udp_array.read().await;
-        if let Some(udp_array) = rlocked.clone() {
-            return Ok(udp_array);
-        }
-
-        drop(rlocked);
        let mut wlocked = self.udp_array.write().await;
-        if let Some(udp_array) = wlocked.clone() {
-            return Ok(udp_array);
-        }
-
-        let udp_array = Arc::new(UdpSocketArray::new_with_context(
-            UDP_ARRAY_SIZE_FOR_HARD_SYM,
-            self.runtime.clone(),
-            self.runtime.socket_context(),
-        ));
+        let udp_array = wlocked
+            .get_or_insert_with(|| {
+                Arc::new(UdpSocketArray::new_with_context(
+                    UDP_ARRAY_SIZE_FOR_HARD_SYM,
+                    self.runtime.clone(),
+                    self.runtime.socket_context(),
+                ))
+            })
+            .clone();
+        // Successful candidates leave the array even if transport admission fails.
+        // Refill only those slots so other candidates retain their NAT mappings.
        udp_array.start().await?;
-        wlocked.replace(udp_array.clone());
        Ok(udp_array)
    }

@@ -513,11 +507,10 @@ where
                SendPunchPacketBothEasySym {
                    transaction_id: tid,
                    public_ip: my_public_ip,
-                    dst_port_num: if me_is_incremental {
-                        cur_mapped_addr.port().saturating_add(DST_PORT_OFFSET)
-                    } else {
-                        cur_mapped_addr.port().saturating_sub(DST_PORT_OFFSET)
-                    } as u32,
+                    dst_port_num: apply_peer_easy_sym_port_offset(
+                        cur_mapped_addr.port(),
+                        me_is_incremental,
+                    ) as u32,
                    udp_socket_count: UDP_ARRAY_SIZE_FOR_BOTH_EASY_SYM as u32,
                    wait_time_ms: REMOTE_WAIT_TIME_MS as u32,
                },
@@ -534,6 +527,7 @@ where
            .ok_or(anyhow::anyhow!("remote_mapped_addr is required"))?;

        let now = Instant::now();
+        let punch_packet = new_hole_punch_packet(tid, HOLE_PUNCH_PACKET_BODY_LEN).into_bytes();
        remote_mapped_addr.set_port(apply_peer_easy_sym_port_offset(
            remote_mapped_addr.port(),
            peer_is_incremental,
@@ -546,10 +540,7 @@ where

        while now.elapsed().as_millis() < (REMOTE_WAIT_TIME_MS + 1000).into() {
            udp_array
-                .send_with_all(
-                    &new_hole_punch_packet(tid, HOLE_PUNCH_PACKET_BODY_LEN).into_bytes(),
-                    remote_mapped_addr,
-                )
+                .send_with_all(&punch_packet, remote_mapped_addr)
                .await?;

            crate::foundation::time::sleep(std::time::Duration::from_millis(100)).await;
@@ -571,17 +562,25 @@ where
            );

            for _ in 0..2 {
-                match self
+                // The peer may still be receiving raw HolePunch packets. Keep
+                // probing this candidate until it switches to a session listener.
+                send_from_local(socket.socket.as_ref(), &punch_packet, remote_mapped_addr).await?;
+                let connect = self
                    .runtime
-                    .connect_with_socket(socket.socket.clone(), remote_mapped_addr)
-                    .await
-                {
-                    Ok(socket) => {
-                        return Ok(Some(socket));
-                    }
-                    Err(e) => {
-                        tracing::error!(?e, "failed to connect with socket");
-                        continue;
+                    .connect_with_socket(socket.socket.clone(), remote_mapped_addr);
+                tokio::pin!(connect);
+                loop {
+                    tokio::select! {
+                        result = &mut connect => {
+                            match result {
+                                Ok(socket) => return Ok(Some(socket)),
+                                Err(e) => tracing::error!(?e, "failed to connect with socket"),
+                            }
+                            break;
+                        }
+                        _ = crate::foundation::time::sleep(std::time::Duration::from_millis(200)) => {
+                            send_from_local(socket.socket.as_ref(), &punch_packet, remote_mapped_addr).await?;
+                        }
                    }
                }
            }
@@ -626,6 +625,7 @@ mod tests {
        incoming_notify: tokio::sync::Notify,
        sent: std::sync::Mutex<Vec<(Vec<u8>, SocketAddr)>>,
        drop_syn_packets: AtomicUsize,
+        required_punch_packets: AtomicUsize,
    }

    #[async_trait]
@@ -636,7 +636,20 @@ mod tests {

        async fn send_to(&self, data: &[u8], addr: SocketAddr) -> io::Result<usize> {
            self.sent.lock().unwrap().push((data.to_vec(), addr));
+            if crate::packet::hole_punch_packet_tid(data, HOLE_PUNCH_PACKET_BODY_LEN).is_some()
+                && self.required_punch_packets.load(Ordering::Relaxed) != 0
+            {
+                self.required_punch_packets.fetch_sub(1, Ordering::Relaxed);
+                self.incoming
+                    .lock()
+                    .unwrap()
+                    .push_back((data.to_vec(), addr));
+                self.incoming_notify.notify_one();
+            }
            if data.len() == 16 && data[4] == crate::packet::UdpPacketType::Syn as u8 {
+                if self.required_punch_packets.load(Ordering::Relaxed) != 0 {
+                    return Ok(data.len());
+                }
                if self.drop_syn_packets.load(Ordering::Relaxed) != 0 {
                    self.drop_syn_packets.fetch_sub(1, Ordering::Relaxed);
                    return Ok(data.len());
@@ -672,6 +685,7 @@ mod tests {
        bind_options: tokio::sync::Mutex<Vec<UdpBindOptions>>,
        sockets: std::sync::Mutex<Vec<std::sync::Weak<MockSocket>>>,
        drop_syn_packets: usize,
+        required_punch_packets: usize,
    }

    impl MockRuntime {
@@ -682,6 +696,7 @@ mod tests {
                bind_options: tokio::sync::Mutex::new(Vec::new()),
                sockets: std::sync::Mutex::new(Vec::new()),
                drop_syn_packets: 0,
+                required_punch_packets: 0,
            }
        }
    }
@@ -699,6 +714,7 @@ mod tests {
                incoming_notify: tokio::sync::Notify::new(),
                sent: std::sync::Mutex::new(Vec::new()),
                drop_syn_packets: AtomicUsize::new(self.drop_syn_packets),
+                required_punch_packets: AtomicUsize::new(self.required_punch_packets),
            });
            self.sockets.lock().unwrap().push(Arc::downgrade(&socket));
            Ok(socket)
@@ -1057,6 +1073,18 @@ mod tests {
        let hard_requests = signaling.hard_requests.lock().await;
        assert_eq!(hard_requests.len(), 1);
        assert_eq!(last_port_idx, 42);
+        drop(hard_requests);
+        drop(easy_requests);
+
+        client
+            .do_hole_punching(2, 4, &mut last_port_idx, NatType::SymmetricEasyInc.into())
+            .await
+            .unwrap();
+        assert_eq!(stun.port_mapping_count.load(Ordering::Relaxed), 2);
+        assert_eq!(
+            runtime.bind_count.load(Ordering::Relaxed),
+            UDP_ARRAY_SIZE_FOR_HARD_SYM
+        );
    }

    #[tokio::test]
@@ -1091,9 +1119,104 @@ mod tests {
    }

    #[test]
-    fn both_easy_sym_port_offset_preserves_old_proto_cast_semantics() {
-        assert_eq!(apply_peer_easy_sym_port_offset(65530, true), 14);
-        assert_eq!(apply_peer_easy_sym_port_offset(10, false), 0);
+    fn both_easy_sym_port_offset_stays_in_valid_port_range() {
+        for (base, incremental, expected) in [
+            (10000, true, 10020),
+            (10000, false, 9980),
+            (65516, true, 65535),
+            (65530, true, 65535),
+            (65535, true, 65535),
+            (20, false, 1),
+            (10, false, 1),
+            (1, false, 1),
+        ] {
+            assert_eq!(apply_peer_easy_sym_port_offset(base, incremental), expected);
+        }
+    }
+
+    #[tokio::test]
+    async fn sym_to_cone_replenishes_candidates_after_connected_socket_is_dropped() {
+        let runtime = Arc::new(MockRuntime::new());
+        let client = UdpSymToConePunchClient::new(
+            runtime.clone(),
+            Arc::new(RecordingSignaling::new(0)),
+            Arc::new(MockStunInfoProvider::default()),
+        );
+        let array = client.prepare_udp_array().await.unwrap();
+        let remote = SocketAddr::from(([203, 0, 113, 1], 30000));
+
+        for tid in 0..=UDP_ARRAY_SIZE_FOR_HARD_SYM as u32 {
+            array.add_interest_tid(tid);
+            let socket = runtime
+                .sockets
+                .lock()
+                .unwrap()
+                .last()
+                .unwrap()
+                .upgrade()
+                .unwrap();
+            socket.incoming.lock().unwrap().push_back((
+                new_hole_punch_packet(tid, HOLE_PUNCH_PACKET_BODY_LEN)
+                    .into_bytes()
+                    .to_vec(),
+                remote,
+            ));
+            socket.incoming_notify.notify_one();
+            drop(socket);
+            let punched = tokio::time::timeout(std::time::Duration::from_secs(1), async {
+                loop {
+                    if let Some(socket) = array.try_fetch_punched_socket(tid) {
+                        break socket;
+                    }
+                    tokio::task::yield_now().await;
+                }
+            })
+            .await
+            .unwrap();
+            array.remove_interest_tid(tid);
+            let connected = runtime
+                .connect_with_socket(punched.socket, remote)
+                .await
+                .unwrap();
+            // Model a successful UDP handshake followed by failed peer admission.
+            drop(connected);
+
+            assert!(Arc::ptr_eq(
+                &array,
+                &client.prepare_udp_array().await.unwrap()
+            ));
+            assert_eq!(
+                runtime.bind_count.load(Ordering::Relaxed),
+                UDP_ARRAY_SIZE_FOR_HARD_SYM + tid as usize + 1,
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn both_easy_sym_keeps_probing_hit_until_remote_listener_is_ready() {
+        let mut runtime = MockRuntime::new();
+        // Initial probes produce a reverse hit, but the remote raw receiver
+        // needs more probes before it can accept SYN as a session listener.
+        runtime.required_punch_packets = 6;
+        let runtime = Arc::new(runtime);
+        let client = UdpBothEasySymPunchClient::new(
+            runtime,
+            Arc::new(RecordingSignaling::new(0)),
+            Arc::new(MockStunInfoProvider::default()),
+        );
+        let result = tokio::time::timeout(
+            std::time::Duration::from_secs(2),
+            client.do_hole_punching(
+                2,
+                NatType::SymmetricEasyInc.into(),
+                NatType::SymmetricEasyInc.into(),
+                &mut false,
+            ),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        assert!(result.is_some());
    }

    #[tokio::test]
@@ -163,19 +163,25 @@ where
        }

        let base_port_num = request.base_port_num;
+        anyhow::ensure!(
+            (1..=u16::MAX as u32).contains(&base_port_num),
+            "send_punch_packet_easy_sym invalid base port"
+        );
        let max_port_num = request.max_port_num.max(1);
        let port_start = if request.is_incremental {
            base_port_num.saturating_add(1)
        } else {
-            base_port_num.saturating_sub(max_port_num)
+            base_port_num.saturating_sub(max_port_num).max(1)
        };
        let port_end = if request.is_incremental {
-            base_port_num.saturating_add(max_port_num)
+            base_port_num
+                .saturating_add(max_port_num)
+                .min(u16::MAX as u32)
        } else {
            base_port_num.saturating_sub(1)
        };

-        if port_end <= port_start {
+        if port_end < port_start {
            anyhow::bail!("send_punch_packet_easy_sym invalid port range");
        }

@@ -950,7 +956,6 @@ where

                if let Some(s) = udp_array.try_fetch_punched_socket(transaction_id) {
                    tracing::info!(?s, ?transaction_id, "got punched socket in both easy sym");
-                    assert!(Arc::strong_count(&s.socket) == 1);
                    let Some(port) = s.socket.local_addr().ok().map(|addr| addr.port()) else {
                        tracing::warn!("failed to get local addr from punched socket");
                        continue;
@@ -1646,6 +1651,88 @@ mod tests {
        assert_eq!(sent.len(), 1);
    }

+    #[tokio::test]
+    async fn predictable_sender_uses_valid_boundary_ports() {
+        for (base_port_num, max_port_num, is_incremental, expected_ports) in [
+            (10000, 4, true, vec![10001, 10002, 10003, 10004]),
+            (10000, 4, false, vec![9996, 9997, 9998, 9999]),
+            (65534, 50, true, vec![65535]),
+            (2, 50, false, vec![1]),
+            (65535, 50, true, vec![]),
+            (1, 50, false, vec![]),
+        ] {
+            let public_ips = vec![Ipv4Addr::new(198, 51, 100, 1)];
+            let listener = listener(10007, Vec::new());
+            let socket = listener.socket.clone();
+            let listener_mapped_addr = listener.mapped_addr;
+            let server = UdpHolePunchServer::new(
+                Arc::new(MockRuntime::new(vec![listener])),
+                mock_stun(),
+                Arc::new(MockSink::default()),
+                UdpSymPunchLock::default(),
+            );
+            server.common.select_listener(false, false).await.unwrap();
+            let result = server
+                .send_punch_packet_easy_sym_inner(SendPunchPacketEasySym {
+                    listener_mapped_addr,
+                    public_ips: public_ips.clone(),
+                    transaction_id: 9,
+                    base_port_num,
+                    max_port_num,
+                    is_incremental,
+                })
+                .await;
+            assert_eq!(result.is_ok(), !expected_ports.is_empty());
+            let sent = socket.sent.lock().await;
+            assert_eq!(sent.len(), expected_ports.len() * public_ips.len() * 3 * 2);
+            for ip in &public_ips {
+                for port in &expected_ports {
+                    assert_eq!(
+                        sent.iter()
+                            .filter(|(_, addr)| *addr == SocketAddr::from((*ip, *port)))
+                            .count(),
+                        6
+                    );
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn predictable_sender_keeps_packet_budget_with_multiple_public_ips() {
+        let listener = listener(10007, Vec::new());
+        let socket = listener.socket.clone();
+        let listener_mapped_addr = listener.mapped_addr;
+        let server = UdpHolePunchServer::new(
+            Arc::new(MockRuntime::new(vec![listener])),
+            mock_stun(),
+            Arc::new(MockSink::default()),
+            UdpSymPunchLock::default(),
+        );
+        server.common.select_listener(false, false).await.unwrap();
+        let public_ips = vec![
+            Ipv4Addr::new(198, 51, 100, 1),
+            Ipv4Addr::new(198, 51, 100, 2),
+        ];
+        server
+            .send_punch_packet_easy_sym_inner(SendPunchPacketEasySym {
+                listener_mapped_addr,
+                public_ips: public_ips.clone(),
+                transaction_id: 9,
+                base_port_num: 10000,
+                max_port_num: 50,
+                is_incremental: true,
+            })
+            .await
+            .unwrap();
+
+        let sent = socket.sent.lock().await;
+        assert_eq!(sent.len(), 300);
+        for ip in public_ips {
+            assert!(sent.iter().any(|(_, addr)| addr.ip() == ip));
+        }
+    }
+
    #[tokio::test]
    async fn symmetric_packet_sender_returns_next_port_index() {
        let socket = Arc::new(MockSocket {
@@ -4,7 +4,7 @@ use std::{
    sync::{Arc, Mutex},
 };

-use dashmap::{DashMap, DashSet};
+use dashmap::DashMap;
 use tokio::task::JoinSet;
 use tracing::{Instrument, Level, instrument};

@@ -42,7 +42,6 @@ where
    socket_context: SocketContext,
    tasks: Arc<Mutex<JoinSet<()>>>,

-    interest_tids: Arc<DashSet<u32>>,
    tid_to_socket: Arc<DashMap<u32, Vec<PunchedUdpSocket<R::Socket>>>>,
 }

@@ -68,7 +67,6 @@ where
            socket_context,
            tasks,

-            interest_tids: Arc::new(DashSet::new()),
            tid_to_socket: Arc::new(DashMap::new()),
        }
    }
@@ -80,12 +78,11 @@ where
    pub async fn add_new_socket(&self, socket: Arc<R::Socket>) -> anyhow::Result<()> {
        let socket_map = self.sockets.clone();
        let local_addr = socket.local_addr()?;
-        let interest_tids = self.interest_tids.clone();
        let tid_to_socket = self.tid_to_socket.clone();
        socket_map.insert(local_addr, socket.clone());
        self.tasks.lock().unwrap().spawn(
            async move {
-                let _socket_map_guard = RemoveSocketOnDrop {
+                let socket_map_guard = RemoveSocketOnDrop {
                    sockets: socket_map,
                    local_addr,
                };
@@ -106,17 +103,18 @@ where

                    tracing::debug!(?addr, ?tid, "got udp hole punch packet");

-                    if interest_tids.contains(&tid) {
+                    if let Some(mut punched) = tid_to_socket.get_mut(&tid) {
                        tracing::info!(?addr, ?tid, "got hole punching packet with interest tid");
-                        tid_to_socket
-                            .entry(tid)
-                            .or_default()
-                            .push(PunchedUdpSocket {
-                                socket: socket.clone(),
-                                tid,
-                                remote_addr: addr,
-                            });
-                        break;
+                        // Stop receiving and remove the candidate before handing it off.
+                        // The consumer may immediately reuse the socket or rebind its port.
+                        drop(socket_map_guard);
+                        tracing::debug!(?local_addr, "udp socket recv loop end");
+                        punched.push(PunchedUdpSocket {
+                            socket,
+                            tid,
+                            remote_addr: addr,
+                        });
+                        return;
                    }
                }
                tracing::debug!(?local_addr, "udp socket recv loop end");
@@ -170,7 +168,7 @@ where
    }

    pub fn add_interest_tid(&self, tid: u32) {
-        self.interest_tids.insert(tid);
+        self.tid_to_socket.entry(tid).or_default();
    }

    pub fn add_intreast_tid(&self, tid: u32) {
@@ -178,7 +176,6 @@ where
    }

    pub fn remove_interest_tid(&self, tid: u32) {
-        self.interest_tids.remove(&tid);
        self.tid_to_socket.remove(&tid);
    }

@@ -196,7 +193,6 @@ where
            .field("sockets", &self.sockets.len())
            .field("max_socket_count", &self.max_socket_count)
            .field("started", &self.started())
-            .field("interest_tids", &self.interest_tids.len())
            .field("tid_to_socket", &self.tid_to_socket.len())
            .finish()
    }
@@ -323,6 +319,9 @@ mod tests {
                assert_eq!(punched.tid, tid);
                assert_eq!(punched.remote_addr, remote_addr);
                assert!(Arc::ptr_eq(&punched.socket, &socket));
+                assert!(array.sockets.is_empty());
+                drop(socket);
+                assert_eq!(Arc::strong_count(&punched.socket), 1);
                return;
            }
            tokio::task::yield_now().await;
@@ -331,6 +330,27 @@ mod tests {
        panic!("punched socket was not recorded");
    }

+    #[tokio::test]
+    async fn removing_interest_drops_hits_and_ignores_late_packets() {
+        let array = UdpSocketArray::new(0, Arc::new(MockFactory::new()));
+        let tid = 7;
+        array.add_interest_tid(tid);
+        array.remove_interest_tid(tid);
+        let socket = Arc::new(MockSocket::new(
+            SocketAddr::from(([127, 0, 0, 1], 20000)),
+            vec![(
+                new_hole_punch_packet(tid, HOLE_PUNCH_PACKET_BODY_LEN)
+                    .into_bytes()
+                    .to_vec(),
+                SocketAddr::from(([10, 0, 0, 1], 1234)),
+            )],
+        ));
+        array.add_new_socket(socket).await.unwrap();
+        tokio::task::yield_now().await;
+        assert!(array.tid_to_socket.is_empty());
+        assert_eq!(array.sockets.len(), 1);
+    }
+
    #[tokio::test]
    async fn send_with_all_sends_three_packets_per_socket() {
        let runtime = Arc::new(MockFactory::new());