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 bind(&self, addr: &SockAddr) -> io::Result<()> {
64        self.inner.get_ref().bind(addr)
65    }
66
67    #[cfg(not(unix))]
68    pub fn bind(&self, _addr: &SockAddr) -> io::Result<()> {
69        unsupported()
70    }
71
72    #[cfg(unix)]
73    pub fn set_ttl(&self, addr: IpAddr, ttl: u32) -> io::Result<()> {
74        match addr {
75            IpAddr::V4(_) => self.inner.get_ref().set_ttl_v4(ttl),
76            IpAddr::V6(_) => self.inner.get_ref().set_unicast_hops_v6(ttl),
77        }
78    }
79
80    #[cfg(not(unix))]
81    pub fn set_ttl(&self, _addr: IpAddr, _ttl: u32) -> io::Result<()> {
82        unsupported()
83    }
84
85    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
86    pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> {
87        self.inner.get_ref().bind_device(interface)
88    }
89
90    #[cfg(unix)]
91    pub async fn recv_from(
92        &self,
93        buf: &mut [MaybeUninit<u8>],
94    ) -> io::Result<(usize, Option<SocketAddr>)> {
95        loop {
96            let mut guard = self.inner.readable().await?;
97
98            match guard.try_io(|inner| {
99                inner
100                    .get_ref()
101                    .recv_from(buf)
102                    .map(|(size, addr)| (size, addr.as_socket()))
103            }) {
104                Ok(result) => return result,
105                Err(_would_block) => continue,
106            }
107        }
108    }
109
110    #[cfg(not(unix))]
111    pub async fn recv_from(
112        &self,
113        _buf: &mut [MaybeUninit<u8>],
114    ) -> io::Result<(usize, Option<SocketAddr>)> {
115        unsupported()
116    }
117
118    #[cfg(unix)]
119    pub async fn send_to(&self, buf: &[u8], target: &SockAddr) -> io::Result<usize> {
120        loop {
121            let mut guard = self.inner.writable().await?;
122
123            match guard.try_io(|inner| inner.get_ref().send_to(buf, target)) {
124                Ok(n) => return n,
125                Err(_would_block) => continue,
126            }
127        }
128    }
129
130    #[cfg(not(unix))]
131    pub async fn send_to(&self, _buf: &[u8], _target: &SockAddr) -> io::Result<usize> {
132        unsupported()
133    }
134}
135
136#[cfg(not(unix))]
137fn unsupported<T>() -> io::Result<T> {
138    Err(io::Error::new(
139        io::ErrorKind::Unsupported,
140        "tiny-ping sockets are only implemented on Unix targets",
141    ))
142}