fix(icmp): drive raw sockets with asynchronous readiness (#2756)

Register nonblocking raw ICMP sockets with Tokio instead of creating
a blocking task for every receive. Keep the raw packet layout and
network namespace binding unchanged across native platforms.

Cancel pending sends and receives when the socket closes, without
relying on shutdown of an unconnected raw socket.

Cover receive cancellation, closing multiple waiters, and repeated
loopback echo fanout across 30 raw sockets without blocking workers.

easytier/src/gateway/icmp_proxy.rs

@@ -1,5 +1,4 @@
 use std::{
-    mem::MaybeUninit,
     net::{IpAddr, Ipv4Addr, SocketAddrV4},
     sync::Arc,
 };
@@ -9,6 +8,8 @@ use easytier_core::{
     socket::SocketContext,
 };
 use socket2::Socket;
+use tokio::net::UdpSocket;
+use tokio_util::sync::CancellationToken;
 
 use crate::common::netns::NetNS;
 
@@ -27,39 +28,43 @@ impl RuntimeIcmpProxyHost {
             Ipv4Addr::UNSPECIFIED,
             0,
         )))?;
+        socket.set_nonblocking(true)?;
         Ok(socket)
     }
 }
 
 #[derive(Debug)]
 struct RuntimeIcmpSocket {
-    socket: Arc<Socket>,
+    socket: UdpSocket,
+    cancel: CancellationToken,
 }
 
 #[async_trait::async_trait]
 impl IcmpProxySocket for RuntimeIcmpSocket {
     async fn send(&self, destination: Ipv4Addr, packet: &[u8]) -> Result<(), ProxyRuntimeError> {
-        self.socket
-            .send_to(packet, &SocketAddrV4::new(destination, 0).into())?;
+        tokio::select! {
+            biased;
+            _ = self.cancel.cancelled() => return Err(std::io::Error::from(std::io::ErrorKind::NotConnected).into()),
+            result = self.socket.send_to(packet, SocketAddrV4::new(destination, 0)) => { result?; }
+        }
         Ok(())
     }
 
     async fn recv(&self) -> Result<(IpAddr, Vec<u8>), ProxyRuntimeError> {
-        let socket = self.socket.clone();
-        tokio::task::spawn_blocking(move || {
-            let mut buffer = vec![0_u8; 8192];
-            let uninitialized: &mut [MaybeUninit<u8>] =
-                unsafe { std::mem::transmute(&mut buffer[..]) };
-            let (length, peer_ip) = socket_recv(&socket, uninitialized)?;
-            buffer.truncate(length);
-            Ok((peer_ip, buffer))
-        })
-        .await
-        .map_err(|error| ProxyRuntimeError::Other(error.into()))?
+        let mut buffer = vec![0_u8; 8192];
+        let (length, peer) = tokio::select! {
+            biased;
+            // The proxy receive loop terminates on an empty packet. A permanent
+            // error would be retried immediately and could prevent task abortion.
+            _ = self.cancel.cancelled() => return Ok((IpAddr::V4(Ipv4Addr::UNSPECIFIED), Vec::new())),
+            result = self.socket.recv_from(&mut buffer) => result?,
+        };
+        buffer.truncate(length);
+        Ok((peer.ip(), buffer))
     }
 
     fn close(&self) {
-        let _ = self.socket.shutdown(std::net::Shutdown::Both);
+        self.cancel.cancel();
     }
 }
 
@@ -72,20 +77,136 @@ impl IcmpProxyHost for RuntimeIcmpProxyHost {
         let socket = Self::create_raw_socket(&context).inspect_err(|error| {
             tracing::warn!(?error, "create ICMP socket failed");
         })?;
+        // Tokio only supplies readiness and recvfrom/sendto here; the underlying
+        // socket remains raw ICMP, including the IPv4 header on received packets.
         Ok(Arc::new(RuntimeIcmpSocket {
-            socket: Arc::new(socket),
+            socket: UdpSocket::from_std(socket.into())?,
+            cancel: CancellationToken::new(),
         }))
     }
 }
 
