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

compio_net/
tcp.rs
