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    /// Sends some data to the socket from the buffer, returning the original
223    /// buffer and quantity of data sent.
224    pub async fn send<T: IoBuf>(&self, buffer: T) -> BufResult<usize, T> {
225        self.inner.send(buffer, 0).await
226    }
227
228    /// Sends some data to the socket from the buffer, returning the original
229    /// buffer and quantity of data sent.
230    pub async fn send_vectored<T: IoVectoredBuf>(&self, buffer: T) -> BufResult<usize, T> {
231        self.inner.send_vectored(buffer, 0).await
232    }
233
234    /// Receives a single datagram message on the socket. On success, returns
235    /// the number of bytes received and the origin.
236    pub async fn recv_from<T: IoBufMut>(&self, buffer: T) -> BufResult<(usize, SocketAddr), T> {
237        self.inner
238            .recv_from(buffer, 0)
239            .await
240            .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
241    }
242
243    /// Receives a single datagram message on the socket. On success, returns
244    /// the number of bytes received and the origin.
245    pub async fn recv_from_vectored<T: IoVectoredBufMut>(
246        &self,
247        buffer: T,
248    ) -> BufResult<(usize, SocketAddr), T> {
249        self.inner
250            .recv_from_vectored(buffer, 0)
251            .await
252            .map_res(|(n, addr)| (n, addr.as_socket().expect("should be SocketAddr")))
253    }
254
255    /// Receives a single datagram message and ancillary data on the socket. On
256    /// success, returns the number of bytes received and the origin.
257    pub async fn recv_msg<T: IoBufMut, C: IoBufMut>(
258        &self,
259        buffer: T,
260        control: C,
261    ) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
262        self.inner
263            .recv_msg(buffer, control, 0)
264            .await
265            .map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
266    }
267
268    /// Receives a single datagram message and ancillary data on the socket. On
269    /// success, returns the number of bytes received and the origin.
270    pub async fn recv_msg_vectored<T: IoVectoredBufMut, C: IoBufMut>(
271        &self,
272        buffer: T,
273        control: C,
274    ) -> BufResult<(usize, usize, SocketAddr), (T, C)> {
275        self.inner
276            .recv_msg_vectored(buffer, control, 0)
277            .await
278            .map_res(|(n, m, addr)| (n, m, addr.as_socket().expect("should be SocketAddr")))
279    }
280
281    /// Sends data on the socket to the given address. On success, returns the
282    /// number of bytes sent.
283    pub async fn send_to<T: IoBuf>(
284        &self,
285        buffer: T,
286        addr: impl ToSocketAddrsAsync,
287    ) -> BufResult<usize, T> {
288        super::first_addr_buf(addr, buffer, |addr, buffer| async move {
289            self.inner.send_to(buffer, &SockAddr::from(addr), 0).await
290        })
291        .await
292    }
293
294    /// Sends data on the socket to the given address. On success, returns the
295    /// number of bytes sent.
296    pub async fn send_to_vectored<T: IoVectoredBuf>(
297        &self,
298        buffer: T,
299        addr: impl ToSocketAddrsAsync,
300    ) -> BufResult<usize, T> {
301        super::first_addr_buf(addr, buffer, |addr, buffer| async move {
302            self.inner
303                .send_to_vectored(buffer, &SockAddr::from(addr), 0)
304                .await
305        })
306        .await
307    }
308
309    /// Sends data on the socket to the given address accompanied by ancillary
310    /// data. On success, returns the number of bytes sent.
311    pub async fn send_msg<T: IoBuf, C: IoBuf>(
312        &self,
313        buffer: T,
314        control: C,
315        addr: impl ToSocketAddrsAsync,
316    ) -> BufResult<usize, (T, C)> {
317        super::first_addr_buf(
318            addr,
319            (buffer, control),
320            |addr, (buffer, control)| async move {
321                self.inner
322                    .send_msg(buffer, control, &SockAddr::from(addr), 0)
323                    .await
324            },
325        )
326        .await
327    }
328
329    /// Sends data on the socket to the given address accompanied by ancillary
330    /// data. On success, returns the number of bytes sent.
331    pub async fn send_msg_vectored<T: IoVectoredBuf, C: IoBuf>(
332        &self,
333        buffer: T,
334        control: C,
335        addr: impl ToSocketAddrsAsync,
336    ) -> BufResult<usize, (T, C)> {
337        super::first_addr_buf(
338            addr,
339            (buffer, control),
340            |addr, (buffer, control)| async move {
341                self.inner
342                    .send_msg_vectored(buffer, control, &SockAddr::from(addr), 0)
343                    .await
344            },
345        )
346        .await
347    }
348
349    /// Gets a socket option.
350    ///
351    /// # Safety
352    ///
353    /// The caller must ensure `T` is the correct type for `level` and `name`.
354    pub unsafe fn get_socket_option<T: Copy>(&self, level: i32, name: i32) -> io::Result<T> {
355        unsafe { self.inner.get_socket_option(level, name) }
356    }
357
358    /// Sets a socket option.
359    ///
360    /// # Safety
361    ///
362    /// The caller must ensure `T` is the correct type for `level` and `name`.
363    pub unsafe fn set_socket_option<T: Copy>(
364        &self,
365        level: i32,
366        name: i32,
367        value: &T,
368    ) -> io::Result<()> {
369        unsafe { self.inner.set_socket_option(level, name, value) }
370    }
371}
372
373impl_raw_fd!(UdpSocket, socket2::Socket, inner, socket);
compio_net/udp.rs

compio_net/
udp.rs
