nuclei/net/
tcp.rs

1//! TCP listeners and streams.
2
3use core::future::{Future, poll_fn};
4use core::pin::Pin;
5use core::task::{Context, Poll};
6use std::io;
7use std::net::{SocketAddr, TcpListener, TcpStream};
8
9use futures_core::Stream;
10
11use super::addr::{ToSocketAddrs, resolve};
12use super::none_worked;
13use std::time::Duration;
14
15use crate::io::IoObject;
16use crate::io::raw::Owned;
17use crate::io::source::write::Peer;
18use crate::io::source::{Addr, Kind};
19use crate::io::{Handle, sys};
20use crate::rt::counters::out_of_descriptors;
21use crate::time::{Sleep, sleep};
22
23impl Handle<TcpListener> {
24    /// A listener bound to the first of the addresses `addr` names that
25    /// it can bind.
26    ///
27    /// # Errors
28    ///
29    /// The last address's error when none could be bound, or the lookup's.
30    pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<Handle<TcpListener>> {
31        let mut last = None;
32        for addr in resolve(&addr).await? {
33            match TcpListener::bind(addr) {
34                Ok(listener) => return Handle::adopt(listener, Kind::TcpListener),
35                Err(err) => last = Some(err),
36            }
37        }
38        Err(none_worked(last))
39    }
40
41    /// The next connection, with its peer's address.
42    ///
43    /// On macOS an accept that finds the process out of descriptors closes
44    /// the connection it took off the backlog, as the system's own does; the
45    /// next ones fail with the same error and take none until a descriptor
46    /// is free.
47    ///
48    /// # Errors
49    ///
50    /// As `accept(2)`: a connection reset before it was taken, or the
51    /// process out of descriptors.
52    pub async fn accept(&self) -> io::Result<(Handle<TcpStream>, SocketAddr)> {
53        poll_fn(|cx| self.poll_accept(cx)).await
54    }
55
56    /// Polls for the next connection, with its peer's address: `Pending`
57    /// with `cx`'s waker kept until one comes.
58    ///
59    /// # Errors
60    ///
61    /// As [`accept`](Self::accept).
62    pub fn poll_accept(
63        &self,
64        cx: &mut Context<'_>,
65    ) -> Poll<io::Result<(Handle<TcpStream>, SocketAddr)>> {
66        let source = self.source();
67        source.poll_accept(cx).map(|accepted| {
68            let (fd, addr) = accepted?;
69            Ok((stream(fd, source), inet(addr)?))
70        })
71    }
72
73    /// A connection, without its address: on io_uring the listener's
74    /// accept stays armed for every connection that comes, taking them
75    /// ahead for this handle.
76    fn poll_accept_stream(&self, cx: &mut Context<'_>) -> Poll<io::Result<Handle<TcpStream>>> {
77        let source = self.source();
78        source
79            .poll_accept_stream(cx)
80            .map(|accepted| accepted.map(|fd| stream(fd, source)))
81    }
82
83    /// The connections as they come, without their addresses; the stream
84    /// never ends, and an error is one of its items.
85    pub fn incoming(
86        &self,
87    ) -> impl Stream<Item = io::Result<Handle<TcpStream>>> + Send + Unpin + '_ {
88        Incoming { listener: self }
89    }
90
91    /// The connections as they come, from a stream that owns its own copy of
92    /// the listener, so it may outlive this handle. A connection reset before
93    /// it was taken is skipped, and so is the process having no descriptor
94    /// left to accept into: the stream waits, with a growing pause, until
95    /// there is one, and the connections wait in the listener's backlog (on
96    /// macOS but for the one the first accept to find none took, which the
97    /// system closes, as [`accept`](Self::accept) says).
98    /// Any other error ends the stream, and
99    /// [`TcpStreamGenerator::take_error`] says what it was.
100    ///
101    /// On io_uring the stream keeps one accept armed for every connection
102    /// that comes, so it takes connections ahead of its polls; those it took
103    /// and never yielded close when it drops, as a listener's own backlog
104    /// does when the last copy of the listener closes.
105    ///
106    /// # Errors
107    ///
108    /// When the listener's descriptor cannot be copied.
109    pub async fn accept_multi(&self) -> io::Result<TcpStreamGenerator> {
110        let listener = self.get_ref().try_clone()?;
111        Ok(TcpStreamGenerator {
112            listener: Handle::adopt(listener, Kind::TcpListener)?,
113            error: None,
114            done: false,
115            backoff: None,
116            delay: OUT_OF_DESCRIPTORS_DELAY,
117        })
118    }
119}
120
121/// A connection `listener` accepted, set up already.
122fn stream(fd: crate::io::raw::Owned, listener: &crate::io::Source) -> Handle<TcpStream> {
123    #[cfg(any(target_os = "linux", target_os = "android"))]
124    if listener.passes_inq() {
125        return Handle::from_nonblocking_counting(TcpStream::from(fd), Kind::Tcp);
126    }
127    #[cfg(not(any(target_os = "linux", target_os = "android")))]
128    let _ = listener;
129    Handle::from_nonblocking(TcpStream::from(fd), Kind::Tcp)
130}
131
132fn inet(addr: Addr) -> io::Result<SocketAddr> {
133    match addr {
134        Addr::Inet(addr) => Ok(addr),
135        #[cfg(unix)]
136        Addr::Unix(_) => Err(io::Error::new(
137            io::ErrorKind::InvalidData,
138            "a TCP listener's peer has no internet address",
139        )),
140    }
141}
142
143struct Incoming<'a> {
144    listener: &'a Handle<TcpListener>,
145}
146
147impl Stream for Incoming<'_> {
148    type Item = io::Result<Handle<TcpStream>>;
149
150    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
151        self.listener
152            .poll_accept(cx)
153            .map(|accepted| Some(accepted.map(|(stream, _)| stream)))
154    }
155}
156
157/// The connections of a listener, as [`Handle::accept_multi`] makes them.
158#[derive(Debug)]
159pub struct TcpStreamGenerator {
160    listener: Handle<TcpListener>,
161    error: Option<io::Error>,
162    done: bool,
163    backoff: Option<Sleep>,
164    delay: Duration,
165}
166
167/// The first wait of a generator whose process has no descriptor to accept
168/// into, doubled up to [`OUT_OF_DESCRIPTORS_DELAY_MAX`] while it lasts.
169const OUT_OF_DESCRIPTORS_DELAY: Duration = Duration::from_millis(10);
170const OUT_OF_DESCRIPTORS_DELAY_MAX: Duration = Duration::from_secs(1);
171
172impl TcpStreamGenerator {
173    /// Why the stream ended, once it has.
174    pub fn take_error(&mut self) -> Option<io::Error> {
175        self.error.take()
176    }
177}
178
179impl Stream for TcpStreamGenerator {
180    type Item = Handle<TcpStream>;
181
182    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
183        let this = &mut *self;
184        while !this.done {
185            if let Some(backoff) = this.backoff.as_mut() {
186                if Pin::new(backoff).poll(cx).is_pending() {
187                    return Poll::Pending;
188                }
189                this.backoff = None;
190            }
191            match this.listener.poll_accept_stream(cx) {
192                Poll::Pending => return Poll::Pending,
193                Poll::Ready(Ok(stream)) => {
194                    this.delay = OUT_OF_DESCRIPTORS_DELAY;
195                    return Poll::Ready(Some(stream));
196                }
197                Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::ConnectionAborted => {}
198                Poll::Ready(Err(err)) if out_of_descriptors(&err) => {
199                    this.backoff = Some(sleep(this.delay));
200                    this.delay = (this.delay * 2).min(OUT_OF_DESCRIPTORS_DELAY_MAX);
201                }
202                Poll::Ready(Err(err)) => {
203                    this.error = Some(err);
204                    this.done = true;
205                }
206            }
207        }
208        Poll::Ready(None)
209    }
210}
211
212impl Handle<TcpStream> {
213    /// A stream connected to the first of `sock_addrs` that takes the
214    /// connection.
215    ///
216    /// # Errors
217    ///
218    /// The last address's error when none took it, or the lookup's.
219    pub async fn connect<A: ToSocketAddrs>(sock_addrs: A) -> io::Result<Handle<TcpStream>> {
220        let mut last = None;
221        for addr in resolve(&sock_addrs).await? {
222            match connect(addr).await {
223                Ok(stream) => return Ok(stream),
224                Err(err) => last = Some(err),
225            }
226        }
227        Err(none_worked(last))
228    }
229
230    /// Sends bytes; returns how many were taken (into the write-behind, as
231    /// for [`AsyncWrite`]).
232    ///
233    /// [`AsyncWrite`]: futures_io::AsyncWrite
234    ///
235    /// # Errors
236    ///
237    /// As `send(2)`, or the error that stopped earlier bytes.
238    pub async fn send(&self, buf: &[u8]) -> io::Result<usize> {
239        poll_fn(|cx| self.source().poll_write(cx, buf)).await
240    }
241
242    /// Receives bytes; zero at the end of the stream.
243    ///
244    /// # Errors
245    ///
246    /// As `recv(2)`.
247    pub async fn recv(&self, buf: &mut [u8]) -> io::Result<usize> {
248        poll_fn(|cx| self.source().poll_read(cx, buf, false)).await
249    }
250
251    /// Looks at the bytes the next receive returns, without taking them.
252    ///
253    /// # Errors
254    ///
255    /// As `recv(2)` with `MSG_PEEK`.
256    pub async fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
257        poll_fn(|cx| self.source().poll_read(cx, buf, true)).await
258    }
259
260    /// Polls for the bytes the next receive returns, into the unfilled part
261    /// of `buf`, without taking them; `buf` advances by what came.
262    ///
263    /// # Errors
264    ///
265    /// As `recv(2)` with `MSG_PEEK`.
266    pub fn poll_peek(
267        &self,
268        cx: &mut Context<'_>,
269        buf: &mut crate::io::ReadBuf<'_>,
270    ) -> Poll<io::Result<usize>> {
271        let n = core::task::ready!(self.source().poll_read(cx, buf.initialize_unfilled(), true))?;
272        buf.advance(n);
273        Poll::Ready(Ok(n))
274    }
275
276    /// `SO_LINGER`: how long a close waits for unsent data, `None` if it
277    /// does not wait.
278    pub fn linger(&self) -> io::Result<Option<Duration>> {
279        super::opt::linger_time(self.raw_io())
280    }
281
282    /// Sets `SO_LINGER`: with `Some(time)` a close waits up to `time` (whole
283    /// seconds) for unsent data; with `None` it returns at once and the
284    /// system sends what is left.
285    pub fn set_linger(&self, linger: Option<Duration>) -> io::Result<()> {
286        super::opt::set_linger_time(self.raw_io(), linger)
287    }
288
289    /// Sets `SO_LINGER` to zero: the close drops unsent data and resets the
290    /// connection.
291    pub fn set_zero_linger(&self) -> io::Result<()> {
292        super::opt::set_linger_time(self.raw_io(), Some(Duration::ZERO))
293    }
294}
295
296/// A new non-blocking TCP socket, IPv6 or IPv4.
297pub(super) fn new_socket(v6: bool) -> io::Result<Owned> {
298    #[cfg(unix)]
299    return sys::socket(
300        if v6 { libc::AF_INET6 } else { libc::AF_INET },
301        libc::SOCK_STREAM,
302    );
303    #[cfg(windows)]
304    return sys::tcp_socket(if v6 {
305        crate::syscore::windows::AF_INET6
306    } else {
307        crate::syscore::windows::AF_INET
308    });
309}
310
311/// One connect: a non-blocking socket, its connect, and the driver seeing
312/// it end.
313async fn connect(addr: SocketAddr) -> io::Result<Handle<TcpStream>> {
314    connect_socket(new_socket(addr.is_ipv6())?, true, addr).await
315}
316
317/// The connect of `socket`, non-blocking already or not, to `addr`.
318pub(super) async fn connect_socket(
319    socket: Owned,
320    nonblocking: bool,
321    addr: SocketAddr,
322) -> io::Result<Handle<TcpStream>> {
323    let stream = if nonblocking {
324        Handle::from_nonblocking(TcpStream::from(socket), Kind::Tcp)
325    } else {
326        Handle::adopt(TcpStream::from(socket), Kind::Tcp)?
327    };
328    poll_fn(|cx| stream.source().poll_connect(cx, Peer::Inet(addr))).await?;
329    Ok(stream)
330}
nuclei/net/tcp.rs

nuclei/net/
tcp.rs
