diff --git a/tokio-tcp/src/stream.rs b/tokio-tcp/src/stream.rs index 397906f29..358a19fa7 100644 --- a/tokio-tcp/src/stream.rs +++ b/tokio-tcp/src/stream.rs @@ -2,6 +2,7 @@ use crate::split::{ split, split_mut, TcpStreamReadHalf, TcpStreamReadHalfMut, TcpStreamWriteHalf, TcpStreamWriteHalfMut, }; +use async_util::future::poll_fn; use bytes::{Buf, BufMut}; use futures_core::ready; use iovec::IoVec; @@ -341,6 +342,22 @@ impl TcpStream { } } + /// 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 + /// + /// ``` + /// unimplemented!(); + /// ``` + pub async fn peek(&mut self, buf: &mut [u8]) -> io::Result { + poll_fn(|cx| self.poll_peek(cx, 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 diff --git a/tokio-tcp/tests/tcp_peek.rs b/tokio-tcp/tests/tcp_peek.rs new file mode 100644 index 000000000..829254ef4 --- /dev/null +++ b/tokio-tcp/tests/tcp_peek.rs @@ -0,0 +1,27 @@ +#![deny(warnings, rust_2018_idioms)] +#![feature(async_await)] + +use std::thread; +use std::{convert::TryInto, io::Write, net}; +use tokio::io::AsyncReadExt; +use tokio::net::TcpStream; +use tokio_test::assert_ok; + +#[tokio::test] +async fn peek() { + let listener = net::TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let t = thread::spawn(move || assert_ok!(listener.accept()).0); + + let left = net::TcpStream::connect(&addr).unwrap(); + let mut right = t.join().unwrap(); + right.write(&[1, 2, 3, 4]).unwrap(); + + let mut left: TcpStream = left.try_into().unwrap(); + let mut buf = [0u8; 16]; + let n = assert_ok!(left.peek(&mut buf).await); + assert_eq!([1, 2, 3, 4], buf[..n]); + + let n = assert_ok!(left.read(&mut buf).await); + assert_eq!([1, 2, 3, 4], buf[..n]); +}