feat: Introduce Read/WriteVolatile traits

With #217, various `read_from` and `write_to` methods across vm-memory
incurred a performance penalty due to an unneccesary copy of all data
read and written. This is due to these functions operating on arbitrary
Read/Write implementations, which cannot generally be assumed to respect
the volatile semantics of guest memory. Additionally, they required
taking slices to guest memory, which was undefined behavior.

This patch series fixes those short-comings by introducing a pair of new
traits, `ReadVolatile` and `WriteVolatile` which are identical to the
standard library's `Read` and `Write` traits, with the difference being
that they operate on `VolatileSlice`s instead of `&mut [u8]`. Using
these traits will therefore not incur the performance penalty of #217.
Additionally, this will make using the vm-memory library easier for
users familiar with std::io, as the APIs will be more familiar (in
particular, the unintuitive "swapping" of reader/writer src and dst of
the old read_from/write_to APIs is resolved).

Please note that `impl<T: AsRawFd> ReadVolatile for T` is not be
possible, as the orphan rules will note that "upstream crates may add a
new impl of trait std::os::fd::AsRawFd for type &[u8] in future
versions".

Signed-off-by: Patrick Roy <roypat@amazon.co.uk>

Author: roypat

src/bytes.rs

@@ -280,6 +280,9 @@ pub trait Bytes<A> {
     /// * `addr` - Begin writing at this address.
     /// * `src` - Copy from `src` into the container.
     /// * `count` - Copy `count` bytes from `src` into the container.
+    #[deprecated(
+        note = "Use `.read_from_volatile` or the functions of the `ReadVolatile` trait instead"
+    )]
     fn read_from<F>(&self, addr: A, src: &mut F, count: usize) -> Result<usize, Self::E>
     where
         F: Read;
@@ -295,6 +298,9 @@ pub trait Bytes<A> {
     /// * `addr` - Begin writing at this address.
     /// * `src` - Copy from `src` into the container.
     /// * `count` - Copy exactly `count` bytes from `src` into the container.
+    #[deprecated(
+        note = "Use `.read_exact_from_volatile` or the functions of the `ReadVolatile` trait instead"
+    )]
     fn read_exact_from<F>(&self, addr: A, src: &mut F, count: usize) -> Result<(), Self::E>
     where
         F: Read;
@@ -307,6 +313,9 @@ pub trait Bytes<A> {
     /// * `addr` - Begin reading from this address.
     /// * `dst` - Copy from the container to `dst`.
     /// * `count` - Copy `count` bytes from the container to `dst`.
+    #[deprecated(
+        note = "Use `.write_to_volatile` or the functions of the `WriteVolatile` trait instead"
+    )]
     fn write_to<F>(&self, addr: A, dst: &mut F, count: usize) -> Result<usize, Self::E>
     where
         F: Write;
@@ -322,6 +331,9 @@ pub trait Bytes<A> {
     /// * `addr` - Begin reading from this address.
     /// * `dst` - Copy from the container to `dst`.
     /// * `count` - Copy exactly `count` bytes from the container to `dst`.
+    #[deprecated(
+        note = "Use `.write_all_to_volatile` or the functions of the `WriteVolatile` trait instead"
+    )]
     fn write_all_to<F>(&self, addr: A, dst: &mut F, count: usize) -> Result<(), Self::E>
     where
         F: Write;

src/guest_memory.rs

