Refactor Socket::peek

Author: Thomasdezeeuw

src/socket.rs

@@ -15,11 +15,6 @@ use std::net::{self, Ipv4Addr, Ipv6Addr, Shutdown};
 use std::os::unix::net::{UnixDatagram, UnixListener, UnixStream};
 use std::time::Duration;
 
-#[cfg(all(unix, feature = "all", not(target_os = "redox")))]
-use libc::MSG_OOB;
-#[cfg(all(windows, feature = "all"))]
-use winapi::um::winsock2::MSG_OOB;
-
 use crate::sys;
 #[cfg(not(target_os = "redox"))]
 use crate::RecvFlags;
@@ -303,7 +298,7 @@ impl Socket {
     /// [`out_of_band_inline`]: Socket::out_of_band_inline
     #[cfg(all(feature = "all", not(target_os = "redox")))]
     pub fn recv_out_of_band(&self, buf: &mut [u8]) -> io::Result<usize> {
-        self.recv_with_flags(buf, MSG_OOB)
+        self.recv_with_flags(buf, sys::MSG_OOB)
     }
 
     /// Identical to [`recv`] but allows for specification of arbitrary flags to
@@ -351,7 +346,7 @@ impl Socket {
     /// 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.inner().peek(buf)
+        self.recv_with_flags(buf, sys::MSG_PEEK)
     }
 
     /// Receives data from the socket. On success, returns the number of bytes
@@ -433,7 +428,7 @@ impl Socket {
     /// [`out_of_band_inline`]: #method.out_of_band_inline
     #[cfg(all(feature = "all", not(target_os = "redox")))]
     pub fn send_out_of_band(&self, buf: &[u8]) -> io::Result<usize> {
-        self.inner().send(buf, MSG_OOB)
+        self.inner().send(buf, sys::MSG_OOB)
     }
 
     /// Sends data on the socket to the given address. On success, returns the

src/sys/unix.rs

@@ -42,6 +42,13 @@ pub(crate) use libc::{IPPROTO_ICMP, IPPROTO_ICMPV6, IPPROTO_TCP, IPPROTO_UDP};
 pub(crate) use libc::{
     sa_family_t, sockaddr, sockaddr_in, sockaddr_in6, sockaddr_storage, socklen_t,
 };
+// Used in `RecvFlags`.
+#[cfg(not(target_os = "redox"))]
+pub(crate) use libc::MSG_TRUNC;
+// Used in `Socket`.
+#[cfg(all(unix, feature = "all", not(target_os = "redox")))]
+pub(crate) use libc::MSG_OOB;
+pub(crate) use libc::MSG_PEEK;
 
 cfg_if::cfg_if! {
     if #[cfg(any(target_os = "dragonfly", target_os = "freebsd",
@@ -82,10 +89,6 @@ macro_rules! syscall {
     }};
 }
 
-// Re-export message flags for the RecvFlags struct.
-#[cfg(not(target_os = "redox"))]
-pub(crate) use libc::MSG_TRUNC;
-
 /// Maximum size of a buffer passed to system call like `recv` and `send`.
 #[cfg(not(target_vendor = "apple"))]
 const MAX_BUF_LEN: usize = <ssize_t>::max_value() as usize;
@@ -550,16 +553,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
-        let n = syscall!(recv(
-            self.fd,
-            buf.as_mut_ptr() as *mut c_void,
-            cmp::min(buf.len(), max_len()),
-            libc::MSG_PEEK,
-        ))?;
-        Ok(n as usize)
-    }
-
     pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SockAddr)> {
         self.recv_from(buf, libc::MSG_PEEK)
     }

src/sys/windows.rs

@@ -33,7 +33,6 @@ use winapi::um::winsock2::{self as sock, u_long, SD_BOTH, SD_RECEIVE, SD_SEND};
 
 use crate::{RecvFlags, SockAddr, Type};
 
-const MSG_PEEK: c_int = 0x2;
 const SIO_KEEPALIVE_VALS: DWORD = 0x98000004;
 
 pub use winapi::ctypes::c_int;
@@ -62,6 +61,10 @@ pub(crate) use winapi::shared::ws2def::{
 };
 pub(crate) use winapi::shared::ws2ipdef::SOCKADDR_IN6_LH as sockaddr_in6;
 pub(crate) use winapi::um::ws2tcpip::socklen_t;
+// Used in `Socket`.
+#[cfg(all(windows, feature = "all"))]
+pub(crate) use winapi::um::winsock2::MSG_OOB;
+pub(crate) use winapi::um::winsock2::MSG_PEEK;
 
 /// Maximum size of a buffer passed to system call like `recv` and `send`.
 const MAX_BUF_LEN: usize = <c_int>::max_value() as usize;
@@ -375,24 +378,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
-        unsafe {
-            let n = {
-                sock::recv(
-                    self.socket,
-                    buf.as_mut_ptr() as *mut c_char,
-                    clamp(buf.len()),
-                    MSG_PEEK,
-                )
-            };
-            match n {
-                sock::SOCKET_ERROR if sock::WSAGetLastError() == sock::WSAESHUTDOWN as i32 => Ok(0),
-                sock::SOCKET_ERROR => Err(last_error()),
-                n => Ok(n as usize),
-            }
-        }
-    }
-
     pub fn peek_from(&self, buf: &mut [u8]) -> io::Result<(usize, SockAddr)> {
         self.recv_from(buf, MSG_PEEK)
     }