diff --git a/docs/core-architecture.md b/docs/core-architecture.md index bcb1a780..f2acf670 100644 --- a/docs/core-architecture.md +++ b/docs/core-architecture.md @@ -179,6 +179,21 @@ wire codecs. It does not own socket I/O or connection policy. - UDP session multiplexing, classification, and lifecycle; - in-process Ring sockets. +UDP sessions expose a scoped hole-punch source observer keyed by transaction +ID. Listeners carry a registration handle for their existing receive loop, which +reports the first matching source; dropping the observer removes the registration. +Cone clients optionally prepare that transaction in the listener-selection RPC; +the response echoes its ID only after registration. The registration retains +early UDP sources until the cone-sending RPC claims it. Unclaimed registrations +expire after 30 seconds through the existing listener cleanup loop, covering +local socket and STUN preparation without a separate receive or sending task. +Clients omitting the optional transaction retain registration at the sending RPC; +older servers omit the echoed ID and do not confirm early registration. +Connectivity owns preparation expiry, the active RPC lifetime, candidate selection +and bounded sending to the signaled and observed addresses. +Cone clients start the existing UDP handshake without requiring a reverse +hole-punch packet, and the session layer remains the sole receiver. + `host/` is the single home of Host capability seams: - DNS and DNS record resolution; diff --git a/easytier-core/src/connectivity/hole_punch/udp/binding.rs b/easytier-core/src/connectivity/hole_punch/udp/binding.rs index 72691a32..e14713fa 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/binding.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/binding.rs @@ -331,6 +331,7 @@ where }; let layer = self.session_layer(socket.clone()); + let hole_punch_observers = layer.hole_punch_observers(); let conn_counter = Arc::new(CoreUdpPunchConnCounter { layer: Arc::downgrade(&layer), }); @@ -338,6 +339,7 @@ where Ok(UdpPunchListener { socket, + hole_punch_observers, mapped_addr: resolved.mapped_addr, conn_counter, acceptor, diff --git a/easytier-core/src/connectivity/hole_punch/udp/client.rs b/easytier-core/src/connectivity/hole_punch/udp/client.rs index a40fce5b..9de579a9 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/client.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/client.rs @@ -16,7 +16,7 @@ use crate::{ config::PeerId, connectivity::stun::StunInfoProvider, packet::{HOLE_PUNCH_PACKET_BODY_LEN, new_hole_punch_packet}, - socket::udp::{UdpBindOptions, VirtualUdpSocketFactory}, + socket::udp::{UdpBindOptions, VirtualUdpSocket}, }; use super::{ @@ -62,17 +62,29 @@ where tracing::info!(?dst_peer_id, "start hole punching"); let tid = rand::random(); - let udp_array = UdpSocketArray::new_with_context(1, runtime.clone(), runtime.socket_context()); - let resp = signaling .select_punch_listener( dst_peer_id, SelectPunchListener { force_new: false, prefer_port_mapping: true, + transaction_id: Some(tid), }, ) .await?; + if resp + .transaction_id + .is_some_and(|registered| registered != tid) + { + return Err(UdpHolePunchSignalError::RemoteRejected( + "listener registered a different hole-punch transaction".into(), + ) + .into()); + } + tracing::debug!( + source_observation_registered = resp.transaction_id == Some(tid), + "selected hole-punch listener" + ); let remote_mapped_addr = resp.listener_mapped_addr; let local_socket = UdpHolePunchRuntime::bind_udp( @@ -96,14 +108,12 @@ where "hole punch got remote listener" ); - udp_array.add_new_socket(local_socket).await?; - udp_array.add_intreast_tid(tid); let punch_packet = new_hole_punch_packet(tid, HOLE_PUNCH_PACKET_BODY_LEN).into_bytes(); - send_from_local(&udp_array, &punch_packet, remote_mapped_addr).await?; + send_from_local(local_socket.as_ref(), &punch_packet, remote_mapped_addr).await?; let signaling_for_task = signaling.clone(); - let punch_task = AbortOnDropHandle::new(tokio::spawn(async move { + let _punch_task = AbortOnDropHandle::new(tokio::spawn(async move { if let Err(e) = signaling_for_task .send_punch_packet_cone( dst_peer_id, @@ -122,33 +132,27 @@ where } })); - let mut finish_time: Option = None; - while finish_time.is_none() || finish_time.as_ref().unwrap().elapsed().as_millis() < 1000 { - crate::foundation::time::sleep(std::time::Duration::from_millis(200)).await; - - if finish_time.is_none() && punch_task.is_finished() { - finish_time = Some(Instant::now()); - } - - let Some(socket) = udp_array.try_fetch_punched_socket(tid) else { - tracing::debug!("no punched socket found, send some more hole punch packets"); - send_from_local(&udp_array, &punch_packet, remote_mapped_addr).await?; - continue; - }; - - tracing::debug!(?socket, ?tid, "punched socket found, try connect with it"); - - for _ in 0..2 { - match runtime - .connect_with_socket(socket.socket.clone(), remote_mapped_addr) - .await - { - Ok(socket) => { - tracing::info!(?socket, "hole punched"); - return Ok(Some(socket)); + // The session layer owns receive while connecting. A reachable listener can + // reply to SYN's actual source even when the STUN candidate differs. + for _ in 0..2 { + let connect = runtime.connect_with_socket(local_socket.clone(), remote_mapped_addr); + tokio::pin!(connect); + loop { + tokio::select! { + result = &mut connect => { + match result { + Ok(socket) => { + tracing::info!(?socket, "hole punched"); + return Ok(Some(socket)); + } + Err(e) => { + tracing::error!(?e, "failed to connect with socket"); + break; + } + } } - Err(e) => { - tracing::error!(?e, "failed to connect with socket"); + _ = crate::foundation::time::sleep(std::time::Duration::from_millis(200)) => { + send_from_local(local_socket.as_ref(), &punch_packet, remote_mapped_addr).await?; } } } @@ -157,17 +161,20 @@ where Ok(None) } -async fn send_from_local( - udp_array: &UdpSocketArray, +async fn send_from_local( + socket: &S, punch_packet: &[u8], remote_mapped_addr: SocketAddr, ) -> UdpHolePunchClientResult<()> where - R: VirtualUdpSocketFactory, + S: VirtualUdpSocket, { - udp_array - .send_with_all(punch_packet, remote_mapped_addr) - .await?; + for _ in 0..3 { + socket + .send_to(punch_packet, remote_mapped_addr) + .await + .map_err(anyhow::Error::from)?; + } Ok(()) } @@ -353,6 +360,7 @@ where SelectPunchListener { force_new: false, prefer_port_mapping: true, + transaction_id: None, }, ) .await?; @@ -614,6 +622,10 @@ mod tests { struct MockSocket { local_addr: SocketAddr, + incoming: std::sync::Mutex, SocketAddr)>>, + incoming_notify: tokio::sync::Notify, + sent: std::sync::Mutex, SocketAddr)>>, + drop_syn_packets: AtomicUsize, } #[async_trait] @@ -622,12 +634,35 @@ mod tests { Ok(self.local_addr) } - async fn send_to(&self, _data: &[u8], _addr: SocketAddr) -> io::Result { - Ok(0) + async fn send_to(&self, data: &[u8], addr: SocketAddr) -> io::Result { + self.sent.lock().unwrap().push((data.to_vec(), addr)); + if data.len() == 16 && data[4] == crate::packet::UdpPacketType::Syn as u8 { + if self.drop_syn_packets.load(Ordering::Relaxed) != 0 { + self.drop_syn_packets.fetch_sub(1, Ordering::Relaxed); + return Ok(data.len()); + } + let conn_id = u32::from_le_bytes(data[..4].try_into().unwrap()); + let magic = u64::from_le_bytes(data[8..].try_into().unwrap()); + self.incoming.lock().unwrap().push_back(( + crate::socket::udp::new_sack_packet(conn_id, magic) + .into_bytes() + .to_vec(), + addr, + )); + self.incoming_notify.notify_one(); + } + Ok(data.len()) } - async fn recv_from(&self, _buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { - std::future::pending().await + async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { + loop { + if let Some((data, addr)) = self.incoming.lock().unwrap().pop_front() { + let len = data.len().min(buf.len()); + buf[..len].copy_from_slice(&data[..len]); + return Ok((len, addr)); + } + self.incoming_notify.notified().await; + } } } @@ -635,6 +670,8 @@ mod tests { bind_count: AtomicUsize, resolve_count: AtomicUsize, bind_options: tokio::sync::Mutex>, + sockets: std::sync::Mutex>>, + drop_syn_packets: usize, } impl MockRuntime { @@ -643,6 +680,8 @@ mod tests { bind_count: AtomicUsize::new(0), resolve_count: AtomicUsize::new(0), bind_options: tokio::sync::Mutex::new(Vec::new()), + sockets: std::sync::Mutex::new(Vec::new()), + drop_syn_packets: 0, } } } @@ -654,18 +693,25 @@ mod tests { async fn bind_udp(&self, options: UdpBindOptions) -> anyhow::Result> { self.bind_options.lock().await.push(options); let bind_idx = self.bind_count.fetch_add(1, Ordering::Relaxed); - Ok(Arc::new(MockSocket { + let socket = Arc::new(MockSocket { local_addr: SocketAddr::from(([127, 0, 0, 1], 10000 + bind_idx as u16)), - })) + incoming: std::sync::Mutex::new(std::collections::VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), + sent: std::sync::Mutex::new(Vec::new()), + drop_syn_packets: AtomicUsize::new(self.drop_syn_packets), + }); + self.sockets.lock().unwrap().push(Arc::downgrade(&socket)); + Ok(socket) } async fn resolve_udp_public_addr( &self, - _socket: Arc, + socket: Arc, ) -> anyhow::Result { + assert_eq!(socket.local_addr()?.port(), 10000); self.resolve_count.fetch_add(1, Ordering::Relaxed); Ok(super::super::UdpResolvedPublicAddr { - mapped_addr: SocketAddr::from(([203, 0, 113, 1], 10000)), + mapped_addr: SocketAddr::from(([203, 0, 113, 1], 40000)), port_mapping_lease: None, }) } @@ -686,10 +732,12 @@ mod tests { async fn connect_with_socket( &self, - _socket: Arc, - _remote: SocketAddr, + socket: Arc, + remote: SocketAddr, ) -> anyhow::Result { - unimplemented!("not used by cone client tests") + let layer = Arc::new(crate::socket::udp::UdpSessionLayer::new(socket)); + let session = layer.connect(remote).await?; + Ok(UdpPunchSocket::new(session, remote, layer)) } } @@ -797,6 +845,8 @@ mod tests { } struct RecordingSignaling { + listener_requests: tokio::sync::Mutex>, + confirm_transaction_id: bool, easy_requests: tokio::sync::Mutex>, hard_requests: tokio::sync::Mutex>, both_requests: tokio::sync::Mutex>, @@ -804,9 +854,100 @@ mod tests { next_port_index: u32, } + #[tokio::test] + async fn cone_connects_without_reverse_hole_punch_and_retries_dropped_syn() { + for (drop_syn_packets, confirm_transaction_id) in + [(0, true), (1, true), (0, false), (1, false)] + { + let mut runtime = MockRuntime::new(); + runtime.drop_syn_packets = drop_syn_packets; + let runtime = Arc::new(runtime); + let mut signaling = RecordingSignaling::new(0); + signaling.confirm_transaction_id = confirm_transaction_id; + let signaling = Arc::new(signaling); + let result = tokio::time::timeout( + std::time::Duration::from_secs(1), + punch_cone_to_cone(runtime.clone(), signaling.clone(), 2), + ) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(runtime.bind_count.load(Ordering::Relaxed), 1); + assert_eq!(runtime.resolve_count.load(Ordering::Relaxed), 1); + let tid = signaling.listener_requests.lock().await[0] + .transaction_id + .unwrap(); + let socket = runtime.sockets.lock().unwrap()[0].upgrade().unwrap(); + let sent = socket.sent.lock().unwrap(); + let syns = sent + .iter() + .filter(|(packet, _)| packet[4] == crate::packet::UdpPacketType::Syn as u8) + .count(); + assert_eq!(syns, drop_syn_packets + 1); + assert!(sent.iter().any(|(packet, _)| { + crate::packet::hole_punch_packet_tid(packet, HOLE_PUNCH_PACKET_BODY_LEN) + == Some(tid) + })); + assert!( + sent.iter() + .all(|(_, addr)| *addr == SocketAddr::from(([127, 0, 0, 1], 30000))) + ); + drop(sent); + drop(result); + } + } + + #[tokio::test] + async fn cancelling_cone_connect_releases_socket_and_receive_task() { + let mut runtime = MockRuntime::new(); + runtime.drop_syn_packets = usize::MAX; + let runtime = Arc::new(runtime); + let task_runtime = runtime.clone(); + let task = tokio::spawn(async move { + punch_cone_to_cone(task_runtime, Arc::new(RecordingSignaling::new(0)), 2).await + }); + let socket = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let socket = runtime + .sockets + .lock() + .unwrap() + .first() + .and_then(std::sync::Weak::upgrade); + if let Some(socket) = socket + && socket + .sent + .lock() + .unwrap() + .iter() + .any(|(packet, _)| packet[4] == crate::packet::UdpPacketType::Syn as u8) + { + break socket; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let weak = Arc::downgrade(&socket); + drop(socket); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + } + impl RecordingSignaling { fn new(next_port_index: u32) -> Self { Self { + listener_requests: tokio::sync::Mutex::new(Vec::new()), + confirm_transaction_id: true, easy_requests: tokio::sync::Mutex::new(Vec::new()), hard_requests: tokio::sync::Mutex::new(Vec::new()), both_requests: tokio::sync::Mutex::new(Vec::new()), @@ -832,10 +973,16 @@ mod tests { async fn select_punch_listener( &self, _dst_peer_id: PeerId, - _request: SelectPunchListener, + request: SelectPunchListener, ) -> Result { + let transaction_id = self + .confirm_transaction_id + .then_some(request.transaction_id) + .flatten(); + self.listener_requests.lock().await.push(request); Ok(super::super::SelectPunchListenerResponse { listener_mapped_addr: SocketAddr::from(([127, 0, 0, 1], 30000)), + transaction_id, }) } diff --git a/easytier-core/src/connectivity/hole_punch/udp/rpc.rs b/easytier-core/src/connectivity/hole_punch/udp/rpc.rs index e5ea1d4a..7b011e84 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/rpc.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/rpc.rs @@ -56,6 +56,7 @@ fn select_listener_request_to_rpc(request: SelectPunchListener) -> SelectPunchLi SelectPunchListenerRequest { force_new: request.force_new, prefer_port_mapping: request.prefer_port_mapping, + transaction_id: request.transaction_id, } } @@ -63,6 +64,7 @@ fn select_listener_request_from_rpc(input: SelectPunchListenerRequest) -> Select SelectPunchListener { force_new: input.force_new, prefer_port_mapping: input.prefer_port_mapping, + transaction_id: input.transaction_id, } } @@ -70,6 +72,7 @@ fn select_listener_response_from_rpc( response: SelectPunchListenerResponse, ) -> Result { Ok(CoreSelectPunchListenerResponse { + transaction_id: response.transaction_id, listener_mapped_addr: SocketAddr::from( response .listener_mapped_addr @@ -83,6 +86,7 @@ fn select_listener_response_to_rpc( ) -> SelectPunchListenerResponse { SelectPunchListenerResponse { listener_mapped_addr: Some(response.listener_mapped_addr.into()), + transaction_id: response.transaction_id, } } @@ -553,8 +557,62 @@ where mod tests { use std::{future, net::Ipv4Addr, time::Duration}; + use prost::Message; + use super::*; + #[derive(Clone, PartialEq, Message)] + struct LegacySelectPunchListenerRequest { + #[prost(bool, tag = "1")] + force_new: bool, + #[prost(bool, tag = "2")] + prefer_port_mapping: bool, + } + + #[derive(Clone, PartialEq, Message)] + struct LegacySelectPunchListenerResponse { + #[prost(message, optional, tag = "1")] + listener_mapped_addr: Option, + } + + #[test] + fn listener_registration_fields_are_compatible_with_legacy_wire_messages() { + let request = select_listener_request_to_rpc(SelectPunchListener { + force_new: true, + prefer_port_mapping: false, + transaction_id: Some(0), + }); + assert_eq!( + SelectPunchListenerRequest::decode(request.encode_to_vec().as_slice()) + .unwrap() + .transaction_id, + Some(0) + ); + let legacy_request = + LegacySelectPunchListenerRequest::decode(request.encode_to_vec().as_slice()).unwrap(); + assert!(legacy_request.force_new); + let request = + SelectPunchListenerRequest::decode(legacy_request.encode_to_vec().as_slice()).unwrap(); + assert_eq!( + select_listener_request_from_rpc(request).transaction_id, + None + ); + + let mapped_addr = "198.51.100.1:31001".parse().unwrap(); + let response = select_listener_response_to_rpc(CoreSelectPunchListenerResponse { + listener_mapped_addr: mapped_addr, + transaction_id: Some(0), + }); + let legacy_response = + LegacySelectPunchListenerResponse::decode(response.encode_to_vec().as_slice()).unwrap(); + let response = + SelectPunchListenerResponse::decode(legacy_response.encode_to_vec().as_slice()) + .unwrap(); + let response = select_listener_response_from_rpc(response).unwrap(); + assert_eq!(response.listener_mapped_addr, mapped_addr); + assert_eq!(response.transaction_id, None); + } + #[test] fn outbound_rpc_controllers_preserve_timeouts() { assert_eq!(cone_controller().timeout_ms, 4000); @@ -568,14 +626,17 @@ mod tests { let domain_request = SelectPunchListener { force_new: true, prefer_port_mapping: false, + transaction_id: Some(0), }; let rpc_request = select_listener_request_to_rpc(domain_request.clone()); assert!(rpc_request.force_new); assert!(!rpc_request.prefer_port_mapping); + assert_eq!(rpc_request.transaction_id, Some(0)); assert_eq!( select_listener_request_from_rpc(SelectPunchListenerRequest { force_new: true, prefer_port_mapping: false, + transaction_id: Some(0), }), domain_request ); @@ -583,6 +644,7 @@ mod tests { let mapped_addr: SocketAddr = "198.51.100.1:31001".parse().unwrap(); let core_response = CoreSelectPunchListenerResponse { listener_mapped_addr: mapped_addr, + transaction_id: Some(0), }; let rpc_response = select_listener_response_to_rpc(core_response.clone()); assert_eq!( diff --git a/easytier-core/src/connectivity/hole_punch/udp/runtime.rs b/easytier-core/src/connectivity/hole_punch/udp/runtime.rs index 5be63eb3..e1fd3480 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/runtime.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/runtime.rs @@ -17,7 +17,10 @@ use crate::{ peers::PeerConnectionOrigin, socket::{ ListenerConnectionCounter, SocketContext, - udp::{UdpBindOptions, UdpSession, VirtualUdpSocket, VirtualUdpSocketFactory}, + udp::{ + UdpBindOptions, UdpHolePunchObserverRegistry, UdpSession, VirtualUdpSocket, + VirtualUdpSocketFactory, + }, }, tunnel::Tunnel, }; @@ -73,6 +76,7 @@ fn udp_url(addr: SocketAddr) -> url::Url { pub struct UdpPunchListener { pub socket: Arc, + pub hole_punch_observers: UdpHolePunchObserverRegistry, pub mapped_addr: SocketAddr, pub conn_counter: Arc, pub acceptor: Box, @@ -88,11 +92,13 @@ pub struct UdpResolvedPublicAddr { pub struct SelectPunchListener { pub force_new: bool, pub prefer_port_mapping: bool, + pub transaction_id: Option, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct SelectPunchListenerResponse { pub listener_mapped_addr: SocketAddr, + pub transaction_id: Option, } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/easytier-core/src/connectivity/hole_punch/udp/server.rs b/easytier-core/src/connectivity/hole_punch/udp/server.rs index 082266ee..4946cf53 100644 --- a/easytier-core/src/connectivity/hole_punch/udp/server.rs +++ b/easytier-core/src/connectivity/hole_punch/udp/server.rs @@ -35,6 +35,9 @@ use super::{ }; const MAX_K1_FOR_RANDOM_HARD_SYM: u32 = 180; +// Selection precedes local socket/STUN preparation. Match the listener's +// selection lifetime so abandoned preparations expire without per-RPC tasks. +const PREPARED_CONE_PUNCH_TIMEOUT: Duration = Duration::from_secs(30); pub struct SelectedUdpPunchListener { /// Kept for listener-selection side effects and tests; production callers @@ -266,8 +269,16 @@ where UdpHolePunchSignalError::RemoteRejected("no listener available".into()) })?; + if let Some(transaction_id) = request.transaction_id { + self.common + .prepare_hole_punch(&selected.mapped_addr, transaction_id) + .await + .map_err(Self::anyhow_to_signal_error)?; + } + Ok(SelectPunchListenerResponse { listener_mapped_addr: selected.mapped_addr, + transaction_id: request.transaction_id, }) } @@ -276,17 +287,12 @@ where request: SendPunchPacketCone, ) -> Result<(), UdpHolePunchSignalError> { let _admission = self.admit().await?; - let listener = self + let (listener, observer) = self .common - .find_listener(&request.listener_mapped_addr) + .observe_hole_punch(&request.listener_mapped_addr, request.transaction_id) .await - .ok_or_else(|| { - UdpHolePunchSignalError::RemoteRejected( - "send_punch_packet_for_cone failed to find listener".into(), - ) - })?; - - send_cone_hole_punch_packets(listener, &request) + .map_err(Self::anyhow_to_signal_error)?; + send_cone_hole_punch_packets(listener, &request, observer) .await .map_err(Self::anyhow_to_signal_error) } @@ -386,6 +392,7 @@ where let mut index = 0; while index < listeners.len() { let listener = &listeners[index]; + listener.hole_punch_observers.expire_prepared(); let active = listener.last_active_time.load().elapsed().as_secs() < 40 || listener.last_select_time.load().elapsed().as_secs() < 30; if active { @@ -473,6 +480,36 @@ where Some(listener.get_socket()) } + async fn prepare_hole_punch( + &self, + addr: &SocketAddr, + transaction_id: u32, + ) -> anyhow::Result<()> { + let listeners = self.listeners.lock().await; + let listener = listeners + .iter() + .find(|listener| listener.mapped_addr == *addr && listener.running.load()) + .ok_or_else(|| anyhow::anyhow!("failed to find hole-punch listener"))?; + listener + .hole_punch_observers + .prepare(transaction_id, PREPARED_CONE_PUNCH_TIMEOUT)?; + Ok(()) + } + + async fn observe_hole_punch( + &self, + addr: &SocketAddr, + transaction_id: u32, + ) -> anyhow::Result<(Arc, crate::socket::udp::UdpHolePunchObserver)> { + let listeners = self.listeners.lock().await; + let listener = listeners + .iter() + .find(|listener| listener.mapped_addr == *addr && listener.running.load()) + .ok_or_else(|| anyhow::anyhow!("send_punch_packet_for_cone failed to find listener"))?; + let observer = listener.hole_punch_observers.observe(transaction_id)?; + Ok((listener.get_socket(), observer)) + } + pub async fn select_listener( &self, force_new_listener: bool, @@ -572,6 +609,7 @@ where struct UdpPunchListenerRecord { socket: Arc, + hole_punch_observers: crate::socket::udp::UdpHolePunchObserverRegistry, tasks: Mutex>, connection_tasks: Arc>>, running: Arc>, @@ -595,6 +633,7 @@ where { let UdpPunchListener { socket, + hole_punch_observers, mapped_addr, conn_counter, mut acceptor, @@ -646,6 +685,7 @@ where Self { socket, + hole_punch_observers, tasks: Mutex::new(tasks), connection_tasks, running, @@ -726,6 +766,7 @@ where pub async fn send_cone_hole_punch_packets( udp: Arc, request: &SendPunchPacketCone, + mut observer: crate::socket::udp::UdpHolePunchObserver, ) -> anyhow::Result<()> where S: VirtualUdpSocket + 'static, @@ -738,21 +779,47 @@ where ); } - for _ in 0..request.packet_batch_count { - tracing::info!(?request, "sending hole punching packet"); - + let udp = udp.as_ref(); + let send_batch = |dest_addr| async move { for _ in 0..request.packet_count_per_batch { let udp_packet = new_hole_punch_packet(request.transaction_id, HOLE_PUNCH_PACKET_BODY_LEN); - if let Err(err) = udp - .send_to(&udp_packet.into_bytes(), request.dest_addr) - .await - { - tracing::error!(?err, "failed to send hole punch packet to dest addr"); + if let Err(err) = udp.send_to(&udp_packet.into_bytes(), dest_addr).await { + tracing::error!( + ?err, + ?dest_addr, + "failed to send hole punch packet to dest addr" + ); } } - crate::foundation::time::sleep(Duration::from_millis(request.packet_interval_ms as u64)) - .await; + }; + let mut observed_addr = observer.source_addr(); + for _ in 0..request.packet_batch_count { + tracing::info!(?request, "sending hole punching packet"); + + send_batch(request.dest_addr).await; + if let Some(addr) = observed_addr.filter(|addr| *addr != request.dest_addr) { + send_batch(addr).await; + } + let interval = crate::foundation::time::sleep(Duration::from_millis( + request.packet_interval_ms as u64, + )); + tokio::pin!(interval); + if observed_addr.is_none() { + tokio::select! { + _ = &mut interval => {} + addr = observer.wait_for_source_addr() => { + let addr = addr?; + observed_addr = Some(addr); + if addr != request.dest_addr { + send_batch(addr).await; + } + interval.await; + } + } + } else { + interval.await; + } } Ok(()) @@ -971,6 +1038,8 @@ mod tests { local_addr: SocketAddr, sent: tokio::sync::Mutex, SocketAddr)>>, fail_next_send: AtomicUsize, + incoming: StdMutex, SocketAddr)>>, + incoming_notify: tokio::sync::Notify, } #[async_trait] @@ -988,8 +1057,15 @@ mod tests { Ok(data.len()) } - async fn recv_from(&self, _buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { - std::future::pending().await + async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> { + loop { + if let Some((data, addr)) = self.incoming.lock().unwrap().pop_front() { + let len = data.len().min(buf.len()); + buf[..len].copy_from_slice(&data[..len]); + return Ok((len, addr)); + } + self.incoming_notify.notified().await; + } } } @@ -1006,6 +1082,7 @@ mod tests { struct MockAcceptor { sockets: VecDeque, + _layer: Arc>, } #[async_trait] @@ -1039,6 +1116,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 0)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), })) } @@ -1130,16 +1209,22 @@ mod tests { } fn listener(port: u16, sockets: Vec) -> UdpPunchListener { + let socket = Arc::new(MockSocket { + local_addr: SocketAddr::from(([127, 0, 0, 1], port)), + sent: tokio::sync::Mutex::new(Vec::new()), + fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), + }); + let layer = Arc::new(crate::socket::udp::UdpSessionLayer::new(socket.clone())); UdpPunchListener { - socket: Arc::new(MockSocket { - local_addr: SocketAddr::from(([127, 0, 0, 1], port)), - sent: tokio::sync::Mutex::new(Vec::new()), - fail_next_send: AtomicUsize::new(0), - }), + socket, + hole_punch_observers: layer.hole_punch_observers(), mapped_addr: SocketAddr::from(([203, 0, 113, 1], port)), conn_counter: Arc::new(MockCounter::default()), acceptor: Box::new(MockAcceptor { sockets: sockets.into(), + _layer: layer, }), port_mapping_lease: None, } @@ -1206,6 +1291,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 10001)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let remote_addr = SocketAddr::from(([127, 0, 0, 1], 20001)); let session = @@ -1231,12 +1318,248 @@ mod tests { assert_eq!(sink.server_tunnels.load(Ordering::Relaxed), 1); } + #[tokio::test] + async fn selected_listener_observes_first_packet_before_cone_rpc_arrives() { + let runtime = Arc::new(MockRuntime::new(vec![listener(10002, Vec::new())])); + let server = UdpHolePunchServer::new( + runtime, + mock_stun(), + Arc::new(MockSink::default()), + UdpSymPunchLock::default(), + ); + server.start().await; + let selected = server + .select_punch_listener(SelectPunchListener { + force_new: false, + prefer_port_mapping: false, + transaction_id: Some(0), + }) + .await + .unwrap(); + assert_eq!(selected.transaction_id, Some(0)); + let socket = server + .common + .find_listener(&selected.listener_mapped_addr) + .await + .unwrap(); + let source = SocketAddr::from(([198, 51, 100, 1], 33140)); + socket.incoming.lock().unwrap().push_back(( + new_hole_punch_packet(0, HOLE_PUNCH_PACKET_BODY_LEN) + .into_bytes() + .to_vec(), + source, + )); + socket.incoming_notify.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + while !socket.incoming.lock().unwrap().is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let dest_addr = SocketAddr::from(([198, 51, 100, 1], 33105)); + server + .send_punch_packet_cone(SendPunchPacketCone { + listener_mapped_addr: selected.listener_mapped_addr, + dest_addr, + transaction_id: 0, + packet_count_per_batch: 2, + packet_batch_count: 1, + packet_interval_ms: 0, + }) + .await + .unwrap(); + let sent = socket.sent.lock().await; + assert_eq!(sent.len(), 4); + assert_eq!( + sent.iter().filter(|(_, addr)| *addr == dest_addr).count(), + 2 + ); + assert_eq!(sent.iter().filter(|(_, addr)| *addr == source).count(), 2); + drop(sent); + server.stop().await; + } + + #[tokio::test(start_paused = true)] + async fn cleanup_expires_abandoned_listener_registration() { + let runtime = Arc::new(MockRuntime::new(vec![listener(10002, Vec::new())])); + let server = UdpHolePunchServer::new( + runtime, + mock_stun(), + Arc::new(MockSink::default()), + UdpSymPunchLock::default(), + ); + server.start().await; + let selected = server + .select_punch_listener(SelectPunchListener { + force_new: false, + prefer_port_mapping: false, + transaction_id: Some(9), + }) + .await + .unwrap(); + assert!( + server + .common + .prepare_hole_punch(&selected.listener_mapped_addr, 9) + .await + .is_err() + ); + // Let the existing periodic cleanup run as the paused clock advances. + crate::foundation::time::sleep(PREPARED_CONE_PUNCH_TIMEOUT + Duration::from_secs(5)).await; + server + .common + .prepare_hole_punch(&selected.listener_mapped_addr, 9) + .await + .unwrap(); + server.stop().await; + } + + #[tokio::test] + async fn cone_sender_adds_observed_source_during_last_batch_and_deduplicates() { + let runtime = Arc::new(MockRuntime::new(vec![listener(10002, Vec::new())])); + let server = UdpHolePunchServer::new( + runtime, + mock_stun(), + Arc::new(MockSink::default()), + UdpSymPunchLock::default(), + ); + server.start().await; + let selected = server + .select_punch_listener(SelectPunchListener { + force_new: false, + prefer_port_mapping: false, + transaction_id: None, + }) + .await + .unwrap(); + let socket = server + .common + .find_listener(&selected.listener_mapped_addr) + .await + .unwrap(); + let request = SendPunchPacketCone { + listener_mapped_addr: selected.listener_mapped_addr, + dest_addr: SocketAddr::from(([198, 51, 100, 1], 33105)), + transaction_id: 9, + packet_count_per_batch: 2, + packet_batch_count: 1, + packet_interval_ms: 100, + }; + for observed_addr in [ + SocketAddr::from(([198, 51, 100, 1], 33140)), + request.dest_addr, + ] { + socket.sent.lock().await.clear(); + let send = server.send_punch_packet_cone(request.clone()); + let observe = async { + loop { + if socket.sent.lock().await.len() == 2 { + break; + } + tokio::task::yield_now().await; + } + socket.incoming.lock().unwrap().push_back(( + new_hole_punch_packet(request.transaction_id, HOLE_PUNCH_PACKET_BODY_LEN) + .into_bytes() + .to_vec(), + observed_addr, + )); + socket.incoming_notify.notify_one(); + }; + let (result, _) = tokio::time::timeout(Duration::from_secs(2), async { + tokio::join!(send, observe) + }) + .await + .unwrap(); + result.unwrap(); + let sent = socket.sent.lock().await; + let expected_count = if observed_addr == request.dest_addr { + 2 + } else { + 4 + }; + assert_eq!(sent.len(), expected_count); + assert_eq!( + sent.iter() + .filter(|(_, addr)| *addr == request.dest_addr) + .count(), + 2 + ); + for (packet, _) in sent.iter() { + assert_eq!( + crate::packet::hole_punch_packet_tid(packet, HOLE_PUNCH_PACKET_BODY_LEN), + Some(request.transaction_id) + ); + } + if observed_addr != request.dest_addr { + assert_eq!( + sent.iter() + .filter(|(_, addr)| *addr == observed_addr) + .count(), + 2 + ); + } + } + let (_, observer) = server + .common + .observe_hole_punch(&request.listener_mapped_addr, request.transaction_id) + .await + .unwrap(); + assert_eq!(observer.source_addr(), None); + drop(observer); + server.stop().await; + } + + #[tokio::test] + async fn cancelling_cone_sender_unregisters_observed_source() { + let socket = Arc::new(MockSocket { + local_addr: SocketAddr::from(([127, 0, 0, 1], 10002)), + sent: tokio::sync::Mutex::new(Vec::new()), + fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), + }); + let layer = crate::socket::udp::UdpSessionLayer::new(socket.clone()); + let request = SendPunchPacketCone { + listener_mapped_addr: SocketAddr::from(([203, 0, 113, 1], 10002)), + dest_addr: SocketAddr::from(([198, 51, 100, 1], 33105)), + transaction_id: 9, + packet_count_per_batch: 2, + packet_batch_count: 5, + packet_interval_ms: 400, + }; + let observer = layer + .hole_punch_observers() + .observe(request.transaction_id) + .unwrap(); + let send_socket = socket.clone(); + let task = tokio::spawn(async move { + send_cone_hole_punch_packets(send_socket, &request, observer).await + }); + while socket.sent.lock().await.is_empty() { + tokio::task::yield_now().await; + } + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert_eq!( + layer + .hole_punch_observers() + .observe(9) + .unwrap() + .source_addr(), + None + ); + } + #[tokio::test] async fn cone_packet_sender_keeps_old_batch_shape() { let socket = Arc::new(MockSocket { local_addr: SocketAddr::from(([127, 0, 0, 1], 10002)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let request = SendPunchPacketCone { listener_mapped_addr: SocketAddr::from(([203, 0, 113, 1], 10002)), @@ -1247,7 +1570,12 @@ mod tests { packet_interval_ms: 0, }; - send_cone_hole_punch_packets(socket.clone(), &request) + let layer = crate::socket::udp::UdpSessionLayer::new(socket.clone()); + let observer = layer + .hole_punch_observers() + .observe(request.transaction_id) + .unwrap(); + send_cone_hole_punch_packets(socket.clone(), &request, observer) .await .unwrap(); @@ -1262,6 +1590,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 10004)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let request = SendPunchPacketCone { listener_mapped_addr: SocketAddr::from(([203, 0, 113, 1], 10004)), @@ -1272,7 +1602,12 @@ mod tests { packet_interval_ms: 0, }; - let err = send_cone_hole_punch_packets(socket.clone(), &request) + let layer = crate::socket::udp::UdpSessionLayer::new(socket.clone()); + let observer = layer + .hole_punch_observers() + .observe(request.transaction_id) + .unwrap(); + let err = send_cone_hole_punch_packets(socket.clone(), &request, observer) .await .unwrap_err(); @@ -1286,6 +1621,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 10005)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(1), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let request = SendPunchPacketCone { listener_mapped_addr: SocketAddr::from(([203, 0, 113, 1], 10005)), @@ -1296,7 +1633,12 @@ mod tests { packet_interval_ms: 0, }; - send_cone_hole_punch_packets(socket.clone(), &request) + let layer = crate::socket::udp::UdpSessionLayer::new(socket.clone()); + let observer = layer + .hole_punch_observers() + .observe(request.transaction_id) + .unwrap(); + send_cone_hole_punch_packets(socket.clone(), &request, observer) .await .unwrap(); @@ -1310,6 +1652,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 10003)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(0), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let next_idx = send_symmetric_hole_punch_packet( &[10, 11, 12], @@ -1335,6 +1679,8 @@ mod tests { local_addr: SocketAddr::from(([127, 0, 0, 1], 10006)), sent: tokio::sync::Mutex::new(Vec::new()), fail_next_send: AtomicUsize::new(1), + incoming: StdMutex::new(VecDeque::new()), + incoming_notify: tokio::sync::Notify::new(), }); let err = send_symmetric_hole_punch_packet( diff --git a/easytier-core/src/socket/udp/layer.rs b/easytier-core/src/socket/udp/layer.rs index ae2ab934..f9dc4be2 100644 --- a/easytier-core/src/socket/udp/layer.rs +++ b/easytier-core/src/socket/udp/layer.rs @@ -25,12 +25,13 @@ use super::{ session::{ ClassifiedUdpSessionAccept, ClassifiedUdpSessionAccepts, ClassifiedUdpSessionKey, ClassifiedUdpSessionRegistry, PendingUdpSessionConnect, PendingUdpSessionConnects, - UdpConnectControl, UdpSession, UdpSessionClose, UdpSessionCodec, UdpSessionConnectError, - UdpSessionConnectRequest, UdpSessionConnector, UdpSessionDatagram, UdpSessionEnqueuePolicy, - UdpSessionKey, UdpSessionKind, UdpSessionLayerControl, UdpSessionProtocol, - UdpSessionRegistry, close_all_classified_udp_sessions, close_all_udp_sessions, - close_classified_udp_session, close_udp_session, create_udp_session_rings, - dispatch_payload_to_session, udp_session_registry_entry, + UdpConnectControl, UdpHolePunchObserverMap, UdpHolePunchObserverRegistry, UdpSession, + UdpSessionClose, UdpSessionCodec, UdpSessionConnectError, UdpSessionConnectRequest, + UdpSessionConnector, UdpSessionDatagram, UdpSessionEnqueuePolicy, UdpSessionKey, + UdpSessionKind, UdpSessionLayerControl, UdpSessionProtocol, UdpSessionRegistry, + close_all_classified_udp_sessions, close_all_udp_sessions, close_classified_udp_session, + close_udp_session, create_udp_session_rings, dispatch_payload_to_session, + udp_session_registry_entry, }, virtual_socket::{ MAX_UDP_SESSION_DATAGRAM_SIZE, NoopUdpSessionStunResponder, PreferredIpv6Source, @@ -49,6 +50,7 @@ pub struct UdpSessionLayer { classified_sessions: Arc, pub(super) classified_accepts: Arc, pub(super) pending_connects: Arc, + hole_punch_observers: UdpHolePunchObserverRegistry, mux_accepted_rx: TokioMutex>, _control_rx: TokioMutex>, session_shutdown_tx: watch::Sender, @@ -90,15 +92,17 @@ where let classified_sessions = Arc::new(DashMap::new()); let classified_accepts = create_classified_udp_session_accepts(); let pending_connects = Arc::new(DashMap::new()); + let hole_punch_observers = Arc::new(DashMap::new()); let (mux_accepted_tx, mux_accepted_rx) = mpsc::channel(UDP_SESSION_QUEUE_CAPACITY); let (control_tx, control_rx) = mpsc::channel(UDP_SESSION_QUEUE_CAPACITY); - let (session_shutdown_tx, _) = watch::channel(false); + let (session_shutdown_tx, session_shutdown_rx) = watch::channel(false); let recv_task = tokio::spawn(udp_session_layer_recv_task( socket.clone(), sessions.clone(), classified_sessions.clone(), classified_accepts.clone(), pending_connects.clone(), + hole_punch_observers.clone(), mux_accepted_tx, control_tx, stun_responder.clone(), @@ -112,6 +116,10 @@ where classified_sessions, classified_accepts, pending_connects, + hole_punch_observers: UdpHolePunchObserverRegistry { + observers: hole_punch_observers, + shutdown: session_shutdown_rx, + }, mux_accepted_rx: TokioMutex::new(mux_accepted_rx), _control_rx: TokioMutex::new(control_rx), session_shutdown_tx, @@ -123,6 +131,10 @@ where self.socket.local_addr() } + pub fn hole_punch_observers(&self) -> UdpHolePunchObserverRegistry { + self.hole_punch_observers.clone() + } + pub fn active_session_count(&self) -> usize { self.sessions.len() } @@ -355,6 +367,7 @@ impl Drop for UdpSessionLayer { fn drop(&mut self) { let _ = self.session_shutdown_tx.send(true); self.pending_connects.clear(); + self.hole_punch_observers.observers.clear(); close_all_udp_sessions(&self.sessions); close_all_classified_udp_sessions(&self.classified_sessions); self.recv_task.abort(); @@ -437,6 +450,7 @@ pub(super) async fn udp_session_layer_recv_task( classified_sessions: Arc, classified_accepts: Arc, pending_connects: Arc, + hole_punch_observers: Arc, mux_accepted: mpsc::Sender, control: mpsc::Sender, stun_responder: Arc, @@ -453,6 +467,7 @@ pub(super) async fn udp_session_layer_recv_task( tracing::debug!(?err, "udp session recv loop stopped"); let _ = session_shutdown_tx.send(true); pending_connects.clear(); + hole_punch_observers.clear(); close_all_udp_sessions(&sessions); close_all_classified_udp_sessions(&classified_sessions); break; @@ -519,6 +534,7 @@ pub(super) async fn udp_session_layer_recv_task( &socket, &sessions, &pending_connects, + &hole_punch_observers, &mux_accepted, &control, &control_permits, @@ -570,6 +586,7 @@ fn dispatch_easy_tier_udp_datagram( socket: &Arc, sessions: &Arc, pending_connects: &Arc, + hole_punch_observers: &Arc, mux_accepted: &mpsc::Sender, control: &mpsc::Sender, control_permits: &Arc, @@ -600,7 +617,20 @@ where dispatch_sack_packet(sessions, pending_connects, remote_addr, conn_id, &packet) } EasyTierUdpPacketKind::HolePunch => { - dispatch_hole_punch_packet(pending_connects, remote_addr) + let mut observed = false; + if packet.udp_payload().len() == crate::packet::HOLE_PUNCH_PACKET_BODY_LEN as usize + && let Some(observer) = hole_punch_observers.get(&conn_id) + { + observer.source.send_if_modified(|source| { + if source.is_some() { + return false; + } + *source = Some(remote_addr); + true + }); + observed = true; + } + dispatch_hole_punch_packet(pending_connects, remote_addr) || observed } EasyTierUdpPacketKind::V4HolePunch => dispatch_v4_hole_punch_control( socket.clone(), diff --git a/easytier-core/src/socket/udp/mod.rs b/easytier-core/src/socket/udp/mod.rs index dd708c98..88239f29 100644 --- a/easytier-core/src/socket/udp/mod.rs +++ b/easytier-core/src/socket/udp/mod.rs @@ -20,9 +20,10 @@ pub use packet::{ parse_udp_session_datagram, }; pub use session::{ - UdpSession, UdpSessionConnectError, UdpSessionConnectRequest, UdpSessionConnector, - UdpSessionKind, UdpSessionLayerControl, UdpSessionListenRequest, UdpSessionListener, - UdpSessionProtocol, UdpSessionRecvMeta, UdpSessionSocket, + UdpHolePunchObserver, UdpHolePunchObserverRegistry, UdpSession, UdpSessionConnectError, + UdpSessionConnectRequest, UdpSessionConnector, UdpSessionKind, UdpSessionLayerControl, + UdpSessionListenRequest, UdpSessionListener, UdpSessionProtocol, UdpSessionRecvMeta, + UdpSessionSocket, }; pub(crate) use session::{ UdpSessionCleanup, UdpSessionCodec, UdpSessionDatagram, UdpSessionOutbound, diff --git a/easytier-core/src/socket/udp/session.rs b/easytier-core/src/socket/udp/session.rs index 4bc3f523..4b170292 100644 --- a/easytier-core/src/socket/udp/session.rs +++ b/easytier-core/src/socket/udp/session.rs @@ -1,7 +1,7 @@ use std::{ io, net::{IpAddr, SocketAddr, SocketAddrV4, SocketAddrV6}, - sync::{Arc, Mutex as StdMutex}, + sync::{Arc, Mutex as StdMutex, Weak}, }; use async_trait::async_trait; @@ -15,6 +15,7 @@ use tokio::{ }; use crate::{ + foundation::time::{Duration, Instant}, packet::{UdpPacketType, ZCPacket, ZCPacketType}, socket::ring::{RingSocket, RingSocketReceiver, RingSocketSendError, RingSocketSender}, }; @@ -201,6 +202,135 @@ pub(super) type ClassifiedUdpSessionRegistry = pub(super) type ClassifiedUdpSessionAccepts = DashMap>; pub(super) type PendingUdpSessionConnects = DashMap; +pub(super) type UdpHolePunchObserverMap = DashMap; + +#[derive(Debug)] +pub(super) struct UdpHolePunchObserverEntry { + pub(super) source: Arc>>, + // Only unclaimed listener-selection registrations expire here. Once an RPC + // takes the observer, its guard owns cleanup through completion/cancellation. + expires_at: Option, +} + +/// Registers observers on an existing UDP session layer without another receive task. +#[derive(Debug, Clone)] +pub struct UdpHolePunchObserverRegistry { + pub(super) observers: Arc, + pub(super) shutdown: watch::Receiver, +} + +impl UdpHolePunchObserverRegistry { + fn ensure_running(&self) -> io::Result<()> { + if *self.shutdown.borrow() { + return Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "UDP listener closed", + )); + } + Ok(()) + } + + /// Keep the first source even if the UDP packet arrives before the sending RPC. + pub fn prepare(&self, transaction_id: u32, timeout: Duration) -> io::Result<()> { + self.ensure_running()?; + match self.observers.entry(transaction_id) { + dashmap::mapref::entry::Entry::Vacant(entry) => { + let (sender, _) = watch::channel(None); + entry.insert(UdpHolePunchObserverEntry { + source: Arc::new(sender), + expires_at: Some(Instant::now() + timeout), + }); + Ok(()) + } + dashmap::mapref::entry::Entry::Occupied(_) => Err(io::Error::new( + io::ErrorKind::AlreadyExists, + "UDP hole-punch transaction already observed", + )), + } + } + + pub fn expire_prepared(&self) { + let now = Instant::now(); + self.observers + .retain(|_, entry| entry.expires_at.is_none_or(|deadline| deadline > now)); + } + + pub fn observe(&self, transaction_id: u32) -> io::Result { + self.ensure_running()?; + let sender = match self.observers.entry(transaction_id) { + dashmap::mapref::entry::Entry::Occupied(mut entry) => { + let Some(deadline) = entry.get_mut().expires_at.take() else { + return Err(io::Error::new( + io::ErrorKind::AlreadyExists, + "UDP hole-punch transaction already observed", + )); + }; + if deadline <= Instant::now() { + entry.remove(); + return Err(io::Error::new( + io::ErrorKind::TimedOut, + "UDP hole-punch registration expired", + )); + } + entry.get().source.clone() + } + dashmap::mapref::entry::Entry::Vacant(entry) => { + // Older clients select a listener without preparing a transaction. + let (sender, _) = watch::channel(None); + let sender = Arc::new(sender); + entry.insert(UdpHolePunchObserverEntry { + source: sender.clone(), + expires_at: None, + }); + sender + } + }; + Ok(UdpHolePunchObserver { + transaction_id, + observers: self.observers.clone(), + sender: Arc::downgrade(&sender), + source: sender.subscribe(), + }) + } +} + +/// Observes the first valid source for one active hole-punch transaction. +/// Dropping the observer unregisters it from the socket's receive loop. +#[derive(Debug)] +pub struct UdpHolePunchObserver { + pub(super) transaction_id: u32, + pub(super) observers: Arc, + pub(super) sender: Weak>>, + pub(super) source: watch::Receiver>, +} + +impl UdpHolePunchObserver { + pub fn source_addr(&self) -> Option { + *self.source.borrow() + } + + pub async fn wait_for_source_addr(&mut self) -> io::Result { + loop { + if let Some(addr) = self.source_addr() { + return Ok(addr); + } + self.source.changed().await.map_err(|_| { + io::Error::new( + io::ErrorKind::UnexpectedEof, + "UDP hole-punch observer closed", + ) + })?; + } + } +} + +impl Drop for UdpHolePunchObserver { + fn drop(&mut self) { + self.observers.remove_if(&self.transaction_id, |_, entry| { + self.sender.ptr_eq(&Arc::downgrade(&entry.source)) + }); + } +} #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub(super) struct ClassifiedUdpSessionKey { diff --git a/easytier-core/src/socket/udp/tests.rs b/easytier-core/src/socket/udp/tests.rs index 557854de..db3070e5 100644 --- a/easytier-core/src/socket/udp/tests.rs +++ b/easytier-core/src/socket/udp/tests.rs @@ -477,6 +477,182 @@ impl AutoSackVirtualUdpSocket { } } +#[tokio::test] +async fn hole_punch_observer_matches_socket_transaction_and_packet_shape() { + let local_addr = SocketAddr::from(([127, 0, 0, 1], 12000)); + let peer_addr = SocketAddr::from(([198, 51, 100, 1], 33140)); + let other_addr = SocketAddr::from(([198, 51, 100, 1], 33105)); + let socket = Arc::new(AutoSackVirtualUdpSocket::new(local_addr)); + let other_socket = Arc::new(AutoSackVirtualUdpSocket::new(local_addr)); + let layer = UdpSessionLayer::new(socket.clone()); + let other_layer = UdpSessionLayer::new(other_socket.clone()); + let mut observer = layer.hole_punch_observers().observe(7).unwrap(); + let other_observer = other_layer.hole_punch_observers().observe(7).unwrap(); + let mut concurrent_observer = layer.hole_punch_observers().observe(8).unwrap(); + + for (tid, body_len, source) in [ + (9, crate::packet::HOLE_PUNCH_PACKET_BODY_LEN, other_addr), + (7, 32, other_addr), + (8, crate::packet::HOLE_PUNCH_PACKET_BODY_LEN, other_addr), + (7, crate::packet::HOLE_PUNCH_PACKET_BODY_LEN, peer_addr), + (7, crate::packet::HOLE_PUNCH_PACKET_BODY_LEN, other_addr), + ] { + socket.incoming.lock().unwrap().push_back(( + crate::packet::new_hole_punch_packet(tid, body_len) + .into_bytes() + .to_vec(), + source, + )); + } + socket.incoming_notify.notify_one(); + + assert_eq!( + tokio::time::timeout(Duration::from_secs(1), observer.wait_for_source_addr()) + .await + .unwrap() + .unwrap(), + peer_addr, + ); + assert_eq!( + concurrent_observer.wait_for_source_addr().await.unwrap(), + other_addr + ); + assert_eq!(observer.source_addr(), Some(peer_addr)); + assert_eq!(other_observer.source_addr(), None); + assert!(socket.sent().is_empty()); + + assert_eq!( + layer.hole_punch_observers().observe(7).unwrap_err().kind(), + io::ErrorKind::AlreadyExists + ); + drop(observer); + assert_eq!( + layer + .hole_punch_observers() + .observe(7) + .unwrap() + .source_addr(), + None + ); +} + +#[tokio::test] +async fn prepared_hole_punch_keeps_source_until_rpc_claims_it() { + let socket = Arc::new(AutoSackVirtualUdpSocket::new(SocketAddr::from(( + [127, 0, 0, 1], + 12000, + )))); + let layer = UdpSessionLayer::new(socket.clone()); + let registry = layer.hole_punch_observers(); + registry.prepare(0, Duration::from_secs(30)).unwrap(); + let peer_addr = SocketAddr::from(([198, 51, 100, 1], 33140)); + socket.incoming.lock().unwrap().push_back(( + crate::packet::new_hole_punch_packet(0, crate::packet::HOLE_PUNCH_PACKET_BODY_LEN) + .into_bytes() + .to_vec(), + peer_addr, + )); + socket.incoming_notify.notify_one(); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if *registry.observers.get(&0).unwrap().source.borrow() == Some(peer_addr) { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + let observer = registry.observe(0).unwrap(); + assert_eq!(observer.source_addr(), Some(peer_addr)); + assert_eq!( + registry.observe(0).unwrap_err().kind(), + io::ErrorKind::AlreadyExists + ); + drop(observer); + assert!(registry.observers.is_empty()); +} + +#[tokio::test(start_paused = true)] +async fn prepared_hole_punch_expires_without_removing_active_rpc() { + let socket = Arc::new(AutoSackVirtualUdpSocket::new(SocketAddr::from(( + [127, 0, 0, 1], + 12000, + )))); + let layer = UdpSessionLayer::new(socket); + let registry = layer.hole_punch_observers(); + registry.prepare(7, Duration::from_secs(1)).unwrap(); + registry.prepare(8, Duration::from_secs(1)).unwrap(); + registry.prepare(9, Duration::from_secs(1)).unwrap(); + let observer = registry.observe(8).unwrap(); + + tokio::time::advance(Duration::from_secs(1)).await; + assert_eq!( + registry.observe(7).unwrap_err().kind(), + io::ErrorKind::TimedOut + ); + registry.expire_prepared(); + assert!(!registry.observers.contains_key(&9)); + assert_eq!( + registry.observe(8).unwrap_err().kind(), + io::ErrorKind::AlreadyExists + ); + registry.prepare(9, Duration::from_secs(1)).unwrap(); + drop(observer); + assert!(!registry.observers.contains_key(&8)); +} + +#[tokio::test] +async fn dropping_hole_punch_observer_preserves_udp_handshake_receive() { + let socket = Arc::new(AutoSackVirtualUdpSocket::new(SocketAddr::from(( + [127, 0, 0, 1], + 12000, + )))); + let layer = UdpSessionLayer::new(socket); + let observer = layer.hole_punch_observers().observe(7).unwrap(); + drop(observer); + let peer_addr = SocketAddr::from(([198, 51, 100, 1], 33140)); + let session = layer.connect(peer_addr).await.unwrap(); + assert_eq!(session.peer_addr().unwrap(), peer_addr); +} + +#[tokio::test] +async fn hole_punch_observer_registry_does_not_keep_listener_alive() { + let socket = Arc::new(AutoSackVirtualUdpSocket::new(SocketAddr::from(( + [127, 0, 0, 1], + 12000, + )))); + let layer = UdpSessionLayer::new(socket.clone()); + let registry = layer.hole_punch_observers(); + let mut observer = registry.observe(7).unwrap(); + registry.prepare(8, Duration::from_secs(30)).unwrap(); + drop(layer); + assert_eq!( + registry.observe(8).unwrap_err().kind(), + io::ErrorKind::BrokenPipe + ); + assert_eq!( + registry + .prepare(9, Duration::from_secs(30)) + .unwrap_err() + .kind(), + io::ErrorKind::BrokenPipe + ); + assert_eq!( + observer.wait_for_source_addr().await.unwrap_err().kind(), + io::ErrorKind::UnexpectedEof + ); + drop(observer); + for _ in 0..10 { + if Arc::strong_count(&socket) == 1 { + return; + } + tokio::task::yield_now().await; + } + assert_eq!(Arc::strong_count(&socket), 1); +} + #[async_trait] impl VirtualUdpSocket for AutoSackVirtualUdpSocket { fn local_addr(&self) -> io::Result { @@ -1301,6 +1477,7 @@ async fn udp_session_recv_loop_error_closes_registered_sessions() { classified_sessions.clone(), classified_accepts, pending_connects.clone(), + Arc::new(DashMap::new()), mux_accepted_tx, control_tx, Arc::new(NoopUdpSessionStunResponder), @@ -1624,6 +1801,7 @@ async fn udp_session_recv_loop_does_not_wait_for_stun_responder() { classified_sessions, classified_accepts, pending_connects, + Arc::new(DashMap::new()), mux_accepted_tx, control_tx, stun_responder.clone(), diff --git a/easytier-proto/proto/peer_rpc.proto b/easytier-proto/proto/peer_rpc.proto index e3e4f75a..59eb086c 100644 --- a/easytier-proto/proto/peer_rpc.proto +++ b/easytier-proto/proto/peer_rpc.proto @@ -204,10 +204,14 @@ service DirectConnectorRpc { message SelectPunchListenerRequest { bool force_new = 1; bool prefer_port_mapping = 2; + // Register source observation before returning the listener address. + optional uint32 transaction_id = 3; } message SelectPunchListenerResponse { common.SocketAddr listener_mapped_addr = 1; + // Echoed only after source observation is registered; absent on older peers. + optional uint32 transaction_id = 2; } message SendPunchPacketConeRequest {