From a40be0857c7bf48e39f815417b0b5293cd8ed1aa Mon Sep 17 00:00:00 2001
From: Tyler Julian <tjulian@uber.com>
Date: Tue, 10 Jan 2017 19:11:56 -0800
Subject: 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.
---
 src/libstd/sys_common/net.rs | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

(limited to 'src/libstd/sys_common')

diff --git a/src/libstd/sys_common/net.rs b/src/libstd/sys_common/net.rs
index 10ad61f4c80..3cdeb511945 100644
--- a/src/libstd/sys_common/net.rs
+++ b/src/libstd/sys_common/net.rs
@@ -91,7 +91,7 @@ fn sockname<F>(f: F) -> io::Result<SocketAddr>
     }
 }
 
-fn sockaddr_to_addr(storage: &c::sockaddr_storage,
+pub fn sockaddr_to_addr(storage: &c::sockaddr_storage,
                     len: usize) -> io::Result<SocketAddr> {
     match storage.ss_family as c_int {
         c::AF_INET => {
@@ -222,6 +222,10 @@ impl TcpStream {
         self.inner.timeout(c::SO_SNDTIMEO)
     }
 
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.peek(buf)
+    }
+
     pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
         self.inner.read(buf)
     }
@@ -441,17 +445,11 @@ impl UdpSocket {
     }
 
     pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
-        let mut storage: c::sockaddr_storage = unsafe { mem::zeroed() };
-        let mut addrlen = mem::size_of_val(&storage) as c::socklen_t;
-        let len = cmp::min(buf.len(), <wrlen_t>::max_value() as usize) as wrlen_t;
+        self.inner.recv_from(buf)
+    }
 
-        let n = cvt(unsafe {
-            c::recvfrom(*self.inner.as_inner(),
-                        buf.as_mut_ptr() as *mut c_void,
-                        len, 0,
-                        &mut storage as *mut _ as *mut _, &mut addrlen)
-        })?;
-        Ok((n as usize, sockaddr_to_addr(&storage, addrlen as usize)?))
+    pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
+        self.inner.peek_from(buf)
     }
 
     pub fn send_to(&self, buf: &[u8], dst: &SocketAddr) -> io::Result<usize> {
@@ -578,6 +576,10 @@ impl UdpSocket {
         self.inner.read(buf)
     }
 
+    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
+        self.inner.peek(buf)
+    }
+
     pub fn send(&self, buf: &[u8]) -> io::Result<usize> {
         let len = cmp::min(buf.len(), <wrlen_t>::max_value() as usize) as wrlen_t;
         let ret = cvt(unsafe {
-- 
cgit 1.4.1-3-g733a5