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 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}