Skip to main content

compio_net/
tcp.rs

1use std::{
2    future::Future,
3    io,
4    net::SocketAddr,
5    pin::Pin,
6    task::{Context, Poll},
7    time::Duration,
8};
9
10use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
11use compio_driver::{
12    BufferRef, SharedFd, impl_raw_fd,
13    op::{RecvFlags, RecvMsgMultiResult, SendFlags, SendMsgZc, SendVectoredZc, SendZc},
14};
15use compio_io::{
16    AsyncRead, AsyncReadManaged, AsyncReadMulti, AsyncWrite, AsyncWriteZerocopy,
17    ancillary::{
18        AsyncReadAncillary, AsyncReadAncillaryManaged, AsyncReadAncillaryMulti,
19        AsyncWriteAncillary, AsyncWriteAncillaryZerocopy, ReturnFlags,
20    },
21    util::Splittable,
22};
23use compio_runtime::{Runtime, fd::PollFd};
24use futures_util::{Stream, StreamExt, stream::FusedStream};
25use socket2::{Protocol, SockAddr, Socket as Socket2, Type};
26
27use crate::{
28    Extract, Incoming, MSG_NOSIGNAL, ReadHalf, Socket, ToSocketAddrsAsync, WriteHalf, Zerocopy,
29};
30
31/// A TCP socket server, listening for connections.
32///
33/// You can accept a new connection by using the
34/// [`accept`](`TcpListener::accept`) method.
35///
36/// # Examples
37///
38/// ```
39/// use std::net::SocketAddr;
40///
41/// use compio_io::{AsyncReadExt, AsyncWriteExt};
42/// use compio_net::{TcpListener, TcpStream};
43/// use socket2::SockAddr;
44///
45/// # compio_runtime::Runtime::new().unwrap().block_on(async move {
46/// let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
47///
48/// let addr = listener.local_addr().unwrap();
49///
50/// let tx_fut = TcpStream::connect(&addr);
51///
52/// let rx_fut = listener.accept();
53///
54/// let (mut tx, (mut rx, _)) = futures_util::try_join!(tx_fut, rx_fut).unwrap();
55///
56/// tx.write_all("test").await.0.unwrap();
57///
58/// let (_, buf) = rx.read_exact(Vec::with_capacity(4)).await.unwrap();
59///
60/// assert_eq!(buf, b"test");
61/// # });
62/// ```
63#[derive(Debug, Clone)]
64pub struct TcpListener {
65    inner: Socket,
66}
67
68impl TcpListener {
69    /// Creates a new `TcpListener`, which will be bound to the specified
70    /// address.
71    ///
72    /// The returned listener is ready for accepting connections.
73    ///
74    /// Binding with a port number of 0 will request that the OS assigns a port
75    /// to this listener.
76    ///
77    /// It enables the `SO_REUSEADDR` option by default.
78    ///
79    /// To configure the socket before binding, you can use the [`TcpSocket`]
80    /// type.
81    pub async fn bind(addr: impl ToSocketAddrsAsync) -> io::Result<Self> {
82        super::each_addr(addr, |addr| async move {
83            let sa = SockAddr::from(addr);
84            let socket = Socket::new(sa.domain(), Type::STREAM, Some(Protocol::TCP)).await?;
85            socket.socket.set_reuse_address(true)?;
86            socket.bind(&sa).await?;
87            socket.listen(128).await?;
88            Ok(Self { inner: socket })
89        })
90        .await
91    }
92
93    /// Creates new TcpListener from a [`std::net::TcpListener`].
94    pub fn from_std(stream: std::net::TcpListener) -> io::Result<Self> {
95        if Runtime::with_current(|r| r.driver_type().is_polling()) {
96            stream.set_nonblocking(true)?;
97        }
98
99        Ok(Self {
100            inner: Socket::from_socket2(Socket2::from(stream))?,
101        })
102    }
103
104    /// Close the socket. If the returned future is dropped before polling, the
105    /// socket won't be closed.
106    ///
107    /// See [`TcpStream::close`] for more details.
108    ///
109    /// [`TcpStream::close`]: crate::tcp::TcpStream::close
110    pub fn close(self) -> impl Future<Output = io::Result<()>> {
111        self.inner.close()
112    }
113
114    /// Accepts a new incoming connection from this listener.
115    ///
116    /// This function will yield once a new TCP connection is established. When
117    /// established, the corresponding [`TcpStream`] and the remote peer's
118    /// address will be returned.
119    pub async fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> {
120        let (socket, addr) = self.inner.accept().await?;
121        let stream = TcpStream { inner: socket };
122        Ok((stream, addr.as_socket().expect("should be SocketAddr")))
123    }
124
125    /// Accepts a new incoming connection from this listener using the provided
126    /// socket.
127    #[cfg(windows)]
128    pub async fn accept_with(&self, sock: TcpSocket) -> io::Result<(TcpStream, SocketAddr)> {
129        let (socket, addr) = self.inner.accept_with(sock.inner).await?;
130        let stream = TcpStream { inner: socket };
131        Ok((stream, addr.as_socket().expect("should be SocketAddr")))
132    }
133
134    /// Returns a stream of incoming connections to this listener.
135    pub fn incoming(&self) -> TcpIncoming<'_> {
136        TcpIncoming {
137            inner: self.inner.incoming(),
138        }
139    }
140
141    /// Returns the local address that this listener is bound to.
142    ///
143    /// This can be useful, for example, when binding to port 0 to
144    /// figure out which port was actually bound.
145    ///
146    /// # Examples
147    ///
148    /// ```
149    /// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
150    ///
151    /// use compio_net::TcpListener;
152    /// use socket2::SockAddr;
153    ///
154    /// # compio_runtime::Runtime::new().unwrap().block_on(async {
155    /// let listener = TcpListener::bind("127.0.0.1:8080").await.unwrap();
156    ///
157    /// let addr = listener.local_addr().expect("Couldn't get local address");
158    /// assert_eq!(
159    ///     addr,
160    ///     SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080))
161    /// );
162    /// # });
163    /// ```
164    pub fn local_addr(&self) -> io::Result<SocketAddr> {
165        self.inner
166            .local_addr()
167            .map(|addr| addr.as_socket().expect("should be SocketAddr"))
168    }
169
170    /// Returns the value of the `SO_ERROR` option.
171    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
172        self.inner.socket.take_error()
173    }
174
175    /// Gets the value of the `IP_TTL` option for this socket.
176    ///
177    /// For more information about this option, see [`set_ttl_v4`].
178    ///
179    /// [`set_ttl_v4`]: method@Self::set_ttl_v4
180    pub fn ttl_v4(&self) -> io::Result<u32> {
181        self.inner.socket.ttl_v4()
182    }
183
184    /// Sets the value for the `IP_TTL` option on this socket.
185    ///
186    /// This value sets the time-to-live field that is used in every packet sent
187    /// from this socket.
188    pub fn set_ttl_v4(&self, ttl: u32) -> io::Result<()> {
189        self.inner.socket.set_ttl_v4(ttl)
190    }
191}
192
193impl_raw_fd!(TcpListener, socket2::Socket, inner, socket);
194
195/// A stream of incoming TCP connections.
196pub struct TcpIncoming<'a> {
197    inner: Incoming<'a>,
198}
199
200impl Stream for TcpIncoming<'_> {
201    type Item = io::Result<TcpStream>;
202
203    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
204        let this = self.get_mut();
205        this.inner.poll_next_unpin(cx).map(|res| {
206            res.map(|res| {
207                let socket = res?;
208                Ok(TcpStream { inner: socket })
209            })
210        })
211    }
212}
213
214impl FusedStream for TcpIncoming<'_> {
215    fn is_terminated(&self) -> bool {
216        self.inner.is_terminated()
217    }
218}
219
220/// A TCP stream between a local and a remote socket.
221///
222/// A TCP stream can either be created by connecting to an endpoint, via the
223/// `connect` method, or by accepting a connection from a listener.
224///
225/// # Examples
226///
227/// ```no_run
228/// use std::net::SocketAddr;
229///
230/// use compio_io::AsyncWrite;
231/// use compio_net::TcpStream;
232///
233/// # compio_runtime::Runtime::new().unwrap().block_on(async {
234/// // Connect to a peer
235/// let mut stream = TcpStream::connect("127.0.0.1:8080").await.unwrap();
236///
237/// // Write some data.
238/// stream.write("hello world!").await.unwrap();
239/// # })
240/// ```
241#[derive(Debug, Clone)]
242pub struct TcpStream {
243    inner: Socket,
244}
245
246impl TcpStream {
247    /// Opens a TCP connection to a remote host.
248    ///
249    /// To configure the socket before connecting, you can use the [`TcpSocket`]
250    /// type.
251    pub async fn connect(addr: impl ToSocketAddrsAsync) -> io::Result<Self> {
252        use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
253
254        super::each_addr(addr, |addr| async move {
255            let addr2 = SockAddr::from(addr);
256            let socket = Socket::new(addr2.domain(), Type::STREAM, Some(Protocol::TCP)).await?;
257            if cfg!(windows) {
258                let bind_addr = if addr.is_ipv4() {
259                    SockAddr::from(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))
260                } else if addr.is_ipv6() {
261                    SockAddr::from(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, 0, 0, 0))
262                } else {
263                    return Err(io::Error::new(
264                        io::ErrorKind::AddrNotAvailable,
265                        "Unsupported address domain.",
266                    ));
267                };
268                socket.bind(&bind_addr).await?;
269            };
270            socket.connect_async(&addr2).await?;
271            Ok(Self { inner: socket })
272        })
273        .await
274    }
275
276    /// Creates new TcpStream from a [`std::net::TcpStream`].
277    pub fn from_std(stream: std::net::TcpStream) -> io::Result<Self> {
278        if Runtime::with_current(|r| r.driver_type().is_polling()) {
279            stream.set_nonblocking(true)?;
280        }
281
282        Ok(Self {
283            inner: Socket::from_socket2(Socket2::from(stream))?,
284        })
285    }
286
287    /// Close the socket. If the returned future is dropped before polling, the
288    /// socket won't be closed.
289    ///
290    /// As the socket is clonable, users can call `close` on a clone, but the
291    /// future will never complete until all clones are dropped. Some
292    /// operations may keep a strong reference to the socket, so the future
293    /// may never complete if there are pending operations.
294    ///
295    /// It's OK to drop the socket directly without calling `close`, but the
296    /// socket may not be closed immediately.
297    pub fn close(self) -> impl Future<Output = io::Result<()>> {
298        self.inner.close()
299    }
300
301    /// Returns the socket address of the remote peer of this TCP connection.
302    pub fn peer_addr(&self) -> io::Result<SocketAddr> {
303        self.inner
304            .peer_addr()
305            .map(|addr| addr.as_socket().expect("should be SocketAddr"))
306    }
307
308    /// Returns the socket address of the local half of this TCP connection.
309    pub fn local_addr(&self) -> io::Result<SocketAddr> {
310        self.inner
311            .local_addr()
312            .map(|addr| addr.as_socket().expect("should be SocketAddr"))
313    }
314
315    /// Returns the value of the `SO_ERROR` option.
316    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
317        self.inner.socket.take_error()
318    }
319
320    /// Splits a [`TcpStream`] into a read half and a write half, which can be
321    /// used to read and write the stream concurrently.
322    ///
323    /// This method is more efficient than
324    /// [`into_split`](TcpStream::into_split), but the halves cannot
325    /// be moved into independently spawned tasks.
326    pub fn split(&self) -> (ReadHalf<'_, Self>, WriteHalf<'_, Self>) {
327        crate::split(self)
328    }
329
330    /// Splits a [`TcpStream`] into a read half and a write half, which can be
331    /// used to read and write the stream concurrently.
332    ///
333    /// Unlike [`split`](TcpStream::split), the owned halves can be moved to
334    /// separate tasks.
335    pub fn into_split(self) -> (Self, Self) {
336        (self.clone(), self)
337    }
338
339    /// Create [`PollFd`] from inner socket.
340    pub fn to_poll_fd(&self) -> io::Result<PollFd<Socket2>> {
341        self.inner.to_poll_fd()
342    }
343
344    /// Create [`PollFd`] from inner socket.
345    pub fn into_poll_fd(self) -> io::Result<PollFd<Socket2>> {
346        self.inner.into_poll_fd()
347    }
348
349    /// Close the connection of the socket, and reuse it to create a new
350    /// connection. This method is useful when the socket is created by
351    /// [`TcpListener::accept`], and will be reused in
352    /// [`TcpListener::accept_with`] to accept a new connection.
353    #[cfg(windows)]
354    pub async fn disconnect(self) -> io::Result<TcpSocket> {
355        self.inner.disconnect().await?;
356        Ok(TcpSocket { inner: self.inner })
357    }
358
359    /// Gets the value of the `TCP_NODELAY` option on this socket.
360    ///
361    /// For more information about this option, see
362    /// [`TcpStream::set_nodelay`].
363    pub fn nodelay(&self) -> io::Result<bool> {
364        self.inner.socket.tcp_nodelay()
365    }
366
367    /// Sets the value of the TCP_NODELAY option on this socket.
368    ///
369    /// If set, this option disables the Nagle algorithm. This means
370    /// that segments are always sent as soon as possible, even if
371    /// there is only a small amount of data. When not set, data is
372    /// buffered until there is a sufficient amount to send out,
373    /// thereby avoiding the frequent sending of small packets.
374    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
375        self.inner.socket.set_tcp_nodelay(nodelay)
376    }
377
378    /// Gets the value of the `TCP_QUICKACK` option on this socket.
379    ///
380    /// For more information about this option, see [`TcpStream::set_quickack`].
381    #[cfg(any(
382        target_os = "linux",
383        target_os = "android",
384        target_os = "fuchsia",
385        target_os = "cygwin",
386    ))]
387    pub fn quickack(&self) -> io::Result<bool> {
388        self.inner.socket.tcp_quickack()
389    }
390
391    /// Enable or disable `TCP_QUICKACK`.
392    ///
393    /// This flag causes Linux to eagerly send `ACK`s rather than delaying them.
394    /// Linux may reset this flag after further operations on the socket.
395    #[cfg(any(
396        target_os = "linux",
397        target_os = "android",
398        target_os = "fuchsia",
399        target_os = "cygwin",
400    ))]
401    pub fn set_quickack(&self, quickack: bool) -> io::Result<()> {
402        self.inner.socket.set_tcp_quickack(quickack)
403    }
404
405    /// Reads the linger duration for this socket by getting the `SO_LINGER`
406    /// option.
407    pub fn linger(&self) -> io::Result<Option<Duration>> {
408        self.inner.socket.linger()
409    }
410
411    /// Sets a linger duration of zero on this socket by setting the `SO_LINGER`
412    /// option.
413    pub fn set_zero_linger(&self) -> io::Result<()> {
414        self.inner.socket.set_linger(Some(Duration::ZERO))
415    }
416
417    /// Gets the value of the `IP_TTL` option for this socket.
418    ///
419    /// For more information about this option, see [`set_ttl_v4`].
420    ///
421    /// [`set_ttl_v4`]: TcpStream::set_ttl_v4
422    pub fn ttl_v4(&self) -> io::Result<u32> {
423        self.inner.socket.ttl_v4()
424    }
425
426    /// Sets the value for the `IP_TTL` option on this socket.
427    ///
428    /// This value sets the time-to-live field that is used in every packet sent
429    /// from this socket.
430    pub fn set_ttl_v4(&self, ttl: u32) -> io::Result<()> {
431        self.inner.socket.set_ttl_v4(ttl)
432    }
433
434    /// Sends out-of-band data on this socket.
435    ///
436    /// Out-of-band data is sent with the `MSG_OOB` flag.
437    pub async fn send_out_of_band<T: IoBuf>(&self, buf: T) -> BufResult<usize, T> {
438        #[cfg(unix)]
439        use libc::MSG_OOB;
440        #[cfg(windows)]
441        use windows_sys::Win32::Networking::WinSock::MSG_OOB;
442
443        self.inner
444            .send(
445                buf,
446                SendFlags::from_bits_retain(MSG_OOB as _) | MSG_NOSIGNAL,
447            )
448            .await
449    }
450
451    /// Peeks at data from this socket without consuming it
452    ///
453    /// ## Platform-specific
454    /// * Windows: this method may work, but is not ensured by Microsoft.
455    pub async fn peek<T: IoBufMut>(&self, buffer: T) -> BufResult<usize, T> {
456        self.inner.recv(buffer, RecvFlags::PEEK).await
457    }
458}
459
460impl AsyncRead for TcpStream {
461    #[inline]
462    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
463        (&*self).read(buf).await
464    }
465
466    #[inline]
467    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
468        (&*self).read_vectored(buf).await
469    }
470}
471
472impl AsyncRead for &TcpStream {
473    #[inline]
474    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
475        self.inner.recv(buf, RecvFlags::empty()).await
476    }
477
478    #[inline]
479    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
480        self.inner.recv_vectored(buf, RecvFlags::empty()).await
481    }
482}
483
484impl AsyncReadManaged for TcpStream {
485    type Buffer = BufferRef;
486
487    async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
488        (&*self).read_managed(len).await
489    }
490}
491
492impl AsyncReadManaged for &TcpStream {
493    type Buffer = BufferRef;
494
495    async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
496        self.inner.recv_managed(len, RecvFlags::empty()).await
497    }
498}
499
500impl AsyncReadMulti for TcpStream {
501    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
502        self.inner.recv_multi(len, RecvFlags::empty())
503    }
504}
505
506impl AsyncReadMulti for &TcpStream {
507    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
508        self.inner.recv_multi(len, RecvFlags::empty())
509    }
510}
511
512impl AsyncReadAncillary for TcpStream {
513    #[inline]
514    async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
515        &mut self,
516        buffer: T,
517        control: C,
518    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
519        (&*self).read_with_ancillary(buffer, control).await
520    }
521
522    #[inline]
523    async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
524        &mut self,
525        buffer: T,
526        control: C,
527    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
528        (&*self).read_vectored_with_ancillary(buffer, control).await
529    }
530}
531
532impl AsyncReadAncillary for &TcpStream {
533    #[inline]
534    async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
535        &mut self,
536        buffer: T,
537        control: C,
538    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
539        self.inner
540            .recv_msg(buffer, control, RecvFlags::empty())
541            .await
542            .map_res(|(res, len, _addr, flags)| (res, len, flags))
543    }
544
545    #[inline]
546    async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
547        &mut self,
548        buffer: T,
549        control: C,
550    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
551        self.inner
552            .recv_msg_vectored(buffer, control, RecvFlags::empty())
553            .await
554            .map_res(|(res, len, _addr, flags)| (res, len, flags))
555    }
556}
557
558impl AsyncReadAncillaryManaged for TcpStream {
559    #[inline]
560    async fn read_managed_with_ancillary<C: IoBufMut>(
561        &mut self,
562        len: usize,
563        control: C,
564    ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
565        (&*self).read_managed_with_ancillary(len, control).await
566    }
567}
568
569impl AsyncReadAncillaryManaged for &TcpStream {
570    #[inline]
571    async fn read_managed_with_ancillary<C: IoBufMut>(
572        &mut self,
573        len: usize,
574        control: C,
575    ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
576        self.inner
577            .recv_msg_managed(len, control, RecvFlags::empty())
578            .await
579            .map(|res| res.map(|(res, len, _addr, flags)| (res, len, flags)))
580    }
581}
582
583impl AsyncReadAncillaryMulti for TcpStream {
584    type Return = RecvMsgMultiResult;
585
586    #[inline]
587    fn read_multi_with_ancillary(
588        &mut self,
589        control_len: usize,
590    ) -> impl Stream<Item = io::Result<Self::Return>> {
591        self.inner.recv_msg_multi(control_len, RecvFlags::empty())
592    }
593}
594
595impl AsyncReadAncillaryMulti for &TcpStream {
596    type Return = RecvMsgMultiResult;
597
598    #[inline]
599    fn read_multi_with_ancillary(
600        &mut self,
601        control_len: usize,
602    ) -> impl Stream<Item = io::Result<Self::Return>> {
603        self.inner.recv_msg_multi(control_len, RecvFlags::empty())
604    }
605}
606
607impl AsyncWrite for TcpStream {
608    #[inline]
609    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
610        (&*self).write(buf).await
611    }
612
613    #[inline]
614    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
615        (&*self).write_vectored(buf).await
616    }
617
618    #[inline]
619    async fn flush(&mut self) -> io::Result<()> {
620        (&*self).flush().await
621    }
622
623    #[inline]
624    async fn shutdown(&mut self) -> io::Result<()> {
625        (&*self).shutdown().await
626    }
627}
628
629impl AsyncWrite for &TcpStream {
630    #[inline]
631    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
632        self.inner.send(buf, MSG_NOSIGNAL).await
633    }
634
635    #[inline]
636    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
637        self.inner.send_vectored(buf, MSG_NOSIGNAL).await
638    }
639
640    #[inline]
641    async fn flush(&mut self) -> io::Result<()> {
642        Ok(())
643    }
644
645    #[inline]
646    async fn shutdown(&mut self) -> io::Result<()> {
647        self.inner.shutdown().await
648    }
649}
650
651impl AsyncWriteAncillary for TcpStream {
652    #[inline]
653    async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
654        &mut self,
655        buffer: T,
656        control: C,
657    ) -> BufResult<usize, (T, C)> {
658        (&*self).write_with_ancillary(buffer, control).await
659    }
660
661    #[inline]
662    async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
663        &mut self,
664        buffer: T,
665        control: C,
666    ) -> BufResult<usize, (T, C)> {
667        (&*self)
668            .write_vectored_with_ancillary(buffer, control)
669            .await
670    }
671}
672
673impl AsyncWriteAncillary for &TcpStream {
674    #[inline]
675    async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
676        &mut self,
677        buffer: T,
678        control: C,
679    ) -> BufResult<usize, (T, C)> {
680        self.inner
681            .send_msg(buffer, control, None, MSG_NOSIGNAL)
682            .await
683    }
684
685    #[inline]
686    async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
687        &mut self,
688        buffer: T,
689        control: C,
690    ) -> BufResult<usize, (T, C)> {
691        self.inner
692            .send_msg_vectored(buffer, control, None, MSG_NOSIGNAL)
693            .await
694    }
695}
696
697impl AsyncWriteZerocopy for TcpStream {
698    type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<Socket2>>>;
699    type VectoredBufferReadyFuture<T: IoVectoredBuf> =
700        Zerocopy<SendVectoredZc<T, SharedFd<Socket2>>>;
701
702    async fn write_zerocopy<T: IoBuf>(
703        &mut self,
704        buf: T,
705    ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
706        self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
707    }
708
709    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
710        &mut self,
711        buf: T,
712    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
713        self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
714    }
715}
716
717impl AsyncWriteZerocopy for &TcpStream {
718    type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<Socket2>>>;
719    type VectoredBufferReadyFuture<T: IoVectoredBuf> =
720        Zerocopy<SendVectoredZc<T, SharedFd<Socket2>>>;
721
722    async fn write_zerocopy<T: IoBuf>(
723        &mut self,
724        buf: T,
725    ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
726        self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
727    }
728
729    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
730        &mut self,
731        buf: T,
732    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
733        self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
734    }
735}
736
737impl AsyncWriteAncillaryZerocopy for TcpStream {
738    type BufferReadyFuture<T: IoBuf, C: IoBuf> =
739        Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<Socket2>>>, T, C>;
740    type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
741        Zerocopy<SendMsgZc<T, C, SharedFd<Socket2>>>;
742
743    async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
744        &mut self,
745        buf: T,
746        control: C,
747    ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
748        self.inner
749            .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
750            .await
751    }
752
753    async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
754        &mut self,
755        buf: T,
756        control: C,
757    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
758        self.inner
759            .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
760            .await
761    }
762}
763
764impl AsyncWriteAncillaryZerocopy for &TcpStream {
765    type BufferReadyFuture<T: IoBuf, C: IoBuf> =
766        Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<Socket2>>>, T, C>;
767    type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
768        Zerocopy<SendMsgZc<T, C, SharedFd<Socket2>>>;
769
770    async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
771        &mut self,
772        buf: T,
773        control: C,
774    ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
775        self.inner
776            .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
777            .await
778    }
779
780    async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
781        &mut self,
782        buf: T,
783        control: C,
784    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
785        self.inner
786            .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
787            .await
788    }
789}
790
791impl Splittable for TcpStream {
792    type ReadHalf = Self;
793    type WriteHalf = Self;
794
795    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
796        self.into_split()
797    }
798}
799
800impl<'a> Splittable for &'a TcpStream {
801    type ReadHalf = ReadHalf<'a, TcpStream>;
802    type WriteHalf = WriteHalf<'a, TcpStream>;
803
804    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
805        crate::split(self)
806    }
807}
808
809impl<'a> Splittable for &'a mut TcpStream {
810    type ReadHalf = ReadHalf<'a, TcpStream>;
811    type WriteHalf = WriteHalf<'a, TcpStream>;
812
813    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
814        crate::split(self)
815    }
816}
817
818impl_raw_fd!(TcpStream, socket2::Socket, inner, socket);
819
820/// A TCP socket that has not yet been converted to a [`TcpStream`] or
821/// [`TcpListener`].
822#[derive(Debug)]
823pub struct TcpSocket {
824    inner: Socket,
825}
826
827impl TcpSocket {
828    /// Creates a new socket configured for IPv4.
829    pub async fn new_v4() -> io::Result<TcpSocket> {
830        TcpSocket::new(socket2::Domain::IPV4).await
831    }
832
833    /// Creates a new socket configured for IPv6.
834    pub async fn new_v6() -> io::Result<TcpSocket> {
835        TcpSocket::new(socket2::Domain::IPV6).await
836    }
837
838    async fn new(domain: socket2::Domain) -> io::Result<TcpSocket> {
839        let inner =
840            Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)).await?;
841        Ok(TcpSocket { inner })
842    }
843
844    /// Sets value for the `SO_KEEPALIVE` option on this socket.
845    pub fn set_keepalive(&self, keepalive: bool) -> io::Result<()> {
846        self.inner.socket.set_keepalive(keepalive)
847    }
848
849    /// Gets the value of the `SO_KEEPALIVE` option on this socket.
850    pub fn keepalive(&self) -> io::Result<bool> {
851        self.inner.socket.keepalive()
852    }
853
854    /// Allows the socket to bind to an in-use address.
855    pub fn set_reuseaddr(&self, reuseaddr: bool) -> io::Result<()> {
856        self.inner.socket.set_reuse_address(reuseaddr)
857    }
858
859    /// Retrieves the value set for `SO_REUSEADDR` on this socket.
860    pub fn reuseaddr(&self) -> io::Result<bool> {
861        self.inner.socket.reuse_address()
862    }
863
864    /// Allows the socket to bind to an in-use port. Only available for
865    /// supported unix systems.
866    #[cfg(all(
867        unix,
868        not(target_os = "solaris"),
869        not(target_os = "illumos"),
870        not(target_os = "cygwin"),
871    ))]
872    pub fn set_reuseport(&self, reuseport: bool) -> io::Result<()> {
873        self.inner.socket.set_reuse_port(reuseport)
874    }
875
876    /// Allows the socket to bind to an in-use port. Only available for
877    /// supported unix systems.
878    #[cfg(all(
879        unix,
880        not(target_os = "solaris"),
881        not(target_os = "illumos"),
882        not(target_os = "cygwin"),
883    ))]
884    pub fn reuseport(&self) -> io::Result<bool> {
885        self.inner.socket.reuse_port()
886    }
887
888    /// Sets the size of the TCP send buffer on this socket.
889    ///
890    /// On most operating systems, this sets the `SO_SNDBUF` socket option.
891    pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> {
892        self.inner.socket.set_send_buffer_size(size as usize)
893    }
894
895    /// Returns the size of the TCP send buffer for this socket.
896    ///
897    /// On most operating systems, this is the value of the `SO_SNDBUF` socket
898    /// option.
899    pub fn send_buffer_size(&self) -> io::Result<u32> {
900        self.inner.socket.send_buffer_size().map(|n| n as u32)
901    }
902
903    /// Sets the size of the TCP receive buffer on this socket.
904    ///
905    /// On most operating systems, this sets the `SO_RCVBUF` socket option.
906    pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> {
907        self.inner.socket.set_recv_buffer_size(size as usize)
908    }
909
910    /// Returns the size of the TCP receive buffer for this socket.
911    ///
912    /// On most operating systems, this is the value of the `SO_RCVBUF` socket
913    /// option.
914    pub fn recv_buffer_size(&self) -> io::Result<u32> {
915        self.inner.socket.recv_buffer_size().map(|n| n as u32)
916    }
917
918    /// Sets a linger duration of zero on this socket by setting the `SO_LINGER`
919    /// option.
920    pub fn set_zero_linger(&self) -> io::Result<()> {
921        self.inner.socket.set_linger(Some(Duration::ZERO))
922    }
923
924    /// Reads the linger duration for this socket by getting the `SO_LINGER`
925    /// option.
926    pub fn linger(&self) -> io::Result<Option<Duration>> {
927        self.inner.socket.linger()
928    }
929
930    /// Sets the value of the `TCP_NODELAY` option on this socket.
931    ///
932    /// If set, this option disables the Nagle algorithm. This means that
933    /// segments are always sent as soon as possible, even if there is only
934    /// a small amount of data. When not set, data is buffered until there
935    /// is a sufficient amount to send out, thereby avoiding the frequent
936    /// sending of small packets.
937    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
938        self.inner.socket.set_tcp_nodelay(nodelay)
939    }
940
941    /// Gets the value of the `TCP_NODELAY` option on this socket.
942    ///
943    /// For more information about this option, see [`set_nodelay`].
944    ///
945    /// [`set_nodelay`]: TcpSocket::set_nodelay
946    pub fn nodelay(&self) -> io::Result<bool> {
947        self.inner.socket.tcp_nodelay()
948    }
949
950    /// Gets the value of the `IPV6_TCLASS` option for this socket.
951    ///
952    /// For more information about this option, see [`set_tclass_v6`].
953    ///
954    /// [`set_tclass_v6`]: Self::set_tclass_v6
955    #[cfg(any(
956        target_os = "android",
957        target_os = "dragonfly",
958        target_os = "freebsd",
959        target_os = "fuchsia",
960        target_os = "linux",
961        target_os = "macos",
962        target_os = "netbsd",
963        target_os = "openbsd",
964        target_os = "cygwin",
965    ))]
966    pub fn tclass_v6(&self) -> io::Result<u32> {
967        self.inner.socket.tclass_v6()
968    }
969
970    /// Sets the value for the `IPV6_TCLASS` option on this socket.
971    ///
972    /// Specifies the traffic class field that is used in every packet
973    /// sent from this socket.
974    ///
975    /// # Note
976    ///
977    /// This may not have any effect on IPv4 sockets.
978    #[cfg(any(
979        target_os = "android",
980        target_os = "dragonfly",
981        target_os = "freebsd",
982        target_os = "fuchsia",
983        target_os = "linux",
984        target_os = "macos",
985        target_os = "netbsd",
986        target_os = "openbsd",
987        target_os = "cygwin",
988    ))]
989    pub fn set_tclass_v6(&self, tclass: u32) -> io::Result<()> {
990        self.inner.socket.set_tclass_v6(tclass)
991    }
992
993    /// Gets the value of the `IP_TOS` option for this socket.
994    ///
995    /// For more information about this option, see [`set_tos_v4`].
996    ///
997    /// [`set_tos_v4`]: Self::set_tos_v4
998    #[cfg(not(any(
999        target_os = "fuchsia",
1000        target_os = "redox",
1001        target_os = "solaris",
1002        target_os = "illumos",
1003        target_os = "haiku"
1004    )))]
1005    pub fn tos_v4(&self) -> io::Result<u32> {
1006        self.inner.socket.tos_v4()
1007    }
1008
1009    /// Sets the value for the `IP_TOS` option on this socket.
1010    ///
1011    /// This value sets the type-of-service field that is used in every packet
1012    /// sent from this socket.
1013    ///
1014    /// # Note
1015    ///
1016    /// - This may not have any effect on IPv6 sockets.
1017    /// - On Windows, `IP_TOS` is only supported on [Windows 8+ or
1018    ///   Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options)
1019    #[cfg(not(any(
1020        target_os = "fuchsia",
1021        target_os = "redox",
1022        target_os = "solaris",
1023        target_os = "illumos",
1024        target_os = "haiku"
1025    )))]
1026    pub fn set_tos_v4(&self, tos: u32) -> io::Result<()> {
1027        self.inner.socket.set_tos_v4(tos)
1028    }
1029
1030    /// Gets the value for the `SO_BINDTODEVICE` option on this socket
1031    ///
1032    /// Returns the interface name of the device to which this socket is bound.
1033    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))]
1034    pub fn device(&self) -> io::Result<Option<Vec<u8>>> {
1035        self.inner.socket.device()
1036    }
1037
1038    /// Sets the value for the `SO_BINDTODEVICE` option on this socket
1039    ///
1040    /// If a socket is bound to an interface, only packets received from that
1041    /// particular interface are processed by the socket. Note that this only
1042    /// works for some socket types, particularly `AF_INET` sockets.
1043    ///
1044    /// If `interface` is `None` or an empty string it removes the binding.
1045    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
1046    pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> {
1047        self.inner.socket.bind_device(interface)
1048    }
1049
1050    /// Gets the local address of this socket.
1051    pub fn local_addr(&self) -> io::Result<SocketAddr> {
1052        Ok(self
1053            .inner
1054            .local_addr()?
1055            .as_socket()
1056            .expect("should be SocketAddr"))
1057    }
1058
1059    /// Returns the value of the `SO_ERROR` option.
1060    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
1061        self.inner.socket.take_error()
1062    }
1063
1064    /// Binds the socket to the given address.
1065    pub async fn bind(&self, addr: SocketAddr) -> io::Result<()> {
1066        self.inner.bind(&addr.into()).await
1067    }
1068
1069    /// Establishes a TCP connection with a peer at the specified socket
1070    /// address.
1071    ///
1072    /// The [`TcpSocket`] is consumed. Once the connection is established, a
1073    /// connected [`TcpStream`] is returned. If the connection fails, the
1074    /// encountered error is returned.
1075    ///
1076    /// On Windows, the socket should be bound to an address before connecting.
1077    pub async fn connect(self, addr: SocketAddr) -> io::Result<TcpStream> {
1078        self.inner.connect_async(&addr.into()).await?;
1079        Ok(TcpStream { inner: self.inner })
1080    }
1081
1082    /// Converts the socket into a `TcpListener`.
1083    ///
1084    /// `backlog` defines the maximum number of pending connections that are
1085    /// queued by the operating system at any given time. Connections are
1086    /// removed from the queue with [`TcpListener::accept`]. When the queue
1087    /// is full, the operating system will start rejecting connections.
1088    pub async fn listen(self, backlog: i32) -> io::Result<TcpListener> {
1089        self.inner.listen(backlog).await?;
1090        Ok(TcpListener { inner: self.inner })
1091    }
1092
1093    /// Converts a [`std::net::TcpStream`] into a [`TcpSocket`]. The provided
1094    /// socket must not have been connected prior to calling this function. This
1095    /// function is typically used together with crates such as [`socket2`] to
1096    /// configure socket options that are not available on [`TcpSocket`].
1097    pub fn from_std_stream(stream: std::net::TcpStream) -> io::Result<TcpSocket> {
1098        if Runtime::with_current(|r| r.driver_type().is_polling()) {
1099            stream.set_nonblocking(true)?;
1100        }
1101
1102        Ok(Self {
1103            inner: Socket::from_socket2(Socket2::from(stream))?,
1104        })
1105    }
1106}
1107
1108impl_raw_fd!(TcpSocket, socket2::Socket, inner, socket);