Add VolatileRead/VolatileWrite traits

These are essentially clones of the Read/Write traits from the standard
library, but instead of operating on &[u8] and &mut [u8], they operate
on VolatileSlices. We cannot use the stdlib traits to operate on guest
memory due to unsoundness concerns, see rust-vmm/vm-memory#217 and 228.

A default implementation is provided for `File` and `UnixStream` (as
these are the types for which firecracker needs them). I have done
experiements with instead providing a blanket implementation for `T:
AsRawFd`, however ran into some problems with trait coherence rules, as
such a blanket implementation makes the implementation of
Read/WriteVolatile for &[u8]/&mut [u8] impossible.
This allows us to also potentially choose different implementations for
different kind of `AsRawFd`s (e.g. using `recv` for sockets). This can
all be revised later though.

These traits would fit nicely into rust-vmm, and I'll put out feelers
about interest regarding these traits after our release tasks are done.

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

src/utils/src/vm_memory.rs

@@ -5,17 +5,17 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
-use std::io::Error as IoError;
+use std::io::{Error as IoError, ErrorKind};
 use std::os::unix::io::AsRawFd;
 
 use vm_memory::bitmap::AtomicBitmap;
-pub use vm_memory::bitmap::Bitmap;
+pub use vm_memory::bitmap::{Bitmap, BitmapSlice};
 use vm_memory::mmap::{check_file_offset, NewBitmap};
 pub use vm_memory::mmap::{MmapRegionBuilder, MmapRegionError};
 pub use vm_memory::{
     address, Address, ByteValued, Bytes, Error, FileOffset, GuestAddress, GuestMemory,
     GuestMemoryError, GuestMemoryRegion, GuestUsize, MemoryRegionAddress, MmapRegion,
-    VolatileMemory, VolatileMemoryError,
+    VolatileMemory, VolatileMemoryError, VolatileSlice,
 };
 
 pub type GuestMemoryMmap = vm_memory::GuestMemoryMmap<Option<AtomicBitmap>>;
@@ -142,6 +142,212 @@ pub fn mark_dirty_mem(mem: &GuestMemoryMmap, addr: GuestAddress, len: usize) {
     });
 }
 
