diff --git a/easytier-core/src/connectivity/hole_punch/udp/client.rs b/easytier-core/src/connectivity/hole_punch/udp/client.rs index 9de579a9..b0d2d1b6 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/client.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/client.rs @@ -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>> { - 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, 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 { 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>, sockets: std::sync::Mutex>>, 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] diff --git a/easytier-core/src/connectivity/hole_punch/udp/server.rs b/easytier-core/src/connectivity/hole_punch/udp/server.rs index 4946cf53..4af0fbe6 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/server.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/server.rs @@ -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 { diff --git a/easytier-core/src/connectivity/hole_punch/udp/socket_array.rs b/easytier-core/src/connectivity/hole_punch/udp/socket_array.rs index 7d07c8c2..36f741ab 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/socket_array.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/socket_array.rs @@ -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>>, - interest_tids: Arc>, tid_to_socket: Arc>>>, } @@ -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) -> 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());