Ensure non-empty buffers for large vectored I/O

thaliaarchi · thaliaarchi · commit d4bc926beed2 · 2025-03-24T17:44:17.000-07:00
`readv` and `writev` are constrained by a platform-specific upper bound
on the number of buffers which can be passed. Currently, `read_vectored`
and `write_vectored` implementations simply truncate to this limit when
larger. However, when the only non-empty buffers are at indices above
this limit, they will erroneously return `Ok(0)`.

Instead, slice the buffers starting at the first non-empty buffer. This
trades a conditional move for a branch, so it's barely a penalty in the
common case.

The new method `limit_slices` on `IoSlice` and `IoSliceMut` may be
generally useful to users like `advance_slices` is, but I have left it
as `pub(crate)` for now.
diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
@@ -297,6 +297,7 @@
 #[cfg(test)]
 mod tests;
 
+use core::intrinsics;
 #[unstable(feature = "read_buf", issue = "78485")]
 pub use core::io::{BorrowedBuf, BorrowedCursor};
 use core::slice::memchr;
@@ -1388,6 +1389,24 @@ impl<'a> IoSliceMut<'a> {
         }
     }
 
+    /// Limits a slice of buffers to at most `n` buffers.
+    ///
+    /// When the slice contains over `n` buffers, ensure that at least one
+    /// non-empty buffer is in the truncated slice, if there is one.
+    #[inline]
+    pub(crate) fn limit_slices(bufs: &mut &mut [IoSliceMut<'a>], n: usize) {
+        if intrinsics::unlikely(bufs.len() > n) {
+            for (i, buf) in bufs.iter().enumerate() {
+                if !buf.is_empty() {
+                    let len = cmp::min(bufs.len() - i, n);
+                    *bufs = &mut take(bufs)[i..i + len];
+                    return;
+                }
+            }
+            *bufs = &mut take(bufs)[..0];
+        }
+    }
+
     /// Get the underlying bytes as a mutable slice with the original lifetime.
     ///
     /// # Examples
@@ -1549,6 +1568,24 @@ impl<'a> IoSlice<'a> {
         }
     }
 
+    /// Limits a slice of buffers to at most `n` buffers.
+    ///
+    /// When the slice contains over `n` buffers, ensure that at least one
+    /// non-empty buffer is in the truncated slice, if there is one.
+    #[inline]
+    pub(crate) fn limit_slices(bufs: &mut &[IoSlice<'a>], n: usize) {
+        if intrinsics::unlikely(bufs.len() > n) {
+            for (i, buf) in bufs.iter().enumerate() {
+                if !buf.is_empty() {
+                    let len = cmp::min(bufs.len() - i, n);
+                    *bufs = &bufs[i..i + len];
+                    return;
+                }
+            }
+            *bufs = &bufs[..0];
+        }
+    }
+
     /// Get the underlying bytes as a slice with the original lifetime.
     ///
     /// This doesn't borrow from `self`, so is less restrictive than calling
diff --git a/library/std/src/sys/net/connection/socket/solid.rs b/library/std/src/sys/net/connection/socket/solid.rs
@@ -9,7 +9,7 @@ use crate::os::solid::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, Owne
 use crate::sys::abi;
 use crate::sys_common::{FromInner, IntoInner};
 use crate::time::Duration;
-use crate::{cmp, mem, ptr, str};
+use crate::{mem, ptr, str};
 
 pub(super) mod netc {
     pub use crate::sys::abi::sockets::*;
@@ -222,13 +222,10 @@ impl Socket {
         self.recv_with_flags(buf, 0)
     }
 
-    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+    pub fn read_vectored(&self, mut bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        IoSliceMut::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
-            netc::readv(
-                self.as_raw_fd(),
-                bufs.as_ptr() as *const netc::iovec,
-                cmp::min(bufs.len(), max_iov()) as c_int,
-            )
+            netc::readv(self.as_raw_fd(), bufs.as_ptr() as *const netc::iovec, bufs.len() as c_int)
         })?;
         Ok(ret as usize)
     }
@@ -267,13 +264,10 @@ impl Socket {
         self.recv_from_with_flags(buf, MSG_PEEK)
     }
 
-    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+    pub fn write_vectored(&self, mut bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        IoSlice::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
-            netc::writev(
-                self.as_raw_fd(),
-                bufs.as_ptr() as *const netc::iovec,
-                cmp::min(bufs.len(), max_iov()) as c_int,
-            )
+            netc::writev(self.as_raw_fd(), bufs.as_ptr() as *const netc::iovec, bufs.len() as c_int)
         })?;
         Ok(ret as usize)
     }
