Ensure non-empty buffers for large vectored I/O

`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.

Author: thaliaarchi

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,25 @@ 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.
+    #[allow(dead_code)] // Not used on all platforms
+    #[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 +1569,25 @@ 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.
+    #[allow(dead_code)] // Not used on all platforms
+    #[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

library/std/src/sys/fd/hermit.rs

@@ -1,6 +1,5 @@
 #![unstable(reason = "not public", issue = "none", feature = "fd")]
 
-use crate::cmp;
 use crate::io::{self, BorrowedCursor, IoSlice, IoSliceMut, Read, SeekFrom};
 use crate::os::hermit::hermit_abi;
 use crate::os::hermit::io::{AsFd, AsRawFd, BorrowedFd, FromRawFd, IntoRawFd, OwnedFd, RawFd};
@@ -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)

library/std/src/sys/fd/unix.rs

@@ -108,12 +108,13 @@ impl FileDesc {
         target_os = "vita",
         target_os = "nuttx"
     )))]
-    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 {
             libc::readv(
                 self.as_raw_fd(),
                 bufs.as_mut_ptr() as *mut libc::iovec as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
             )
         })?;
         Ok(ret as usize)
@@ -198,12 +199,17 @@ impl FileDesc {
         target_os = "netbsd",
         target_os = "openbsd", // OpenBSD 2.7
     ))]
-    pub fn read_vectored_at(&self, bufs: &mut [IoSliceMut<'_>], offset: u64) -> io::Result<usize> {
+    pub fn read_vectored_at(
+        &self,
+        mut bufs: &mut [IoSliceMut<'_>],
+        offset: u64,
+    ) -> io::Result<usize> {
+        IoSliceMut::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             libc::preadv(
                 self.as_raw_fd(),
                 bufs.as_mut_ptr() as *mut libc::iovec as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
                 offset as _,
             )
         })?;
@@ -235,7 +241,11 @@ impl FileDesc {
     // passing 64-bits parameters to syscalls, so we fallback to the default
     // implementation if `preadv` is not available.
     #[cfg(all(target_os = "android", target_pointer_width = "64"))]
-    pub fn read_vectored_at(&self, bufs: &mut [IoSliceMut<'_>], offset: u64) -> io::Result<usize> {
+    pub fn read_vectored_at(
+        &self,
+        mut bufs: &mut [IoSliceMut<'_>],
+        offset: u64,
+    ) -> io::Result<usize> {
         syscall!(
             fn preadv(
                 fd: libc::c_int,
@@ -245,11 +255,12 @@ impl FileDesc {
             ) -> isize;
         );
 
+        IoSliceMut::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             preadv(
                 self.as_raw_fd(),
                 bufs.as_mut_ptr() as *mut libc::iovec as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
                 offset as _,
             )
         })?;
@@ -260,7 +271,11 @@ impl FileDesc {
     // FIXME(#115199): Rust currently omits weak function definitions
     // and its metadata from LLVM IR.
     #[no_sanitize(cfi)]
-    pub fn read_vectored_at(&self, bufs: &mut [IoSliceMut<'_>], offset: u64) -> io::Result<usize> {
+    pub fn read_vectored_at(
+        &self,
+        mut bufs: &mut [IoSliceMut<'_>],
+        offset: u64,
+    ) -> io::Result<usize> {
         weak!(
             fn preadv64(
                 fd: libc::c_int,
@@ -272,11 +287,12 @@ impl FileDesc {
 
         match preadv64.get() {
             Some(preadv) => {
+                IoSliceMut::limit_slices(&mut bufs, max_iov());
                 let ret = cvt(unsafe {
                     preadv(
                         self.as_raw_fd(),
                         bufs.as_mut_ptr() as *mut libc::iovec as *const libc::iovec,
-                        cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                        bufs.len() as libc::c_int,
                         offset as _,
                     )
                 })?;
@@ -296,7 +312,11 @@ impl FileDesc {
     // These versions may be newer than the minimum supported versions of OS's we support so we must
     // use "weak" linking.
     #[cfg(target_vendor = "apple")]
-    pub fn read_vectored_at(&self, bufs: &mut [IoSliceMut<'_>], offset: u64) -> io::Result<usize> {
+    pub fn read_vectored_at(
+        &self,
+        mut bufs: &mut [IoSliceMut<'_>],
+        offset: u64,
+    ) -> io::Result<usize> {
         weak!(
             fn preadv(
                 fd: libc::c_int,
@@ -308,11 +328,12 @@ impl FileDesc {
 
         match preadv.get() {
             Some(preadv) => {
+                IoSliceMut::limit_slices(&mut bufs, max_iov());
                 let ret = cvt(unsafe {
                     preadv(
                         self.as_raw_fd(),
                         bufs.as_mut_ptr() as *mut libc::iovec as *const libc::iovec,
-                        cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                        bufs.len() as libc::c_int,
                         offset as _,
                     )
                 })?;
@@ -339,12 +360,13 @@ impl FileDesc {
         target_os = "vita",
         target_os = "nuttx"
     )))]
-    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 {
             libc::writev(
                 self.as_raw_fd(),
                 bufs.as_ptr() as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
             )
         })?;
         Ok(ret as usize)
@@ -408,12 +430,13 @@ impl FileDesc {
         target_os = "netbsd",
         target_os = "openbsd", // OpenBSD 2.7
     ))]
-    pub fn write_vectored_at(&self, bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+    pub fn write_vectored_at(&self, mut bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+        IoSlice::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             libc::pwritev(
                 self.as_raw_fd(),
                 bufs.as_ptr() as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
                 offset as _,
             )
         })?;
@@ -445,7 +468,7 @@ impl FileDesc {
     // passing 64-bits parameters to syscalls, so we fallback to the default
     // implementation if `pwritev` is not available.
     #[cfg(all(target_os = "android", target_pointer_width = "64"))]
-    pub fn write_vectored_at(&self, bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+    pub fn write_vectored_at(&self, mut bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
         syscall!(
             fn pwritev(
                 fd: libc::c_int,
@@ -455,19 +478,20 @@ impl FileDesc {
             ) -> isize;
         );
 
+        IoSlice::limit_slices(&mut bufs, max_iov());
         let ret = cvt(unsafe {
             pwritev(
                 self.as_raw_fd(),
                 bufs.as_ptr() as *const libc::iovec,
-                cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                bufs.len() as libc::c_int,
                 offset as _,
             )
         })?;
         Ok(ret as usize)
     }
 
     #[cfg(all(target_os = "android", target_pointer_width = "32"))]
-    pub fn write_vectored_at(&self, bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+    pub fn write_vectored_at(&self, mut bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
         weak!(
             fn pwritev64(
                 fd: libc::c_int,
@@ -479,11 +503,12 @@ impl FileDesc {
 
         match pwritev64.get() {
             Some(pwritev) => {
+                IoSlice::limit_slices(&mut bufs, max_iov());
                 let ret = cvt(unsafe {
                     pwritev(
                         self.as_raw_fd(),
                         bufs.as_ptr() as *const libc::iovec,
-                        cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                        bufs.len() as libc::c_int,
                         offset as _,
                     )
                 })?;
@@ -503,7 +528,7 @@ impl FileDesc {
     // These versions may be newer than the minimum supported versions of OS's we support so we must
     // use "weak" linking.
     #[cfg(target_vendor = "apple")]
-    pub fn write_vectored_at(&self, bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+    pub fn write_vectored_at(&self, mut bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
         weak!(
             fn pwritev(
                 fd: libc::c_int,
@@ -515,11 +540,12 @@ impl FileDesc {
 
         match pwritev.get() {
             Some(pwritev) => {
+                IoSlice::limit_slices(&mut bufs, max_iov());
                 let ret = cvt(unsafe {
                     pwritev(
                         self.as_raw_fd(),
                         bufs.as_ptr() as *const libc::iovec,
-                        cmp::min(bufs.len(), max_iov()) as libc::c_int,
+                        bufs.len() as libc::c_int,
                         offset as _,
                     )
                 })?;

library/std/src/sys/fd/unix/tests.rs

@@ -1,12 +1,23 @@
 use core::mem::ManuallyDrop;
 
-use super::FileDesc;
+use super::{FileDesc, max_iov};
 use crate::io::IoSlice;
 use crate::os::unix::io::FromRawFd;
 
 #[test]
 fn limit_vector_count() {
+    const IOV_MAX: usize = max_iov();
+
     let stdout = ManuallyDrop::new(unsafe { FileDesc::from_raw_fd(1) });
-    let bufs = (0..1500).map(|_| IoSlice::new(&[])).collect::<Vec<_>>();
-    assert!(stdout.write_vectored(&bufs).is_ok());
+    let mut bufs = vec![IoSlice::new(&[]); IOV_MAX * 2 + 1];
+    assert_eq!(stdout.write_vectored(&bufs).unwrap(), 0);
+
+    // The slice of buffers is truncated to IOV_MAX buffers. However, since the
+    // first IOV_MAX buffers are all empty, it is sliced starting at the first
+    // non-empty buffer to avoid erroneously returning Ok(0). In this case, that
+    // starts with the b"hello" buffer and ends just before the b"world!"
+    // buffer.
+    bufs[IOV_MAX] = IoSlice::new(b"hello");
+    bufs[IOV_MAX * 2] = IoSlice::new(b"world!");
+    assert_eq!(stdout.write_vectored(&bufs).unwrap(), b"hello".len())
 }

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)
     }