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