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
452impl AsyncRead for TcpStream {
453    #[inline]
454    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
455        (&*self).read(buf).await
456    }
457
458    #[inline]
459    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
460        (&*self).read_vectored(buf).await
461    }
462}
463
464impl AsyncRead for &TcpStream {
465    #[inline]
466    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
467        self.inner.recv(buf, RecvFlags::empty()).await
468    }
469
470    #[inline]
471    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
472        self.inner.recv_vectored(buf, RecvFlags::empty()).await
473    }
474}
475
476impl AsyncReadManaged for TcpStream {
477    type Buffer = BufferRef;
478
479    async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
480        (&*self).read_managed(len).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.inner.recv_managed(len, RecvFlags::empty()).await
489    }
490}
491
492impl AsyncReadMulti for TcpStream {
493    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
494        self.inner.recv_multi(len, RecvFlags::empty())
495    }
496}
497
498impl AsyncReadMulti for &TcpStream {
499    fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
500        self.inner.recv_multi(len, RecvFlags::empty())
501    }
502}
503
504impl AsyncReadAncillary for TcpStream {
505    #[inline]
506    async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
507        &mut self,
508        buffer: T,
509        control: C,
510    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
511        (&*self).read_with_ancillary(buffer, control).await
512    }
513
514    #[inline]
515    async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
516        &mut self,
517        buffer: T,
518        control: C,
519    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
520        (&*self).read_vectored_with_ancillary(buffer, control).await
521    }
522}
523
524impl AsyncReadAncillary for &TcpStream {
525    #[inline]
526    async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
527        &mut self,
528        buffer: T,
529        control: C,
530    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
531        self.inner
532            .recv_msg(buffer, control, RecvFlags::empty())
533            .await
534            .map_res(|(res, len, _addr, flags)| (res, len, flags))
535    }
536
537    #[inline]
538    async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
539        &mut self,
540        buffer: T,
541        control: C,
542    ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
543        self.inner
544            .recv_msg_vectored(buffer, control, RecvFlags::empty())
545            .await
546            .map_res(|(res, len, _addr, flags)| (res, len, flags))
547    }
548}
549
550impl AsyncReadAncillaryManaged for TcpStream {
551    #[inline]
552    async fn read_managed_with_ancillary<C: IoBufMut>(
553        &mut self,
554        len: usize,
555        control: C,
556    ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
557        (&*self).read_managed_with_ancillary(len, control).await
558    }
559}
560
561impl AsyncReadAncillaryManaged for &TcpStream {
562    #[inline]
563    async fn read_managed_with_ancillary<C: IoBufMut>(
564        &mut self,
565        len: usize,
566        control: C,
567    ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
568        self.inner
569            .recv_msg_managed(len, control, RecvFlags::empty())
570            .await
571            .map(|res| res.map(|(res, len, _addr, flags)| (res, len, flags)))
572    }
573}
574
575impl AsyncReadAncillaryMulti for TcpStream {
576    type Return = RecvMsgMultiResult;
577
578    #[inline]
579    fn read_multi_with_ancillary(
580        &mut self,
581        control_len: usize,
582    ) -> impl Stream<Item = io::Result<Self::Return>> {
583        self.inner.recv_msg_multi(control_len, RecvFlags::empty())
584    }
585}
586
587impl AsyncReadAncillaryMulti for &TcpStream {
588    type Return = RecvMsgMultiResult;
589
590    #[inline]
591    fn read_multi_with_ancillary(
592        &mut self,
593        control_len: usize,
594    ) -> impl Stream<Item = io::Result<Self::Return>> {
595        self.inner.recv_msg_multi(control_len, RecvFlags::empty())
596    }
597}
598
599impl AsyncWrite for TcpStream {
600    #[inline]
601    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
602        (&*self).write(buf).await
603    }
604
605    #[inline]
606    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
607        (&*self).write_vectored(buf).await
608    }
609
610    #[inline]
611    async fn flush(&mut self) -> io::Result<()> {
612        (&*self).flush().await
613    }
614
615    #[inline]
616    async fn shutdown(&mut self) -> io::Result<()> {
617        (&*self).shutdown().await
618    }
619}
620
621impl AsyncWrite for &TcpStream {
622    #[inline]
623    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
624        self.inner.send(buf, MSG_NOSIGNAL).await
625    }
626
627    #[inline]
628    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
629        self.inner.send_vectored(buf, MSG_NOSIGNAL).await
630    }
631
632    #[inline]
633    async fn flush(&mut self) -> io::Result<()> {
634        Ok(())
635    }
636
637    #[inline]
638    async fn shutdown(&mut self) -> io::Result<()> {
639        self.inner.shutdown().await
640    }
641}
642
643impl AsyncWriteAncillary for TcpStream {
644    #[inline]
645    async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
646        &mut self,
647        buffer: T,
648        control: C,
649    ) -> BufResult<usize, (T, C)> {
650        (&*self).write_with_ancillary(buffer, control).await
651    }
652
653    #[inline]
654    async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
655        &mut self,
656        buffer: T,
657        control: C,
658    ) -> BufResult<usize, (T, C)> {
659        (&*self)
660            .write_vectored_with_ancillary(buffer, control)
661            .await
662    }
663}
664
665impl AsyncWriteAncillary for &TcpStream {
666    #[inline]
667    async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
668        &mut self,
669        buffer: T,
670        control: C,
671    ) -> BufResult<usize, (T, C)> {
672        self.inner
673            .send_msg(buffer, control, None, MSG_NOSIGNAL)
674            .await
675    }
676
677    #[inline]
678    async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
679        &mut self,
680        buffer: T,
681        control: C,
682    ) -> BufResult<usize, (T, C)> {
683        self.inner
684            .send_msg_vectored(buffer, control, None, MSG_NOSIGNAL)
685            .await
686    }
687}
688
689impl AsyncWriteZerocopy for TcpStream {
690    type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<Socket2>>>;
691    type VectoredBufferReadyFuture<T: IoVectoredBuf> =
692        Zerocopy<SendVectoredZc<T, SharedFd<Socket2>>>;
693
694    async fn write_zerocopy<T: IoBuf>(
695        &mut self,
696        buf: T,
697    ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
698        self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
699    }
700
701    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
702        &mut self,
703        buf: T,
704    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
705        self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
706    }
707}
708
709impl AsyncWriteZerocopy for &TcpStream {
710    type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<Socket2>>>;
711    type VectoredBufferReadyFuture<T: IoVectoredBuf> =
712        Zerocopy<SendVectoredZc<T, SharedFd<Socket2>>>;
713
714    async fn write_zerocopy<T: IoBuf>(
715        &mut self,
716        buf: T,
717    ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
718        self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
719    }
720
721    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
722        &mut self,
723        buf: T,
724    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
725        self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
726    }
727}
728
729impl AsyncWriteAncillaryZerocopy for TcpStream {
730    type BufferReadyFuture<T: IoBuf, C: IoBuf> =
731        Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<Socket2>>>, T, C>;
732    type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
733        Zerocopy<SendMsgZc<T, C, SharedFd<Socket2>>>;
734
735    async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
736        &mut self,
737        buf: T,
738        control: C,
739    ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
740        self.inner
741            .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
742            .await
743    }
744
745    async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
746        &mut self,
747        buf: T,
748        control: C,
749    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
750        self.inner
751            .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
752            .await
753    }
754}
755
756impl AsyncWriteAncillaryZerocopy for &TcpStream {
757    type BufferReadyFuture<T: IoBuf, C: IoBuf> =
758        Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<Socket2>>>, T, C>;
759    type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
760        Zerocopy<SendMsgZc<T, C, SharedFd<Socket2>>>;
761
762    async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
763        &mut self,
764        buf: T,
765        control: C,
766    ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
767        self.inner
768            .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
769            .await
770    }
771
772    async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
773        &mut self,
774        buf: T,
775        control: C,
776    ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
777        self.inner
778            .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
779            .await
780    }
781}
782
783impl Splittable for TcpStream {
784    type ReadHalf = Self;
785    type WriteHalf = Self;
786
787    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
788        self.into_split()
789    }
790}
791
792impl<'a> Splittable for &'a TcpStream {
793    type ReadHalf = ReadHalf<'a, TcpStream>;
794    type WriteHalf = WriteHalf<'a, TcpStream>;
795
796    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
797        crate::split(self)
798    }
799}
800
801impl<'a> Splittable for &'a mut TcpStream {
802    type ReadHalf = ReadHalf<'a, TcpStream>;
803    type WriteHalf = WriteHalf<'a, TcpStream>;
804
805    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
806        crate::split(self)
807    }
808}
809
810impl_raw_fd!(TcpStream, socket2::Socket, inner, socket);
811
812/// A TCP socket that has not yet been converted to a [`TcpStream`] or
813/// [`TcpListener`].
814#[derive(Debug)]
815pub struct TcpSocket {
816    inner: Socket,
817}
818
819impl TcpSocket {
820    /// Creates a new socket configured for IPv4.
821    pub async fn new_v4() -> io::Result<TcpSocket> {
822        TcpSocket::new(socket2::Domain::IPV4).await
823    }
824
825    /// Creates a new socket configured for IPv6.
826    pub async fn new_v6() -> io::Result<TcpSocket> {
827        TcpSocket::new(socket2::Domain::IPV6).await
828    }
829
830    async fn new(domain: socket2::Domain) -> io::Result<TcpSocket> {
831        let inner =
832            Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP)).await?;
833        Ok(TcpSocket { inner })
834    }
835
836    /// Sets value for the `SO_KEEPALIVE` option on this socket.
837    pub fn set_keepalive(&self, keepalive: bool) -> io::Result<()> {
838        self.inner.socket.set_keepalive(keepalive)
839    }
840
841    /// Gets the value of the `SO_KEEPALIVE` option on this socket.
842    pub fn keepalive(&self) -> io::Result<bool> {
843        self.inner.socket.keepalive()
844    }
845
846    /// Allows the socket to bind to an in-use address.
847    pub fn set_reuseaddr(&self, reuseaddr: bool) -> io::Result<()> {
848        self.inner.socket.set_reuse_address(reuseaddr)
849    }
850
851    /// Retrieves the value set for `SO_REUSEADDR` on this socket.
852    pub fn reuseaddr(&self) -> io::Result<bool> {
853        self.inner.socket.reuse_address()
854    }
855
856    /// Allows the socket to bind to an in-use port. Only available for
857    /// supported unix systems.
858    #[cfg(all(
859        unix,
860        not(target_os = "solaris"),
861        not(target_os = "illumos"),
862        not(target_os = "cygwin"),
863    ))]
864    pub fn set_reuseport(&self, reuseport: bool) -> io::Result<()> {
865        self.inner.socket.set_reuse_port(reuseport)
866    }
867
868    /// Allows the socket to bind to an in-use port. Only available for
869    /// supported unix systems.
870    #[cfg(all(
871        unix,
872        not(target_os = "solaris"),
873        not(target_os = "illumos"),
874        not(target_os = "cygwin"),
875    ))]
876    pub fn reuseport(&self) -> io::Result<bool> {
877        self.inner.socket.reuse_port()
878    }
879
880    /// Sets the size of the TCP send buffer on this socket.
881    ///
882    /// On most operating systems, this sets the `SO_SNDBUF` socket option.
883    pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> {
884        self.inner.socket.set_send_buffer_size(size as usize)
885    }
886
887    /// Returns the size of the TCP send buffer for this socket.
888    ///
889    /// On most operating systems, this is the value of the `SO_SNDBUF` socket
890    /// option.
891    pub fn send_buffer_size(&self) -> io::Result<u32> {
892        self.inner.socket.send_buffer_size().map(|n| n as u32)
893    }
894
895    /// Sets the size of the TCP receive buffer on this socket.
896    ///
897    /// On most operating systems, this sets the `SO_RCVBUF` socket option.
898    pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> {
899        self.inner.socket.set_recv_buffer_size(size as usize)
900    }
901
902    /// Returns the size of the TCP receive buffer for this socket.
903    ///
904    /// On most operating systems, this is the value of the `SO_RCVBUF` socket
905    /// option.
906    pub fn recv_buffer_size(&self) -> io::Result<u32> {
907        self.inner.socket.recv_buffer_size().map(|n| n as u32)
908    }
909
910    /// Sets a linger duration of zero on this socket by setting the `SO_LINGER`
911    /// option.
912    pub fn set_zero_linger(&self) -> io::Result<()> {
913        self.inner.socket.set_linger(Some(Duration::ZERO))
914    }
915
916    /// Reads the linger duration for this socket by getting the `SO_LINGER`
917    /// option.
918    pub fn linger(&self) -> io::Result<Option<Duration>> {
919        self.inner.socket.linger()
920    }
921
922    /// Sets the value of the `TCP_NODELAY` option on this socket.
923    ///
924    /// If set, this option disables the Nagle algorithm. This means that
925    /// segments are always sent as soon as possible, even if there is only
926    /// a small amount of data. When not set, data is buffered until there
927    /// is a sufficient amount to send out, thereby avoiding the frequent
928    /// sending of small packets.
929    pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> {
930        self.inner.socket.set_tcp_nodelay(nodelay)
931    }
932
933    /// Gets the value of the `TCP_NODELAY` option on this socket.
934    ///
935    /// For more information about this option, see [`set_nodelay`].
936    ///
937    /// [`set_nodelay`]: TcpSocket::set_nodelay
938    pub fn nodelay(&self) -> io::Result<bool> {
939        self.inner.socket.tcp_nodelay()
940    }
941
942    /// Gets the value of the `IPV6_TCLASS` option for this socket.
943    ///
944    /// For more information about this option, see [`set_tclass_v6`].
945    ///
946    /// [`set_tclass_v6`]: Self::set_tclass_v6
947    #[cfg(any(
948        target_os = "android",
949        target_os = "dragonfly",
950        target_os = "freebsd",
951        target_os = "fuchsia",
952        target_os = "linux",
953        target_os = "macos",
954        target_os = "netbsd",
955        target_os = "openbsd",
956        target_os = "cygwin",
957    ))]
958    pub fn tclass_v6(&self) -> io::Result<u32> {
959        self.inner.socket.tclass_v6()
960    }
961
962    /// Sets the value for the `IPV6_TCLASS` option on this socket.
963    ///
964    /// Specifies the traffic class field that is used in every packet
965    /// sent from this socket.
966    ///
967    /// # Note
968    ///
969    /// This may not have any effect on IPv4 sockets.
970    #[cfg(any(
971        target_os = "android",
972        target_os = "dragonfly",
973        target_os = "freebsd",
974        target_os = "fuchsia",
975        target_os = "linux",
976        target_os = "macos",
977        target_os = "netbsd",
978        target_os = "openbsd",
979        target_os = "cygwin",
980    ))]
981    pub fn set_tclass_v6(&self, tclass: u32) -> io::Result<()> {
982        self.inner.socket.set_tclass_v6(tclass)
983    }
984
985    /// Gets the value of the `IP_TOS` option for this socket.
986    ///
987    /// For more information about this option, see [`set_tos_v4`].
988    ///
989    /// [`set_tos_v4`]: Self::set_tos_v4
990    #[cfg(not(any(
991        target_os = "fuchsia",
992        target_os = "redox",
993        target_os = "solaris",
994        target_os = "illumos",
995        target_os = "haiku"
996    )))]
997    pub fn tos_v4(&self) -> io::Result<u32> {
998        self.inner.socket.tos_v4()
999    }
1000
1001    /// Sets the value for the `IP_TOS` option on this socket.
1002    ///
1003    /// This value sets the type-of-service field that is used in every packet
1004    /// sent from this socket.
1005    ///
1006    /// # Note
1007    ///
1008    /// - This may not have any effect on IPv6 sockets.
1009    /// - On Windows, `IP_TOS` is only supported on [Windows 8+ or
1010    ///   Windows Server 2012+.](https://docs.microsoft.com/en-us/windows/win32/winsock/ipproto-ip-socket-options)
1011    #[cfg(not(any(
1012        target_os = "fuchsia",
1013        target_os = "redox",
1014        target_os = "solaris",
1015        target_os = "illumos",
1016        target_os = "haiku"
1017    )))]
1018    pub fn set_tos_v4(&self, tos: u32) -> io::Result<()> {
1019        self.inner.socket.set_tos_v4(tos)
1020    }
1021
1022    /// Gets the value for the `SO_BINDTODEVICE` option on this socket
1023    ///
1024    /// Returns the interface name of the device to which this socket is bound.
1025    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux",))]
1026    pub fn device(&self) -> io::Result<Option<Vec<u8>>> {
1027        self.inner.socket.device()
1028    }
1029
1030    /// Sets the value for the `SO_BINDTODEVICE` option on this socket
1031    ///
1032    /// If a socket is bound to an interface, only packets received from that
1033    /// particular interface are processed by the socket. Note that this only
1034    /// works for some socket types, particularly `AF_INET` sockets.
1035    ///
1036    /// If `interface` is `None` or an empty string it removes the binding.
1037    #[cfg(any(target_os = "android", target_os = "fuchsia", target_os = "linux"))]
1038    pub fn bind_device(&self, interface: Option<&[u8]>) -> io::Result<()> {
1039        self.inner.socket.bind_device(interface)
1040    }
1041
1042    /// Gets the local address of this socket.
1043    pub fn local_addr(&self) -> io::Result<SocketAddr> {
1044        Ok(self
1045            .inner
1046            .local_addr()?
1047            .as_socket()
1048            .expect("should be SocketAddr"))
1049    }
1050
1051    /// Returns the value of the `SO_ERROR` option.
1052    pub fn take_error(&self) -> io::Result<Option<io::Error>> {
1053        self.inner.socket.take_error()
1054    }
1055
1056    /// Binds the socket to the given address.
1057    pub async fn bind(&self, addr: SocketAddr) -> io::Result<()> {
1058        self.inner.bind(&addr.into()).await
1059    }
1060
1061    /// Establishes a TCP connection with a peer at the specified socket
1062    /// address.
1063    ///
1064    /// The [`TcpSocket`] is consumed. Once the connection is established, a
1065    /// connected [`TcpStream`] is returned. If the connection fails, the
1066    /// encountered error is returned.
1067    ///
1068    /// On Windows, the socket should be bound to an address before connecting.
1069    pub async fn connect(self, addr: SocketAddr) -> io::Result<TcpStream> {
1070        self.inner.connect_async(&addr.into()).await?;
1071        Ok(TcpStream { inner: self.inner })
1072    }
1073
1074    /// Converts the socket into a `TcpListener`.
1075    ///
1076    /// `backlog` defines the maximum number of pending connections that are
1077    /// queued by the operating system at any given time. Connections are
1078    /// removed from the queue with [`TcpListener::accept`]. When the queue
1079    /// is full, the operating system will start rejecting connections.
1080    pub async fn listen(self, backlog: i32) -> io::Result<TcpListener> {
1081        self.inner.listen(backlog).await?;
1082        Ok(TcpListener { inner: self.inner })
1083    }
1084
1085    /// Converts a [`std::net::TcpStream`] into a [`TcpSocket`]. The provided
1086    /// socket must not have been connected prior to calling this function. This
1087    /// function is typically used together with crates such as [`socket2`] to
1088    /// configure socket options that are not available on [`TcpSocket`].
1089    pub fn from_std_stream(stream: std::net::TcpStream) -> io::Result<TcpSocket> {
1090        if Runtime::with_current(|r| r.driver_type().is_polling()) {
1091            stream.set_nonblocking(true)?;
1092        }
1093
1094        Ok(Self {
1095            inner: Socket::from_socket2(Socket2::from(stream))?,
1096        })
1097    }
1098}
1099
1100impl_raw_fd!(TcpSocket, socket2::Socket, inner, socket);