libstd/net: Add `peek` APIs to UdpSocket and TcpStream

These methods enable socket reads without side-effects. That is,
repeated calls to peek() return identical data. This is accomplished
by providing the POSIX flag MSG_PEEK to the underlying socket read
operations.

This also moves the current implementation of recv_from out of the
platform-independent sys_common and into respective sys/windows and
sys/unix implementations. This allows for more platform-dependent
implementations.

2 files changed, 151 insertions, 0 deletions

diff --git a/src/libstd/net/tcp.rs b/src/libstd/net/tcp.rs
index ed1f08f9c90..ba6160cc723 100644
--- a/src/libstd/net/tcp.rs
+++ b/src/libstd/net/tcp.rs
@@ -296,6 +296,29 @@ impl TcpStream {
         self.0.write_timeout()
     }
 
+    /// Receives data on the socket from the remote adress 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
+    /// #![feature(peek)]
+    /// use std::net::TcpStream;
+    ///
+    /// let stream = TcpStream::connect("127.0.0.1:8000")
+    ///                        .expect("couldn't bind to address");
+    /// let mut buf = [0; 10];
+    /// let len = stream.peek(&mut buf).expect("peek failed");
+    /// ```
+    #[unstable(feature = "peek", issue = "38980")]
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.0.peek(buf)
+    }
+
     /// Sets the value of the `TCP_NODELAY` option on this socket.
     ///
     /// If set, this option disables the Nagle algorithm. This means that
@@ -1405,4 +1428,35 @@ mod tests {
             Err(e) => panic!("unexpected error {}", e),
         }
     }
+
+    #[test]
+    fn peek() {
+        each_ip(&mut |addr| {
+            let (txdone, rxdone) = channel();
+
+            let srv = t!(TcpListener::bind(&addr));
+            let _t = thread::spawn(move|| {
+                let mut cl = t!(srv.accept()).0;
+                cl.write(&[1,3,3,7]).unwrap();
+                t!(rxdone.recv());
+            });
+
+            let mut c = t!(TcpStream::connect(&addr));
+            let mut b = [0; 10];
+            for _ in 1..3 {
+                let len = c.peek(&mut b).unwrap();
+                assert_eq!(len, 4);
+            }
+            let len = c.read(&mut b).unwrap();
+            assert_eq!(len, 4);
+
+            t!(c.set_nonblocking(true));
+            match c.peek(&mut b) {
+                Ok(_) => panic!("expected error"),
+                Err(ref e) if e.kind() == ErrorKind::WouldBlock => {}
+                Err(e) => panic!("unexpected error {}", e),
+            }
+            t!(txdone.send(()));
+        })
+    }
 }
diff --git a/src/libstd/net/udp.rs b/src/libstd/net/udp.rs
index f8a5ec0b379..2f28f475dc8 100644
--- a/src/libstd/net/udp.rs
+++ b/src/libstd/net/udp.rs
@@ -83,6 +83,30 @@ impl UdpSocket {
         self.0.recv_from(buf)
     }
 
+    /// Receives data from the socket, without removing it from the queue.
+    ///
+    /// Successive calls return the same data. This is accomplished by passing
+    /// `MSG_PEEK` as a flag to the underlying `recvfrom` system call.
+    ///
+    /// On success, returns the number of bytes peeked and the address from
+    /// whence the data came.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// #![feature(peek)]
+    /// use std::net::UdpSocket;
+    ///
+    /// let socket = UdpSocket::bind("127.0.0.1:34254").expect("couldn't bind to address");
+    /// let mut buf = [0; 10];
+    /// let (number_of_bytes, src_addr) = socket.peek_from(&mut buf)
+    ///                                         .expect("Didn't receive data");
+    /// ```
+    #[unstable(feature = "peek", issue = "38980")]
+    pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        self.0.peek_from(buf)
+    }
+
     /// Sends data on the socket to the given address. On success, returns the
     /// number of bytes written.
     ///
@@ -579,6 +603,37 @@ impl UdpSocket {
         self.0.recv(buf)
     }
 
+    /// Receives data on the socket from the remote adress 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.
+    ///
+    /// # Errors
+    ///
+    /// This method will fail if the socket is not connected. The `connect` method
+    /// will connect this socket to a remote address.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// #![feature(peek)]
+    /// use std::net::UdpSocket;
+    ///
+    /// let socket = UdpSocket::bind("127.0.0.1:34254").expect("couldn't bind to address");
+    /// socket.connect("127.0.0.1:8080").expect("connect function failed");
+    /// let mut buf = [0; 10];
+    /// match socket.peek(&mut buf) {
+    ///     Ok(received) => println!("received {} bytes", received),
+    ///     Err(e) => println!("peek function failed: {:?}", e),
+    /// }
+    /// ```
+    #[unstable(feature = "peek", issue = "38980")]
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.0.peek(buf)
+    }
+
     /// Moves this UDP socket into or out of nonblocking mode.
     ///
     /// On Unix this corresponds to calling fcntl, and on Windows this
@@ -870,6 +925,48 @@ mod tests {
     }
 
     #[test]
+    fn connect_send_peek_recv() {
+        each_ip(&mut |addr, _| {
+            let socket = t!(UdpSocket::bind(&addr));
+            t!(socket.connect(addr));
+
+            t!(socket.send(b"hello world"));
+
+            for _ in 1..3 {
+                let mut buf = [0; 11];
+                let size = t!(socket.peek(&mut buf));
+                assert_eq!(b"hello world", &buf[..]);
+                assert_eq!(size, 11);
+            }
+
+            let mut buf = [0; 11];
+            let size = t!(socket.recv(&mut buf));
+            assert_eq!(b"hello world", &buf[..]);
+            assert_eq!(size, 11);
+        })
+    }
+
+    #[test]
+    fn peek_from() {
+        each_ip(&mut |addr, _| {
+            let socket = t!(UdpSocket::bind(&addr));
+            t!(socket.send_to(b"hello world", &addr));
+
+            for _ in 1..3 {
+                let mut buf = [0; 11];
+                let (size, _) = t!(socket.peek_from(&mut buf));
+                assert_eq!(b"hello world", &buf[..]);
+                assert_eq!(size, 11);
+            }
+
+            let mut buf = [0; 11];
+            let (size, _) = t!(socket.recv_from(&mut buf));
+            assert_eq!(b"hello world", &buf[..]);
+            assert_eq!(size, 11);
+        })
+    }
+
+    #[test]
     fn ttl() {
         let ttl = 100;