Skip to main content

nex_socket/udp/
async_impl.rs

1use crate::udp::UdpConfig;
2use socket2::{Domain, Protocol, Socket, Type as SockType};
3use std::io;
4use std::net::{SocketAddr, UdpSocket as StdUdpSocket};
5use tokio::net::UdpSocket;
6
7/// Asynchronous UDP socket built on top of Tokio.
8#[derive(Debug)]
9pub struct AsyncUdpSocket {
10    inner: UdpSocket,
11}
12
13impl AsyncUdpSocket {
14    /// Create an asynchronous UDP socket from the given configuration.
15    pub fn from_config(config: &UdpConfig) -> io::Result<Self> {
16        let socket = Socket::new(
17            config.socket_family.to_domain(),
18            config.socket_type.to_sock_type(),
19            Some(Protocol::UDP),
20        )?;
21
22        socket.set_nonblocking(true)?;
23
24        // Set socket options based on configuration
25        if let Some(flag) = config.reuseaddr {
26            socket.set_reuse_address(flag)?;
27        }
28        #[cfg(any(
29            target_os = "android",
30            target_os = "dragonfly",
31            target_os = "freebsd",
32            target_os = "fuchsia",
33            target_os = "ios",
34            target_os = "linux",
35            target_os = "macos",
36            target_os = "netbsd",
37            target_os = "openbsd",
38            target_os = "tvos",
39            target_os = "visionos",
40            target_os = "watchos"
41        ))]
42        if let Some(flag) = config.reuseport {
43            socket.set_reuse_port(flag)?;
44        }
45        if let Some(flag) = config.broadcast {
46            socket.set_broadcast(flag)?;
47        }
48        if let Some(ttl) = config.ttl {
49            socket.set_ttl(ttl)?;
50        }
51        if let Some(hoplimit) = config.hoplimit {
52            socket.set_unicast_hops_v6(hoplimit)?;
53        }
54        if let Some(timeout) = config.read_timeout {
55            socket.set_read_timeout(Some(timeout))?;
56        }
57        if let Some(timeout) = config.write_timeout {
58            socket.set_write_timeout(Some(timeout))?;
59        }
60        if let Some(size) = config.recv_buffer_size {
61            socket.set_recv_buffer_size(size)?;
62        }
63        if let Some(size) = config.send_buffer_size {
64            socket.set_send_buffer_size(size)?;
65        }
66        if let Some(tos) = config.tos {
67            socket.set_tos(tos)?;
68        }
69        #[cfg(any(
70            target_os = "android",
71            target_os = "dragonfly",
72            target_os = "freebsd",
73            target_os = "fuchsia",
74            target_os = "ios",
75            target_os = "linux",
76            target_os = "macos",
77            target_os = "netbsd",
78            target_os = "openbsd",
79            target_os = "tvos",
80            target_os = "visionos",
81            target_os = "watchos"
82        ))]
83        if let Some(tclass) = config.tclass_v6 {
84            socket.set_tclass_v6(tclass)?;
85        }
86        if let Some(only_v6) = config.only_v6 {
87            socket.set_only_v6(only_v6)?;
88        }
89        if let Some(on) = config.recv_pktinfo {
90            crate::udp::set_recv_pktinfo(&socket, config.socket_family, on)?;
91        }
92
93        // Linux: optional interface name
94        #[cfg(any(target_os = "linux", target_os = "android", target_os = "fuchsia"))]
95        if let Some(iface) = &config.bind_device {
96            socket.bind_device(Some(iface.as_bytes()))?;
97        }
98
99        // bind to the specified address if provided
100        if let Some(addr) = config.bind_addr {
101            socket.bind(&addr.into())?;
102        }
103
104        #[cfg(windows)]
105        let std_socket = unsafe {
106            use std::os::windows::io::{FromRawSocket, IntoRawSocket};
107            StdUdpSocket::from_raw_socket(socket.into_raw_socket())
108        };
109        #[cfg(unix)]
110        let std_socket = unsafe {
111            use std::os::fd::{FromRawFd, IntoRawFd};
112            StdUdpSocket::from_raw_fd(socket.into_raw_fd())
113        };
114
115        let inner = UdpSocket::from_std(std_socket)?;
116
117        Ok(Self { inner })
118    }
119
120    /// Create a socket of arbitrary type (DGRAM or RAW).
121    pub fn new(domain: Domain, sock_type: SockType) -> io::Result<Self> {
122        let socket = Socket::new(domain, sock_type, Some(Protocol::UDP))?;
123        socket.set_nonblocking(true)?;
124
125        #[cfg(windows)]
126        let std_socket = unsafe {
127            use std::os::windows::io::{FromRawSocket, IntoRawSocket};
128            StdUdpSocket::from_raw_socket(socket.into_raw_socket())
129        };
130        #[cfg(unix)]
131        let std_socket = unsafe {
132            use std::os::fd::{FromRawFd, IntoRawFd};
133            StdUdpSocket::from_raw_fd(socket.into_raw_fd())
134        };
135
136        let inner = UdpSocket::from_std(std_socket)?;
137
138        Ok(Self { inner })
139    }
140
141    /// Convenience constructor for IPv4 DGRAM.
142    pub fn v4_dgram() -> io::Result<Self> {
143        Self::new(Domain::IPV4, SockType::DGRAM)
144    }
145
146    /// Convenience constructor for IPv6 DGRAM.
147    pub fn v6_dgram() -> io::Result<Self> {
148        Self::new(Domain::IPV6, SockType::DGRAM)
149    }
150
151    /// IPv4 RAW UDP. Requires administrator privileges.
152    pub fn raw_v4() -> io::Result<Self> {
153        Self::new(Domain::IPV4, SockType::RAW)
154    }
155
156    /// IPv6 RAW UDP. Requires administrator privileges.
157    pub fn raw_v6() -> io::Result<Self> {
158        Self::new(Domain::IPV6, SockType::RAW)
159    }
160
161    /// Send data asynchronously.
162    pub async fn send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
163        self.inner.send_to(buf, target).await
164    }
165
166    /// Receive data asynchronously.
167    pub async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
168        self.inner.recv_from(buf).await
169    }
170
171    /// Retrieve the local socket address.
172    pub fn local_addr(&self) -> io::Result<SocketAddr> {
173        self.inner.local_addr()
174    }
175
176    pub fn into_tokio_socket(self) -> io::Result<UdpSocket> {
177        Ok(self.inner)
178    }
179
180    /// Construct from a standard UDP socket.
181    pub fn from_std_socket(socket: StdUdpSocket) -> io::Result<Self> {
182        Ok(Self {
183            inner: UdpSocket::from_std(socket)?,
184        })
185    }
186
187    /// Convert into a standard UDP socket.
188    pub fn into_std_socket(self) -> io::Result<StdUdpSocket> {
189        self.inner.into_std()
190    }
191
192    /// Set IPv4 time-to-live.
193    pub fn set_ttl(&self, ttl: u32) -> io::Result<()> {
194        self.inner.set_ttl(ttl)
195    }
196
197    /// Get IPv4 time-to-live.
198    pub fn ttl(&self) -> io::Result<u32> {
199        self.inner.ttl()
200    }
201
202    /// Set broadcast mode.
203    pub fn set_broadcast(&self, on: bool) -> io::Result<()> {
204        self.inner.set_broadcast(on)
205    }
206
207    /// Get broadcast mode.
208    pub fn broadcast(&self) -> io::Result<bool> {
209        self.inner.broadcast()
210    }
211
212    #[cfg(unix)]
213    pub fn as_raw_fd(&self) -> std::os::unix::io::RawFd {
214        use std::os::fd::AsRawFd;
215        self.inner.as_raw_fd()
216    }
217
218    #[cfg(windows)]
219    pub fn as_raw_socket(&self) -> std::os::windows::io::RawSocket {
220        use std::os::windows::io::AsRawSocket;
221        self.inner.as_raw_socket()
222    }
223}