+pub trait ReadVolatile {
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError>;
+
+    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
+
+        // Doing the trick the stdlib does where in each loop, the given slice is adjusted does not
+        // work for us as VolatileSlice::offset returns an owned copy, but `buf` is a
+        // reference. This means the first iteration of the loop would have to work with a
+        // reference, and all following iterations with an owned copy. This is not (easily)
+        // expressible in rust's type system, which is why we use this approach. Note that the total
+        // number of `.offset` calls are the same, meaning we do not incur a performance impact (as
+        // VolatileSlice::offset is a O(1) operation).
+        let mut bytes_read_so_far = 0;
+
+        while bytes_read_so_far < buf.len() {
+            match self.read_volatile(&mut buf.offset(bytes_read_so_far)?) {
+                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) => bytes_read_so_far += bytes_read,
+                Err(err) => return Err(err),
+            }
+        }
+
+        Ok(())
+    }
+}
+
+pub trait WriteVolatile {
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError>;
+
+    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
+
+        // See comment in ReadVolatile::read_exact_volatile for reasons for this variable.
+        let mut bytes_written_so_far = 0;
+
+        while bytes_written_so_far < buf.len() {
+            match self.write_volatile(&buf.offset(bytes_written_so_far)?) {
+                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) => bytes_written_so_far += bytes_written,
+                Err(err) => return Err(err),
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl ReadVolatile for std::fs::File {
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        read_volatile_raw_fd(self, buf)
+    }
+}
+
+impl ReadVolatile for std::os::unix::net::UnixStream {
+    fn read_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &mut VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        read_volatile_raw_fd(self, buf)
+    }
+}
+
+fn read_volatile_raw_fd(
+    raw_fd: &mut impl AsRawFd,
+    buf: &mut VolatileSlice<impl BitmapSlice>,
+) -> Result<usize, VolatileMemoryError> {
+    let fd = raw_fd.as_raw_fd();
+    let dst = buf.as_ptr().cast::<libc::c_void>();
+
+    buf.bitmap().mark_dirty(0, buf.len());
+
+    // 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 {
+        Err(VolatileMemoryError::IOError(std::io::Error::last_os_error()))
+    } else {
+        Ok(bytes_read.try_into().unwrap())
+    }
+}
+
+impl WriteVolatile for std::fs::File {
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        write_volatile_raw_fd(self, buf)
+    }
+}
+
+impl WriteVolatile for std::os::unix::net::UnixStream {
+    fn write_volatile<B: BitmapSlice>(
+        &mut self,
+        buf: &VolatileSlice<B>,
+    ) -> Result<usize, VolatileMemoryError> {
+        write_volatile_raw_fd(self, buf)
+    }
+}
+
+fn write_volatile_raw_fd(
+    raw_fd: &mut impl AsRawFd,
+    buf: &VolatileSlice<impl BitmapSlice>,
+) -> Result<usize, VolatileMemoryError> {
+    let fd = raw_fd.as_raw_fd();
+    let src = buf.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> {
+        // NOTE: The duality of read <-> write here is correct. This is because we translate a call
+        // "slice.write(buf)" (e.g. write into slice from buf) into "buf.read(slice)" (e.g. read
+        // from buffer into slice). Both express data transfer from the buffer to the slice
+        buf.read(self, 0)
+    }
+
+    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> {
+        // NOTE: the duality of read <-> write here is correct. This is because we translate a call
+        // "slice.read(buf)" (e.g. "read from slice into buf") into "buf.write(slice)" (e.g. write
+        // into buf from slice)
+        buf.write(self, 0)
+    }
+
+    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(|_| ())
+    }
+}
+
 pub mod test_utils {
     use super::*;
 
@@ -193,6 +399,8 @@ pub mod test_utils {
 #[cfg(test)]
 mod tests {
     #![allow(clippy::undocumented_unsafe_blocks)]
+    use std::io::{Read, Seek, Write};
+
     use super::*;
     use crate::get_page_size;
     use crate::tempfile::TempFile;
@@ -444,4 +652,147 @@ mod tests {
                 .unwrap();
         }
     }
+
+    #[test]
+    fn test_read_volatile() {
+        let test_cases = [
+            ([1u8, 2].as_ref(), [1u8, 2, 0, 0, 0]),
+            ([1, 2, 3, 4].as_ref(), [1, 2, 3, 4, 0]),
+            // ensure we don't have a buffer overrun
+            ([5, 6, 7, 8, 9].as_ref(), [5, 6, 7, 8, 0]),
+        ];
+
+        for (mut input, output) in test_cases {
+            // ---- Test ReadVolatile for &[u8] ----
+            //
+            // Test read_volatile for &[u8] works
+            let mut memory = vec![0u8; 5];
+
+            assert_eq!(
+                input
+                    .read_volatile(&mut VolatileSlice::from(&mut memory[..4]))
+                    .unwrap(),
+                input.len().min(4)
+            );
+            assert_eq!(&memory, &output);
+
+            // Test read_exact_volatile for &[u8] works
+            let mut memory = vec![0u8; 5];
+            let result = input.read_exact_volatile(&mut VolatileSlice::from(&mut memory[..4]));
+
+            // read_exact fails if there are not enough bytes in input to completely fill
+            // memory[..4]
+            if input.len() < 4 {
+                result.unwrap_err();
+                assert_eq!(memory, vec![0u8; 5]);
+            } else {
+                result.unwrap();
+                assert_eq!(&memory, &output);
+            }
+
+            // ---- Test ReadVolatile for File ----
+
+            let mut temp_file = TempFile::new().unwrap().into_file();
+            temp_file.write_all(input).unwrap();
+            temp_file.rewind().unwrap();
+
+            // Test read_volatile for File works
+            let mut memory = vec![0u8; 5];
+
+            assert_eq!(
+                temp_file
+                    .read_volatile(&mut VolatileSlice::from(&mut memory[..4]))
+                    .unwrap(),
+                input.len().min(4)
+            );
+            assert_eq!(&memory, &output);
+
+            temp_file.rewind().unwrap();
+
+            // Test read_exact_volatile for File works
+            let mut memory = vec![0u8; 5];
+
+            let read_exact_result =
+                temp_file.read_exact_volatile(&mut VolatileSlice::from(&mut memory[..4]));
+
+            if input.len() < 4 {
+                read_exact_result.unwrap_err();
+            } else {
+                read_exact_result.unwrap();
+            }
+            assert_eq!(&memory, &output);
+        }
+    }
+
+    #[test]
+    fn test_write_volatile() {
+        let test_cases = [
+            (vec![1u8, 2], [1u8, 2, 0, 0, 0]),
+            (vec![1, 2, 3, 4], [1, 2, 3, 4, 0]),
+            // ensure we don't have a buffer overrun
+            (vec![5, 6, 7, 8, 9], [5, 6, 7, 8, 0]),
+        ];
+
+        for (mut input, output) in test_cases {
+            // ---- Test WriteVolatile for &mut [u8] ----
+            //
+            // Test write_volatile for &mut [u8] works
+            let mut memory = vec![0u8; 5];
+
+            assert_eq!(
+                (&mut memory[..4])
+                    .write_volatile(&VolatileSlice::from(input.as_mut_slice()))
+                    .unwrap(),
+                input.len().min(4)
+            );
+            assert_eq!(&memory, &output);
+
+            // Test write_all_volatile for &mut [u8] works
+            let mut memory = vec![0u8; 5];
+
+            let result =
+                (&mut memory[..4]).write_all_volatile(&VolatileSlice::from(input.as_mut_slice()));
+
+            if input.len() > 4 {
+                result.unwrap_err();
+                // This quirky behavior of writing to the slice even in the case of failure is also
+                // exhibited by the stdlib
+                assert_eq!(&memory, &output);
+            } else {
+                result.unwrap();
+                assert_eq!(&memory, &output);
+            }
+
+            // ---- Test ẂriteVolatile for File works
+            // Test write_volatile for File works
+            let mut temp_file = TempFile::new().unwrap().into_file();
+
+            temp_file
+                .write_volatile(&VolatileSlice::from(input.as_mut_slice()))
+                .unwrap();
+            temp_file.rewind().unwrap();
+
+            let mut written = vec![0u8; input.len()];
+            temp_file.read_exact(written.as_mut_slice()).unwrap();
+
+            assert_eq!(input, written);
+            // check no excess bytes were written to the file
+            assert_eq!(temp_file.read(&mut [0u8]).unwrap(), 0);
+
+            // Test write_all_volatile for File works
+            let mut temp_file = TempFile::new().unwrap().into_file();
+
+            temp_file
+                .write_all_volatile(&VolatileSlice::from(input.as_mut_slice()))
+                .unwrap();
+            temp_file.rewind().unwrap();
+
+            let mut written = vec![0u8; input.len()];
+            temp_file.read_exact(written.as_mut_slice()).unwrap();
+
+            assert_eq!(input, written);
+            // check no excess bytes were written to the file
+            assert_eq!(temp_file.read(&mut [0u8]).unwrap(), 0);
+        }
+    }
 }