surge_ping/
client.rs

1#[cfg(unix)]
2use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
3#[cfg(windows)]
4use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket};
5
6use std::{
7    collections::HashMap,
8    io,
9    net::{IpAddr, SocketAddr},
10    sync::atomic::{AtomicBool, Ordering},
11    sync::Arc,
12    time::Instant,
13};
14
15use parking_lot::Mutex;
16use socket2::{Domain, Protocol, Socket, Type as SockType};
17use tokio::{
18    net::UdpSocket,
19    sync::oneshot,
20    task::{self, JoinHandle},
21};
22use tracing::debug;
23
24use crate::{
25    config::Config,
26    icmp::{icmpv4::Icmpv4Packet, icmpv6::Icmpv6Packet},
27    IcmpPacket, PingIdentifier, PingSequence, Pinger, SurgeError, ICMP,
28};
29
30// Check, if the platform's socket operates with ICMP packets in a casual way
31#[macro_export]
32macro_rules! is_linux_icmp_socket {
33    ($sock_type:expr) => {
34        if ($sock_type == socket2::Type::DGRAM
35            && cfg!(not(any(target_os = "linux", target_os = "android"))))
36            || $sock_type == socket2::Type::RAW
37        {
38            false
39        } else {
40            true
41        }
42    };
43}
44
45#[derive(Clone)]
46pub struct AsyncSocket {
47    inner: Arc<UdpSocket>,
48    sock_type: SockType,
49}
50
51impl AsyncSocket {
52    pub fn new(config: &Config) -> io::Result<Self> {
53        let (sock_type, socket) = Self::create_socket(config)?;
54
55        socket.set_nonblocking(true)?;
56        if let Some(sock_addr) = &config.bind {
57            socket.bind(sock_addr)?;
58        }
59        #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
60        if let Some(interface) = &config.interface {
61            socket.bind_device(Some(interface.as_bytes()))?;
62        }
63        #[cfg(any(
64            target_os = "ios",
65            target_os = "visionos",
66            target_os = "macos",
67            target_os = "tvos",
68            target_os = "watchos",
69            target_os = "illumos",
70            target_os = "solaris",
71            target_os = "linux",
72            target_os = "android",
73        ))]
74        {
75            if config.interface_index.is_some() {
76                match config.kind {
77                    ICMP::V4 => socket.bind_device_by_index_v4(config.interface_index)?,
78                    ICMP::V6 => socket.bind_device_by_index_v6(config.interface_index)?,
79                }
80            }
81        }
82        if let Some(ttl) = config.ttl {
83            match config.kind {
84                ICMP::V4 => socket.set_ttl_v4(ttl)?,
85                ICMP::V6 => socket.set_unicast_hops_v6(ttl)?,
86            }
87        }
88        #[cfg(target_os = "freebsd")]
89        if let Some(fib) = config.fib {
90            socket.set_fib(fib)?;
91        }
92        #[cfg(windows)]
93        let socket = UdpSocket::from_std(unsafe {
94            std::net::UdpSocket::from_raw_socket(socket.into_raw_socket())
95        })?;
96        #[cfg(unix)]
97        let socket =
98            UdpSocket::from_std(unsafe { std::net::UdpSocket::from_raw_fd(socket.into_raw_fd()) })?;
99        Ok(Self {
100            inner: Arc::new(socket),
101            sock_type,
102        })
103    }
104
105    fn create_socket(config: &Config) -> io::Result<(SockType, Socket)> {
106        let (domain, proto) = match config.kind {
107            ICMP::V4 => (Domain::IPV4, Some(Protocol::ICMPV4)),
108            ICMP::V6 => (Domain::IPV6, Some(Protocol::ICMPV6)),
109        };
110
111        let first_err = match Socket::new(domain, config.sock_type_hint, proto) {
112            Ok(sock) => return Ok((config.sock_type_hint, sock)),
113            Err(err) => err,
114        };
115
116        let fallback_type = if config.sock_type_hint == SockType::DGRAM {
117            SockType::RAW
118        } else {
119            SockType::DGRAM
120        };
121
122        debug!(
123            "error opening {:?} type socket, trying {:?}: {:?}",
124            config.sock_type_hint, fallback_type, first_err
125        );
126
127        match Socket::new(domain, fallback_type, proto) {
128            Ok(sock) => Ok((fallback_type, sock)),
129            Err(_second_err) => {
130                #[cfg(all(target_os = "linux", any(target_arch = "x86", target_arch = "x86_64", target_arch = "arm", target_arch = "aarch64")))]
131                {
132                    if config.sock_type_hint == SockType::DGRAM
133                        && first_err.kind() == io::ErrorKind::PermissionDenied
134                    {
135                        return Err(io::Error::new(
136                            io::ErrorKind::PermissionDenied,
137                            format!(
138                                "Permission denied creating ICMP socket. On Linux, you may need to:\n\
139                                1. Enable non-privileged ICMP: `sudo sysctl -w net.ipv{}.ping_group_range=\"0 2147483647\"`\n\
140                                2. Run with sudo or set CAP_NET_RAW: `sudo setcap cap_net_raw+ep <binary>`\n\
141                                Original error: {}",
142                                if config.kind == ICMP::V4 { "4" } else { "6" },
143                                first_err
144                            ),
145                        ));
146                    }
147                }
148
149                Err(first_err)
150            }
151        }
152    }
153
154    pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
155        self.inner.recv_from(buf).await
156    }
157
158    pub async fn send_to(&self, buf: &mut [u8], target: &SocketAddr) -> io::Result<usize> {
159        self.inner.send_to(buf, target).await
160    }
161
162    pub fn local_addr(&self) -> io::Result<SocketAddr> {
163        self.inner.local_addr()
164    }
165
166    pub fn get_type(&self) -> SockType {
167        self.sock_type
168    }
169
170    #[cfg(unix)]
171    pub fn get_native_sock(&self) -> RawFd {
172        self.inner.as_raw_fd()
173    }
174
175    #[cfg(windows)]
176    pub fn get_native_sock(&self) -> RawSocket {
177        self.inner.as_raw_socket()
178    }
179}
180
181#[derive(PartialEq, Eq, Hash)]
182struct ReplyToken(IpAddr, Option<PingIdentifier>, PingSequence);
183
184pub(crate) struct Reply {
185    pub timestamp: Instant,
186    pub packet: IcmpPacket,
187}
188
189#[derive(Clone)]
190pub(crate) struct ReplyMap {
191    inner: Arc<Mutex<HashMap<ReplyToken, oneshot::Sender<Reply>>>>,
192    alive: Arc<AtomicBool>,
193}
194
195impl Default for ReplyMap {
196    fn default() -> Self {
197        Self {
198            inner: Arc::new(Mutex::new(HashMap::new())),
199            alive: Arc::new(AtomicBool::new(true)),
200        }
201    }
202}
203
204impl ReplyMap {
205    /// Register to wait for a reply from host with ident and sequence number.
206    /// If there is already someone waiting for this specific reply then an
207    /// error is returned.
208    pub fn new_waiter(
209        &self,
210        host: IpAddr,
211        ident: Option<PingIdentifier>,
212        seq: PingSequence,
213    ) -> Result<oneshot::Receiver<Reply>, SurgeError> {
214        if !self.alive.load(Ordering::Relaxed) {
215            return Err(SurgeError::ClientDestroyed);
216        }
217        let (tx, rx) = oneshot::channel();
218        if self
219            .inner
220            .lock()
221            .insert(ReplyToken(host, ident, seq), tx)
222            .is_some()
223        {
224            return Err(SurgeError::IdenticalRequests { host, ident, seq });
225        }
226        Ok(rx)
227    }
228
229    /// Remove a waiter.
230    pub(crate) fn remove(
231        &self,
232        host: IpAddr,
233        ident: Option<PingIdentifier>,
234        seq: PingSequence,
235    ) -> Option<oneshot::Sender<Reply>> {
236        self.inner.lock().remove(&ReplyToken(host, ident, seq))
237    }
238
239    /// Mark the client as destroyed. This is called when the Client is dropped.
240    pub(crate) fn mark_destroyed(&self) {
241        self.alive.store(false, Ordering::Relaxed);
242    }
243}
244
245///
246/// If you want to pass the `Client` in the task, please wrap it with `Arc`: `Arc<Client>`.
247/// and can realize the simultaneous ping of multiple addresses when only one `socket` is created.
248///
249#[derive(Clone)]
250pub struct Client {
251    socket: AsyncSocket,
252    reply_map: ReplyMap,
253    recv: Arc<JoinHandle<()>>,
254}
255
256impl Drop for Client {
257    fn drop(&mut self) {
258        // Mark the reply_map as destroyed so any pending or new ping operations
259        // will fail with ClientDestroyed error instead of timing out.
260        self.reply_map.mark_destroyed();
261        // The client may pass through multiple tasks, so need to judge whether the number of references is 1.
262        if Arc::strong_count(&self.recv) <= 1 {
263            self.recv.abort();
264        }
265    }
266}
267
268impl Client {
269    /// A client is generated according to the configuration. In fact, a `AsyncSocket` is wrapped inside,
270    /// and you can clone to any `task` at will.
271    pub fn new(config: &Config) -> io::Result<Self> {
272        let socket = AsyncSocket::new(config)?;
273        let reply_map = ReplyMap::default();
274        let recv = task::spawn(recv_task(socket.clone(), reply_map.clone()));
275        Ok(Self {
276            socket,
277            reply_map,
278            recv: Arc::new(recv),
279        })
280    }
281
282    /// Create a `Pinger` instance, you can make special configuration for this instance.
283    pub async fn pinger(&self, host: IpAddr, ident: PingIdentifier) -> Pinger {
284        Pinger::new(host, ident, self.socket.clone(), self.reply_map.clone())
285    }
286
287    /// Expose the underlying socket, if user wants to modify any options on it
288    pub fn get_socket(&self) -> AsyncSocket {
289        self.socket.clone()
290    }
291}
292
293async fn recv_task(socket: AsyncSocket, reply_map: ReplyMap) {
294    let mut buf = [0; 2048];
295    loop {
296        if let Ok((sz, addr)) = socket.recv_from(&mut buf).await {
297            let timestamp = Instant::now();
298            let message = &buf[..sz];
299            let local_addr = socket.local_addr().unwrap().ip();
300            let packet = {
301                let result = match addr.ip() {
302                    IpAddr::V4(src_addr) => {
303                        let local_addr_ip4 = match local_addr {
304                            IpAddr::V4(local_addr_ip4) => local_addr_ip4,
305                            _ => continue,
306                        };
307
308                        Icmpv4Packet::decode(message, socket.sock_type, src_addr, local_addr_ip4)
309                            .map(IcmpPacket::V4)
310                    }
311                    IpAddr::V6(src_addr) => {
312                        Icmpv6Packet::decode(message, src_addr).map(IcmpPacket::V6)
313                    }
314                };
315                match result {
316                    Ok(packet) => packet,
317                    Err(err) => {
318                        debug!("error decoding ICMP packet: {:?}", err);
319                        continue;
320                    }
321                }
322            };
323
324            let ident = if is_linux_icmp_socket!(socket.get_type()) {
325                None
326            } else {
327                Some(packet.get_identifier())
328            };
329
330            if let Some(waiter) = reply_map.remove(addr.ip(), ident, packet.get_sequence()) {
331                // If send fails the receiving end has closed. Nothing to do.
332                let _ = waiter.send(Reply { timestamp, packet });
333            } else {
334                debug!("no one is waiting for ICMP packet ({:?})", packet);
335            }
336        }
337    }
338}
surge_ping/client.rs

surge_ping/
client.rs
