fix(hole-punch): connect UDP when STUN mapping differs (#2680)

* fix(hole-punch): connect UDP without reverse probe gating

Start cone UDP handshakes on the same punched socket without waiting
for reverse HolePunch packets. Keep outgoing probes while the existing
session layer retries SYN, allowing a reachable listener to reply to the
actual source when the STUN candidate uses a different mapping.

Register source observers on the listener's existing receive loop for
the active cone RPC. Send bounded batches to the signaled endpoint and
the first matching observed source, deduplicating equal endpoints.
Unregister observers on completion, cancellation, or listener shutdown.

Cover missing reverse probes, dropped SYN, late observed sources,
transaction isolation, and cleanup. Document socket and connectivity
ownership of observation and reply policy.

* fix(hole-punch): prepare source observation during listener selection

Carry an optional transaction ID in SelectPunchListener and echo it only
after registering the listener's existing UDP receive loop. Preserve the
first source arriving before SendPunchPacketCone, then transfer cleanup
ownership to that RPC without clearing the observed address.

Keep prepared and active registrations in the same transaction table.
Expire abandoned preparations after 30 seconds through listener cleanup;
active observers still unregister on completion or cancellation.
Preserve legacy requests without a transaction ID and tolerate older
responses without a registration acknowledgement.

Cover early packets, transaction ID zero, preparation expiry, active RPC
cleanup, legacy protobuf interoperability, and SYN retries against peers
with and without registration confirmation. Document the lifetime and
compatibility contract in the core architecture.

Validation: 820 core tests; workspace fmt and full-feature clippy
passed.

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;

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,

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<Instant> = 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<R>(
-    udp_array: &UdpSocketArray<R>,
+async fn send_from_local<S>(
+    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<std::collections::VecDeque<(Vec<u8>, SocketAddr)>>,
+        incoming_notify: tokio::sync::Notify,
+        sent: std::sync::Mutex<Vec<(Vec<u8>, 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<usize> {
-            Ok(0)
+        async fn send_to(&self, data: &[u8], addr: SocketAddr) -> io::Result<usize> {
+            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<Vec<UdpBindOptions>>,
+        sockets: std::sync::Mutex<Vec<std::sync::Weak<MockSocket>>>,
+        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<Arc<Self::Socket>> {
             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<Self::Socket>,
+            socket: Arc<Self::Socket>,
         ) -> anyhow::Result<super::super::UdpResolvedPublicAddr> {
+            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<Self::Socket>,
-            _remote: SocketAddr,
+            socket: Arc<Self::Socket>,
+            remote: SocketAddr,
         ) -> anyhow::Result<UdpPunchSocket> {
-            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<Vec<SelectPunchListener>>,
+        confirm_transaction_id: bool,
         easy_requests: tokio::sync::Mutex<Vec<SendPunchPacketEasySym>>,
         hard_requests: tokio::sync::Mutex<Vec<SendPunchPacketHardSym>>,
         both_requests: tokio::sync::Mutex<Vec<SendPunchPacketBothEasySym>>,
@@ -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<super::super::SelectPunchListenerResponse, UdpHolePunchSignalError> {
+            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,
             })
         }
 

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<CoreSelectPunchListenerResponse, UdpHolePunchSignalError> {
     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<crate::proto::common::SocketAddr>,
+    }
+
+    #[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!(

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<S> {
     pub socket: Arc<S>,
+    pub hole_punch_observers: UdpHolePunchObserverRegistry,
     pub mapped_addr: SocketAddr,
     pub conn_counter: Arc<dyn ListenerConnectionCounter>,
     pub acceptor: Box<dyn UdpPunchAcceptor>,
@@ -88,11 +92,13 @@ pub struct UdpResolvedPublicAddr {
 pub struct SelectPunchListener {
     pub force_new: bool,
     pub prefer_port_mapping: bool,
+    pub transaction_id: Option<u32>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct SelectPunchListenerResponse {
     pub listener_mapped_addr: SocketAddr,
+    pub transaction_id: Option<u32>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]

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<S> {
     /// 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<R::Socket>, 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<S> {
     socket: Arc<S>,
+    hole_punch_observers: crate::socket::udp::UdpHolePunchObserverRegistry,
     tasks: Mutex<JoinSet<()>>,
     connection_tasks: Arc<Mutex<JoinSet<()>>>,
     running: Arc<AtomicCell<bool>>,
@@ -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<S>(
     udp: Arc<S>,
     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<Vec<(Vec<u8>, SocketAddr)>>,
         fail_next_send: AtomicUsize,
+        incoming: StdMutex<VecDeque<(Vec<u8>, 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<UdpPunchSocket>,
+        _layer: Arc<crate::socket::udp::UdpSessionLayer<MockSocket>>,
     }
 
     #[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<UdpPunchSocket>) -> UdpPunchListener<MockSocket> {
+        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(

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<S, R = NoopUdpSessionStunResponder> {
     classified_sessions: Arc<ClassifiedUdpSessionRegistry>,
     pub(super) classified_accepts: Arc<ClassifiedUdpSessionAccepts>,
     pub(super) pending_connects: Arc<PendingUdpSessionConnects>,
+    hole_punch_observers: UdpHolePunchObserverRegistry,
     mux_accepted_rx: TokioMutex<mpsc::Receiver<UdpSession>>,
     _control_rx: TokioMutex<mpsc::Receiver<UdpSessionLayerControl>>,
     session_shutdown_tx: watch::Sender<bool>,
@@ -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<S, R> Drop for UdpSessionLayer<S, R> {
     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<S, R>(
     classified_sessions: Arc<ClassifiedUdpSessionRegistry>,
     classified_accepts: Arc<ClassifiedUdpSessionAccepts>,
     pending_connects: Arc<PendingUdpSessionConnects>,
+    hole_punch_observers: Arc<UdpHolePunchObserverMap>,
     mux_accepted: mpsc::Sender<UdpSession>,
     control: mpsc::Sender<UdpSessionLayerControl>,
     stun_responder: Arc<R>,
@@ -453,6 +467,7 @@ pub(super) async fn udp_session_layer_recv_task<S, R>(
                 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<S, R>(
                     &socket,
                     &sessions,
                     &pending_connects,
+                    &hole_punch_observers,
                     &mux_accepted,
                     &control,
                     &control_permits,
@@ -570,6 +586,7 @@ fn dispatch_easy_tier_udp_datagram<S>(
     socket: &Arc<S>,
     sessions: &Arc<UdpSessionRegistry>,
     pending_connects: &Arc<PendingUdpSessionConnects>,
+    hole_punch_observers: &Arc<UdpHolePunchObserverMap>,
     mux_accepted: &mpsc::Sender<UdpSession>,
     control: &mpsc::Sender<UdpSessionLayerControl>,
     control_permits: &Arc<Semaphore>,
@@ -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(),

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,

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<UdpSessionProtocol, Arc<ClassifiedUdpSessionAccept>>;
 pub(super) type PendingUdpSessionConnects = DashMap<u32, PendingUdpSessionConnect>;
+pub(super) type UdpHolePunchObserverMap = DashMap<u32, UdpHolePunchObserverEntry>;
+
+#[derive(Debug)]
+pub(super) struct UdpHolePunchObserverEntry {
+    pub(super) source: Arc<watch::Sender<Option<SocketAddr>>>,
+    // Only unclaimed listener-selection registrations expire here. Once an RPC
+    // takes the observer, its guard owns cleanup through completion/cancellation.
+    expires_at: Option<Instant>,
+}
+
+/// Registers observers on an existing UDP session layer without another receive task.
+#[derive(Debug, Clone)]
+pub struct UdpHolePunchObserverRegistry {
+    pub(super) observers: Arc<UdpHolePunchObserverMap>,
+    pub(super) shutdown: watch::Receiver<bool>,
+}
+
+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<UdpHolePunchObserver> {
+        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<UdpHolePunchObserverMap>,
+    pub(super) sender: Weak<watch::Sender<Option<SocketAddr>>>,
+    pub(super) source: watch::Receiver<Option<SocketAddr>>,
+}
+
+impl UdpHolePunchObserver {
+    pub fn source_addr(&self) -> Option<SocketAddr> {
+        *self.source.borrow()
+    }
+
+    pub async fn wait_for_source_addr(&mut self) -> io::Result<SocketAddr> {
+        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 {

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<SocketAddr> {
@@ -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(),

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 {