Skip to main content

tiny_ping/
socket.rs

1use std::mem::MaybeUninit;
2use std::net::SocketAddr;
3#[cfg(unix)]
4use std::sync::Arc;
5use std::{io, net::IpAddr};
6
7use socket2::SockAddr;
8#[cfg(unix)]
9use socket2::{Domain, Protocol, Socket, Type};
10#[cfg(unix)]
11use tokio::io::unix::AsyncFd;
12
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14pub enum SocketType {
15    Raw,
16    Dgram,
17}
18
19#[derive(Debug, Clone)]
20pub struct AsyncSocket {
21    #[cfg(unix)]
22    inner: Arc<AsyncFd<Socket>>,
23    socket_type: SocketType,
24}
25
26impl AsyncSocket {
27    #[cfg(unix)]
28    pub fn new(addr: IpAddr, socket_type: SocketType) -> io::Result<AsyncSocket> {
29        let ty = match socket_type {
30            SocketType::Raw => Type::RAW,
31            SocketType::Dgram => Type::DGRAM,
32        };
33        let socket = match addr {
34            IpAddr::V4(_) => Socket::new(Domain::IPV4, ty, Some(Protocol::ICMPV4))?,
35            IpAddr::V6(_) => Socket::new(Domain::IPV6, ty, Some(Protocol::ICMPV6))?,
36        };
37
38        // TODO: Type filtering,
39        // https://tools.ietf.org/html/rfc3542#section-3.2. Currently blocked
40        // on https://github.com/rust-lang/socket2/issues/199
41
42        // TODO: Get access to the hop limits
43        // https://tools.ietf.org/html/rfc3542#section-4, to show the TTL for
44        // ICMPv6.
45
46        socket.set_nonblocking(true)?;
47        Ok(AsyncSocket {
48            inner: Arc::new(AsyncFd::new(socket)?),
49            socket_type,
50        })
51    }
52
53    #[cfg(not(unix))]
54    pub fn new(_addr: IpAddr, socket_type: SocketType) -> io::Result<AsyncSocket> {
55        Ok(AsyncSocket { socket_type })
56    }
57
58    pub fn socket_type(&self) -> SocketType {
59        self.socket_type
60    }
61
62    #[cfg(unix)]
63    pub fn set_ttl(&self, addr: IpAddr, ttl: u32) -> io::Result<()> {
64        match addr {
65            IpAddr::V4(_) => self.inner.get_ref().set_ttl_v4(ttl),
66            IpAddr::V6(_) => self.inner.get_ref().set_unicast_hops_v6(ttl),
67        }
68    }
69
70    #[cfg(not(unix))]
71    pub fn set_ttl(&self, _addr: IpAddr, _ttl: u32) -> io::Result<()> {
72        unsupported()
73    }
74
75    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
76    pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> {
77        self.inner.get_ref().bind_device(interface)
78    }
79
80    #[cfg(unix)]
81    pub async fn recv_from(
82        &self,
83        buf: &mut [MaybeUninit<u8>],
84    ) -> io::Result<(usize, Option<SocketAddr>)> {
85        loop {
86            let mut guard = self.inner.readable().await?;
87
88            match guard.try_io(|inner| {
89                inner
90                    .get_ref()
91                    .recv_from(buf)
92                    .map(|(size, addr)| (size, addr.as_socket()))
93            }) {
94                Ok(result) => return result,
95                Err(_would_block) => continue,
96            }
97        }
98    }
99
100    #[cfg(not(unix))]
101    pub async fn recv_from(
102        &self,
103        _buf: &mut [MaybeUninit<u8>],
104    ) -> io::Result<(usize, Option<SocketAddr>)> {
105        unsupported()
106    }
107
108    #[cfg(unix)]
109    pub async fn send_to(&self, buf: &[u8], target: &SockAddr) -> io::Result<usize> {
110        loop {
111            let mut guard = self.inner.writable().await?;
112
113            match guard.try_io(|inner| inner.get_ref().send_to(buf, target)) {
114                Ok(n) => return n,
115                Err(_would_block) => continue,
116            }
117        }
118    }
119
120    #[cfg(not(unix))]
121    pub async fn send_to(&self, _buf: &[u8], _target: &SockAddr) -> io::Result<usize> {
122        unsupported()
123    }
124}
125
126#[cfg(not(unix))]
127fn unsupported<T>() -> io::Result<T> {
128    Err(io::Error::new(
129        io::ErrorKind::Unsupported,
130        "tiny-ping sockets are only implemented on Unix targets",
131    ))
132}