Ensure non-empty buffers for large vectored I/O

library/std/src/io/mod.rs

@@ -1626,6 +1626,76 @@ impl<'a> Deref for IoSlice<'a> {
     }
 }
 
+/// Limits a slice of buffers to at most `n` buffers and ensures that it has at
+/// least one buffer, even if empty.
+///
+/// When the slice contains over `n` buffers, ensure that at least one non-empty
+/// buffer is in the truncated slice, if there is one.
+///
+/// For example, [POSIX writev] requires that `bufs.len()` is in `1..=IOV_MAX`.
+///
+/// [POSIX writev]: https://pubs.opengroup.org/onlinepubs/9799919799/functions/writev.html
+#[allow(unused_macros)] // Not used on all platforms
+pub(crate) macro limit_slices(&mut $bufs:ident, $n:expr) {
+    // Rebind $bufs, so the lifetime does not need to live as long as the
+    // function parameter, and shadow it so the macro caller gets the result.
+    let mut $bufs: &[IoSlice<'_>] = $bufs;
+    let n: usize = $n;
+
+    let empty = &[IoSlice::new(&[])];
+
+    if $bufs.len() > n || $bufs.is_empty() {
+        crate::hint::cold_path();
+        'fixup: {
+            // Find the first non-empty buffer and take up to `n` buffers after it.
+            for (i, buf) in $bufs.iter().enumerate() {
+                if !buf.is_empty() {
+                    let len = cmp::min($bufs.len() - i, n);
+                    $bufs = &$bufs[i..i + len];
+                    break 'fixup;
+                }
+            }
+            // If no non-empty buffer is found, use a single empty buffer.
+            $bufs = empty;
+        }
+    }
+}
+
+/// Limits a slice of buffers to at most `n` buffers and ensures that it has at
+/// least one buffer, even if empty.
+///
+/// When the slice contains over `n` buffers, ensure that at least one non-empty
+/// buffer is in the truncated slice, if there is one.
+///
+/// For example, [POSIX readv] requires that `bufs.len()` is in `1..=IOV_MAX`.
+///
+/// [POSIX readv]: https://pubs.opengroup.org/onlinepubs/9799919799/functions/readv.html
+#[allow(unused_macros)] // Not used on all platforms
+pub(crate) macro limit_slices_mut(&mut $bufs:ident, $n:expr) {
+    // Rebind $bufs, so the lifetime does not need to live as long as the
+    // function parameter, and shadow it so the macro caller gets the result.
+    let mut $bufs: &mut [IoSliceMut<'_>] = $bufs;
+    let n: usize = $n;
+
+    let empty = &mut [IoSliceMut::new(&mut [])];
+
+    if $bufs.len() > n || $bufs.is_empty() {
+        crate::hint::cold_path();
+        'fixup: {
+            // Find the first non-empty buffer and take up to `n` buffers after it.
+            for (i, buf) in $bufs.iter().enumerate() {
+                if !buf.is_empty() {
+                    let len = cmp::min($bufs.len() - i, n);
+                    $bufs = &mut $bufs[i..i + len];
+                    break 'fixup;
+                }
+            }
+            // If no non-empty buffer is found, use a single empty buffer.
+            $bufs = empty;
+        }
+    }
+}
+
 /// A trait for objects which are byte-oriented sinks.
 ///
 /// Implementors of the `Write` trait are sometimes called 'writers'.

library/std/src/lib.rs

@@ -280,6 +280,7 @@
 #![feature(cfg_target_thread_local)]
 #![feature(cfi_encoding)]
 #![feature(char_max_len)]
+#![feature(cold_path)]
 #![feature(const_trait_impl)]
 #![feature(core_float_math)]
 #![feature(decl_macro)]

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};
@@ -39,11 +38,12 @@ impl FileDesc {
     }
 
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        io::limit_slices_mut!(&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)
@@ -66,11 +66,12 @@ impl FileDesc {
     }
 
     pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        io::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

@@ -126,11 +126,12 @@ impl FileDesc {
         target_os = "nuttx"
     )))]
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        io::limit_slices_mut!(&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)
@@ -221,11 +222,12 @@ impl FileDesc {
         target_os = "openbsd", // OpenBSD 2.7
     ))]
     pub fn read_vectored_at(&self, bufs: &mut [IoSliceMut<'_>], offset: u64) -> io::Result<usize> {
+        io::limit_slices_mut!(&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 _,
             )
         })?;
@@ -267,11 +269,12 @@ impl FileDesc {
             ) -> isize;
         );
 
+        io::limit_slices_mut!(&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 _,
             )
         })?;
@@ -291,11 +294,12 @@ impl FileDesc {
 
         match preadv64.get() {
             Some(preadv) => {
+                io::limit_slices_mut!(&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 _,
                     )
                 })?;
@@ -327,11 +331,12 @@ impl FileDesc {
 
         match preadv.get() {
             Some(preadv) => {
+                io::limit_slices_mut!(&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 _,
                     )
                 })?;
@@ -359,11 +364,12 @@ impl FileDesc {
         target_os = "nuttx"
     )))]
     pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        io::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)
@@ -427,11 +433,12 @@ impl FileDesc {
         target_os = "openbsd", // OpenBSD 2.7
     ))]
     pub fn write_vectored_at(&self, bufs: &[IoSlice<'_>], offset: u64) -> io::Result<usize> {
+        io::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 _,
             )
         })?;
@@ -473,11 +480,12 @@ impl FileDesc {
             ) -> isize;
         );
 
+        io::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 _,
             )
         })?;
@@ -497,11 +505,12 @@ impl FileDesc {
 
         match pwritev64.get() {
             Some(pwritev) => {
+                io::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 _,
                     )
                 })?;
@@ -533,11 +542,12 @@ impl FileDesc {
 
         match pwritev.get() {
             Some(pwritev) => {
+                io::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,32 @@
 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 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())
+}
+
+#[test]
+fn empty_vector() {
+    let stdin = ManuallyDrop::new(unsafe { FileDesc::from_raw_fd(0) });
+    assert_eq!(stdin.read_vectored(&mut []).unwrap(), 0);
+
     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());
+    assert_eq!(stdout.write_vectored(&[]).unwrap(), 0);
 }

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::*;
@@ -223,12 +223,9 @@ impl Socket {
     }
 
     pub fn read_vectored(&self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
+        io::limit_slices_mut!(&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)
     }
@@ -268,12 +265,9 @@ impl Socket {
     }
 
     pub fn write_vectored(&self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
+        io::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)
     }