tokio_ping/
ping.rs

1use std::io;
2
3use std::collections::HashMap;
4use std::net::{IpAddr, SocketAddr};
5use std::sync::Arc;
6use std::time::{Instant, Duration};
7
8use futures::{Async, Future, Poll, Stream};
9use futures::sync::oneshot;
10use rand::random;
11use parking_lot::Mutex;
12use socket2::{Domain, Protocol, Type};
13
14use tokio_executor::spawn;
15use tokio_reactor::Handle;
16use tokio_timer::Delay;
17
18use errors::{Error, ErrorKind};
19use packet::{IpV4Packet, IpV4Protocol};
20use packet::{ICMP_HEADER_SIZE, IcmpV4, IcmpV6, EchoRequest, EchoReply};
21use socket::{Socket, Send};
22
23const DEFAULT_TIMEOUT: u64 = 2;
24const TOKEN_SIZE: usize = 24;
25const ECHO_REQUEST_BUFFER_SIZE: usize = ICMP_HEADER_SIZE + TOKEN_SIZE;
26type Token = [u8; TOKEN_SIZE];
27type EchoRequestBuffer = [u8; ECHO_REQUEST_BUFFER_SIZE];
28
29#[derive(Clone)]
30struct PingState {
31    inner: Arc<Mutex<HashMap<Token, oneshot::Sender<Instant>>>>,
32}
33
34impl PingState {
35    fn new() -> Self {
36        Self {
37            inner: Arc::new(Mutex::new(HashMap::new())),
38        }
39    }
40
41    fn insert(&self, key: Token, value: oneshot::Sender<Instant>) {
42        self.inner.lock().insert(key, value);
43    }
44
45    fn remove(&self, key: &[u8]) -> Option<oneshot::Sender<Instant>> {
46        self.inner.lock().remove(key)
47    }
48}
49
50/// Represent a future that resolves into ping response time, resolves into `None` if timed out.
51#[must_use = "futures do nothing unless polled"]
52pub struct PingFuture {
53    inner: PingFutureKind
54}
55
56enum PingFutureKind {
57    Normal(NormalPingFutureKind),
58    PacketEncodeError,
59    InvalidProtocol,
60}
61
62struct NormalPingFutureKind {
63    start_time: Instant,
64    state: PingState,
65    token: Token,
66    delay: Delay,
67    send: Option<Send<EchoRequestBuffer>>,
68    receiver: oneshot::Receiver<Instant>,
69}
70
71impl Future for PingFuture {
72    type Item = Option<Duration>;
73    type Error = Error;
74
75    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
76        match self.inner {
77            PingFutureKind::Normal(ref mut normal) => {
78                let mut swap_send = false;
79                if let Some(ref mut send) = normal.send {
80                    match send.poll() {
81                        Ok(Async::NotReady) => (),
82                        Ok(Async::Ready(_)) => swap_send = true,
83                        Err(_) => return Err(ErrorKind::InternalError.into()),
84                    }
85                }
86
87                if swap_send {
88                    normal.send = None;
89                }
90
91                match normal.receiver.poll() {
92                    Ok(Async::NotReady) => (),
93                    Ok(Async::Ready(stop_time)) => {
94                        return Ok(Async::Ready(Some(stop_time - normal.start_time)))
95                    },
96                    Err(_) => return Err(ErrorKind::InternalError.into()),
97                }
98
99                match normal.delay.poll() {
100                    Ok(Async::NotReady) => (),
101                    Ok(Async::Ready(_)) => return Ok(Async::Ready(None)),
102                    Err(_) => return Err(ErrorKind::InternalError.into()),
103                }
104            }
105            PingFutureKind::InvalidProtocol => {
106                return Err(ErrorKind::InvalidProtocol.into())
107            }
108            PingFutureKind::PacketEncodeError => {
109                return Err(ErrorKind::InternalError.into())
110            }
111        }
112        Ok(Async::NotReady)
113    }
114}
115
116impl Drop for PingFuture {
117    fn drop(&mut self) {
118        match self.inner {
119            PingFutureKind::Normal(ref normal) => {
120                normal.state.remove(&normal.token);
121            }
122            | PingFutureKind::InvalidProtocol
123            | PingFutureKind::PacketEncodeError => (),
124        }
125    }
126}
127
128/// Ping the same host several times.
129pub struct PingChain {
130    pinger: Pinger,
131    hostname: IpAddr,
132    ident: Option<u16>,
133    seq_cnt: Option<u16>,
134    timeout: Option<Duration>,
135}
136
137impl PingChain {
138    fn new(pinger: Pinger, hostname: IpAddr) -> Self {
139        Self {
140            pinger: pinger,
141            hostname: hostname,
142            ident: None,
143            seq_cnt: None,
144            timeout: None,
145        }
146    }
147
148    /// Set ICMP ident. Default value is randomized.
149    pub fn ident(mut self, ident: u16) -> Self {
150        self.ident = Some(ident);
151        self
152    }
153
154    /// Set ICMP seq_cnt, this value will be incremented by one for every `send`.
155    /// Default value is 0.
156    pub fn seq_cnt(mut self, seq_cnt: u16) -> Self {
157        self.seq_cnt = Some(seq_cnt);
158        self
159    }
160
161    /// Set ping timeout. Default timeout is two seconds.
162    pub fn timeout(mut self, timeout: Duration) -> Self {
163        self.timeout = Some(timeout);
164        self
165    }
166
167    /// Send ICMP request and wait for response.
168    pub fn send(&mut self) -> PingFuture {
169        let ident = match self.ident {
170            Some(ident) => ident,
171            None => {
172                let ident = random();
173                self.ident = Some(ident);
174                ident
175            }
176        };
177
178        let seq_cnt = match self.seq_cnt {
179            Some(seq_cnt) => {
180                self.seq_cnt = Some(seq_cnt.wrapping_add(1));
181                seq_cnt
182            }
183            None => {
184                self.seq_cnt = Some(1);
185                0
186            }
187        };
188
189        let timeout = match self.timeout {
190            Some(timeout) => timeout,
191            None => {
192                let timeout = Duration::from_secs(DEFAULT_TIMEOUT);
193                self.timeout = Some(timeout);
194                timeout
195            }
196        };
197
198        self.pinger.ping(self.hostname, ident, seq_cnt, timeout)
199    }
200
201    /// Create infinite stream of ping response times.
202    pub fn stream(self) -> PingChainStream {
203        PingChainStream::new(self)
204    }
205}
206
207/// Stream of sequential ping response times, iterates `None` if timed out.
208pub struct PingChainStream {
209    chain: PingChain,
210    future: Option<PingFuture>,
211}
212
213impl PingChainStream {
214    fn new(chain: PingChain) -> Self {
215        Self {
216            chain: chain,
217            future: None,
218        }
219    }
220}
221
222impl Stream for PingChainStream {
223    type Item = Option<Duration>;
224    type Error = Error;
225
226    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
227        let mut future = self.future.take().unwrap_or_else(|| self.chain.send());
228
229        match future.poll() {
230            Ok(Async::Ready(item)) => Ok(Async::Ready(Some(item))),
231            Ok(Async::NotReady) => {
232                self.future = Some(future);
233                Ok(Async::NotReady)
234            },
235            Err(err) => Err(err),
236        }
237    }
238}
239
240/// ICMP packets sender and receiver.
241#[derive(Clone)]
242pub struct Pinger {
243    inner: Arc<PingInner>,
244}
245
246struct PingInner {
247    sockets: Sockets,
248    state: PingState,
249    _v4_finalize: Option<oneshot::Sender<()>>,
250    _v6_finalize: Option<oneshot::Sender<()>>,
251}
252
253enum Sockets {
254    V4(Socket),
255    V6(Socket),
256    Both { v4: Socket, v6: Socket },
257}
258
259impl Sockets {
260    fn new(handle: &Handle) -> io::Result<Self> {
261        let mb_v4socket = Socket::new(Domain::ipv4(), Type::raw(), Protocol::icmpv4(), handle);
262        let mb_v6socket = Socket::new(Domain::ipv6(), Type::raw(), Protocol::icmpv6(), handle);
263        match (mb_v4socket, mb_v6socket) {
264            (Ok(v4_socket), Ok(v6_socket)) => Ok(Sockets::Both {
265                v4: v4_socket,
266                v6: v6_socket,
267            }),
268            (Ok(v4_socket), Err(_)) => Ok(Sockets::V4(v4_socket)),
269            (Err(_), Ok(v6_socket)) => Ok(Sockets::V6(v6_socket)),
270            (Err(err), Err(_)) => Err(err),
271        }
272    }
273
274    fn v4(&self) -> Option<&Socket> {
275        match *self {
276            Sockets::V4(ref socket) => Some(socket),
277            Sockets::Both { ref v4, .. } => Some(v4),
278            Sockets::V6(_) => None,
279        }
280    }
281
282    fn v6(&self) -> Option<&Socket> {
283        match *self {
284            Sockets::V4(_) => None,
285            Sockets::Both { ref v6, .. } => Some(v6),
286            Sockets::V6(ref socket) => Some(socket),
287        }
288    }
289}
290
291impl Pinger {
292    /// Create new `Pinger` instance, will fail if unable to create both IPv4 and IPv6 sockets.
293    pub fn new() -> impl Future<Item=Self, Error=Error> {
294        ::futures::future::lazy(||
295            Self::with_handle(&Handle::default()).map_err(From::from)
296        )
297
298    }
299
300    fn with_handle(handle: &Handle) -> io::Result<Self> {
301        let sockets = Sockets::new(handle)?;
302
303        let state = PingState::new();
304
305        let v4_finalize = if let Some(v4_socket) = sockets.v4() {
306            let (s, r) = oneshot::channel();
307            let receiver =
308                Receiver::<IcmpV4>::new(v4_socket.clone(), state.clone());
309            spawn(receiver.select(r.map_err(|_| ())).then(|_| Ok(())));
310            Some(s)
311        } else {
312            None
313        };
314
315        let v6_finalize = if let Some(v6_socket) = sockets.v6() {
316            let (s, r) = oneshot::channel();
317            let receiver =
318                Receiver::<IcmpV6>::new(v6_socket.clone(), state.clone());
319            spawn(receiver.select(r.map_err(|_| ())).then(|_| Ok(())));
320            Some(s)
321        } else {
322            None
323        };
324
325        let inner = PingInner {
326            sockets: sockets,
327            state: state,
328            _v4_finalize: v4_finalize,
329            _v6_finalize: v6_finalize,
330        };
331
332        Ok(Self {
333            inner: Arc::new(inner),
334        })
335    }
336
337    /// Ping the same host several times.
338    pub fn chain(&self, hostname: IpAddr) -> PingChain {
339        PingChain::new(self.clone(), hostname)
340    }
341
342    /// Send ICMP request and wait for response.
343    pub fn ping(
344        &self,
345        hostname: IpAddr,
346        ident: u16,
347        seq_cnt: u16,
348        timeout: Duration,
349    ) -> PingFuture {
350        let (sender, receiver) = oneshot::channel();
351
352        let deadline = Instant::now() + timeout;
353
354        let token = random();
355        self.inner.state.insert(token, sender);
356
357        let dest = SocketAddr::new(hostname, 0);
358        let mut buffer = [0; ECHO_REQUEST_BUFFER_SIZE];
359
360        let request = EchoRequest {
361            ident: ident,
362            seq_cnt: seq_cnt,
363            payload: &token,
364        };
365
366        let (encode_result, mb_socket) = {
367            if dest.is_ipv4() {
368                (
369                    request.encode::<IcmpV4>(&mut buffer[..]),
370                    self.inner.sockets.v4().cloned()
371                )
372            } else {
373                (
374                    request.encode::<IcmpV6>(&mut buffer[..]),
375                    self.inner.sockets.v6().cloned(),
376                )
377            }
378        };
379
380
381        let socket = match mb_socket {
382            Some(socket) => socket,
383            None => {
384                return PingFuture {
385                    inner: PingFutureKind::InvalidProtocol
386                }
387            }
388        };
389
390        if let Err(_) = encode_result {
391            return PingFuture {
392                inner: PingFutureKind::PacketEncodeError
393            }
394        }
395
396        let send_future = socket.send_to(buffer, &dest);
397
398        PingFuture {
399            inner: PingFutureKind::Normal(NormalPingFutureKind {
400                start_time: Instant::now(),
401                state: self.inner.state.clone(),
402                token: token,
403                delay: Delay::new(deadline),
404                send: Some(send_future),
405                receiver: receiver,
406            })
407        }
408    }
409}
410
411struct Receiver<Message> {
412    socket: Socket,
413    state: PingState,
414    buffer: [u8; 2048],
415    _phantom: ::std::marker::PhantomData<Message>,
416}
417
418trait ParseReply {
419    fn reply_payload(data: &[u8]) -> Option<&[u8]>;
420}
421
422impl ParseReply for IcmpV4 {
423    fn reply_payload(data: &[u8]) -> Option<&[u8]> {
424        if let Ok(ipv4_packet) = IpV4Packet::decode(data) {
425            if ipv4_packet.protocol != IpV4Protocol::Icmp {
426                return None;
427            }
428
429            if let Ok(reply) = EchoReply::decode::<IcmpV4>(ipv4_packet.data) {
430                return Some(reply.payload);
431            }
432        }
433        None
434    }
435}
436
437impl ParseReply for IcmpV6 {
438    fn reply_payload(data: &[u8]) -> Option<&[u8]> {
439        if let Ok(reply) = EchoReply::decode::<IcmpV6>(data) {
440            return Some(reply.payload);
441        }
442        None
443    }
444}
445
446impl<Proto> Receiver<Proto> {
447    fn new(socket: Socket, state: PingState) -> Self {
448        Self {
449            socket: socket,
450            state: state,
451            buffer: [0; 2048],
452            _phantom: ::std::marker::PhantomData,
453        }
454    }
455}
456
457impl<Message: ParseReply> Future for Receiver<Message> {
458    type Item = ();
459    type Error = ();
460
461    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
462        match self.socket.recv(&mut self.buffer) {
463            Ok(Async::Ready(bytes)) => {
464                if let Some(payload) = Message::reply_payload(&self.buffer[..bytes]) {
465                    let now = Instant::now();
466                    if let Some(sender) = self.state.remove(payload) {
467                        sender.send(now).unwrap_or_default()
468                    }
469                }
470                self.poll()
471            }
472            Ok(Async::NotReady) => Ok(Async::NotReady),
473            Err(_) => Err(()),
474        }
475    }
476}
tokio_ping/ping.rs

tokio_ping/
ping.rs
