Refactor Socket::send_to_vectored

Author: Thomasdezeeuw

src/sockaddr.rs

@@ -59,6 +59,12 @@ impl SockAddr {
         &self.storage as *const _ as *const _
     }
 
+    /// Returns a raw pointer to the address storage.
+    #[cfg(unix)]
+    pub(crate) fn as_storage_ptr(&self) -> *const sockaddr_storage {
+        &self.storage
+    }
+
     /// Returns this address as a `SocketAddr` if it is in the `AF_INET` (IP v4)
     /// or `AF_INET6` (IP v6) family, otherwise returns `None`.
     pub fn as_std(&self) -> Option<SocketAddr> {

src/socket.rs

@@ -465,21 +465,27 @@ impl Socket {
         sys::send_to(self.inner, buf, addr, flags)
     }
 
-    /// Identical to [`send_with_flags`] but writes from a slice of buffers.
+    /// Send data to a peer listening on `addr`. Returns the amount of bytes
+    /// written.
+    #[cfg(not(target_os = "redox"))]
+    pub fn send_to_vectored(&self, bufs: &[IoSlice<'_>], addr: &SockAddr) -> io::Result<usize> {
+        self.send_to_vectored_with_flags(bufs, addr, 0)
+    }
+
+    /// Identical to [`send_to_vectored`] but allows for specification of
+    /// arbitrary flags to the underlying `sendmsg`/`WSASendTo` call.
     ///
-    /// [`send_with_flags`]: #method.send_with_flags
+    /// [`send_to_vectored`]: Socket::send_to_vectored
     #[cfg(not(target_os = "redox"))]
-    pub fn send_to_vectored(
+    pub fn send_to_vectored_with_flags(
         &self,
         bufs: &[IoSlice<'_>],
         addr: &SockAddr,
         flags: i32,
     ) -> io::Result<usize> {
-        self.inner().send_to_vectored(bufs, flags, addr)
+        sys::send_to_vectored(self.inner, bufs, addr, flags)
     }
 
-    // ================================================
-
     /// Gets the value of the `IP_TTL` option for this socket.
     ///
     /// For more information about this option, see [`set_ttl`][link].

src/sys/unix.rs

@@ -493,20 +493,7 @@ pub(crate) fn send_vectored(
     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)
+    sendmsg(fd, ptr::null(), 0, bufs, flags)
 }
 
 pub(crate) fn send_to(
@@ -526,6 +513,40 @@ pub(crate) fn send_to(
     .map(|n| n as usize)
 }
 
+#[cfg(not(target_os = "redox"))]
+pub(crate) fn send_to_vectored(
+    fd: SysSocket,
+    bufs: &[IoSlice<'_>],
+    addr: &SockAddr,
+    flags: c_int,
+) -> io::Result<usize> {
+    sendmsg(fd, addr.as_storage_ptr(), addr.len(), bufs, flags)
+}
+
+/// Returns the (bytes received, sending address len, `RecvFlags`).
+fn sendmsg(
+    fd: SysSocket,
+    msg_name: *const sockaddr_storage,
+    msg_namelen: socklen_t,
+    bufs: &[IoSlice<'_>],
+    flags: c_int,
+) -> io::Result<usize> {
+    let mut msg = libc::msghdr {
+        // 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_name: (msg_name as *mut sockaddr_storage).cast(),
+        msg_namelen,
+        // Safety: Same as above about `*const` -> `*mut`.
+        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.
@@ -661,27 +682,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    #[cfg(not(target_os = "redox"))]
-    pub fn send_to_vectored(
-        &self,
-        bufs: &[IoSlice<'_>],
-        flags: c_int,
-        addr: &SockAddr,
-    ) -> io::Result<usize> {
-        let mut msg = libc::msghdr {
-            msg_name: addr.as_ptr() as *mut c_void,
-            msg_namelen: addr.len(),
-            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)
-    }
-
     pub fn ttl(&self) -> io::Result<u32> {
         unsafe {
             let raw: c_int = self.getsockopt(libc::IPPROTO_IP, libc::IP_TTL)?;

src/sys/windows.rs

@@ -437,6 +437,8 @@ pub(crate) fn send_vectored(
             //
             // Tracking issue: https://github.com/rust-lang/socket2-rs/issues/129.
             //
+            // NOTE: `send_to_vectored` has the same problem.
+            //
             // [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,
@@ -472,6 +474,32 @@ pub(crate) fn send_to(
     .map(|n| n as usize)
 }
 
+pub(crate) fn send_to_vectored(
+    socket: SysSocket,
+    bufs: &[IoSlice<'_>],
+    addr: &SockAddr,
+    flags: c_int,
+) -> io::Result<usize> {
+    let mut nsent = 0;
+    syscall!(
+        WSASendTo(
+            socket,
+            // FIXME: Same problem as in `send_vectored`.
+            bufs.as_ptr() as *mut _,
+            bufs.len().min(DWORD::MAX as usize) as DWORD,
+            &mut nsent,
+            flags as DWORD,
+            addr.as_ptr(),
+            addr.len(),
+            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();
@@ -525,34 +553,6 @@ pub struct Socket {
 }
 
 impl Socket {
-    pub fn send_to_vectored(
-        &self,
-        bufs: &[IoSlice<'_>],
-        flags: c_int,
-        addr: &SockAddr,
-    ) -> io::Result<usize> {
-        let mut nsent = 0;
-        let ret = unsafe {
-            sock::WSASendTo(
-                self.socket,
-                bufs.as_ptr() as *mut WSABUF,
-                bufs.len().min(DWORD::MAX as usize) as DWORD,
-                &mut nsent,
-                flags as DWORD,
-                addr.as_ptr(),
-                addr.len(),
-                std::ptr::null_mut(),
-                None,
-            )
-        };
-        match ret {
-            0 => Ok(nsent as usize),
-            _ => Err(last_error()),
-        }
-    }
-
-    // ================================================
-
     pub fn ttl(&self) -> io::Result<u32> {
         unsafe {
             let raw: c_int = self.getsockopt(IPPROTO_IP, IP_TTL)?;

src/tests.rs

@@ -215,7 +215,6 @@ fn send_from_recv_to_vectored() {
                 IoSlice::new(b"menswear"),
             ],
             &addr_b,
-            0,
         )
         .unwrap();
     assert_eq!(sent, 18);