compio_net/udp.rs
1use std::{future::Future, io, net::SocketAddr};
2
3use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
4use compio_driver::impl_raw_fd;
5use compio_runtime::{BorrowedBuffer, BufferPool};
6use socket2::{Protocol, SockAddr, Socket as Socket2, Type};
7
8use crate::{Socket, SocketOpts, ToSocketAddrsAsync};
9
10/// A UDP socket.
11///
12/// UDP is "connectionless", unlike TCP. Meaning, regardless of what address
13/// you've bound to, a `UdpSocket` is free to communicate with many different
14/// remotes. There are basically two main ways to use `UdpSocket`:
15///
16/// * one to many: [`bind`](`UdpSocket::bind`) and use
17/// [`send_to`](`UdpSocket::send_to`) and
18/// [`recv_from`](`UdpSocket::recv_from`) to communicate with many different
19/// addresses
20/// * one to one: [`connect`](`UdpSocket::connect`) and associate with a single
21/// address, using [`send`](`UdpSocket::send`) and [`recv`](`UdpSocket::recv`)
22/// to communicate only with that remote address
23///
24/// # Examples
25/// Bind and connect a pair of sockets and send a packet:
26///
27/// ```
28/// use std::net::SocketAddr;
29///
30/// use compio_net::UdpSocket;
31///
32/// # compio_runtime::Runtime::new().unwrap().block_on(async {
33/// let first_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
34/// let second_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
35///
36/// // bind sockets
37/// let mut socket = UdpSocket::bind(first_addr).await.unwrap();
38/// let first_addr = socket.local_addr().unwrap();
39/// let mut other_socket = UdpSocket::bind(second_addr).await.unwrap();
40/// let second_addr = other_socket.local_addr().unwrap();
41///
42/// // connect sockets
43/// socket.connect(second_addr).await.unwrap();
44/// other_socket.connect(first_addr).await.unwrap();
45///
46/// let buf = Vec::with_capacity(12);
47///
48/// // write data
49/// socket.send("Hello world!").await.unwrap();
50///
51/// // read data
52/// let (n_bytes, buf) = other_socket.recv(buf).await.unwrap();
53///
54/// assert_eq!(n_bytes, buf.len());
55/// assert_eq!(buf, b"Hello world!");
56/// # });
57/// ```
58/// Send and receive packets without connecting:
59///
60/// ```
61/// use std::net::SocketAddr;
62///
63/// use compio_net::UdpSocket;
64/// use socket2::SockAddr;
65///
66/// # compio_runtime::Runtime::new().unwrap().block_on(async {
67/// let first_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
68/// let second_addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
69///
70/// // bind sockets
71/// let mut socket = UdpSocket::bind(first_addr).await.unwrap();
72/// let first_addr = socket.local_addr().unwrap();
73/// let mut other_socket = UdpSocket::bind(second_addr).await.unwrap();
74/// let second_addr = other_socket.local_addr().unwrap();
75///
76/// let buf = Vec::with_capacity(32);
77///
78/// // write data
79/// socket.send_to("hello world", second_addr).await.unwrap();
80///
81/// // read data
82/// let ((n_bytes, addr), buf) = other_socket.recv_from(buf).await.unwrap();
83///
84/// assert_eq!(addr, first_addr);
85/// assert_eq!(n_bytes, buf.len());
86/// assert_eq!(buf, b"hello world");
87/// # });
88/// ```
89#[derive(Debug, Clone)]
90pub struct UdpSocket {
91 inner: Socket,
92}
93
94impl UdpSocket {
95 /// Creates a new UDP socket and attempt to bind it to the addr provided.
96 pub async fn bind(addr: impl ToSocketAddrsAsync) -> io::Result<Self> {
97 Self::bind_with_options(addr, &SocketOpts::default()).await
98 }
99
100 /// Creates a new UDP socket with [`SocketOpts`] and attempt to bind it to
101 /// the addr provided.
102 pub async fn bind_with_options(
103 addr: impl ToSocketAddrsAsync,
104 opts: &SocketOpts,
105 ) -> io::Result<Self> {
106 super::each_addr(addr, |addr| async move {
107 let socket =
108 Socket::bind(&SockAddr::from(addr), Type::DGRAM, Some(Protocol::UDP)).await?;
109 opts.setup_socket(&socket)?;
110 Ok(Self { inner: socket })
111 })
112 .await
113 }
114
115 /// Connects this UDP socket to a remote address, allowing the `send` and
116 /// `recv` to be used to send data and also applies filters to only
117 /// receive data from the specified address.
118 ///
119 /// Note that usually, a successful `connect` call does not specify
120 /// that there is a remote server listening on the port, rather, such an
121 /// error would only be detected after the first send.
122 pub async fn connect(&self, addr: impl ToSocketAddrsAsync) -> io::Result<()> {
123 super::each_addr(addr, |addr| async move {
124 self.inner.connect(&SockAddr::from(addr))
125 })
126 .await
127 }
128
129 /// Creates new UdpSocket from a std::net::UdpSocket.
130 pub fn from_std(socket: std::net::UdpSocket) -> io::Result<Self> {
131 Ok(Self {
132 inner: Socket::from_socket2(Socket2::from(socket))?,
133 })
134 }
135
136 /// Close the socket. If the returned future is dropped before polling, the
137 /// socket won't be closed.
138 pub fn close(self) -> impl Future<Output = io::Result<()>> {
139 self.inner.close()
140 }
141
142 /// Returns the socket address of the remote peer this socket was connected
143 /// to.
144 ///
145 /// # Examples
146 ///
147 /// ```no_run
148 /// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
149 ///
150 /// use compio_net::UdpSocket;
151 /// use socket2::SockAddr;
152 ///
153 /// # compio_runtime::Runtime::new().unwrap().block_on(async {
154 /// let socket = UdpSocket::bind("127.0.0.1:34254")
155 /// .await
156 /// .expect("couldn't bind to address");
157 /// socket
158 /// .connect("192.168.0.1:41203")
159 /// .await
160 /// .expect("couldn't connect to address");
161 /// assert_eq!(
162 /// socket.peer_addr().unwrap(),
163 /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(192, 168, 0, 1), 41203))
164 /// );
165 /// # });
166 /// ```
167 pub fn peer_addr(&self) -> io::Result<SocketAddr> {
168 self.inner
169 .peer_addr()
170 .map(|addr| addr.as_socket().expect("should be SocketAddr"))
171 }
172
173 /// Returns the local address that this socket is bound to.
174 ///
175 /// # Example
176 ///
177 /// ```
178 /// use std::net::SocketAddr;
179 ///
180 /// use compio_net::UdpSocket;
181 /// use socket2::SockAddr;
182 ///
183 /// # compio_runtime::Runtime::new().unwrap().block_on(async {
184 /// let addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
185 /// let sock = UdpSocket::bind(&addr).await.unwrap();
186 /// // the address the socket is bound to
187 /// let local_addr = sock.local_addr().unwrap();
188 /// assert_eq!(local_addr, addr);
189 /// # });
190 /// ```
191 pub fn local_addr(&self) -> io::Result<SocketAddr> {
192 self.inner
193 .local_addr()
194 .map(|addr| addr.as_socket().expect("should be SocketAddr"))
195 }
196
197 /// Receives a packet of data from the socket into the buffer, returning the
198 /// original buffer and quantity of data received.
199 pub async fn recv<T: IoBufMut>(&self, buffer: T) -> BufResult<usize, T> {
200 self.inner.recv(buffer, 0).await
201 }
202
203 /// Receives a packet of data from the socket into the buffer, returning the
204 /// original buffer and quantity of data received.
205 pub async fn recv_vectored<T: IoVectoredBufMut>(&self, buffer: T) -> BufResult<usize, T> {
206 self.inner.recv_vectored(buffer, 0).await
207 }
208
209 /// Read some bytes from this source with [`BufferPool`] and return
210 /// a [`BorrowedBuffer`].
211 ///
212 /// If `len` == 0, will use [`BufferPool`] inner buffer size as the max len,
213 /// if `len` > 0, `min(len, inner buffer size)` will be the read max len
214 pub async fn recv_managed<'a>(
215 &self,
216 buffer_pool: &'a BufferPool,
217 len: usize,
218 ) -> io::Result<BorrowedBuffer<'a>> {
219 self.inner.recv_managed(buffer_pool, len, 0).await
220 }
221
222 /// Read some bytes from this source with [`BufferPool`] and return
223 /// a [`BorrowedBuffer`] with the sender address.
224 ///
225 /// If `len` == 0, will use [`BufferPool`] inner buffer size as the max len,
226 /// if `len` > 0, `min(len, inner buffer size)` will be the read max len
227 pub async fn recv_from_managed<'a>(
228 &self,
229 buffer_pool: &'a BufferPool,
230 len: usize,
231 ) -> io::Result<(BorrowedBuffer<'a>, SocketAddr)> {
232 self.inner
233 .recv_from_managed(buffer_pool, len, 0)
234 .await
235 .map(|(buffer, addr)| (buffer, addr.as_socket().expect("should be SocketAddr")))
236 }
237
238 /// Sends some data to the socket from the buffer, returning the original
239 /// buffer and quantity of data sent.
240 pub async fn send<T: IoBuf>(&self, buffer: T) -> BufResult<usize, T> {
241 self.inner.send(buffer, 0).await
242 }
243
244 /// Sends some data to the socket from the buffer, returning the original
245 /// buffer and quantity of data sent.
246 pub async fn send_vectored<T: IoVectoredBuf>(&self, buffer: T) -> BufResult<usize, T> {
247 self.inner.send_vectored(buffer, 0).await
248 }
249
250 /// Receives a single datagram message on the socket. On success, returns
251 /// the number of bytes received and the origin.
252 pub async fn recv_from<T: IoBufMut>(&self, buffer: T) -> BufResult<(usize, SocketAddr), T> {
253 self.inner
254 .recv_from(buffer, 0)
255 .await
256 .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
257 }
258
259 /// Receives a single datagram message on the socket. On success, returns
260 /// the number of bytes received and the origin.
261 pub async fn recv_from_vectored<T: IoVectoredBufMut>(
262 &self,
263 buffer: T,
264 ) -> BufResult<(usize, SocketAddr), T> {
265 self.inner
266 .recv_from_vectored(buffer, 0)
267 .await
268 .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
269 }
270
271 /// Receives a single datagram message and ancillary data on the socket. On
272 /// success, returns the number of bytes received and the origin.
273 pub async fn recv_msg<T: IoBufMut, C: IoBufMut>(
274 &self,
275 buffer: T,
276 control: C,
277 ) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
278 self.inner
279 .recv_msg(buffer, control, 0)
280 .await
281 .map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
282 }
283
284 /// Receives a single datagram message and ancillary data on the socket. On
285 /// success, returns the number of bytes received and the origin.
286 pub async fn recv_msg_vectored<T: IoVectoredBufMut, C: IoBufMut>(
287 &self,
288 buffer: T,
289 control: C,
290 ) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
291 self.inner
292 .recv_msg_vectored(buffer, control, 0)
293 .await
294 .map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
295 }
296
297 /// Sends data on the socket to the given address. On success, returns the
298 /// number of bytes sent.
299 pub async fn send_to<T: IoBuf>(
300 &self,
301 buffer: T,
302 addr: impl ToSocketAddrsAsync,
303 ) -> BufResult<usize, T> {
304 super::first_addr_buf(addr, buffer, |addr, buffer| async move {
305 self.inner.send_to(buffer, &SockAddr::from(addr), 0).await
306 })
307 .await
308 }
309
310 /// Sends data on the socket to the given address. On success, returns the
311 /// number of bytes sent.
312 pub async fn send_to_vectored<T: IoVectoredBuf>(
313 &self,
314 buffer: T,
315 addr: impl ToSocketAddrsAsync,
316 ) -> BufResult<usize, T> {
317 super::first_addr_buf(addr, buffer, |addr, buffer| async move {
318 self.inner
319 .send_to_vectored(buffer, &SockAddr::from(addr), 0)
320 .await
321 })
322 .await
323 }
324
325 /// Sends data on the socket to the given address accompanied by ancillary
326 /// data. On success, returns the number of bytes sent.
327 pub async fn send_msg<T: IoBuf, C: IoBuf>(
328 &self,
329 buffer: T,
330 control: C,
331 addr: impl ToSocketAddrsAsync,
332 ) -> BufResult<usize, (T, C)> {
333 super::first_addr_buf(
334 addr,
335 (buffer, control),
336 |addr, (buffer, control)| async move {
337 self.inner
338 .send_msg(buffer, control, &SockAddr::from(addr), 0)
339 .await
340 },
341 )
342 .await
343 }
344
345 /// Sends data on the socket to the given address accompanied by ancillary
346 /// data. On success, returns the number of bytes sent.
347 pub async fn send_msg_vectored<T: IoVectoredBuf, C: IoBuf>(
348 &self,
349 buffer: T,
350 control: C,
351 addr: impl ToSocketAddrsAsync,
352 ) -> BufResult<usize, (T, C)> {
353 super::first_addr_buf(
354 addr,
355 (buffer, control),
356 |addr, (buffer, control)| async move {
357 self.inner
358 .send_msg_vectored(buffer, control, &SockAddr::from(addr), 0)
359 .await
360 },
361 )
362 .await
363 }
364
365 /// Gets a socket option.
366 ///
367 /// # Safety
368 ///
369 /// The caller must ensure `T` is the correct type for `level` and `name`.
370 pub unsafe fn get_socket_option<T: Copy>(&self, level: i32, name: i32) -> io::Result<T> {
371 unsafe { self.inner.get_socket_option(level, name) }
372 }
373
374 /// Sets a socket option.
375 ///
376 /// # Safety
377 ///
378 /// The caller must ensure `T` is the correct type for `level` and `name`.
379 pub unsafe fn set_socket_option<T: Copy>(
380 &self,
381 level: i32,
382 name: i32,
383 value: &T,
384 ) -> io::Result<()> {
385 unsafe { self.inner.set_socket_option(level, name, value) }
386 }
387}
388
389impl_raw_fd!(UdpSocket, socket2::Socket, inner, socket);