sys_util: Add ReadWriteAtVolatile trait and `*vectored` functions

Add the FileReadWriteAtVolatile trait, which is basically the same as
the FileReadWriteVolatile trait but additionally takes an offest.  This
is only useful for types that are seekable and can allow concurrent
operations on the same underlying type.

Also add `*_vectored` versions of all the functions.  These match the
`*_vectored` functions in the standard library and can reduce the number
of system calls needed to read or write a whole buffer.

Implement both traits for `&mut T` if `T` implements them.

Change the trait implementation for `File` to a macro so that we can
also implement it for `GuestMemory`.

BUG=b:136128319
TEST=unit tests

Change-Id: I3d8eb7bba17fe3247e18649b1b04e21a91a841e2
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/1724229
Auto-Submit: Chirantan Ekbote <chirantan@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Daniel Verkamp <dverkamp@chromium.org>
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>

2 files changed, 343 insertions, 28 deletions

diff --git a/sys_util/src/file_traits.rs b/sys_util/src/file_traits.rs
index c0c7a9b..d35bc4f 100644
--- a/sys_util/src/file_traits.rs
+++ b/sys_util/src/file_traits.rs
@@ -8,7 +8,12 @@ use std::os::unix::io::{AsRawFd, RawFd};
 
 use data_model::VolatileSlice;
 
-use libc::{c_void, read, write};
+use libc::{
+    c_int, c_void, off64_t, pread64, preadv64, pwrite64, pwritev64, read, readv, size_t, write,
+    writev,
+};
+
+use crate::SharedMemory;
 
 /// A trait for flushing the contents of a file to disk.
 /// This is equivalent to File's `sync_all` method, but
