1 files changed, 163 insertions, 0 deletions
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..cce50f6ab3
--- /dev/null
+++ b/third_party/rust/tokio/src/net/tcp/split.rs
@@ -0,0 +1,163 @@
+//! `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};
+use crate::net::TcpStream;
+
+use bytes::Buf;
+use std::io;
+use std::mem::MaybeUninit;
+use std::net::Shutdown;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// Read half of a `TcpStream`.
+#[derive(Debug)]
+pub struct ReadHalf<'a>(&'a TcpStream);
+
+/// Write half of a `TcpStream`.
+///
+/// Note that in the `AsyncWrite` implemenation of `TcpStreamWriteHalf`,
+/// `poll_shutdown` actually shuts down the TCP stream in the write direction.
+#[derive(Debug)]
+pub struct WriteHalf<'a>(&'a TcpStream);
+
+pub(crate) fn split(stream: &mut TcpStream) -> (ReadHalf<'_>, WriteHalf<'_>) {
+    (ReadHalf(&*stream), WriteHalf(&*stream))
+}
+
+impl ReadHalf<'_> {
+    /// 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.
+    ///
+    /// See the [`TcpStream::poll_peek`] level documenation for more details.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use tokio::io;
+    /// 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];
+    ///
+    ///     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 [u8]) -> Poll<io::Result<usize>> {
+        self.0.poll_peek2(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 documenation for more details.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use tokio::net::TcpStream;
+    /// use tokio::prelude::*;
+    /// 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(())
+    /// }
+    /// ```
+    ///
+    /// [`TcpStream::peek`]: TcpStream::peek
+    pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        poll_fn(|cx| self.poll_peek(cx, buf)).await
+    }
+}
+
+impl AsyncRead for ReadHalf<'_> {
+    unsafe fn prepare_uninitialized_buffer(&self, _: &mut [MaybeUninit<u8>]) -> bool {
+        false
+    }
+
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        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_buf<B: Buf>(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut B,
+    ) -> Poll<io::Result<usize>> {
+        self.0.poll_write_buf_priv(cx, buf)
+    }
+
+    #[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(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
+    }
+}