Refactor Socket::send_vectored

Author: Thomasdezeeuw

src/socket.rs

@@ -422,12 +422,19 @@ impl Socket {
         sys::send(self.inner, buf, flags)
     }
 
-    /// Identical to [`send_with_flags`] but writes from a slice of buffers.
+    /// Send data to the connected peer. Returns the amount of bytes written.
+    #[cfg(not(target_os = "redox"))]
+    pub fn send_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        self.send_vectored_with_flags(bufs, 0)
+    }
+
+    /// Identical to [`send_vectored`] but allows for specification of arbitrary
+    /// flags to the underlying `sendmsg`/`WSASend` call.
     ///
-    /// [`send_with_flags`]: #method.send_with_flags
+    /// [`send_vectored`]: Socket::send_vectored
     #[cfg(not(target_os = "redox"))]
-    pub fn send_vectored(&self, bufs: &[IoSlice<'_>], flags: i32) -> io::Result<usize> {
-        self.inner().send_vectored(bufs, flags)
+    pub fn send_vectored_with_flags(&self, bufs: &[IoSlice<'_>], flags: i32) -> io::Result<usize> {
+        sys::send_vectored(self.inner, bufs, flags)
     }
 
     /// Sends out-of-band (OOB) data on the socket to connected peer
@@ -909,6 +916,7 @@ impl Read for Socket {
         self.recv(buf)
     }
 
+    #[cfg(not(target_os = "redox"))]
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.recv_vectored(bufs).map(|(n, _)| n)
     }
@@ -919,6 +927,7 @@ impl<'a> Read for &'a Socket {
         self.recv(buf)
     }
 
+    #[cfg(not(target_os = "redox"))]
     fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
         self.recv_vectored(bufs).map(|(n, _)| n)
     }

src/sys/unix.rs

@@ -471,7 +471,7 @@ fn recvmsg(
         msg_controllen: 0,
         msg_flags: 0,
     };
-    syscall!(recvmsg(fd, &mut msg as *mut _, flags))
+    syscall!(recvmsg(fd, &mut msg, flags))
         .map(|n| (n as usize, msg.msg_namelen, RecvFlags(msg.msg_flags)))
 }
 
@@ -485,6 +485,28 @@ pub(crate) fn send(fd: SysSocket, buf: &[u8], flags: c_int) -> io::Result<usize>
     .map(|n| n as usize)
 }
 
+#[cfg(not(target_os = "redox"))]
+pub(crate) fn send_vectored(
+    fd: SysSocket,
+    bufs: &[IoSlice<'_>],
+    flags: c_int,
+) -> io::Result<usize> {
+    let mut msg = libc::msghdr {
+        msg_name: ptr::null_mut(),
+        msg_namelen: 0,
+        // Safety: we're creating a `*mut` pointer from a reference, which is UB
+        // once actually used. However the OS should not write to it in the
+        // `sendmsg` system call.
+        msg_iov: bufs.as_ptr() as *mut _,
+        msg_iovlen: min(bufs.len(), IovLen::MAX as usize) as IovLen,
+        msg_control: ptr::null_mut(),
+        msg_controllen: 0,
+        msg_flags: 0,
+    };
+
+    syscall!(sendmsg(fd, &mut msg, flags)).map(|n| n as usize)
+}
+
 /// Unix only API.
 impl crate::Socket {
     /// Accept a new incoming connection from this listener.
@@ -632,22 +654,6 @@ impl Socket {
         Ok(n as usize)
     }
 
-    #[cfg(not(target_os = "redox"))]
-    pub fn send_vectored(&self, bufs: &[IoSlice<'_>], flags: c_int) -> io::Result<usize> {
-        let mut msg = libc::msghdr {
-            msg_name: std::ptr::null_mut(),
-            msg_namelen: 0,
-            msg_iov: bufs.as_ptr() as *mut libc::iovec,
-            msg_iovlen: bufs.len().min(IovLen::MAX as usize) as IovLen,
-            msg_control: std::ptr::null_mut(),
-            msg_controllen: 0,
-            msg_flags: 0,
-        };
-
-        let n = syscall!(sendmsg(self.fd, &mut msg as *mut libc::msghdr, flags))?;
-        Ok(n as usize)
-    }
-
     #[cfg(not(target_os = "redox"))]
     pub fn send_to_vectored(
         &self,

src/sys/windows.rs

@@ -421,6 +421,39 @@ pub(crate) fn send(socket: SysSocket, buf: &[u8], flags: c_int) -> io::Result<us
     .map(|n| n as usize)
 }
 
+pub(crate) fn send_vectored(
+    socket: SysSocket,
+    bufs: &[IoSlice<'_>],
+    flags: c_int,
+) -> io::Result<usize> {
+    let mut nsent = 0;
+    syscall!(
+        WSASend(
+            socket,
+            // FIXME: From the `WSASend` docs [1]:
+            // > For a Winsock application, once the WSASend function is called,
+            // > the system owns these buffers and the application may not
+            // > access them.
+            //
+            // So what we're doing is actually UB as `bufs` needs to be `&mut
+            // [IoSlice<'_>]`.
+            //
+            // Tracking issue: https://github.com/rust-lang/socket2-rs/issues/129.
+            //
+            // [1] https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-wsasend
+            bufs.as_ptr() as *mut _,
+            min(bufs.len(), DWORD::max_value() as usize) as DWORD,
+            &mut nsent,
+            flags as DWORD,
+            std::ptr::null_mut(),
+            None,
+        ),
+        PartialEq::eq,
+        sock::SOCKET_ERROR
+    )
+    .map(|_| nsent as usize)
+}
+
 /// Caller must ensure `T` is the correct type for `opt` and `val`.
 unsafe fn getsockopt<T>(socket: SysSocket, opt: c_int, val: c_int) -> io::Result<T> {
     let mut payload: MaybeUninit<T> = MaybeUninit::uninit();
@@ -494,25 +527,6 @@ impl Socket {
         }
     }
 
-    pub fn send_vectored(&self, bufs: &[IoSlice<'_>], flags: c_int) -> io::Result<usize> {
-        let mut nsent = 0;
-        let ret = unsafe {
-            sock::WSASend(
-                self.socket,
-                bufs.as_ptr() as *mut WSABUF,
-                bufs.len().min(DWORD::MAX as usize) as DWORD,
-                &mut nsent,
-                flags as DWORD,
-                std::ptr::null_mut(),
-                None,
-            )
-        };
-        match ret {
-            0 => Ok(nsent as usize),
-            _ => Err(last_error()),
-        }
-    }
-
     pub fn send_to_vectored(
         &self,
         bufs: &[IoSlice<'_>],

src/tests.rs

@@ -157,15 +157,12 @@ fn send_recv_vectored() {
     let (socket_a, socket_b) = udp_pair_connected();
 
     let sent = socket_a
-        .send_vectored(
-            &[
-                IoSlice::new(b"the"),
-                IoSlice::new(b"weeknight"),
-                IoSlice::new(b"would"),
-                IoSlice::new(b"yellow"),
-            ],
-            0,
-        )
+        .send_vectored(&[
+            IoSlice::new(b"the"),
+            IoSlice::new(b"weeknight"),
+            IoSlice::new(b"would"),
+            IoSlice::new(b"yellow"),
+        ])
         .unwrap();
     assert_eq!(sent, 23);