tiny_ping/
ping.rs

1use std::{
2    collections::HashMap,
3    mem::MaybeUninit,
4    net::{IpAddr, SocketAddr},
5    sync::Arc,
6    time::{Duration, Instant},
7};
8
9use parking_lot::Mutex;
10use tokio::task;
11use tokio::time::timeout;
12
13use crate::error::{Error, Result};
14use crate::icmp::{EchoReply, EchoRequest};
15use crate::unix::AsyncSocket;
16
17type Token = (u16, u16);
18
19#[derive(Debug, Clone)]
20struct Cache {
21    inner: Arc<Mutex<HashMap<Token, Instant>>>,
22}
23
24impl Cache {
25    fn new() -> Cache {
26        Cache {
27            inner: Arc::new(Mutex::new(HashMap::new())),
28        }
29    }
30
31    fn insert(&self, ident: u16, seq_cnt: u16, time: Instant) {
32        self.inner.lock().insert((ident, seq_cnt), time);
33    }
34
35    fn remove(&self, ident: u16, seq_cnt: u16) -> Option<Instant> {
36        self.inner.lock().remove(&(ident, seq_cnt))
37    }
38}
39
40/// A Ping struct represents the state of one particular ping instance.
41#[derive(Debug, Clone)]
42pub struct Pinger {
43    host: IpAddr,
44    ident: u16,
45    size: usize,
46    timeout: Duration,
47    socket: AsyncSocket,
48    cache: Cache,
49}
50
51impl Pinger {
52    /// Creates a new Ping instance from `IpAddr`.
53    pub fn new(host: IpAddr) -> Result<Pinger> {
54        Ok(Pinger {
55            host,
56            ident: 1,
57            size: 56,
58            timeout: Duration::from_secs(2),
59            socket: AsyncSocket::new(host)?,
60            cache: Cache::new(),
61        })
62    }
63
64    /// Sets the value for the `SO_BINDTODEVICE` option on this socket.
65    ///
66    /// If a socket is bound to an interface, only packets received from that
67    /// particular interface are processed by the socket. Note that this only
68    /// works for some socket types, particularly `AF_INET` sockets.
69    ///
70    /// If `interface` is `None` or an empty string it removes the binding.
71    ///
72    /// This function is only available on Fuchsia and Linux.
73    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
74    pub fn bind_device(&mut self, interface: Option<&[u8]>) -> Result<&mut Pinger> {
75        self.socket.bind_device(interface)?;
76        Ok(self)
77    }
78
79    /// Set the identification of ICMP.
80    pub fn ident(&mut self, val: u16) -> &mut Pinger {
81        self.ident = val;
82        self
83    }
84
85    /// Set the packet size.(default: 56)
86    pub fn size(&mut self, size: usize) -> &mut Pinger {
87        self.size = size;
88        self
89    }
90
91    /// The timeout of each Ping, in seconds. (default: 2s)
92    pub fn timeout(&mut self, timeout: Duration) -> &mut Pinger {
93        self.timeout = timeout;
94        self
95    }
96
97    async fn recv_reply(&self, seq_cnt: u16) -> Result<(EchoReply, Duration)> {
98        let mut buffer = [MaybeUninit::new(0); 2048];
99        loop {
100            let size = self.socket.recv(&mut buffer).await?;
101            let buf = unsafe { assume_init(&buffer[..size]) };
102            match EchoReply::decode(self.host, buf) {
103                Ok(reply) => {
104                    // check reply ident is same
105                    if reply.identifier == self.ident && reply.sequence == seq_cnt {
106                        if let Some(ins) = self.cache.remove(self.ident, seq_cnt) {
107                            return Ok((reply, Instant::now() - ins));
108                        }
109                    }
110                    continue;
111                }
112                Err(Error::NotEchoReply) => continue,
113                Err(Error::NotV6EchoReply) => continue,
114                Err(Error::OtherICMP) => continue,
115                Err(e) => {
116                    return Err(e);
117                }
118            }
119        }
120    }
121
122    /// Send Ping request with sequence number.
123    pub async fn ping(&self, seq_cnt: u16) -> Result<(EchoReply, Duration)> {
124        let sender = self.socket.clone();
125        let mut packet = EchoRequest::new(self.host, self.ident, seq_cnt, self.size).encode()?;
126        let sock_addr = SocketAddr::new(self.host, 0);
127        let ident = self.ident;
128        let cache = self.cache.clone();
129        task::spawn(async move {
130            let _size = sender
131                .send_to(&mut packet, &sock_addr.into())
132                .await
133                .expect("socket send packet error");
134            cache.insert(ident, seq_cnt, Instant::now());
135        });
136
137        match timeout(self.timeout, self.recv_reply(seq_cnt))
138            .await
139            .map_err(|err| {
140                self.cache.remove(ident, seq_cnt);
141                err
142            }) {
143            Ok(rez) => rez,
144            Err(_) => Err(Error::Timeout),
145        }
146    }
147}
148
149/// Assume the `buf`fer to be initialised.
150// TODO: replace with `MaybeUninit::slice_assume_init_ref` once stable.
151unsafe fn assume_init(buf: &[MaybeUninit<u8>]) -> &[u8] {
152    &*(buf as *const [MaybeUninit<u8>] as *const [u8])
153}
tiny_ping/ping.rs

tiny_ping/
ping.rs