diff --git a/library/std/src/sys/net/connection/socket/windows.rs b/library/std/src/sys/net/connection/socket/windows.rs
@@ -299,8 +299,6 @@ impl Socket {
     }
 
     fn recv_with_flags(&self, mut buf: BorrowedCursor<'_>, flags: c_int) -> io::Result<()> {
-        // On unix when a socket is shut down all further reads return 0, so we
-        // do the same on windows to map a shut down socket to returning EOF.
         let length = cmp::min(buf.capacity(), i32::MAX as usize) as i32;
         let result =
             unsafe { c::recv(self.as_raw(), buf.as_mut().as_mut_ptr() as *mut _, length, flags) };
@@ -309,6 +307,9 @@ impl Socket {
             c::SOCKET_ERROR => {
                 let error = unsafe { c::WSAGetLastError() };
 
+                // On Unix when a socket is shut down, all further reads return
+                // 0, so we do the same on Windows to map a shut down socket to
+                // returning EOF.
                 if error == c::WSAESHUTDOWN {
                     Ok(())
                 } else {
@@ -332,17 +333,15 @@ impl Socket {
         self.recv_with_flags(buf, 0)
     }
 
-    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
-        // On unix when a socket is shut down all further reads return 0, so we
-        // do the same on windows to map a shut down socket to returning EOF.
-        let length = cmp::min(bufs.len(), u32::MAX as usize) as u32;
+    pub fn read_vectored(&self, mut bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        IoSliceMut::limit_slices(&mut bufs, u32::MAX as usize);
         let mut nread = 0;
         let mut flags = 0;
         let result = unsafe {
             c::WSARecv(
                 self.as_raw(),
                 bufs.as_mut_ptr() as *mut c::WSABUF,
-                length,
+                bufs.len() as u32,
                 &mut nread,
                 &mut flags,
                 ptr::null_mut(),
@@ -355,6 +354,9 @@ impl Socket {
             _ => {
                 let error = unsafe { c::WSAGetLastError() };
 
+                // On Unix when a socket is shut down, all further reads return
+                // 0, so we do the same on Windows to map a shut down socket to
+                // returning EOF.
                 if error == c::WSAESHUTDOWN {
                     Ok(0)
                 } else {
@@ -384,8 +386,6 @@ impl Socket {
         let mut addrlen = size_of_val(&storage) as netc::socklen_t;
         let length = cmp::min(buf.len(), <wrlen_t>::MAX as usize) as wrlen_t;
 
-        // On unix when a socket is shut down all further reads return 0, so we
-        // do the same on windows to map a shut down socket to returning EOF.
         let result = unsafe {
             c::recvfrom(
                 self.as_raw(),
@@ -401,6 +401,9 @@ impl Socket {
             c::SOCKET_ERROR => {
                 let error = unsafe { c::WSAGetLastError() };
 
+                // On Unix when a socket is shut down, all further reads return
+                // 0, so we do the same on Windows to map a shut down socket to
+                // returning EOF.
                 if error == c::WSAESHUTDOWN {
                     Ok((0, unsafe { socket_addr_from_c(&storage, addrlen as usize)? }))
                 } else {
@@ -419,14 +422,14 @@ impl Socket {
         self.recv_from_with_flags(buf, c::MSG_PEEK)
     }
 
-    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
-        let length = cmp::min(bufs.len(), u32::MAX as usize) as u32;
+    pub fn write_vectored(&self, mut bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        IoSlice::limit_slices(&mut bufs, u32::MAX as usize);
         let mut nwritten = 0;
         let result = unsafe {
             c::WSASend(
                 self.as_raw(),
                 bufs.as_ptr() as *const c::WSABUF as *mut _,
-                length,
+                bufs.len() as u32,
                 &mut nwritten,
                 0,
                 ptr::null_mut(),
diff --git a/library/std/src/sys/pal/hermit/fd.rs b/library/std/src/sys/pal/hermit/fd.rs
@@ -1,7 +1,6 @@
 #![unstable(reason = "not public", issue = "none", feature = "fd")]
 
 use super::hermit_abi;
-use crate::cmp;
 use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut, Read, SeekFrom};
 use crate::os::hermit::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
 use crate::sys::{cvt, unsupported};
@@ -38,12 +37,13 @@ impl FileDesc {
         Ok(())
     }
 
-    pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+    pub fn read_vectored(&self, mut bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        IoSliceMut::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             hermit_abi::readv(
                 self.as_raw_fd(),
                 bufs.as_mut_ptr() as *mut hermit_abi::iovec as *const hermit_abi::iovec,
-                cmp::min(bufs.len(), max_iov()),
+                bufs.len(),
             )
         })?;
         Ok(ret as usize)
@@ -65,12 +65,13 @@ impl FileDesc {
         Ok(result as usize)
     }
 
-    pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+    pub fn write_vectored(&self, mut bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        IoSlice::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             hermit_abi::writev(
                 self.as_raw_fd(),
                 bufs.as_ptr() as *const hermit_abi::iovec,
-                cmp::min(bufs.len(), max_iov()),
+                bufs.len(),
             )
         })?;
         Ok(ret as usize)
diff --git a/library/std/src/sys/pal/unix/fd.rs b/library/std/src/sys/pal/unix/fd.rs
