diff options
Diffstat (limited to 'third_party/rust/tokio/src/net/tcp')
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/listener.rs | 397 | ||||
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/mod.rs | 14 | ||||
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/socket.rs | 690 | ||||
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/split.rs | 401 | ||||
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/split_owned.rs | 485 | ||||
| -rw-r--r-- | third_party/rust/tokio/src/net/tcp/stream.rs | 1310 |
6 files changed, 3297 insertions, 0 deletions
diff --git a/third_party/rust/tokio/src/net/tcp/listener.rs b/third_party/rust/tokio/src/net/tcp/listener.rs new file mode 100644 index 0000000000..8aecb21aaa --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/listener.rs @@ -0,0 +1,397 @@ +use crate::io::{Interest, PollEvented}; +use crate::net::tcp::TcpStream; +use crate::net::{to_socket_addrs, ToSocketAddrs}; + +use std::convert::TryFrom; +use std::fmt; +use std::io; +use std::net::{self, SocketAddr}; +use std::task::{Context, Poll}; + +cfg_net! { + /// A TCP socket server, listening for connections. + /// + /// You can accept a new connection by using the [`accept`](`TcpListener::accept`) + /// method. + /// + /// A `TcpListener` can be turned into a `Stream` with [`TcpListenerStream`]. + /// + /// [`TcpListenerStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.TcpListenerStream.html + /// + /// # Errors + /// + /// Note that accepting a connection can lead to various errors and not all + /// of them are necessarily fatal ‒ for example having too many open file + /// descriptors or the other side closing the connection while it waits in + /// an accept queue. These would terminate the stream if not handled in any + /// way. + /// + /// # Examples + /// + /// Using `accept`: + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// async fn process_socket<T>(socket: T) { + /// # drop(socket); + /// // do work with socket here + /// } + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:8080").await?; + /// + /// loop { + /// let (socket, _) = listener.accept().await?; + /// process_socket(socket).await; + /// } + /// } + /// ``` + pub struct TcpListener { + io: PollEvented<mio::net::TcpListener>, + } +} + +impl TcpListener { + /// Creates a new TcpListener, which will be bound to the specified address. + /// + /// The returned listener is ready for accepting connections. + /// + /// Binding with a port number of 0 will request that the OS assigns a port + /// to this listener. The port allocated can be queried via the `local_addr` + /// method. + /// + /// The address type can be any implementor of the [`ToSocketAddrs`] trait. + /// If `addr` yields multiple addresses, bind will be attempted with each of + /// the addresses until one succeeds and returns the listener. If none of + /// the addresses succeed in creating a listener, the error returned from + /// the last attempt (the last address) is returned. + /// + /// This function sets the `SO_REUSEADDR` option on the socket. + /// + /// To configure the socket before binding, you can use the [`TcpSocket`] + /// type. + /// + /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs + /// [`TcpSocket`]: struct@crate::net::TcpSocket + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:2345").await?; + /// + /// // use the listener + /// + /// # let _ = listener; + /// Ok(()) + /// } + /// ``` + pub async fn bind<A: ToSocketAddrs>(addr: A) -> io::Result<TcpListener> { + let addrs = to_socket_addrs(addr).await?; + + let mut last_err = None; + + for addr in addrs { + match TcpListener::bind_addr(addr) { + Ok(listener) => return Ok(listener), + Err(e) => last_err = Some(e), + } + } + + Err(last_err.unwrap_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "could not resolve to any address", + ) + })) + } + + fn bind_addr(addr: SocketAddr) -> io::Result<TcpListener> { + let listener = mio::net::TcpListener::bind(addr)?; + TcpListener::new(listener) + } + + /// Accepts a new incoming connection from this listener. + /// + /// This function will yield once a new TCP connection is established. When + /// established, the corresponding [`TcpStream`] and the remote peer's + /// address will be returned. + /// + /// # Cancel safety + /// + /// This method is cancel safe. If the method is used as the event in a + /// [`tokio::select!`](crate::select) statement and some other branch + /// completes first, then it is guaranteed that no new connections were + /// accepted by this method. + /// + /// [`TcpStream`]: struct@crate::net::TcpStream + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:8080").await?; + /// + /// match listener.accept().await { + /// Ok((_socket, addr)) => println!("new client: {:?}", addr), + /// Err(e) => println!("couldn't get client: {:?}", e), + /// } + /// + /// Ok(()) + /// } + /// ``` + pub async fn accept(&self) -> io::Result<(TcpStream, SocketAddr)> { + let (mio, addr) = self + .io + .registration() + .async_io(Interest::READABLE, || self.io.accept()) + .await?; + + let stream = TcpStream::new(mio)?; + Ok((stream, addr)) + } + + /// Polls to accept a new incoming connection to this listener. + /// + /// If there is no connection to accept, `Poll::Pending` is returned and the + /// current task will be notified by a waker. Note that on multiple calls + /// to `poll_accept`, only the `Waker` from the `Context` passed to the most + /// recent call is scheduled to receive a wakeup. + pub fn poll_accept(&self, cx: &mut Context<'_>) -> Poll<io::Result<(TcpStream, SocketAddr)>> { + loop { + let ev = ready!(self.io.registration().poll_read_ready(cx))?; + + match self.io.accept() { + Ok((io, addr)) => { + let io = TcpStream::new(io)?; + return Poll::Ready(Ok((io, addr))); + } + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.io.registration().clear_readiness(ev); + } + Err(e) => return Poll::Ready(Err(e)), + } + } + } + + /// Creates new `TcpListener` from a `std::net::TcpListener`. + /// + /// This function is intended to be used to wrap a TCP listener from the + /// standard library in the Tokio equivalent. The conversion assumes nothing + /// about the underlying listener; it is left up to the user to set it in + /// non-blocking mode. + /// + /// This API is typically paired with the `socket2` crate and the `Socket` + /// type to build up and customize a listener before it's shipped off to the + /// backing event loop. This allows configuration of options like + /// `SO_REUSEPORT`, binding to multiple addresses, etc. + /// + /// # Examples + /// + /// ```rust,no_run + /// use std::error::Error; + /// use tokio::net::TcpListener; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let std_listener = std::net::TcpListener::bind("127.0.0.1:0")?; + /// std_listener.set_nonblocking(true)?; + /// let listener = TcpListener::from_std(std_listener)?; + /// Ok(()) + /// } + /// ``` + /// + /// # Panics + /// + /// This function panics if thread-local runtime is not set. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_std(listener: net::TcpListener) -> io::Result<TcpListener> { + let io = mio::net::TcpListener::from_std(listener); + let io = PollEvented::new(io)?; + Ok(TcpListener { io }) + } + + /// Turns a [`tokio::net::TcpListener`] into a [`std::net::TcpListener`]. + /// + /// The returned [`std::net::TcpListener`] will have nonblocking mode set as + /// `true`. Use [`set_nonblocking`] to change the blocking mode if needed. + /// + /// # Examples + /// + /// ```rust,no_run + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let tokio_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?; + /// let std_listener = tokio_listener.into_std()?; + /// std_listener.set_nonblocking(false)?; + /// Ok(()) + /// } + /// ``` + /// + /// [`tokio::net::TcpListener`]: TcpListener + /// [`std::net::TcpListener`]: std::net::TcpListener + /// [`set_nonblocking`]: fn@std::net::TcpListener::set_nonblocking + pub fn into_std(self) -> io::Result<std::net::TcpListener> { + #[cfg(unix)] + { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + self.io + .into_inner() + .map(|io| io.into_raw_fd()) + .map(|raw_fd| unsafe { std::net::TcpListener::from_raw_fd(raw_fd) }) + } + + #[cfg(windows)] + { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + self.io + .into_inner() + .map(|io| io.into_raw_socket()) + .map(|raw_socket| unsafe { std::net::TcpListener::from_raw_socket(raw_socket) }) + } + } + + pub(crate) fn new(listener: mio::net::TcpListener) -> io::Result<TcpListener> { + let io = PollEvented::new(listener)?; + Ok(TcpListener { io }) + } + + /// Returns the local address that this listener is bound to. + /// + /// This can be useful, for example, when binding to port 0 to figure out + /// which port was actually bound. + /// + /// # Examples + /// + /// ```rust,no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:8080").await?; + /// + /// assert_eq!(listener.local_addr()?, + /// SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080))); + /// + /// Ok(()) + /// } + /// ``` + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.io.local_addr() + } + + /// Gets the value of the `IP_TTL` option for this socket. + /// + /// For more information about this option, see [`set_ttl`]. + /// + /// [`set_ttl`]: method@Self::set_ttl + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:0").await?; + /// + /// listener.set_ttl(100).expect("could not set TTL"); + /// assert_eq!(listener.ttl()?, 100); + /// + /// Ok(()) + /// } + /// ``` + pub fn ttl(&self) -> io::Result<u32> { + self.io.ttl() + } + + /// Sets the value for the `IP_TTL` option on this socket. + /// + /// This value sets the time-to-live field that is used in every packet sent + /// from this socket. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpListener; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let listener = TcpListener::bind("127.0.0.1:0").await?; + /// + /// listener.set_ttl(100).expect("could not set TTL"); + /// + /// Ok(()) + /// } + /// ``` + pub fn set_ttl(&self, ttl: u32) -> io::Result<()> { + self.io.set_ttl(ttl) + } +} + +impl TryFrom<net::TcpListener> for TcpListener { + type Error = io::Error; + + /// Consumes stream, returning the tokio I/O object. + /// + /// This is equivalent to + /// [`TcpListener::from_std(stream)`](TcpListener::from_std). + fn try_from(stream: net::TcpListener) -> Result<Self, Self::Error> { + Self::from_std(stream) + } +} + +impl fmt::Debug for TcpListener { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.io.fmt(f) + } +} + +#[cfg(unix)] +mod sys { + use super::TcpListener; + use std::os::unix::prelude::*; + + impl AsRawFd for TcpListener { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } + } +} + +#[cfg(windows)] +mod sys { + use super::TcpListener; + use std::os::windows::prelude::*; + + impl AsRawSocket for TcpListener { + fn as_raw_socket(&self) -> RawSocket { + self.io.as_raw_socket() + } + } +} diff --git a/third_party/rust/tokio/src/net/tcp/mod.rs b/third_party/rust/tokio/src/net/tcp/mod.rs new file mode 100644 index 0000000000..cb8a8b238b --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/mod.rs @@ -0,0 +1,14 @@ +//! TCP utility types. + +pub(crate) mod listener; + +pub(crate) mod socket; + +mod split; +pub use split::{ReadHalf, WriteHalf}; + +mod split_owned; +pub use split_owned::{OwnedReadHalf, OwnedWriteHalf, ReuniteError}; + +pub(crate) mod stream; +pub(crate) use stream::TcpStream; diff --git a/third_party/rust/tokio/src/net/tcp/socket.rs b/third_party/rust/tokio/src/net/tcp/socket.rs new file mode 100644 index 0000000000..171e240189 --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/socket.rs @@ -0,0 +1,690 @@ +use crate::net::{TcpListener, TcpStream}; + +use std::fmt; +use std::io; +use std::net::SocketAddr; + +#[cfg(unix)] +use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd}; +#[cfg(windows)] +use std::os::windows::io::{AsRawSocket, FromRawSocket, IntoRawSocket, RawSocket}; +use std::time::Duration; + +cfg_net! { + /// A TCP socket that has not yet been converted to a `TcpStream` or + /// `TcpListener`. + /// + /// `TcpSocket` wraps an operating system socket and enables the caller to + /// configure the socket before establishing a TCP connection or accepting + /// inbound connections. The caller is able to set socket option and explicitly + /// bind the socket with a socket address. + /// + /// The underlying socket is closed when the `TcpSocket` value is dropped. + /// + /// `TcpSocket` should only be used directly if the default configuration used + /// by `TcpStream::connect` and `TcpListener::bind` does not meet the required + /// use case. + /// + /// Calling `TcpStream::connect("127.0.0.1:8080")` is equivalent to: + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// let stream = socket.connect(addr).await?; + /// # drop(stream); + /// + /// Ok(()) + /// } + /// ``` + /// + /// Calling `TcpListener::bind("127.0.0.1:8080")` is equivalent to: + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// // On platforms with Berkeley-derived sockets, this allows to quickly + /// // rebind a socket, without needing to wait for the OS to clean up the + /// // previous one. + /// // + /// // On Windows, this allows rebinding sockets which are actively in use, + /// // which allows “socket hijacking”, so we explicitly don't set it here. + /// // https://docs.microsoft.com/en-us/windows/win32/winsock/using-so-reuseaddr-and-so-exclusiveaddruse + /// socket.set_reuseaddr(true)?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// # drop(listener); + /// + /// Ok(()) + /// } + /// ``` + /// + /// Setting socket options not explicitly provided by `TcpSocket` may be done by + /// accessing the `RawFd`/`RawSocket` using [`AsRawFd`]/[`AsRawSocket`] and + /// setting the option with a crate like [`socket2`]. + /// + /// [`RawFd`]: https://doc.rust-lang.org/std/os/unix/io/type.RawFd.html + /// [`RawSocket`]: https://doc.rust-lang.org/std/os/windows/io/type.RawSocket.html + /// [`AsRawFd`]: https://doc.rust-lang.org/std/os/unix/io/trait.AsRawFd.html + /// [`AsRawSocket`]: https://doc.rust-lang.org/std/os/windows/io/trait.AsRawSocket.html + /// [`socket2`]: https://docs.rs/socket2/ + #[cfg_attr(docsrs, doc(alias = "connect_std"))] + pub struct TcpSocket { + inner: socket2::Socket, + } +} + +impl TcpSocket { + /// Creates a new socket configured for IPv4. + /// + /// Calls `socket(2)` with `AF_INET` and `SOCK_STREAM`. + /// + /// # Returns + /// + /// On success, the newly created `TcpSocket` is returned. If an error is + /// encountered, it is returned instead. + /// + /// # Examples + /// + /// Create a new IPv4 socket and start listening. + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// let socket = TcpSocket::new_v4()?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(128)?; + /// # drop(listener); + /// Ok(()) + /// } + /// ``` + pub fn new_v4() -> io::Result<TcpSocket> { + TcpSocket::new(socket2::Domain::IPV4) + } + + /// Creates a new socket configured for IPv6. + /// + /// Calls `socket(2)` with `AF_INET6` and `SOCK_STREAM`. + /// + /// # Returns + /// + /// On success, the newly created `TcpSocket` is returned. If an error is + /// encountered, it is returned instead. + /// + /// # Examples + /// + /// Create a new IPv6 socket and start listening. + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "[::1]:8080".parse().unwrap(); + /// let socket = TcpSocket::new_v6()?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(128)?; + /// # drop(listener); + /// Ok(()) + /// } + /// ``` + pub fn new_v6() -> io::Result<TcpSocket> { + TcpSocket::new(socket2::Domain::IPV6) + } + + fn new(domain: socket2::Domain) -> io::Result<TcpSocket> { + let ty = socket2::Type::STREAM; + #[cfg(any( + target_os = "android", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "fuchsia", + target_os = "illumos", + target_os = "linux", + target_os = "netbsd", + target_os = "openbsd" + ))] + let ty = ty.nonblocking(); + let inner = socket2::Socket::new(domain, ty, Some(socket2::Protocol::TCP))?; + #[cfg(not(any( + target_os = "android", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "fuchsia", + target_os = "illumos", + target_os = "linux", + target_os = "netbsd", + target_os = "openbsd" + )))] + inner.set_nonblocking(true)?; + Ok(TcpSocket { inner }) + } + + /// Allows the socket to bind to an in-use address. + /// + /// Behavior is platform specific. Refer to the target platform's + /// documentation for more details. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.set_reuseaddr(true)?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// # drop(listener); + /// + /// Ok(()) + /// } + /// ``` + pub fn set_reuseaddr(&self, reuseaddr: bool) -> io::Result<()> { + self.inner.set_reuse_address(reuseaddr) + } + + /// Retrieves the value set for `SO_REUSEADDR` on this socket. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.set_reuseaddr(true)?; + /// assert!(socket.reuseaddr().unwrap()); + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// Ok(()) + /// } + /// ``` + pub fn reuseaddr(&self) -> io::Result<bool> { + self.inner.reuse_address() + } + + /// Allows the socket to bind to an in-use port. Only available for unix systems + /// (excluding Solaris & Illumos). + /// + /// Behavior is platform specific. Refer to the target platform's + /// documentation for more details. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.set_reuseport(true)?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// Ok(()) + /// } + /// ``` + #[cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))] + #[cfg_attr( + docsrs, + doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) + )] + pub fn set_reuseport(&self, reuseport: bool) -> io::Result<()> { + self.inner.set_reuse_port(reuseport) + } + + /// Allows the socket to bind to an in-use port. Only available for unix systems + /// (excluding Solaris & Illumos). + /// + /// Behavior is platform specific. Refer to the target platform's + /// documentation for more details. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.set_reuseport(true)?; + /// assert!(socket.reuseport().unwrap()); + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// Ok(()) + /// } + /// ``` + #[cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))] + #[cfg_attr( + docsrs, + doc(cfg(all(unix, not(target_os = "solaris"), not(target_os = "illumos")))) + )] + pub fn reuseport(&self) -> io::Result<bool> { + self.inner.reuse_port() + } + + /// Sets the size of the TCP send buffer on this socket. + /// + /// On most operating systems, this sets the `SO_SNDBUF` socket option. + pub fn set_send_buffer_size(&self, size: u32) -> io::Result<()> { + self.inner.set_send_buffer_size(size as usize) + } + + /// Returns the size of the TCP send buffer for this socket. + /// + /// On most operating systems, this is the value of the `SO_SNDBUF` socket + /// option. + /// + /// Note that if [`set_send_buffer_size`] has been called on this socket + /// previously, the value returned by this function may not be the same as + /// the argument provided to `set_send_buffer_size`. This is for the + /// following reasons: + /// + /// * Most operating systems have minimum and maximum allowed sizes for the + /// send buffer, and will clamp the provided value if it is below the + /// minimum or above the maximum. The minimum and maximum buffer sizes are + /// OS-dependent. + /// * Linux will double the buffer size to account for internal bookkeeping + /// data, and returns the doubled value from `getsockopt(2)`. As per `man + /// 7 socket`: + /// > Sets or gets the maximum socket send buffer in bytes. The + /// > kernel doubles this value (to allow space for bookkeeping + /// > overhead) when it is set using `setsockopt(2)`, and this doubled + /// > value is returned by `getsockopt(2)`. + /// + /// [`set_send_buffer_size`]: #method.set_send_buffer_size + pub fn send_buffer_size(&self) -> io::Result<u32> { + self.inner.send_buffer_size().map(|n| n as u32) + } + + /// Sets the size of the TCP receive buffer on this socket. + /// + /// On most operating systems, this sets the `SO_RCVBUF` socket option. + pub fn set_recv_buffer_size(&self, size: u32) -> io::Result<()> { + self.inner.set_recv_buffer_size(size as usize) + } + + /// Returns the size of the TCP receive buffer for this socket. + /// + /// On most operating systems, this is the value of the `SO_RCVBUF` socket + /// option. + /// + /// Note that if [`set_recv_buffer_size`] has been called on this socket + /// previously, the value returned by this function may not be the same as + /// the argument provided to `set_send_buffer_size`. This is for the + /// following reasons: + /// + /// * Most operating systems have minimum and maximum allowed sizes for the + /// receive buffer, and will clamp the provided value if it is below the + /// minimum or above the maximum. The minimum and maximum buffer sizes are + /// OS-dependent. + /// * Linux will double the buffer size to account for internal bookkeeping + /// data, and returns the doubled value from `getsockopt(2)`. As per `man + /// 7 socket`: + /// > Sets or gets the maximum socket send buffer in bytes. The + /// > kernel doubles this value (to allow space for bookkeeping + /// > overhead) when it is set using `setsockopt(2)`, and this doubled + /// > value is returned by `getsockopt(2)`. + /// + /// [`set_recv_buffer_size`]: #method.set_recv_buffer_size + pub fn recv_buffer_size(&self) -> io::Result<u32> { + self.inner.recv_buffer_size().map(|n| n as u32) + } + + /// Sets the linger duration of this socket by setting the SO_LINGER option. + /// + /// This option controls the action taken when a stream has unsent messages and the stream is + /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the + /// data or until the time expires. + /// + /// If SO_LINGER is not specified, and the socket is closed, the system handles the call in a + /// way that allows the process to continue as quickly as possible. + pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> { + self.inner.set_linger(dur) + } + + /// Reads the linger duration for this socket by getting the `SO_LINGER` + /// option. + /// + /// For more information about this option, see [`set_linger`]. + /// + /// [`set_linger`]: TcpSocket::set_linger + pub fn linger(&self) -> io::Result<Option<Duration>> { + self.inner.linger() + } + + /// Gets the local address of this socket. + /// + /// Will fail on windows if called before `bind`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.bind(addr)?; + /// assert_eq!(socket.local_addr().unwrap().to_string(), "127.0.0.1:8080"); + /// let listener = socket.listen(1024)?; + /// Ok(()) + /// } + /// ``` + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr().and_then(convert_address) + } + + /// Binds the socket to the given address. + /// + /// This calls the `bind(2)` operating-system function. Behavior is + /// platform specific. Refer to the target platform's documentation for more + /// details. + /// + /// # Examples + /// + /// Bind a socket before listening. + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// # drop(listener); + /// + /// Ok(()) + /// } + /// ``` + pub fn bind(&self, addr: SocketAddr) -> io::Result<()> { + self.inner.bind(&addr.into()) + } + + /// Establishes a TCP connection with a peer at the specified socket address. + /// + /// The `TcpSocket` is consumed. Once the connection is established, a + /// connected [`TcpStream`] is returned. If the connection fails, the + /// encountered error is returned. + /// + /// [`TcpStream`]: TcpStream + /// + /// This calls the `connect(2)` operating-system function. Behavior is + /// platform specific. Refer to the target platform's documentation for more + /// details. + /// + /// # Examples + /// + /// Connecting to a peer. + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// let stream = socket.connect(addr).await?; + /// # drop(stream); + /// + /// Ok(()) + /// } + /// ``` + pub async fn connect(self, addr: SocketAddr) -> io::Result<TcpStream> { + if let Err(err) = self.inner.connect(&addr.into()) { + #[cfg(unix)] + if err.raw_os_error() != Some(libc::EINPROGRESS) { + return Err(err); + } + #[cfg(windows)] + if err.kind() != io::ErrorKind::WouldBlock { + return Err(err); + } + } + #[cfg(unix)] + let mio = { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + + let raw_fd = self.inner.into_raw_fd(); + unsafe { mio::net::TcpStream::from_raw_fd(raw_fd) } + }; + + #[cfg(windows)] + let mio = { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + + let raw_socket = self.inner.into_raw_socket(); + unsafe { mio::net::TcpStream::from_raw_socket(raw_socket) } + }; + + TcpStream::connect_mio(mio).await + } + + /// Converts the socket into a `TcpListener`. + /// + /// `backlog` defines the maximum number of pending connections are queued + /// by the operating system at any given time. Connection are removed from + /// the queue with [`TcpListener::accept`]. When the queue is full, the + /// operating-system will start rejecting connections. + /// + /// [`TcpListener::accept`]: TcpListener::accept + /// + /// This calls the `listen(2)` operating-system function, marking the socket + /// as a passive socket. Behavior is platform specific. Refer to the target + /// platform's documentation for more details. + /// + /// # Examples + /// + /// Create a `TcpListener`. + /// + /// ```no_run + /// use tokio::net::TcpSocket; + /// + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let addr = "127.0.0.1:8080".parse().unwrap(); + /// + /// let socket = TcpSocket::new_v4()?; + /// socket.bind(addr)?; + /// + /// let listener = socket.listen(1024)?; + /// # drop(listener); + /// + /// Ok(()) + /// } + /// ``` + pub fn listen(self, backlog: u32) -> io::Result<TcpListener> { + self.inner.listen(backlog as i32)?; + #[cfg(unix)] + let mio = { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + + let raw_fd = self.inner.into_raw_fd(); + unsafe { mio::net::TcpListener::from_raw_fd(raw_fd) } + }; + + #[cfg(windows)] + let mio = { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + + let raw_socket = self.inner.into_raw_socket(); + unsafe { mio::net::TcpListener::from_raw_socket(raw_socket) } + }; + + TcpListener::new(mio) + } + + /// Converts a [`std::net::TcpStream`] into a `TcpSocket`. The provided + /// socket must not have been connected prior to calling this function. This + /// function is typically used together with crates such as [`socket2`] to + /// configure socket options that are not available on `TcpSocket`. + /// + /// [`std::net::TcpStream`]: struct@std::net::TcpStream + /// [`socket2`]: https://docs.rs/socket2/ + /// + /// # Examples + /// + /// ``` + /// use tokio::net::TcpSocket; + /// use socket2::{Domain, Socket, Type}; + /// + /// #[tokio::main] + /// async fn main() -> std::io::Result<()> { + /// + /// let socket2_socket = Socket::new(Domain::IPV4, Type::STREAM, None)?; + /// + /// let socket = TcpSocket::from_std_stream(socket2_socket.into()); + /// + /// Ok(()) + /// } + /// ``` + pub fn from_std_stream(std_stream: std::net::TcpStream) -> TcpSocket { + #[cfg(unix)] + { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + + let raw_fd = std_stream.into_raw_fd(); + unsafe { TcpSocket::from_raw_fd(raw_fd) } + } + + #[cfg(windows)] + { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + + let raw_socket = std_stream.into_raw_socket(); + unsafe { TcpSocket::from_raw_socket(raw_socket) } + } + } +} + +fn convert_address(address: socket2::SockAddr) -> io::Result<SocketAddr> { + match address.as_socket() { + Some(address) => Ok(address), + None => Err(io::Error::new( + io::ErrorKind::InvalidInput, + "invalid address family (not IPv4 or IPv6)", + )), + } +} + +impl fmt::Debug for TcpSocket { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + self.inner.fmt(fmt) + } +} + +#[cfg(unix)] +impl AsRawFd for TcpSocket { + fn as_raw_fd(&self) -> RawFd { + self.inner.as_raw_fd() + } +} + +#[cfg(unix)] +impl FromRawFd for TcpSocket { + /// Converts a `RawFd` to a `TcpSocket`. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socket is in + /// non-blocking mode. + unsafe fn from_raw_fd(fd: RawFd) -> TcpSocket { + let inner = socket2::Socket::from_raw_fd(fd); + TcpSocket { inner } + } +} + +#[cfg(unix)] +impl IntoRawFd for TcpSocket { + fn into_raw_fd(self) -> RawFd { + self.inner.into_raw_fd() + } +} + +#[cfg(windows)] +impl IntoRawSocket for TcpSocket { + fn into_raw_socket(self) -> RawSocket { + self.inner.into_raw_socket() + } +} + +#[cfg(windows)] +impl AsRawSocket for TcpSocket { + fn as_raw_socket(&self) -> RawSocket { + self.inner.as_raw_socket() + } +} + +#[cfg(windows)] +impl FromRawSocket for TcpSocket { + /// Converts a `RawSocket` to a `TcpStream`. + /// + /// # Notes + /// + /// The caller is responsible for ensuring that the socket is in + /// non-blocking mode. + unsafe fn from_raw_socket(socket: RawSocket) -> TcpSocket { + let inner = socket2::Socket::from_raw_socket(socket); + TcpSocket { inner } + } +} diff --git a/third_party/rust/tokio/src/net/tcp/split.rs b/third_party/rust/tokio/src/net/tcp/split.rs new file mode 100644 index 0000000000..0e02928495 --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/split.rs @@ -0,0 +1,401 @@ +//! `TcpStream` split support. +//! +//! A `TcpStream` can be split into a `ReadHalf` and a +//! `WriteHalf` with the `TcpStream::split` method. `ReadHalf` +//! implements `AsyncRead` while `WriteHalf` implements `AsyncWrite`. +//! +//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized +//! split has no associated overhead and enforces all invariants at the type +//! level. + +use crate::future::poll_fn; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; +use crate::net::TcpStream; + +use std::io; +use std::net::{Shutdown, SocketAddr}; +use std::pin::Pin; +use std::task::{Context, Poll}; + +cfg_io_util! { + use bytes::BufMut; +} + +/// Borrowed read half of a [`TcpStream`], created by [`split`]. +/// +/// Reading from a `ReadHalf` is usually done using the convenience methods found on the +/// [`AsyncReadExt`] trait. +/// +/// [`TcpStream`]: TcpStream +/// [`split`]: TcpStream::split() +/// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt +#[derive(Debug)] +pub struct ReadHalf<'a>(&'a TcpStream); + +/// Borrowed write half of a [`TcpStream`], created by [`split`]. +/// +/// Note that in the [`AsyncWrite`] implementation of this type, [`poll_shutdown`] will +/// shut down the TCP stream in the write direction. +/// +/// Writing to an `WriteHalf` is usually done using the convenience methods found +/// on the [`AsyncWriteExt`] trait. +/// +/// [`TcpStream`]: TcpStream +/// [`split`]: TcpStream::split() +/// [`AsyncWrite`]: trait@crate::io::AsyncWrite +/// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown +/// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt +#[derive(Debug)] +pub struct WriteHalf<'a>(&'a TcpStream); + +pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) { + (ReadHalf(&*stream), WriteHalf(&*stream)) +} + +impl ReadHalf<'_> { + /// Attempts to receive data on the socket, without removing that data from + /// the queue, registering the current task for wakeup if data is not yet + /// available. + /// + /// Note that on multiple calls to `poll_peek` or `poll_read`, only the + /// `Waker` from the `Context` passed to the most recent call is scheduled + /// to receive a wakeup. + /// + /// See the [`TcpStream::poll_peek`] level documentation for more details. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, ReadBuf}; + /// use tokio::net::TcpStream; + /// + /// use futures::future::poll_fn; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let mut stream = TcpStream::connect("127.0.0.1:8000").await?; + /// let (mut read_half, _) = stream.split(); + /// let mut buf = [0; 10]; + /// let mut buf = ReadBuf::new(&mut buf); + /// + /// poll_fn(|cx| { + /// read_half.poll_peek(cx, &mut buf) + /// }).await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// [`TcpStream::poll_peek`]: TcpStream::poll_peek + pub fn poll_peek( + &mut self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<usize>> { + self.0.poll_peek(cx, buf) + } + + /// Receives data on the socket from the remote address to which it is + /// connected, without removing that data from the queue. On success, + /// returns the number of bytes peeked. + /// + /// See the [`TcpStream::peek`] level documentation for more details. + /// + /// [`TcpStream::peek`]: TcpStream::peek + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncReadExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; + /// let (mut read_half, _) = stream.split(); + /// + /// let mut b1 = [0; 10]; + /// let mut b2 = [0; 10]; + /// + /// // Peek at the data + /// let n = read_half.peek(&mut b1).await?; + /// + /// // Read the data + /// assert_eq!(n, read_half.read(&mut b2[..n]).await?); + /// assert_eq!(&b1[..n], &b2[..n]); + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`read`] method is defined on the [`AsyncReadExt`] trait. + /// + /// [`read`]: fn@crate::io::AsyncReadExt::read + /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt + pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> { + let mut buf = ReadBuf::new(buf); + poll_fn(|cx| self.poll_peek(cx, &mut buf)).await + } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with `try_read()` or `try_write()`. It + /// can be used to concurrently read / write to the same socket on a single + /// task without splitting the socket. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// This function is also equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.0.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.0.try_read(buf) + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.0.try_read_vectored(bufs) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.0.try_read_buf(buf) + } + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } +} + +impl WriteHalf<'_> { + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with `try_read()` or `try_write()`. It + /// can be used to concurrently read / write to the same socket on a single + /// task without splitting the socket. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.0.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.0.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.0.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.0.try_write_vectored(bufs) + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.0.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.0.local_addr() + } +} + +impl AsyncRead for ReadHalf<'_> { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.0.poll_read_priv(cx, buf) + } +} + +impl AsyncWrite for WriteHalf<'_> { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.0.poll_write_priv(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.0.poll_write_vectored_priv(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.0.is_write_vectored() + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + // tcp flush is a no-op + Poll::Ready(Ok(())) + } + + // `poll_shutdown` on a write half shutdowns the stream in the "write" direction. + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + self.0.shutdown_std(Shutdown::Write).into() + } +} + +impl AsRef<TcpStream> for ReadHalf<'_> { + fn as_ref(&self) -> &TcpStream { + self.0 + } +} + +impl AsRef<TcpStream> for WriteHalf<'_> { + fn as_ref(&self) -> &TcpStream { + self.0 + } +} diff --git a/third_party/rust/tokio/src/net/tcp/split_owned.rs b/third_party/rust/tokio/src/net/tcp/split_owned.rs new file mode 100644 index 0000000000..ef4e7b5361 --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/split_owned.rs @@ -0,0 +1,485 @@ +//! `TcpStream` owned split support. +//! +//! A `TcpStream` can be split into an `OwnedReadHalf` and a `OwnedWriteHalf` +//! with the `TcpStream::into_split` method. `OwnedReadHalf` implements +//! `AsyncRead` while `OwnedWriteHalf` implements `AsyncWrite`. +//! +//! Compared to the generic split of `AsyncRead + AsyncWrite`, this specialized +//! split has no associated overhead and enforces all invariants at the type +//! level. + +use crate::future::poll_fn; +use crate::io::{AsyncRead, AsyncWrite, Interest, ReadBuf, Ready}; +use crate::net::TcpStream; + +use std::error::Error; +use std::net::{Shutdown, SocketAddr}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::{fmt, io}; + +cfg_io_util! { + use bytes::BufMut; +} + +/// Owned read half of a [`TcpStream`], created by [`into_split`]. +/// +/// Reading from an `OwnedReadHalf` is usually done using the convenience methods found +/// on the [`AsyncReadExt`] trait. +/// +/// [`TcpStream`]: TcpStream +/// [`into_split`]: TcpStream::into_split() +/// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt +#[derive(Debug)] +pub struct OwnedReadHalf { + inner: Arc<TcpStream>, +} + +/// Owned write half of a [`TcpStream`], created by [`into_split`]. +/// +/// Note that in the [`AsyncWrite`] implementation of this type, [`poll_shutdown`] will +/// shut down the TCP stream in the write direction. Dropping the write half +/// will also shut down the write half of the TCP stream. +/// +/// Writing to an `OwnedWriteHalf` is usually done using the convenience methods found +/// on the [`AsyncWriteExt`] trait. +/// +/// [`TcpStream`]: TcpStream +/// [`into_split`]: TcpStream::into_split() +/// [`AsyncWrite`]: trait@crate::io::AsyncWrite +/// [`poll_shutdown`]: fn@crate::io::AsyncWrite::poll_shutdown +/// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt +#[derive(Debug)] +pub struct OwnedWriteHalf { + inner: Arc<TcpStream>, + shutdown_on_drop: bool, +} + +pub(crate) fn split_owned(stream: TcpStream) -> (OwnedReadHalf, OwnedWriteHalf) { + let arc = Arc::new(stream); + let read = OwnedReadHalf { + inner: Arc::clone(&arc), + }; + let write = OwnedWriteHalf { + inner: arc, + shutdown_on_drop: true, + }; + (read, write) +} + +pub(crate) fn reunite( + read: OwnedReadHalf, + write: OwnedWriteHalf, +) -> Result<TcpStream, ReuniteError> { + if Arc::ptr_eq(&read.inner, &write.inner) { + write.forget(); + // This unwrap cannot fail as the api does not allow creating more than two Arcs, + // and we just dropped the other half. + Ok(Arc::try_unwrap(read.inner).expect("TcpStream: try_unwrap failed in reunite")) + } else { + Err(ReuniteError(read, write)) + } +} + +/// Error indicating that two halves were not from the same socket, and thus could +/// not be reunited. +#[derive(Debug)] +pub struct ReuniteError(pub OwnedReadHalf, pub OwnedWriteHalf); + +impl fmt::Display for ReuniteError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "tried to reunite halves that are not from the same socket" + ) + } +} + +impl Error for ReuniteError {} + +impl OwnedReadHalf { + /// Attempts to put the two halves of a `TcpStream` back together and + /// recover the original socket. Succeeds only if the two halves + /// originated from the same call to [`into_split`]. + /// + /// [`into_split`]: TcpStream::into_split() + pub fn reunite(self, other: OwnedWriteHalf) -> Result<TcpStream, ReuniteError> { + reunite(self, other) + } + + /// Attempt to receive data on the socket, without removing that data from + /// the queue, registering the current task for wakeup if data is not yet + /// available. + /// + /// Note that on multiple calls to `poll_peek` or `poll_read`, only the + /// `Waker` from the `Context` passed to the most recent call is scheduled + /// to receive a wakeup. + /// + /// See the [`TcpStream::poll_peek`] level documentation for more details. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, ReadBuf}; + /// use tokio::net::TcpStream; + /// + /// use futures::future::poll_fn; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let stream = TcpStream::connect("127.0.0.1:8000").await?; + /// let (mut read_half, _) = stream.into_split(); + /// let mut buf = [0; 10]; + /// let mut buf = ReadBuf::new(&mut buf); + /// + /// poll_fn(|cx| { + /// read_half.poll_peek(cx, &mut buf) + /// }).await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// [`TcpStream::poll_peek`]: TcpStream::poll_peek + pub fn poll_peek( + &mut self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<usize>> { + self.inner.poll_peek(cx, buf) + } + + /// Receives data on the socket from the remote address to which it is + /// connected, without removing that data from the queue. On success, + /// returns the number of bytes peeked. + /// + /// See the [`TcpStream::peek`] level documentation for more details. + /// + /// [`TcpStream::peek`]: TcpStream::peek + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncReadExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// let (mut read_half, _) = stream.into_split(); + /// + /// let mut b1 = [0; 10]; + /// let mut b2 = [0; 10]; + /// + /// // Peek at the data + /// let n = read_half.peek(&mut b1).await?; + /// + /// // Read the data + /// assert_eq!(n, read_half.read(&mut b2[..n]).await?); + /// assert_eq!(&b1[..n], &b2[..n]); + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`read`] method is defined on the [`AsyncReadExt`] trait. + /// + /// [`read`]: fn@crate::io::AsyncReadExt::read + /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt + pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> { + let mut buf = ReadBuf::new(buf); + poll_fn(|cx| self.poll_peek(cx, &mut buf)).await + } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with `try_read()` or `try_write()`. It + /// can be used to concurrently read / write to the same socket on a single + /// task without splitting the socket. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// This function is also equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn readable(&self) -> io::Result<()> { + self.inner.readable().await + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + self.inner.try_read(buf) + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: Self::try_read() + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + self.inner.try_read_vectored(bufs) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.inner.try_read_buf(buf) + } + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } +} + +impl AsyncRead for OwnedReadHalf { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.inner.poll_read_priv(cx, buf) + } +} + +impl OwnedWriteHalf { + /// Attempts to put the two halves of a `TcpStream` back together and + /// recover the original socket. Succeeds only if the two halves + /// originated from the same call to [`into_split`]. + /// + /// [`into_split`]: TcpStream::into_split() + pub fn reunite(self, other: OwnedReadHalf) -> Result<TcpStream, ReuniteError> { + reunite(other, self) + } + + /// Destroys the write half, but don't close the write half of the stream + /// until the read half is dropped. If the read half has already been + /// dropped, this closes the stream. + pub fn forget(mut self) { + self.shutdown_on_drop = false; + drop(self); + } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with `try_read()` or `try_write()`. It + /// can be used to concurrently read / write to the same socket on a single + /// task without splitting the socket. + /// + /// This function is equivalent to [`TcpStream::ready`]. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + self.inner.ready(interest).await + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + pub async fn writable(&self) -> io::Result<()> { + self.inner.writable().await + } + + /// Tries to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + self.inner.try_write(buf) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: Self::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + self.inner.try_write_vectored(bufs) + } + + /// Returns the remote address that this stream is connected to. + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.inner.peer_addr() + } + + /// Returns the local address that this stream is bound to. + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.inner.local_addr() + } +} + +impl Drop for OwnedWriteHalf { + fn drop(&mut self) { + if self.shutdown_on_drop { + let _ = self.inner.shutdown_std(Shutdown::Write); + } + } +} + +impl AsyncWrite for OwnedWriteHalf { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.inner.poll_write_priv(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.inner.poll_write_vectored_priv(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + self.inner.is_write_vectored() + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + // tcp flush is a no-op + Poll::Ready(Ok(())) + } + + // `poll_shutdown` on a write half shutdowns the stream in the "write" direction. + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + let res = self.inner.shutdown_std(Shutdown::Write); + if res.is_ok() { + Pin::into_inner(self).shutdown_on_drop = false; + } + res.into() + } +} + +impl AsRef<TcpStream> for OwnedReadHalf { + fn as_ref(&self) -> &TcpStream { + &*self.inner + } +} + +impl AsRef<TcpStream> for OwnedWriteHalf { + fn as_ref(&self) -> &TcpStream { + &*self.inner + } +} diff --git a/third_party/rust/tokio/src/net/tcp/stream.rs b/third_party/rust/tokio/src/net/tcp/stream.rs new file mode 100644 index 0000000000..ebb67b84d1 --- /dev/null +++ b/third_party/rust/tokio/src/net/tcp/stream.rs @@ -0,0 +1,1310 @@ +use crate::future::poll_fn; +use crate::io::{AsyncRead, AsyncWrite, Interest, PollEvented, ReadBuf, Ready}; +use crate::net::tcp::split::{split, ReadHalf, WriteHalf}; +use crate::net::tcp::split_owned::{split_owned, OwnedReadHalf, OwnedWriteHalf}; +use crate::net::{to_socket_addrs, ToSocketAddrs}; + +use std::convert::TryFrom; +use std::fmt; +use std::io; +use std::net::{Shutdown, SocketAddr}; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +cfg_io_util! { + use bytes::BufMut; +} + +cfg_net! { + /// A TCP stream between a local and a remote socket. + /// + /// A TCP stream can either be created by connecting to an endpoint, via the + /// [`connect`] method, or by [accepting] a connection from a [listener]. A + /// TCP stream can also be created via the [`TcpSocket`] type. + /// + /// Reading and writing to a `TcpStream` is usually done using the + /// convenience methods found on the [`AsyncReadExt`] and [`AsyncWriteExt`] + /// traits. + /// + /// [`connect`]: method@TcpStream::connect + /// [accepting]: method@crate::net::TcpListener::accept + /// [listener]: struct@crate::net::TcpListener + /// [`TcpSocket`]: struct@crate::net::TcpSocket + /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt + /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncWriteExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// // Write some data. + /// stream.write_all(b"hello world!").await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. + /// + /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all + /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt + /// + /// To shut down the stream in the write direction, you can call the + /// [`shutdown()`] method. This will cause the other peer to receive a read of + /// length 0, indicating that no more data will be sent. This only closes + /// the stream in one direction. + /// + /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown + pub struct TcpStream { + io: PollEvented<mio::net::TcpStream>, + } +} + +impl TcpStream { + /// Opens a TCP connection to a remote host. + /// + /// `addr` is an address of the remote host. Anything which implements the + /// [`ToSocketAddrs`] trait can be supplied as the address. If `addr` + /// yields multiple addresses, connect will be attempted with each of the + /// addresses until a connection is successful. If none of the addresses + /// result in a successful connection, the error returned from the last + /// connection attempt (the last address) is returned. + /// + /// To configure the socket before connecting, you can use the [`TcpSocket`] + /// type. + /// + /// [`ToSocketAddrs`]: trait@crate::net::ToSocketAddrs + /// [`TcpSocket`]: struct@crate::net::TcpSocket + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncWriteExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// // Write some data. + /// stream.write_all(b"hello world!").await?; + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`write_all`] method is defined on the [`AsyncWriteExt`] trait. + /// + /// [`write_all`]: fn@crate::io::AsyncWriteExt::write_all + /// [`AsyncWriteExt`]: trait@crate::io::AsyncWriteExt + pub async fn connect<A: ToSocketAddrs>(addr: A) -> io::Result<TcpStream> { + let addrs = to_socket_addrs(addr).await?; + + let mut last_err = None; + + for addr in addrs { + match TcpStream::connect_addr(addr).await { + Ok(stream) => return Ok(stream), + Err(e) => last_err = Some(e), + } + } + + Err(last_err.unwrap_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "could not resolve to any address", + ) + })) + } + + /// Establishes a connection to the specified `addr`. + async fn connect_addr(addr: SocketAddr) -> io::Result<TcpStream> { + let sys = mio::net::TcpStream::connect(addr)?; + TcpStream::connect_mio(sys).await + } + + pub(crate) async fn connect_mio(sys: mio::net::TcpStream) -> io::Result<TcpStream> { + let stream = TcpStream::new(sys)?; + + // Once we've connected, wait for the stream to be writable as + // that's when the actual connection has been initiated. Once we're + // writable we check for `take_socket_error` to see if the connect + // actually hit an error or not. + // + // If all that succeeded then we ship everything on up. + poll_fn(|cx| stream.io.registration().poll_write_ready(cx)).await?; + + if let Some(e) = stream.io.take_error()? { + return Err(e); + } + + Ok(stream) + } + + pub(crate) fn new(connected: mio::net::TcpStream) -> io::Result<TcpStream> { + let io = PollEvented::new(connected)?; + Ok(TcpStream { io }) + } + + /// Creates new `TcpStream` from a `std::net::TcpStream`. + /// + /// This function is intended to be used to wrap a TCP stream from the + /// standard library in the Tokio equivalent. The conversion assumes nothing + /// about the underlying stream; it is left up to the user to set it in + /// non-blocking mode. + /// + /// # Examples + /// + /// ```rust,no_run + /// use std::error::Error; + /// use tokio::net::TcpStream; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let std_stream = std::net::TcpStream::connect("127.0.0.1:34254")?; + /// std_stream.set_nonblocking(true)?; + /// let stream = TcpStream::from_std(std_stream)?; + /// Ok(()) + /// } + /// ``` + /// + /// # Panics + /// + /// This function panics if thread-local runtime is not set. + /// + /// The runtime is usually set implicitly when this function is called + /// from a future driven by a tokio runtime, otherwise runtime can be set + /// explicitly with [`Runtime::enter`](crate::runtime::Runtime::enter) function. + pub fn from_std(stream: std::net::TcpStream) -> io::Result<TcpStream> { + let io = mio::net::TcpStream::from_std(stream); + let io = PollEvented::new(io)?; + Ok(TcpStream { io }) + } + + /// Turns a [`tokio::net::TcpStream`] into a [`std::net::TcpStream`]. + /// + /// The returned [`std::net::TcpStream`] will have nonblocking mode set as `true`. + /// Use [`set_nonblocking`] to change the blocking mode if needed. + /// + /// # Examples + /// + /// ``` + /// use std::error::Error; + /// use std::io::Read; + /// use tokio::net::TcpListener; + /// # use tokio::net::TcpStream; + /// # use tokio::io::AsyncWriteExt; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let mut data = [0u8; 12]; + /// let listener = TcpListener::bind("127.0.0.1:34254").await?; + /// # let handle = tokio::spawn(async { + /// # let mut stream: TcpStream = TcpStream::connect("127.0.0.1:34254").await.unwrap(); + /// # stream.write(b"Hello world!").await.unwrap(); + /// # }); + /// let (tokio_tcp_stream, _) = listener.accept().await?; + /// let mut std_tcp_stream = tokio_tcp_stream.into_std()?; + /// # handle.await.expect("The task being joined has panicked"); + /// std_tcp_stream.set_nonblocking(false)?; + /// std_tcp_stream.read_exact(&mut data)?; + /// # assert_eq!(b"Hello world!", &data); + /// Ok(()) + /// } + /// ``` + /// [`tokio::net::TcpStream`]: TcpStream + /// [`std::net::TcpStream`]: std::net::TcpStream + /// [`set_nonblocking`]: fn@std::net::TcpStream::set_nonblocking + pub fn into_std(self) -> io::Result<std::net::TcpStream> { + #[cfg(unix)] + { + use std::os::unix::io::{FromRawFd, IntoRawFd}; + self.io + .into_inner() + .map(|io| io.into_raw_fd()) + .map(|raw_fd| unsafe { std::net::TcpStream::from_raw_fd(raw_fd) }) + } + + #[cfg(windows)] + { + use std::os::windows::io::{FromRawSocket, IntoRawSocket}; + self.io + .into_inner() + .map(|io| io.into_raw_socket()) + .map(|raw_socket| unsafe { std::net::TcpStream::from_raw_socket(raw_socket) }) + } + } + + /// Returns the local address that this stream is bound to. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.local_addr()?); + /// # Ok(()) + /// # } + /// ``` + pub fn local_addr(&self) -> io::Result<SocketAddr> { + self.io.local_addr() + } + + /// Returns the remote address that this stream is connected to. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.peer_addr()?); + /// # Ok(()) + /// # } + /// ``` + pub fn peer_addr(&self) -> io::Result<SocketAddr> { + self.io.peer_addr() + } + + /// Attempts to receive data on the socket, without removing that data from + /// the queue, registering the current task for wakeup if data is not yet + /// available. + /// + /// Note that on multiple calls to `poll_peek`, `poll_read` or + /// `poll_read_ready`, only the `Waker` from the `Context` passed to the + /// most recent call is scheduled to receive a wakeup. (However, + /// `poll_write` retains a second, independent waker.) + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if data is not yet available. + /// * `Poll::Ready(Ok(n))` if data is available. `n` is the number of bytes peeked. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::io::{self, ReadBuf}; + /// use tokio::net::TcpStream; + /// + /// use futures::future::poll_fn; + /// + /// #[tokio::main] + /// async fn main() -> io::Result<()> { + /// let stream = TcpStream::connect("127.0.0.1:8000").await?; + /// let mut buf = [0; 10]; + /// let mut buf = ReadBuf::new(&mut buf); + /// + /// poll_fn(|cx| { + /// stream.poll_peek(cx, &mut buf) + /// }).await?; + /// + /// Ok(()) + /// } + /// ``` + pub fn poll_peek( + &self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<usize>> { + loop { + let ev = ready!(self.io.registration().poll_read_ready(cx))?; + + let b = unsafe { + &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) + }; + + match self.io.peek(b) { + Ok(ret) => { + unsafe { buf.assume_init(ret) }; + buf.advance(ret); + return Poll::Ready(Ok(ret)); + } + Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + self.io.registration().clear_readiness(ev); + } + Err(e) => return Poll::Ready(Err(e)), + } + } + } + + /// Waits for any of the requested ready states. + /// + /// This function is usually paired with `try_read()` or `try_write()`. It + /// can be used to concurrently read / write to the same socket on a single + /// task without splitting the socket. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read or write that fails with `WouldBlock` or + /// `Poll::Pending`. + /// + /// # Examples + /// + /// Concurrently read and write to the stream on the same task without + /// splitting. + /// + /// ```no_run + /// use tokio::io::Interest; + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// let ready = stream.ready(Interest::READABLE | Interest::WRITABLE).await?; + /// + /// if ready.is_readable() { + /// let mut data = vec![0; 1024]; + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_read(&mut data) { + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// + /// } + /// + /// if ready.is_writable() { + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_write(b"hello world") { + /// Ok(n) => { + /// println!("write {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// } + /// } + /// ``` + pub async fn ready(&self, interest: Interest) -> io::Result<Ready> { + let event = self.io.registration().readiness(interest).await?; + Ok(event.ready) + } + + /// Waits for the socket to become readable. + /// + /// This function is equivalent to `ready(Interest::READABLE)` and is usually + /// paired with `try_read()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to read that fails with `WouldBlock` or + /// `Poll::Pending`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// let mut msg = vec![0; 1024]; + /// + /// loop { + /// // Wait for the socket to be readable + /// stream.readable().await?; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_read(&mut msg) { + /// Ok(n) => { + /// msg.truncate(n); + /// break; + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// println!("GOT = {:?}", msg); + /// Ok(()) + /// } + /// ``` + pub async fn readable(&self) -> io::Result<()> { + self.ready(Interest::READABLE).await?; + Ok(()) + } + + /// Polls for read readiness. + /// + /// If the tcp stream is not currently ready for reading, this method will + /// store a clone of the `Waker` from the provided `Context`. When the tcp + /// stream becomes ready for reading, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_read_ready`, `poll_read` or + /// `poll_peek`, only the `Waker` from the `Context` passed to the most + /// recent call is scheduled to receive a wakeup. (However, + /// `poll_write_ready` retains a second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`readable`] is not feasible. Where possible, using [`readable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the tcp stream is not ready for reading. + /// * `Poll::Ready(Ok(()))` if the tcp stream is ready for reading. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`readable`]: method@Self::readable + pub fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_read_ready(cx).map_ok(|_| ()) + } + + /// Tries to read data from the stream into the provided buffer, returning how + /// many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: TcpStream::readable() + /// [`ready()`]: TcpStream::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// // Wait for the socket to be readable + /// stream.readable().await?; + /// + /// // Creating the buffer **after** the `await` prevents it from + /// // being stored in the async task. + /// let mut buf = [0; 4096]; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_read(&mut buf) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read(&self, buf: &mut [u8]) -> io::Result<usize> { + use std::io::Read; + + self.io + .registration() + .try_io(Interest::READABLE, || (&*self.io).read(buf)) + } + + /// Tries to read data from the stream into the provided buffers, returning + /// how many bytes were read. + /// + /// Data is copied to fill each buffer in order, with the final buffer + /// written to possibly being only partially filled. This method behaves + /// equivalently to a single call to [`try_read()`] with concatenated + /// buffers. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_vectored()` is non-blocking, the buffer does not have to be + /// stored by the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`try_read()`]: TcpStream::try_read() + /// [`readable()`]: TcpStream::readable() + /// [`ready()`]: TcpStream::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io::{self, IoSliceMut}; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// // Wait for the socket to be readable + /// stream.readable().await?; + /// + /// // Creating the buffer **after** the `await` prevents it from + /// // being stored in the async task. + /// let mut buf_a = [0; 512]; + /// let mut buf_b = [0; 1024]; + /// let mut bufs = [ + /// IoSliceMut::new(&mut buf_a), + /// IoSliceMut::new(&mut buf_b), + /// ]; + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_read_vectored(&mut bufs) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_vectored(&self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> { + use std::io::Read; + + self.io + .registration() + .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) + } + + cfg_io_util! { + /// Tries to read data from the stream into the provided buffer, advancing the + /// buffer's internal cursor, returning how many bytes were read. + /// + /// Receives any pending data from the socket but does not wait for new data + /// to arrive. On success, returns the number of bytes read. Because + /// `try_read_buf()` is non-blocking, the buffer does not have to be stored by + /// the async task and can exist entirely on the stack. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: TcpStream::readable() + /// [`ready()`]: TcpStream::ready() + /// + /// # Return + /// + /// If data is successfully read, `Ok(n)` is returned, where `n` is the + /// number of bytes read. `Ok(0)` indicates the stream's read half is closed + /// and will no longer yield data. If the stream is not ready to read data + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// // Wait for the socket to be readable + /// stream.readable().await?; + /// + /// let mut buf = Vec::with_capacity(4096); + /// + /// // Try to read data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_read_buf(&mut buf) { + /// Ok(0) => break, + /// Ok(n) => { + /// println!("read {} bytes", n); + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_read_buf<B: BufMut>(&self, buf: &mut B) -> io::Result<usize> { + self.io.registration().try_io(Interest::READABLE, || { + use std::io::Read; + + let dst = buf.chunk_mut(); + let dst = + unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]) }; + + // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the + // buffer. + let n = (&*self.io).read(dst)?; + + unsafe { + buf.advance_mut(n); + } + + Ok(n) + }) + } + } + + /// Waits for the socket to become writable. + /// + /// This function is equivalent to `ready(Interest::WRITABLE)` and is usually + /// paired with `try_write()`. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Once a readiness event occurs, the method + /// will continue to return immediately until the readiness event is + /// consumed by an attempt to write that fails with `WouldBlock` or + /// `Poll::Pending`. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// // Wait for the socket to be writable + /// stream.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_write(b"hello world") { + /// Ok(n) => { + /// break; + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub async fn writable(&self) -> io::Result<()> { + self.ready(Interest::WRITABLE).await?; + Ok(()) + } + + /// Polls for write readiness. + /// + /// If the tcp stream is not currently ready for writing, this method will + /// store a clone of the `Waker` from the provided `Context`. When the tcp + /// stream becomes ready for writing, `Waker::wake` will be called on the + /// waker. + /// + /// Note that on multiple calls to `poll_write_ready` or `poll_write`, only + /// the `Waker` from the `Context` passed to the most recent call is + /// scheduled to receive a wakeup. (However, `poll_read_ready` retains a + /// second, independent waker.) + /// + /// This function is intended for cases where creating and pinning a future + /// via [`writable`] is not feasible. Where possible, using [`writable`] is + /// preferred, as this supports polling from multiple tasks at once. + /// + /// # Return value + /// + /// The function returns: + /// + /// * `Poll::Pending` if the tcp stream is not ready for writing. + /// * `Poll::Ready(Ok(()))` if the tcp stream is ready for writing. + /// * `Poll::Ready(Err(e))` if an error is encountered. + /// + /// # Errors + /// + /// This function may encounter any standard I/O error except `WouldBlock`. + /// + /// [`writable`]: method@Self::writable + pub fn poll_write_ready(&self, cx: &mut Context<'_>) -> Poll<io::Result<()>> { + self.io.registration().poll_write_ready(cx).map_ok(|_| ()) + } + + /// Try to write a buffer to the stream, returning how many bytes were + /// written. + /// + /// The function will attempt to write the entire contents of `buf`, but + /// only part of the buffer may be written. + /// + /// This function is usually paired with `writable()`. + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// loop { + /// // Wait for the socket to be writable + /// stream.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_write(b"hello world") { + /// Ok(n) => { + /// break; + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_write(&self, buf: &[u8]) -> io::Result<usize> { + use std::io::Write; + + self.io + .registration() + .try_io(Interest::WRITABLE, || (&*self.io).write(buf)) + } + + /// Tries to write several buffers to the stream, returning how many bytes + /// were written. + /// + /// Data is written from each buffer in order, with the final buffer read + /// from possible being only partially consumed. This method behaves + /// equivalently to a single call to [`try_write()`] with concatenated + /// buffers. + /// + /// This function is usually paired with `writable()`. + /// + /// [`try_write()`]: TcpStream::try_write() + /// + /// # Return + /// + /// If data is successfully written, `Ok(n)` is returned, where `n` is the + /// number of bytes written. If the stream is not ready to write data, + /// `Err(io::ErrorKind::WouldBlock)` is returned. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use std::error::Error; + /// use std::io; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// let bufs = [io::IoSlice::new(b"hello "), io::IoSlice::new(b"world")]; + /// + /// loop { + /// // Wait for the socket to be writable + /// stream.writable().await?; + /// + /// // Try to write data, this may still fail with `WouldBlock` + /// // if the readiness event is a false positive. + /// match stream.try_write_vectored(&bufs) { + /// Ok(n) => { + /// break; + /// } + /// Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => { + /// continue; + /// } + /// Err(e) => { + /// return Err(e.into()); + /// } + /// } + /// } + /// + /// Ok(()) + /// } + /// ``` + pub fn try_write_vectored(&self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> { + use std::io::Write; + + self.io + .registration() + .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(bufs)) + } + + /// Tries to read or write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation from the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `TcpStream` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: TcpStream::readable() + /// [`writable()`]: TcpStream::writable() + /// [`ready()`]: TcpStream::ready() + pub fn try_io<R>( + &self, + interest: Interest, + f: impl FnOnce() -> io::Result<R>, + ) -> io::Result<R> { + self.io.registration().try_io(interest, f) + } + + /// Receives data on the socket from the remote address to which it is + /// connected, without removing that data from the queue. On success, + /// returns the number of bytes peeked. + /// + /// Successive calls return the same data. This is accomplished by passing + /// `MSG_PEEK` as a flag to the underlying recv system call. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// use tokio::io::AsyncReadExt; + /// use std::error::Error; + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Box<dyn Error>> { + /// // Connect to a peer + /// let mut stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// let mut b1 = [0; 10]; + /// let mut b2 = [0; 10]; + /// + /// // Peek at the data + /// let n = stream.peek(&mut b1).await?; + /// + /// // Read the data + /// assert_eq!(n, stream.read(&mut b2[..n]).await?); + /// assert_eq!(&b1[..n], &b2[..n]); + /// + /// Ok(()) + /// } + /// ``` + /// + /// The [`read`] method is defined on the [`AsyncReadExt`] trait. + /// + /// [`read`]: fn@crate::io::AsyncReadExt::read + /// [`AsyncReadExt`]: trait@crate::io::AsyncReadExt + pub async fn peek(&self, buf: &mut [u8]) -> io::Result<usize> { + self.io + .registration() + .async_io(Interest::READABLE, || self.io.peek(buf)) + .await + } + + /// Shuts down the read, write, or both halves of this connection. + /// + /// This function will cause all pending and future I/O on the specified + /// portions to return immediately with an appropriate value (see the + /// documentation of `Shutdown`). + pub(super) fn shutdown_std(&self, how: Shutdown) -> io::Result<()> { + self.io.shutdown(how) + } + + /// Gets the value of the `TCP_NODELAY` option on this socket. + /// + /// For more information about this option, see [`set_nodelay`]. + /// + /// [`set_nodelay`]: TcpStream::set_nodelay + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.nodelay()?); + /// # Ok(()) + /// # } + /// ``` + pub fn nodelay(&self) -> io::Result<bool> { + self.io.nodelay() + } + + /// Sets the value of the `TCP_NODELAY` option on this socket. + /// + /// If set, this option disables the Nagle algorithm. This means that + /// segments are always sent as soon as possible, even if there is only a + /// small amount of data. When not set, data is buffered until there is a + /// sufficient amount to send out, thereby avoiding the frequent sending of + /// small packets. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// stream.set_nodelay(true)?; + /// # Ok(()) + /// # } + /// ``` + pub fn set_nodelay(&self, nodelay: bool) -> io::Result<()> { + self.io.set_nodelay(nodelay) + } + + /// Reads the linger duration for this socket by getting the `SO_LINGER` + /// option. + /// + /// For more information about this option, see [`set_linger`]. + /// + /// [`set_linger`]: TcpStream::set_linger + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.linger()?); + /// # Ok(()) + /// # } + /// ``` + pub fn linger(&self) -> io::Result<Option<Duration>> { + socket2::SockRef::from(self).linger() + } + + /// Sets the linger duration of this socket by setting the SO_LINGER option. + /// + /// This option controls the action taken when a stream has unsent messages and the stream is + /// closed. If SO_LINGER is set, the system shall block the process until it can transmit the + /// data or until the time expires. + /// + /// If SO_LINGER is not specified, and the stream is closed, the system handles the call in a + /// way that allows the process to continue as quickly as possible. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// stream.set_linger(None)?; + /// # Ok(()) + /// # } + /// ``` + pub fn set_linger(&self, dur: Option<Duration>) -> io::Result<()> { + socket2::SockRef::from(self).set_linger(dur) + } + + /// Gets the value of the `IP_TTL` option for this socket. + /// + /// For more information about this option, see [`set_ttl`]. + /// + /// [`set_ttl`]: TcpStream::set_ttl + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// println!("{:?}", stream.ttl()?); + /// # Ok(()) + /// # } + /// ``` + pub fn ttl(&self) -> io::Result<u32> { + self.io.ttl() + } + + /// Sets the value for the `IP_TTL` option on this socket. + /// + /// This value sets the time-to-live field that is used in every packet sent + /// from this socket. + /// + /// # Examples + /// + /// ```no_run + /// use tokio::net::TcpStream; + /// + /// # async fn dox() -> Result<(), Box<dyn std::error::Error>> { + /// let stream = TcpStream::connect("127.0.0.1:8080").await?; + /// + /// stream.set_ttl(123)?; + /// # Ok(()) + /// # } + /// ``` + pub fn set_ttl(&self, ttl: u32) -> io::Result<()> { + self.io.set_ttl(ttl) + } + + // These lifetime markers also appear in the generated documentation, and make + // it more clear that this is a *borrowed* split. + #[allow(clippy::needless_lifetimes)] + /// Splits a `TcpStream` into a read half and a write half, which can be used + /// to read and write the stream concurrently. + /// + /// This method is more efficient than [`into_split`], but the halves cannot be + /// moved into independently spawned tasks. + /// + /// [`into_split`]: TcpStream::into_split() + pub fn split<'a>(&'a mut self) -> (ReadHalf<'a>, WriteHalf<'a>) { + split(self) + } + + /// Splits a `TcpStream` into a read half and a write half, which can be used + /// to read and write the stream concurrently. + /// + /// Unlike [`split`], the owned halves can be moved to separate tasks, however + /// this comes at the cost of a heap allocation. + /// + /// **Note:** Dropping the write half will shut down the write half of the TCP + /// stream. This is equivalent to calling [`shutdown()`] on the `TcpStream`. + /// + /// [`split`]: TcpStream::split() + /// [`shutdown()`]: fn@crate::io::AsyncWriteExt::shutdown + pub fn into_split(self) -> (OwnedReadHalf, OwnedWriteHalf) { + split_owned(self) + } + + // == Poll IO functions that takes `&self` == + // + // To read or write without mutable access to the `UnixStream`, combine the + // `poll_read_ready` or `poll_write_ready` methods with the `try_read` or + // `try_write` methods. + + pub(crate) fn poll_read_priv( + &self, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + // Safety: `TcpStream::read` correctly handles reads into uninitialized memory + unsafe { self.io.poll_read(cx, buf) } + } + + pub(super) fn poll_write_priv( + &self, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.io.poll_write(cx, buf) + } + + pub(super) fn poll_write_vectored_priv( + &self, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.io.poll_write_vectored(cx, bufs) + } +} + +impl TryFrom<std::net::TcpStream> for TcpStream { + type Error = io::Error; + + /// Consumes stream, returning the tokio I/O object. + /// + /// This is equivalent to + /// [`TcpStream::from_std(stream)`](TcpStream::from_std). + fn try_from(stream: std::net::TcpStream) -> Result<Self, Self::Error> { + Self::from_std(stream) + } +} + +// ===== impl Read / Write ===== + +impl AsyncRead for TcpStream { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll<io::Result<()>> { + self.poll_read_priv(cx, buf) + } +} + +impl AsyncWrite for TcpStream { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll<io::Result<usize>> { + self.poll_write_priv(cx, buf) + } + + fn poll_write_vectored( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + bufs: &[io::IoSlice<'_>], + ) -> Poll<io::Result<usize>> { + self.poll_write_vectored_priv(cx, bufs) + } + + fn is_write_vectored(&self) -> bool { + true + } + + #[inline] + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + // tcp flush is a no-op + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> { + self.shutdown_std(std::net::Shutdown::Write)?; + Poll::Ready(Ok(())) + } +} + +impl fmt::Debug for TcpStream { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.io.fmt(f) + } +} + +#[cfg(unix)] +mod sys { + use super::TcpStream; + use std::os::unix::prelude::*; + + impl AsRawFd for TcpStream { + fn as_raw_fd(&self) -> RawFd { + self.io.as_raw_fd() + } + } +} + +#[cfg(windows)] +mod sys { + use super::TcpStream; + use std::os::windows::prelude::*; + + impl AsRawSocket for TcpStream { + fn as_raw_socket(&self) -> RawSocket { + self.io.as_raw_socket() + } + } +} |