@@ -856,6 +856,7 @@ impl<T: GuestMemory + ?Sized> Bytes<GuestAddress> for T {
     where
         F: Read,
     {
+        #[allow(deprecated)] // this function itself is deprecated
         let res = self.read_from(addr, src, count)?;
         if res != count {
             return Err(Error::PartialBuffer {
@@ -949,6 +950,7 @@ impl<T: GuestMemory + ?Sized> Bytes<GuestAddress> for T {
     where
         F: Write,
     {
+        #[allow(deprecated)] // this function itself is deprecated
         let res = self.write_to(addr, dst, count)?;
         if res != count {
             return Err(Error::PartialBuffer {

src/io.rs

@@ -0,0 +1,272 @@
+// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+//! Module containing versions of the standard library's [`Read`] and [`Write`] traits compatible
+//! with volatile memory accesses.
+
+use crate::bitmap::BitmapSlice;
+use crate::volatile_memory::copy_slice_impl::{copy_from_volatile_slice, copy_to_volatile_slice};
+use crate::{VolatileMemoryError, VolatileSlice};
+use std::io::ErrorKind;
+use std::os::fd::AsRawFd;
+
+/// A version of the standard library's [`Read`] trait that operates on volatile memory instead of
+/// slices
+///
+/// This trait is needed as rust slices (`&[u8]` and `&mut [u8]`) cannot be used when operating on
+/// guest memory [1].
+///
+/// [1]: https://github.com/rust-vmm/vm-memory/pull/217
+pub trait ReadVolatile {
+    /// Tries to read some bytes into the given [`VolatileSlice`] buffer, returning how many bytes
+    /// were read.
+    ///
+    /// The behavior of implementations should be identical to [`Read::read`]
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError>;
+
+    /// Tries to fill the given [`VolatileSlice`] buffer by reading from `self` returning an error
+    /// if insufficient bytes could be read.
+    ///
+    /// The default implementation is identical to that of [`Read::read_exact`]
+    fn read_exact_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        // Implementation based on https://github.com/rust-lang/rust/blob/7e7483d26e3cec7a44ef00cf7ae6c9c8c918bec6/library/std/src/io/mod.rs#L465
+
+        let mut partial_buf = buf.offset(0)?;
+
+        while !partial_buf.is_empty() {
+            match self.read_volatile(&mut partial_buf) {
+                Err(VolatileMemoryError::IOError(err)) if err.kind() == ErrorKind::Interrupted => {
+                    continue
+                }
+                Ok(0) => {
+                    return Err(VolatileMemoryError::IOError(std::io::Error::new(
+                        ErrorKind::UnexpectedEof,
+                        "failed to fill whole buffer",
+                    )))
+                }
+                Ok(bytes_read) => partial_buf = partial_buf.offset(bytes_read)?,
+                Err(err) => return Err(err),
+            }
+        }
+
+        Ok(())
+    }
+}
+
+/// A version of the standard library's [`Write`] trait that operates on volatile memory instead of
+/// slices
+///
+/// This trait is needed as rust slices (`&[u8]` and `&mut [u8]`) cannot be used when operating on
+/// guest memory [1].
+///
+/// [1]: https://github.com/rust-vmm/vm-memory/pull/217
+pub trait WriteVolatile {
+    /// Tries to write some bytes from the given [`VolatileSlice`] buffer, returning how many bytes
+    /// were written.
+    ///
+    /// The behavior of implementations should be identical to [`Write::write`]
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError>;
+
+    /// Tries write the entire content of the given [`VolatileSlice`] buffer to `self` returning an
+    /// error if not all bytes could be written.
+    ///
+    /// The default implementation is identical to that of [`Write::write_all`]
+    fn write_all_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        // Based on https://github.com/rust-lang/rust/blob/7e7483d26e3cec7a44ef00cf7ae6c9c8c918bec6/library/std/src/io/mod.rs#L1570
+
+        let mut partial_buf = buf.offset(0)?;
+
+        while !partial_buf.is_empty() {
+            match self.write_volatile(&partial_buf) {
+                Err(VolatileMemoryError::IOError(err)) if err.kind() == ErrorKind::Interrupted => {
+                    continue
+                }
+                Ok(0) => {
+                    return Err(VolatileMemoryError::IOError(std::io::Error::new(
+                        ErrorKind::WriteZero,
+                        "failed to write whole buffer",
+                    )))
+                }
+                Ok(bytes_written) => partial_buf = partial_buf.offset(bytes_written)?,
+                Err(err) => return Err(err),
+            }
+        }
+
+        Ok(())
+    }
+}
+
+// We explicitly implement our traits for [`std::fs::File`] and [`std::os::unix::net::UnixStream`]
+// instead of providing blanket implementation for [`AsRawFd`] due to trait coherence limitations: A
+// blanket implementation would prevent us from providing implementations for `&mut [u8]` below, as
+// "an upstream crate could implement AsRawFd for &mut [u8]`.
+
+macro_rules! impl_read_write_volatile_for_raw_fd {
+    ($raw_fd_ty:ty) => {
+        impl ReadVolatile for $raw_fd_ty {
+            fn read_volatile<B: BitmapSlice>(
+                &mut self,
+                buf: &mut VolatileSlice<B>,
+            ) -> Result<usize, VolatileMemoryError> {
+                read_volatile_raw_fd(self, buf)
+            }
+        }
+
+        impl WriteVolatile for $raw_fd_ty {
+            fn write_volatile<B: BitmapSlice>(
+                &mut self,
+                buf: &VolatileSlice<B>,
+            ) -> Result<usize, VolatileMemoryError> {
+                write_volatile_raw_fd(self, buf)
+            }
+        }
+    };
+}
+
+impl_read_write_volatile_for_raw_fd!(std::fs::File);
+impl_read_write_volatile_for_raw_fd!(std::os::unix::net::UnixStream);
+impl_read_write_volatile_for_raw_fd!(std::os::fd::OwnedFd);
+impl_read_write_volatile_for_raw_fd!(std::os::fd::BorrowedFd<'_>);
+
+/// Tries to do a single `read` syscall on the provided file descriptor, storing the data raed in
+/// the given [`VolatileSlice`].
+///
+/// Returns the numbers of bytes read.
+fn read_volatile_raw_fd<Fd: AsRawFd>(
+    raw_fd: &mut Fd,
+    buf: &mut VolatileSlice<impl BitmapSlice>,
+) -> Result<usize, VolatileMemoryError> {
+    let fd = raw_fd.as_raw_fd();
+    let guard = buf.ptr_guard_mut();
+
+    let dst = guard.as_ptr().cast::<libc::c_void>();
+
+    // SAFETY: We got a valid file descriptor from `AsRawFd`. The memory pointed to by `dst` is
+    // valid for writes of length `buf.len() by the invariants upheld by the constructor
+    // of `VolatileSlice`.
+    let bytes_read = unsafe { libc::read(fd, dst, buf.len()) };
+
+    if bytes_read < 0 {
+        // We don't know if a partial read might have happened, so mark everything as dirty
+        buf.bitmap().mark_dirty(0, buf.len());
+
+        Err(VolatileMemoryError::IOError(std::io::Error::last_os_error()))
+    } else {
+        let bytes_read = bytes_read.try_into().unwrap();
+        buf.bitmap().mark_dirty(0, bytes_read);
+        Ok(bytes_read)
+    }
+}
+
+/// Tries to do a single `write` syscall on the provided file descriptor, attempting to write the
+/// data stored in the given [`VolatileSlice`].
+///
+/// Returns the numbers of bytes written.
+fn write_volatile_raw_fd<Fd: AsRawFd>(
+    raw_fd: &mut Fd,
+    buf: &VolatileSlice<impl BitmapSlice>,
+) -> Result<usize, VolatileMemoryError> {
+    let fd = raw_fd.as_raw_fd();
+    let guard = buf.ptr_guard();
+
+    let src = guard.as_ptr().cast::<libc::c_void>();
+
+    // SAFETY: We got a valid file descriptor from `AsRawFd`. The memory pointed to by `src` is
+    // valid for reads of length `buf.len() by the invariants upheld by the constructor
+    // of `VolatileSlice`.
+    let bytes_written = unsafe { libc::write(fd, src, buf.len()) };
+
+    if bytes_written < 0 {
+        Err(VolatileMemoryError::IOError(std::io::Error::last_os_error()))
+    } else {
+        Ok(bytes_written.try_into().unwrap())
+    }
+}
+
+impl WriteVolatile for &mut [u8] {
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        let total = buf.len().min(self.len());
+        let src = buf.subslice(0, total)?;
+
+        // SAFETY:
+        // We check above that `src` is contiguously allocated memory of length `total <= self.len())`.
+        // Furthermore, both src and dst of the call to
+        // copy_from_volatile_slice are valid for reads and writes respectively of length `total`
+        // since total is the minimum of lengths of the memory areas pointed to. The areas do not
+        // overlap, since `dst` is inside guest memory, and buf is a slice (no slices to guest
+        // memory are possible without violating rust's aliasing rules).
+        let written = unsafe { copy_from_volatile_slice(self.as_mut_ptr(), &src, total) };
+
+        // Advance the slice, just like the stdlib: https://doc.rust-lang.org/src/std/io/impls.rs.html#335
+        *self = std::mem::take(self).split_at_mut(written).1;
+
+        Ok(written)
+    }
+
+    fn write_all_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        // Based on https://github.com/rust-lang/rust/blob/f7b831ac8a897273f78b9f47165cf8e54066ce4b/library/std/src/io/impls.rs#L376-L382
+        if self.write_volatile(buf)? == buf.len() {
+            Ok(())
+        } else {
+            Err(VolatileMemoryError::IOError(std::io::Error::new(
+                ErrorKind::WriteZero,
+                "failed to write whole buffer",
+            )))
+        }
+    }
+}
+
+impl ReadVolatile for &[u8] {
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        let total = buf.len().min(self.len());
+        let dst = buf.subslice(0, total)?;
+
+        // SAFETY:
+        // We check above that `dst` is contiguously allocated memory of length `total <= self.len())`.
+        // Furthermore, both src and dst of the call to copy_to_volatile_slice are valid for reads
+        // and writes respectively of length `total` since total is the minimum of lengths of the
+        // memory areas pointed to. The areas do not overlap, since `dst` is inside guest memory,
+        // and buf is a slice (no slices to guest memory are possible without violating rust's aliasing rules).
+        let read = unsafe { copy_to_volatile_slice(&dst, self.as_ptr(), total) };
+
+        // Advance the slice, just like the stdlib: https://doc.rust-lang.org/src/std/io/impls.rs.html#232-310
+        *self = self.split_at(read).1;
+
+        Ok(read)
+    }
+
+    fn read_exact_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<(), VolatileMemoryError> {
+        // Based on https://github.com/rust-lang/rust/blob/f7b831ac8a897273f78b9f47165cf8e54066ce4b/library/std/src/io/impls.rs#L282-L302
+        if buf.len() > self.len() {
+            return Err(VolatileMemoryError::IOError(std::io::Error::new(
+                ErrorKind::UnexpectedEof,
+                "failed to fill whole buffer",
+            )));
+        }
+
+        self.read_volatile(buf).map(|_| ())
+    }
+}

src/lib.rs

@@ -45,6 +45,9 @@ pub use guest_memory::{
     GuestMemoryRegion, GuestUsize, MemoryRegionAddress, Result as GuestMemoryResult,
 };
 
+pub mod io;
+pub use io::{ReadVolatile, WriteVolatile};
+
 #[cfg(all(feature = "backend-mmap", not(feature = "xen"), unix))]
 mod mmap_unix;