surge_ping/
ping.rs

1use std::{
2    net::{IpAddr, SocketAddr},
3    time::{Duration, Instant},
4};
5
6use tokio::time::timeout;
7
8use crate::{
9    client::{AsyncSocket, ReplyMap},
10    error::{Result, SurgeError},
11    icmp::{icmpv4, icmpv6, IcmpPacket, PingIdentifier, PingSequence},
12    is_linux_icmp_socket,
13};
14
15/// A Ping struct represents the state of one particular ping instance.
16pub struct Pinger {
17    pub host: IpAddr,
18    pub ident: Option<PingIdentifier>,
19    scope_id: u32,
20    timeout: Duration,
21    socket: AsyncSocket,
22    reply_map: ReplyMap,
23    last_sequence: Option<PingSequence>,
24}
25
26impl Drop for Pinger {
27    fn drop(&mut self) {
28        if let Some(sequence) = self.last_sequence.take() {
29            // Ensure no reply waiter is left hanging if this pinger is dropped while
30            // waiting for a reply.
31            self.reply_map.remove(self.host, self.ident, sequence);
32        }
33    }
34}
35
36impl Pinger {
37    pub(crate) fn new(
38        host: IpAddr,
39        ident_hint: PingIdentifier,
40        socket: AsyncSocket,
41        response_map: ReplyMap,
42    ) -> Pinger {
43        let ident = if is_linux_icmp_socket!(socket.get_type()) {
44            None
45        } else {
46            Some(ident_hint)
47        };
48
49        Pinger {
50            host,
51            ident,
52            scope_id: 0,
53            timeout: Duration::from_secs(2),
54            socket,
55            reply_map: response_map,
56            last_sequence: None,
57        }
58    }
59
60    /// The scope id of the IPv6 socket address.
61    pub fn scope_id(&mut self, scope_id: u32) -> &mut Pinger {
62        self.scope_id = scope_id;
63        self
64    }
65
66    /// The timeout of each Ping, in seconds. (default: 2s)
67    pub fn timeout(&mut self, timeout: Duration) -> &mut Pinger {
68        self.timeout = timeout;
69        self
70    }
71
72    /// Send Ping request with sequence number.
73    pub async fn ping(
74        &mut self,
75        seq: PingSequence,
76        payload: &[u8],
77    ) -> Result<(IcmpPacket, Duration)> {
78        // Register to wait for a reply.
79        let reply_waiter = self.reply_map.new_waiter(self.host, self.ident, seq)?;
80
81        // Send actual packet
82        if let Err(e) = self.send_ping(seq, payload).await {
83            self.reply_map.remove(self.host, self.ident, seq);
84            return Err(e);
85        }
86
87        let send_time = Instant::now();
88        self.last_sequence = Some(seq);
89
90        // Wait for reply or timeout.
91        match timeout(self.timeout, reply_waiter).await {
92            Ok(Ok(reply)) => Ok((
93                reply.packet,
94                reply.timestamp.saturating_duration_since(send_time),
95            )),
96            Ok(Err(_err)) => Err(SurgeError::NetworkError),
97            Err(_) => {
98                self.reply_map.remove(self.host, self.ident, seq);
99                Err(SurgeError::Timeout { seq })
100            }
101        }
102    }
103
104    /// Send a ping packet (useful, when you don't need a reply).
105    pub async fn send_ping(&self, seq: PingSequence, payload: &[u8]) -> Result<()> {
106        // Create and send ping packet.
107        let mut packet = match self.host {
108            IpAddr::V4(_) => icmpv4::make_icmpv4_echo_packet(
109                self.ident.unwrap_or(PingIdentifier(0)),
110                seq,
111                self.socket.get_type(),
112                payload,
113            )?,
114            IpAddr::V6(_) => icmpv6::make_icmpv6_echo_packet(
115                self.ident.unwrap_or(PingIdentifier(0)),
116                seq,
117                payload,
118            )?,
119        };
120
121        let mut target = SocketAddr::new(self.host, 0);
122        if let SocketAddr::V6(sa) = &mut target {
123            sa.set_scope_id(self.scope_id);
124        }
125
126        self.socket.send_to(&mut packet, &target).await?;
127
128        Ok(())
129    }
130}
surge_ping/ping.rs

surge_ping/
ping.rs