-fn socket_recv(
-    socket: &Socket,
-    buffer: &mut [MaybeUninit<u8>],
-) -> Result<(usize, IpAddr), std::io::Error> {
-    let (size, address) = socket.recv_from(buffer)?;
-    let peer_ip = address
-        .as_socket()
-        .map(|address| address.ip())
-        .unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED));
-    Ok((size, peer_ip))
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn close_ends_receivers_and_rejects_further_sends() {
+        let socket = RuntimeIcmpSocket {
+            socket: UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap(),
+            cancel: CancellationToken::new(),
+        };
+        let mut first = Box::pin(socket.recv());
+        let mut second = Box::pin(socket.recv());
+        assert!(futures::poll!(&mut first).is_pending());
+        assert!(futures::poll!(&mut second).is_pending());
+
+        socket.close();
+        socket.close();
+        tokio::time::timeout(Duration::from_secs(1), async {
+            assert!(first.await.unwrap().1.is_empty());
+            assert!(second.await.unwrap().1.is_empty());
+            assert!(socket.recv().await.unwrap().1.is_empty());
+            assert!(socket.send(Ipv4Addr::LOCALHOST, &[0; 8]).await.is_err());
+        })
+        .await
+        .unwrap();
+    }
+
+    #[tokio::test]
+    async fn cancelled_receive_does_not_consume_the_next_packet() {
+        let socket = RuntimeIcmpSocket {
+            socket: UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap(),
+            cancel: CancellationToken::new(),
+        };
+        let mut receive = Box::pin(socket.recv());
+        assert!(futures::poll!(&mut receive).is_pending());
+        drop(receive);
+
+        let sender = UdpSocket::bind((Ipv4Addr::LOCALHOST, 0)).await.unwrap();
+        sender
+            .send_to(b"after cancellation", socket.socket.local_addr().unwrap())
+            .await
+            .unwrap();
+        let (peer, packet) = tokio::time::timeout(Duration::from_secs(1), socket.recv())
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(peer, IpAddr::V4(Ipv4Addr::LOCALHOST));
+        assert_eq!(packet, b"after cancellation");
+    }
+
+    // Raw ICMP requires CAP_NET_RAW; run this test in the privileged test container.
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn raw_echo_fanout_does_not_start_blocking_threads() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let threads = Arc::new(AtomicUsize::new(0));
+        let started = threads.clone();
+        let runtime = tokio::runtime::Builder::new_current_thread()
+            .enable_all()
+            .on_thread_start(move || {
+                started.fetch_add(1, Ordering::Relaxed);
+            })
+            .build()
+            .unwrap();
+        let result = runtime.block_on(async {
+            tokio::time::timeout(Duration::from_secs(5), async {
+                let mut sockets = Vec::new();
+                for _ in 0..30 {
+                    sockets.push(
+                        RuntimeIcmpProxyHost
+                            .open_icmp_v4(SocketContext::default())
+                            .await
+                            .unwrap(),
+                    );
+                }
+                let identifier = rand::random::<u16>().to_be_bytes();
+                for sequence in 0..16 {
+                    let mut request = vec![8, 0, 0, 0, identifier[0], identifier[1], 0, sequence];
+                    request.extend_from_slice(b"async ICMP");
+                    let mut checksum: u32 = request
+                        .chunks_exact(2)
+                        .map(|word| u32::from(u16::from_be_bytes([word[0], word[1]])))
+                        .sum();
+                    while checksum >> 16 != 0 {
+                        checksum = (checksum & 0xffff) + (checksum >> 16);
+                    }
+                    request[2..4].copy_from_slice(&(!(checksum as u16)).to_be_bytes());
+                    let receivers = futures::future::join_all(sockets.iter().map(|socket| async {
+                        loop {
+                            let (peer, packet) = socket.recv().await.unwrap();
+                            let header_len = usize::from(packet[0] & 0x0f) * 4;
+                            let Some(reply) = packet.get(header_len..) else {
+                                continue;
+                            };
+                            if reply.len() == request.len()
+                                && reply[0] == 0
+                                && reply[4..] == request[4..]
+                            {
+                                assert_eq!(peer, IpAddr::V4(Ipv4Addr::LOCALHOST));
+                                break;
+                            }
+                        }
+                    }));
+                    tokio::join!(receivers, async {
+                        sockets[0]
+                            .send(Ipv4Addr::LOCALHOST, &request)
+                            .await
+                            .unwrap();
+                    });
+                }
+                for socket in sockets {
+                    socket.close();
+                }
+            })
+            .await
+        });
+        runtime.shutdown_timeout(Duration::from_secs(1));
+        result.unwrap();
+        assert_eq!(threads.load(Ordering::Relaxed), 0);
+    }
 }