Skip to main content

nuclei/net/
udp_socket.rs

1//! [`UdpSocketBuilder`]: a UDP socket made and its options set before it binds and becomes a
2//! [`UdpSocket`](super::UdpSocket).
3//!
4//! As [`TcpSocket`](super::TcpSocket), it is its descriptor alone, each setting one system call
5//! on it.
6
7use core::mem::ManuallyDrop;
8use std::io;
9use std::net::SocketAddr;
10
11use super::opt;
12use super::socket::{raw_of, std_from_raw};
13use crate::io::raw::{Owned, Raw};
14use crate::io::source::Kind;
15use crate::io::{Handle, sys};
16
17/// A UDP socket to set up before it binds: what only a socket not bound yet takes
18/// (`SO_REUSEPORT` among sockets that share a port, `IPV6_V6ONLY`, a device), and the rest
19/// before the first datagram.
20///
21/// ```
22/// use nuclei::net::UdpSocketBuilder;
23///
24/// # nuclei::block_on(async {
25/// let builder = UdpSocketBuilder::new_v4()?;
26/// builder.set_reuseaddr(true)?;
27/// builder.set_recv_buffer_size(256 * 1024)?;
28/// let socket = builder.bind("127.0.0.1:0".parse().unwrap())?;
29/// let addr = socket.local_addr()?;
30/// socket.send_to(b"hi", addr).await?;
31/// let mut buf = [0u8; 2];
32/// assert_eq!(socket.recv_from(&mut buf).await?, (2, addr));
33/// # std::io::Result::Ok(())
34/// # }).unwrap();
35/// ```
36pub struct UdpSocketBuilder {
37    socket: Owned,
38}
39
40super::sockopt::socket_options!(UdpSocketBuilder);
41
42impl UdpSocketBuilder {
43    /// A new IPv4 UDP socket.
44    ///
45    /// # Errors
46    ///
47    /// When the system has no socket to give.
48    pub fn new_v4() -> io::Result<UdpSocketBuilder> {
49        Self::new(false)
50    }
51
52    /// A new IPv6 UDP socket.
53    ///
54    /// # Errors
55    ///
56    /// When the system has no socket to give.
57    pub fn new_v6() -> io::Result<UdpSocketBuilder> {
58        Self::new(true)
59    }
60
61    fn new(v6: bool) -> io::Result<UdpSocketBuilder> {
62        #[cfg(unix)]
63        let socket = sys::socket(
64            if v6 { libc::AF_INET6 } else { libc::AF_INET },
65            libc::SOCK_DGRAM,
66        )?;
67        #[cfg(windows)]
68        let socket = sys::udp_socket(if v6 {
69            crate::syscore::windows::AF_INET6
70        } else {
71            crate::syscore::windows::AF_INET
72        })?;
73        Ok(UdpSocketBuilder { socket })
74    }
75
76    fn raw(&self) -> Raw {
77        raw_of(&self.socket)
78    }
79
80    /// The socket, seen as std's for the calls std makes.
81    fn as_std(&self) -> ManuallyDrop<std::net::UdpSocket> {
82        // SAFETY: the socket is ours and outlives the view, which never
83        // closes it.
84        ManuallyDrop::new(unsafe { std_from_raw(self.raw()) })
85    }
86
87    /// The address the socket is bound to; before it binds, the unspecified one on unix, and an
88    /// error (`InvalidInput`) on Windows, as std's `local_addr` answers there.
89    pub fn local_addr(&self) -> io::Result<SocketAddr> {
90        self.as_std().local_addr()
91    }
92
93    /// The socket's pending error (`SO_ERROR`), taken.
94    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
95        self.as_std().take_error()
96    }
97
98    /// Binds the socket to `addr`, and makes it a [`UdpSocket`](super::UdpSocket).
99    ///
100    /// # Errors
101    ///
102    /// As `bind(2)`: the address is in use (shared only by sockets that all set
103    /// `SO_REUSEPORT`, or on Windows `SO_REUSEADDR`), or not this host's.
104    pub fn bind(self, addr: SocketAddr) -> io::Result<super::UdpSocket> {
105        sys::bind_inet(self.raw(), &addr)?;
106        let socket = std::net::UdpSocket::from(self.socket);
107        Ok(Handle::from_nonblocking(socket, Kind::Udp))
108    }
109}
110
111impl core::fmt::Debug for UdpSocketBuilder {
112    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
113        f.debug_struct("UdpSocketBuilder")
114            .field("socket", &self.raw())
115            .finish_non_exhaustive()
116    }
117}
118
119#[cfg(unix)]
120impl std::os::fd::AsRawFd for UdpSocketBuilder {
121    fn as_raw_fd(&self) -> std::os::fd::RawFd {
122        self.raw()
123    }
124}
125
126#[cfg(unix)]
127impl std::os::fd::AsFd for UdpSocketBuilder {
128    fn as_fd(&self) -> std::os::fd::BorrowedFd<'_> {
129        std::os::fd::AsFd::as_fd(&self.socket)
130    }
131}
132
133#[cfg(windows)]
134impl std::os::windows::io::AsRawSocket for UdpSocketBuilder {
135    fn as_raw_socket(&self) -> std::os::windows::io::RawSocket {
136        self.raw()
137    }
138}
139
140#[cfg(windows)]
141impl std::os::windows::io::AsSocket for UdpSocketBuilder {
142    fn as_socket(&self) -> std::os::windows::io::BorrowedSocket<'_> {
143        std::os::windows::io::AsSocket::as_socket(&self.socket)
144    }
145}