@@ -47,6 +52,18 @@ pub trait FileReadWriteVolatile {
     /// success.
     fn read_volatile(&mut self, slice: VolatileSlice) -> Result<usize>;
 
+    /// Like `read_volatile`, except it reads to a slice of buffers. Data is copied to fill each
+    /// buffer in order, with the final buffer written to possibly being only partially filled. This
+    /// method must behave as a single call to `read_volatile` with the buffers concatenated would.
+    /// The default implementation calls `read_volatile` with either the first nonempty buffer
+    /// provided, or returns `Ok(0)` if none exists.
+    fn read_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+        bufs.iter()
+            .find(|b| b.size() > 0)
+            .map(|&b| self.read_volatile(b))
+            .unwrap_or(Ok(0))
+    }
+
     /// Reads bytes from this into the given slice until all bytes in the slice are written, or an
     /// error is returned.
     fn read_exact_volatile(&mut self, mut slice: VolatileSlice) -> Result<()> {
@@ -66,6 +83,18 @@ pub trait FileReadWriteVolatile {
     /// success.
     fn write_volatile(&mut self, slice: VolatileSlice) -> Result<usize>;
 
+    /// Like `write_volatile`, except that it writes from a slice of buffers. Data is copied from
+    /// each buffer in order, with the final buffer read from possibly being only partially
+    /// consumed. This method must behave as a call to `write_volatile` with the buffers
+    /// concatenated would. The default implementation calls `write_volatile` with either the first
+    /// nonempty buffer provided, or returns `Ok(0)` if none exists.
+    fn write_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+        bufs.iter()
+            .find(|b| b.size() > 0)
+            .map(|&b| self.write_volatile(b))
+            .unwrap_or(Ok(0))
+    }
+
     /// Write bytes from the slice to the given file until all the bytes from the slice have been
     /// written, or an error is returned.
     fn write_all_volatile(&mut self, mut slice: VolatileSlice) -> Result<()> {
@@ -82,42 +111,326 @@ pub trait FileReadWriteVolatile {
     }
 }
 
-impl FileReadWriteVolatile for File {
+impl<'a, T: FileReadWriteVolatile> FileReadWriteVolatile for &'a mut T {
     fn read_volatile(&mut self, slice: VolatileSlice) -> Result<usize> {
-        // Safe because only bytes inside the slice are accessed and the kernel is expected to
-        // handle arbitrary memory for I/O.
-        let ret = unsafe {
-            read(
-                self.as_raw_fd(),
-                slice.as_ptr() as *mut c_void,
-                slice.size() as usize,
-            )
-        };
-        if ret >= 0 {
-            Ok(ret as usize)
+        (**self).read_volatile(slice)
+    }
+
+    fn read_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+        (**self).read_vectored_volatile(bufs)
+    }
+
+    fn read_exact_volatile(&mut self, slice: VolatileSlice) -> Result<()> {
+        (**self).read_exact_volatile(slice)
+    }
+
+    fn write_volatile(&mut self, slice: VolatileSlice) -> Result<usize> {
+        (**self).write_volatile(slice)
+    }
+
+    fn write_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+        (**self).write_vectored_volatile(bufs)
+    }
+
+    fn write_all_volatile(&mut self, slice: VolatileSlice) -> Result<()> {
+        (**self).write_all_volatile(slice)
+    }
+}
+
+/// A trait similar to the unix `ReadExt` and `WriteExt` traits, but for volatile memory.
+pub trait FileReadWriteAtVolatile {
+    /// Reads bytes from this file at `offset` into the given slice, returning the number of bytes
+    /// read on success.
+    fn read_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize>;
+
+    /// Like `read_at_volatile`, except it reads to a slice of buffers. Data is copied to fill each
+    /// buffer in order, with the final buffer written to possibly being only partially filled. This
+    /// method must behave as a single call to `read_at_volatile` with the buffers concatenated
+    /// would. The default implementation calls `read_at_volatile` with either the first nonempty
+    /// buffer provided, or returns `Ok(0)` if none exists.
+    fn read_vectored_at_volatile(&mut self, bufs: &[VolatileSlice], offset: u64) -> Result<usize> {
+        if let Some(&slice) = bufs.first() {
+            self.read_at_volatile(slice, offset)
         } else {
-            Err(Error::last_os_error())
+            Ok(0)
         }
     }
 
-    fn write_volatile(&mut self, slice: VolatileSlice) -> Result<usize> {
-        // Safe because only bytes inside the slice are accessed and the kernel is expected to
-        // handle arbitrary memory for I/O.
-        let ret = unsafe {
-            write(
-                self.as_raw_fd(),
-                slice.as_ptr() as *const c_void,
-                slice.size() as usize,
-            )
-        };
-        if ret >= 0 {
-            Ok(ret as usize)
+    /// Reads bytes from this file at `offset` into the given slice until all bytes in the slice are
+    /// read, or an error is returned.
+    fn read_exact_at_volatile(&mut self, mut slice: VolatileSlice, mut offset: u64) -> Result<()> {
+        while slice.size() > 0 {
+            match self.read_at_volatile(slice, offset) {
+                Ok(0) => return Err(Error::from(ErrorKind::UnexpectedEof)),
+                Ok(n) => {
+                    slice = slice.offset(n as u64).unwrap();
+                    offset = offset.checked_add(n as u64).unwrap();
+                }
+                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(())
+    }
+
+    /// Writes bytes from this file at `offset` into the given slice, returning the number of bytes
+    /// written on success.
+    fn write_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize>;
+
+    /// Like `write_at_at_volatile`, except that it writes from a slice of buffers. Data is copied
+    /// from each buffer in order, with the final buffer read from possibly being only partially
+    /// consumed. This method must behave as a call to `write_at_volatile` with the buffers
+    /// concatenated would. The default implementation calls `write_at_volatile` with either the
+    /// first nonempty buffer provided, or returns `Ok(0)` if none exists.
+    fn write_vectored_at_volatile(&mut self, bufs: &[VolatileSlice], offset: u64) -> Result<usize> {
+        if let Some(&slice) = bufs.first() {
+            self.write_at_volatile(slice, offset)
         } else {
-            Err(Error::last_os_error())
+            Ok(0)
         }
     }
+
+    /// Writes bytes from this file at `offset` into the given slice until all bytes in the slice
+    /// are written, or an error is returned.
+    fn write_all_at_volatile(&mut self, mut slice: VolatileSlice, mut offset: u64) -> Result<()> {
+        while slice.size() > 0 {
+            match self.write_at_volatile(slice, offset) {
+                Ok(0) => return Err(Error::from(ErrorKind::WriteZero)),
+                Ok(n) => {
+                    slice = slice.offset(n as u64).unwrap();
+                    offset = offset.checked_add(n as u64).unwrap();
+                }
+                Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
+                Err(e) => return Err(e),
+            }
+        }
+        Ok(())
+    }
 }
 
+impl<'a, T: FileReadWriteAtVolatile> FileReadWriteAtVolatile for &'a mut T {
+    fn read_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize> {
+        (**self).read_at_volatile(slice, offset)
+    }
+
+    fn read_vectored_at_volatile(&mut self, bufs: &[VolatileSlice], offset: u64) -> Result<usize> {
+        (**self).read_vectored_at_volatile(bufs, offset)
+    }
+
+    fn read_exact_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<()> {
+        (**self).read_exact_at_volatile(slice, offset)
+    }
+
+    fn write_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize> {
+        (**self).write_at_volatile(slice, offset)
+    }
+
+    fn write_vectored_at_volatile(&mut self, bufs: &[VolatileSlice], offset: u64) -> Result<usize> {
+        (**self).write_vectored_at_volatile(bufs, offset)
+    }
+
+    fn write_all_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<()> {
+        (**self).write_all_at_volatile(slice, offset)
+    }
+}
+
+macro_rules! volatile_impl {
+    ($ty:ty) => {
+        impl FileReadWriteVolatile for $ty {
+            fn read_volatile(&mut self, slice: VolatileSlice) -> Result<usize> {
+                // Safe because only bytes inside the slice are accessed and the kernel is expected
+                // to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    read(
+                        self.as_raw_fd(),
+                        slice.as_ptr() as *mut c_void,
+                        slice.size() as usize,
+                    )
+                };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn read_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+                let iovecs: Vec<libc::iovec> = bufs
+                    .iter()
+                    .map(|s| libc::iovec {
+                        iov_base: s.as_ptr() as *mut c_void,
+                        iov_len: s.size() as size_t,
+                    })
+                    .collect();
+
+                if iovecs.is_empty() {
+                    return Ok(0);
+                }
+
+                // Safe because only bytes inside the buffers are accessed and the kernel is
+                // expected to handle arbitrary memory for I/O.
+                let ret = unsafe { readv(self.as_raw_fd(), &iovecs[0], iovecs.len() as c_int) };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn write_volatile(&mut self, slice: VolatileSlice) -> Result<usize> {
+                // Safe because only bytes inside the slice are accessed and the kernel is expected
+                // to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    write(
+                        self.as_raw_fd(),
+                        slice.as_ptr() as *const c_void,
+                        slice.size() as usize,
+                    )
+                };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn write_vectored_volatile(&mut self, bufs: &[VolatileSlice]) -> Result<usize> {
+                let iovecs: Vec<libc::iovec> = bufs
+                    .iter()
+                    .map(|s| libc::iovec {
+                        iov_base: s.as_ptr() as *mut c_void,
+                        iov_len: s.size() as size_t,
+                    })
+                    .collect();
+
+                if iovecs.is_empty() {
+                    return Ok(0);
+                }
+
+                // Safe because only bytes inside the buffers are accessed and the kernel is
+                // expected to handle arbitrary memory for I/O.
+                let ret = unsafe { writev(self.as_raw_fd(), &iovecs[0], iovecs.len() as c_int) };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+        }
+
+        impl FileReadWriteAtVolatile for $ty {
+            fn read_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize> {
+                // Safe because only bytes inside the slice are accessed and the kernel is expected
+                // to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    pread64(
+                        self.as_raw_fd(),
+                        slice.as_ptr() as *mut c_void,
+                        slice.size() as usize,
+                        offset as off64_t,
+                    )
+                };
+
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn read_vectored_at_volatile(
+                &mut self,
+                bufs: &[VolatileSlice],
+                offset: u64,
+            ) -> Result<usize> {
+                let iovecs: Vec<libc::iovec> = bufs
+                    .iter()
+                    .map(|s| libc::iovec {
+                        iov_base: s.as_ptr() as *mut c_void,
+                        iov_len: s.size() as size_t,
+                    })
+                    .collect();
+
+                if iovecs.is_empty() {
+                    return Ok(0);
+                }
+
+                // Safe because only bytes inside the buffers are accessed and the kernel is
+                // expected to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    preadv64(
+                        self.as_raw_fd(),
+                        &iovecs[0],
+                        iovecs.len() as c_int,
+                        offset as off64_t,
+                    )
+                };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn write_at_volatile(&mut self, slice: VolatileSlice, offset: u64) -> Result<usize> {
+                // Safe because only bytes inside the slice are accessed and the kernel is expected
+                // to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    pwrite64(
+                        self.as_raw_fd(),
+                        slice.as_ptr() as *const c_void,
+                        slice.size() as usize,
+                        offset as off64_t,
+                    )
+                };
+
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+
+            fn write_vectored_at_volatile(
+                &mut self,
+                bufs: &[VolatileSlice],
+                offset: u64,
+            ) -> Result<usize> {
+                let iovecs: Vec<libc::iovec> = bufs
+                    .iter()
+                    .map(|s| libc::iovec {
+                        iov_base: s.as_ptr() as *mut c_void,
+                        iov_len: s.size() as size_t,
+                    })
+                    .collect();
+
+                if iovecs.is_empty() {
+                    return Ok(0);
+                }
+
+                // Safe because only bytes inside the buffers are accessed and the kernel is
+                // expected to handle arbitrary memory for I/O.
+                let ret = unsafe {
+                    pwritev64(
+                        self.as_raw_fd(),
+                        &iovecs[0],
+                        iovecs.len() as c_int,
+                        offset as off64_t,
+                    )
+                };
+                if ret >= 0 {
+                    Ok(ret as usize)
+                } else {
+                    Err(Error::last_os_error())
+                }
+            }
+        }
+    };
+}
+
+volatile_impl!(File);
+volatile_impl!(SharedMemory);
+
 /// A trait similar to `AsRawFd` but supports an arbitrary number of file descriptors.
 pub trait AsRawFds {
     fn as_raw_fds(&self) -> Vec<RawFd>;
diff --git a/sys_util/src/lib.rs b/sys_util/src/lib.rs
index 697b0e1..7cba7eb 100644
--- a/sys_util/src/lib.rs
+++ b/sys_util/src/lib.rs
@@ -63,7 +63,9 @@ pub use crate::terminal::*;
 pub use crate::timerfd::*;
 pub use poll_token_derive::*;
 
-pub use crate::file_traits::{AsRawFds, FileReadWriteVolatile, FileSetLen, FileSync};
+pub use crate::file_traits::{
+    AsRawFds, FileReadWriteAtVolatile, FileReadWriteVolatile, FileSetLen, FileSync,
+};
 pub use crate::guest_memory::Error as GuestMemoryError;
 pub use crate::mmap::Error as MmapError;
 pub use crate::seek_hole::SeekHole;