add TcpStream::peek

Author: seanmonstar

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 
+* Add `TcpStream::peek` function (#773)
 * Raise minimum Rust version to 1.18.0
 * `Poll`: retry select() when interrupted by a signal
 

src/net/tcp.rs

@@ -343,6 +343,16 @@ impl TcpStream {
         self.sys.take_error()
     }
 
+    /// 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.
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.sys.peek(buf)
+    }
+
     /// Read in a list of buffers all at once.
     ///
     /// This operation will attempt to read bytes from this socket and place

src/sys/fuchsia/net.rs

@@ -128,6 +128,10 @@ impl TcpStream {
         self.io.take_error()
     }
 
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.io.peek(buf)
+    }
+
     pub fn readv(&self, bufs: &mut [&mut IoVec]) -> io::Result<usize> {
         unsafe {
             let slice = iovec::as_os_slice_mut(bufs);

src/sys/unix/tcp.rs

@@ -126,6 +126,10 @@ impl TcpStream {
         self.inner.take_error()
     }
 
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.peek(buf)
+    }
+
     pub fn readv(&self, bufs: &mut [&mut IoVec]) -> io::Result<usize> {
         unsafe {
             let slice = iovec::as_os_slice_mut(bufs);

src/sys/windows/tcp.rs

@@ -230,29 +230,7 @@ impl TcpStream {
         self.imp.inner()
     }
 
-    fn post_register(&self, interest: Ready, me: &mut StreamInner) {
-        if interest.is_readable() {
-            self.imp.schedule_read(me);
-        }
-
-        // At least with epoll, if a socket is registered with an interest in
-        // writing and it's immediately writable then a writable event is
-        // generated immediately, so do so here.
-        if interest.is_writable() {
-            if let State::Empty = me.write {
-                self.imp.add_readiness(me, Ready::writable());
-            }
-        }
-    }
-
-    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
-        match IoVec::from_bytes_mut(buf) {
-            Some(vec) => self.readv(&mut [vec]),
-            None => Ok(0),
-        }
-    }
-
-    pub fn readv(&self, bufs: &mut [&mut IoVec]) -> io::Result<usize> {
+    fn before_read(&self) -> io::Result<MutexGuard<StreamInner>> {
         let mut me = self.inner();
 
         match me.read {
@@ -280,6 +258,47 @@ impl TcpStream {
             State::Ready(()) => {}
         }
 
+        Ok(me)
+    }
+
+    fn post_register(&self, interest: Ready, me: &mut StreamInner) {
+        if interest.is_readable() {
+            self.imp.schedule_read(me);
+        }
+
+        // At least with epoll, if a socket is registered with an interest in
+        // writing and it's immediately writable then a writable event is
+        // generated immediately, so do so here.
+        if interest.is_writable() {
+            if let State::Empty = me.write {
+                self.imp.add_readiness(me, Ready::writable());
+            }
+        }
+    }
+
+    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
+        match IoVec::from_bytes_mut(buf) {
+            Some(vec) => self.readv(&mut [vec]),
+            None => Ok(0),
+        }
+    }
+
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        let mut me = self.before_read()?;
+
+        match (&self.imp.inner.socket).peek(buf) {
+            Ok(n) => Ok(n),
+            Err(e) => {
+                me.read = State::Empty;
+                self.imp.schedule_read(&mut me);
+                Err(e)
+            }
+        }
+    }
+
+    pub fn readv(&self, bufs: &mut [&mut IoVec]) -> io::Result<usize> {
+        let mut me = self.before_read()?;
+
         // TODO: Does WSARecv work on a nonblocking sockets? We ideally want to
         //       call that instead of looping over all the buffers and calling
         //       `recv` on each buffer. I'm not sure though if an overlapped

test/test_tcp.rs

@@ -154,6 +154,61 @@ fn read() {
     t.join().unwrap();
 }
 
+#[test]
+fn peek() {
+    const N: usize = 16 * 1024 * 1024;
+    struct H { amt: usize, socket: TcpStream, shutdown: bool }
+
+    let l = net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = l.local_addr().unwrap();
+
+    let t = thread::spawn(move || {
+        let mut s = l.accept().unwrap().0;
+        let b = [0; 1024];
+        let mut amt = 0;
+        while amt < N {
+            amt += s.write(&b).unwrap();
+        }
+    });
+
+    let poll = Poll::new().unwrap();
+    let s = TcpStream::connect(&addr).unwrap();
+
+    poll.register(&s, Token(1), Ready::readable(), PollOpt::edge()).unwrap();
+
+    let mut events = Events::with_capacity(128);
+
+    let mut h = H { amt: 0, socket: s, shutdown: false };
+    while !h.shutdown {
+        poll.poll(&mut events, None).unwrap();
+
+        for event in &events {
+            assert_eq!(event.token(), Token(1));
+            let mut b = [0; 1024];
+            match h.socket.peek(&mut b) {
+                Ok(_) => (),
+                Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+                    continue
+                },
+                Err(e) => panic!("unexpected error: {:?}", e),
+            }
+
+            loop {
+                if let Some(amt) = h.socket.try_read(&mut b).unwrap() {
+                    h.amt += amt;
+                } else {
+                    break
+                }
+                if h.amt >= N {
+                    h.shutdown = true;
+                    break
+                }
+            }
+        }
+    }
+    t.join().unwrap();
+}
+
 #[test]
 fn read_bufs() {
     const N: usize = 16 * 1024 * 1024;