nuclei/net/
socket.rs

1//! [`TcpSocket`]: a TCP socket made, its options set and bound before it
2//! becomes a [`TcpStream`](super::TcpStream) (`connect`) or a
3//! [`TcpListener`](super::TcpListener) (`listen`).
4//!
5//! A builder is its descriptor and nothing else: no allocation, each
6//! setting one system call on it.
7
8use core::mem::ManuallyDrop;
9use std::io;
10use std::net::SocketAddr;
11use std::time::Duration;
12
13use super::opt;
14use crate::io::raw::{Owned, Raw};
15use crate::io::source::Kind;
16use crate::io::{Handle, sys};
17
18/// A TCP socket to set up before it connects or listens.
19///
20/// ```
21/// use nuclei::net::TcpSocket;
22///
23/// # nuclei::block_on(async {
24/// let socket = TcpSocket::new_v4()?;
25/// socket.set_reuseaddr(true)?;
26/// socket.bind("127.0.0.1:0".parse().unwrap())?;
27/// let addr = socket.local_addr()?;
28/// let listener = socket.listen(128)?;
29///
30/// let client = TcpSocket::new_v4()?;
31/// client.set_nodelay(true)?;
32/// let stream = client.connect(addr).await?;
33/// let (_peer, _) = listener.accept().await?;
34/// assert!(stream.nodelay()?);
35/// # std::io::Result::Ok(())
36/// # }).unwrap();
37/// ```
38pub struct TcpSocket {
39    socket: Owned,
40    /// Made non-blocking by nuclei; one taken from elsewhere is set so when
41    /// it connects or listens.
42    nonblocking: bool,
43}
44
45super::sockopt::socket_options!(TcpSocket);
46
47impl TcpSocket {
48    /// A new IPv4 TCP socket.
49    ///
50    /// # Errors
51    ///
52    /// When the system has no socket to give.
53    pub fn new_v4() -> io::Result<TcpSocket> {
54        Self::new(false)
55    }
56
57    /// A new IPv6 TCP socket.
58    ///
59    /// # Errors
60    ///
61    /// When the system has no socket to give.
62    pub fn new_v6() -> io::Result<TcpSocket> {
63        Self::new(true)
64    }
65
66    fn new(v6: bool) -> io::Result<TcpSocket> {
67        Ok(TcpSocket {
68            socket: super::tcp::new_socket(v6)?,
69            nonblocking: true,
70        })
71    }
72
73    pub(super) fn raw(&self) -> Raw {
74        raw_of(&self.socket)
75    }
76
77    /// The socket, seen as std's stream for the calls std makes.
78    fn as_std(&self) -> ManuallyDrop<std::net::TcpStream> {
79        // SAFETY: the socket is ours and outlives the view, which never
80        // closes it.
81        ManuallyDrop::new(unsafe { std_from_raw(self.raw()) })
82    }
83
84    /// Sets `SO_KEEPALIVE`: the system probes an idle connection.
85    pub fn set_keepalive(&self, keepalive: bool) -> io::Result<()> {
86        opt::set_flag(self.raw(), opt::SOL_SOCKET, opt::SO_KEEPALIVE, keepalive)
87    }
88
89    /// Whether `SO_KEEPALIVE` is set.
90    pub fn keepalive(&self) -> io::Result<bool> {
91        opt::flag(self.raw(), opt::SOL_SOCKET, opt::SO_KEEPALIVE)
92    }
93
94    /// Sets `SO_LINGER`: with `Some(time)` a close waits up to `time` (whole
95    /// seconds) for unsent data; with `None` it returns at once and the
96    /// system sends what is left.
97    pub fn set_linger(&self, linger: Option<Duration>) -> io::Result<()> {
98        opt::set_linger_time(self.raw(), linger)
99    }
100
101    /// Sets `SO_LINGER` to zero: a close drops unsent data and resets the
102    /// connection.
103    pub fn set_zero_linger(&self) -> io::Result<()> {
104        opt::set_linger_time(self.raw(), Some(Duration::ZERO))
105    }
106
107    /// `SO_LINGER`: how long a close waits for unsent data, `None` if it
108    /// does not wait.
109    pub fn linger(&self) -> io::Result<Option<Duration>> {
110        opt::linger_time(self.raw())
111    }
112
113    /// Sets `TCP_NODELAY`: small writes go out at once, not gathered.
114    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
115        self.as_std().set_nodelay(nodelay)
116    }
117
118    /// Whether `TCP_NODELAY` is set.
119    pub fn nodelay(&self) -> io::Result<bool> {
120        self.as_std().nodelay()
121    }
122
123    /// The address the socket is bound to.
124    pub fn local_addr(&self) -> io::Result<SocketAddr> {
125        self.as_std().local_addr()
126    }
127
128    /// The socket's pending error (`SO_ERROR`), taken.
129    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
130        self.as_std().take_error()
131    }
132
133    /// Binds the socket to `addr`.
134    ///
135    /// # Errors
136    ///
137    /// As `bind(2)`: the address is in use, or not this host's.
138    pub fn bind(&self, addr: SocketAddr) -> io::Result<()> {
139        sys::bind_inet(self.raw(), &addr)
140    }
141
142    /// Connects the socket to `addr`, as a stream.
143    ///
144    /// # Errors
145    ///
146    /// As `connect(2)`: refused, unreachable, or timed out.
147    pub async fn connect(self, addr: SocketAddr) -> io::Result<super::TcpStream> {
148        let TcpSocket {
149            socket,
150            nonblocking,
151        } = self;
152        super::tcp::connect_socket(socket, nonblocking, addr).await
153    }
154
155    /// Makes the socket a listener, with room for `backlog` connections not
156    /// accepted yet.
157    ///
158    /// # Errors
159    ///
160    /// As `listen(2)`; a socket not bound gets an address of the system's
161    /// choosing on unix, and an error on Windows.
162    pub fn listen(self, backlog: u32) -> io::Result<super::TcpListener> {
163        sys::listen(self.raw(), backlog)?;
164        let listener = std::net::TcpListener::from(self.socket);
165        if self.nonblocking {
166            Ok(Handle::from_nonblocking(listener, Kind::TcpListener))
167        } else {
168            Handle::adopt(listener, Kind::TcpListener)
169        }
170    }
171
172    /// Takes std's stream over as a socket to set up, connected or not.
173    pub fn from_std_stream(stream: std::net::TcpStream) -> TcpSocket {
174        TcpSocket {
175            socket: Owned::from(stream),
176            nonblocking: false,
177        }
178    }
179}
180
181impl core::fmt::Debug for TcpSocket {
182    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
183        f.debug_struct("TcpSocket")
184            .field("socket", &self.raw())
185            .finish_non_exhaustive()
186    }
187}
188
189#[cfg(unix)]
190pub(super) fn raw_of(socket: &Owned) -> Raw {
191    use std::os::fd::AsRawFd;
192    socket.as_raw_fd()
193}
194
195#[cfg(windows)]
196pub(super) fn raw_of(socket: &Owned) -> Raw {
197    use std::os::windows::io::AsRawSocket;
198    socket.as_raw_socket()
199}
200
201/// std's socket `T` over the socket `raw`.
202///
203/// # Safety
204///
205/// `raw` is an open socket of `T`'s kind, which the caller keeps from being
206/// closed twice.
207#[cfg(unix)]
208pub(super) unsafe fn std_from_raw<T>(raw: Raw) -> T
209where
210    T: std::os::fd::FromRawFd,
211{
212    // SAFETY: the caller's contract.
213    unsafe { T::from_raw_fd(raw) }
214}
215
216/// std's socket `T` over the socket `raw`.
217///
218/// # Safety
219///
220/// `raw` is an open socket of `T`'s kind, which the caller keeps from being
221/// closed twice.
222#[cfg(windows)]
223pub(super) unsafe fn std_from_raw<T>(raw: Raw) -> T
224where
225    T: std::os::windows::io::FromRawSocket,
226{
227    // SAFETY: the caller's contract.
228    unsafe { T::from_raw_socket(raw) }
229}
230
231#[cfg(unix)]
232mod fd {
233    use std::os::fd::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
234
235    use super::TcpSocket;
236
237    impl AsRawFd for TcpSocket {
238        fn as_raw_fd(&self) -> RawFd {
239            self.socket.as_raw_fd()
240        }
241    }
242
243    impl AsFd for TcpSocket {
244        fn as_fd(&self) -> BorrowedFd<'_> {
245            self.socket.as_fd()
246        }
247    }
248
249    impl FromRawFd for TcpSocket {
250        /// # Safety
251        ///
252        /// `fd` is an open TCP socket the caller owns and hands over.
253        unsafe fn from_raw_fd(fd: RawFd) -> TcpSocket {
254            TcpSocket {
255                // SAFETY: the caller's contract.
256                socket: unsafe { OwnedFd::from_raw_fd(fd) },
257                nonblocking: false,
258            }
259        }
260    }
261
262    impl IntoRawFd for TcpSocket {
263        fn into_raw_fd(self) -> RawFd {
264            self.socket.into_raw_fd()
265        }
266    }
267}
268
269#[cfg(windows)]
270mod raw_socket {
271    use std::os::windows::io::{
272        AsRawSocket, AsSocket, BorrowedSocket, FromRawSocket, IntoRawSocket, OwnedSocket, RawSocket,
273    };
274
275    use super::TcpSocket;
276
277    impl AsRawSocket for TcpSocket {
278        fn as_raw_socket(&self) -> RawSocket {
279            self.socket.as_raw_socket()
280        }
281    }
282
283    impl AsSocket for TcpSocket {
284        fn as_socket(&self) -> BorrowedSocket<'_> {
285            self.socket.as_socket()
286        }
287    }
288
289    impl FromRawSocket for TcpSocket {
290        /// # Safety
291        ///
292        /// `socket` is an open, overlapped TCP socket the caller owns and
293        /// hands over.
294        unsafe fn from_raw_socket(socket: RawSocket) -> TcpSocket {
295            TcpSocket {
296                // SAFETY: the caller's contract.
297                socket: unsafe { OwnedSocket::from_raw_socket(socket) },
298                nonblocking: false,
299            }
300        }
301    }
302
303    impl IntoRawSocket for TcpSocket {
304        fn into_raw_socket(self) -> RawSocket {
305            self.socket.into_raw_socket()
306        }
307    }
308}
nuclei/net/socket.rs

nuclei/net/
socket.rs
