IoMemory: Add IOVA-space bitmap

Without an IOMMU, we have direct access to guest physical addresses
(GPAs).  In order to track our writes to guest memory (during
migration), we log them into dirty bitmaps, and a page's bit index is
its GPA divided by the page size.

With an IOMMU, however, we no longer know the GPA, instead we operate on
I/O virtual addresses (IOVAs) and VMM user-space addresses (VUAs).
Here, the dirty bitmap bit index is the IOVA divided by the page size.

`IoMemory` types contain an internal "physical" memory type that
operates on these VUAs (`IoMemory::PhysicalMemory).  Any bitmap
functionality that this internal type may already have (e.g.
`GuestMemoryMmap` does) cannot be used for dirty bitmap tracking with an
IOMMU because they would use the VUA, but we need to use the IOVA, and
this information is not available on that lower layer.

Therefore, `IoMemory` itself needs to support bitmaps separately from
its inner `PhysicalMemory`, which will be used when the IOMMU is in use.
Add an associated `IoMemory::Bitmap` type and add a bitmap object to
`IommuMemory`.  Ensure that writes to memory dirty that bitmap
appropriately:
- In `try_access()`, if write access was requested, dirty the handled
  region of the bitmap after the access is done.
- In `get_slice()`, replace the `VolatileSlice`'s bitmap (which comes
  from the inner `PhysicalMemory`) by the correct slice of our IOVA
  bitmap before returning it.

Signed-off-by: Hanna Czenczek <[email protected]>

Author: XanClic

src/guest_memory.rs

@@ -52,7 +52,7 @@ use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use crate::address::{Address, AddressValue};
-use crate::bitmap::{BitmapSlice, MS};
+use crate::bitmap::{Bitmap, BitmapSlice, BS, MS};
 use crate::bytes::{AtomicAccess, Bytes};
 use crate::io::{ReadVolatile, WriteVolatile};
 #[cfg(feature = "iommu")]
@@ -771,6 +771,11 @@ impl Permissions {
     pub fn allow(&self, access: Self) -> bool {
         *self & access == access
     }
+
+    /// Check whether the permissions `self` include write access.
+    pub fn has_write(&self) -> bool {
+        *self & Permissions::Write == Permissions::Write
+    }
 }
 
 impl std::ops::BitOr for Permissions {
@@ -806,6 +811,8 @@ impl std::ops::BitAnd for Permissions {
 pub trait IoMemory {
     /// Underlying `GuestMemory` type.
     type PhysicalMemory: GuestMemory + ?Sized;
+    /// Dirty bitmap type for tracking writes to the IOVA address space.
+    type Bitmap: Bitmap;
 
     /// Return `true` if `addr..(addr + count)` is accessible with `access`.
     fn check_range(&self, addr: GuestAddress, count: usize, access: Permissions) -> bool;
@@ -828,7 +835,7 @@ pub trait IoMemory {
         addr: GuestAddress,
         count: usize,
         access: Permissions,
-    ) -> Result<impl IoMemorySliceIterator<'a, MS<'a, Self::PhysicalMemory>>>;
+    ) -> Result<impl IoMemorySliceIterator<'a, BS<'a, Self::Bitmap>>>;
 
     /// If this virtual memory is just a plain `GuestMemory` object underneath without an IOMMU
     /// translation layer in between, return that `GuestMemory` object.
@@ -867,6 +874,7 @@ pub trait IoMemorySliceIterator<'a, B: BitmapSlice>:
 /// the same [`GuestMemory`] methods (if available), discarding the `access` parameter.
 impl<M: GuestMemory + ?Sized> IoMemory for M {
     type PhysicalMemory = M;
+    type Bitmap = <M::R as GuestMemoryRegion>::B;
 
     fn check_range(&self, addr: GuestAddress, count: usize, _access: Permissions) -> bool {
         <M as GuestMemory>::check_range(self, addr, count)
@@ -877,7 +885,7 @@ impl<M: GuestMemory + ?Sized> IoMemory for M {
         addr: GuestAddress,
         count: usize,
         _access: Permissions,
-    ) -> Result<impl IoMemorySliceIterator<'a, MS<'a, Self::PhysicalMemory>>> {
+    ) -> Result<impl IoMemorySliceIterator<'a, BS<'a, Self::Bitmap>>> {
         Ok(<M as GuestMemory>::get_slices(self, addr, count))
     }
 

src/iommu.rs

@@ -12,14 +12,18 @@
 //! IOTLB misses require sending a notification to the front-end and awaiting a reply that supplies
 //! the desired mapping.
 
+use crate::bitmap::{self, Bitmap};
 use crate::guest_memory::{
     Error as GuestMemoryError, GuestMemorySliceIterator, IoMemorySliceIterator,
     Result as GuestMemoryResult,
 };
-use crate::{bitmap, GuestAddress, GuestMemory, IoMemory, Permissions, VolatileSlice};
+use crate::{
+    Address, GuestAddress, GuestMemory, GuestMemoryRegion, GuestUsize, IoMemory, Permissions,
+    VolatileSlice,
+};
 use rangemap::RangeMap;
 use std::cmp;
-use std::fmt::Debug;
+use std::fmt::{self, Debug};
 use std::num::Wrapping;
 use std::ops::{Deref, Range};
 use std::sync::Arc;
@@ -182,7 +186,15 @@ pub struct IotlbFails {
 /// The underlying [`GuestMemory`] is basically the physical memory, and the [`Iommu`] translates
 /// the I/O virtual address space that `IommuMemory` provides into that underlying physical address
 /// space.
-#[derive(Debug, Default)]
+///
+/// Note on memory write tracking (“logging”):
+/// - When the IOMMU is disabled ([`IommuMemory::set_iommu_enabled()`]), writes to memory are
+///   tracked by the underlying [`GuestMemory`] in its bitmap(s).
+/// - When it is enabled, they are instead tracked in the [`IommuMemory`]’s dirty bitmap; the
+///   offset in the bitmap is calculated from the write’s IOVA.
+///
+/// Therefore, this type should only be used when this is the desired behavior (IOVA-based memory
+/// write logging when IOMMU is used).
 pub struct IommuMemory<M: GuestMemory, I: Iommu> {
     /// Physical memory
     inner: M,
@@ -191,6 +203,8 @@ pub struct IommuMemory<M: GuestMemory, I: Iommu> {
     /// Whether the IOMMU is even to be used or not; disabling it makes this a pass-through to
     /// `inner`.
     use_iommu: bool,
+    /// Dirty bitmap to use for IOVA accesses
+    bitmap: Arc<<M::R as GuestMemoryRegion>::B>,
 }
 
 impl IommuMapping {
@@ -354,27 +368,36 @@ impl TryFrom<Range<u64>> for IovaRange {
 
 impl<M: GuestMemory, I: Iommu> IommuMemory<M, I> {
     /// Create a new `IommuMemory` instance.
-    pub fn new(inner: M, iommu: I, use_iommu: bool) -> Self {
+    pub fn new(inner: M, iommu: I, use_iommu: bool, bitmap: <Self as IoMemory>::Bitmap) -> Self {
         IommuMemory {
             inner,
             iommu: Arc::new(iommu),
             use_iommu,
+            bitmap: Arc::new(bitmap),
         }
     }
 
     /// Create a new version of `self` with the underlying physical memory replaced.
     ///
-    /// Note that the inner `Arc` reference to the IOMMU is cloned, i.e. both the existing and the
-    /// new `IommuMemory` object will share an IOMMU instance.  (The `use_iommu` flag however is
-    /// copied, so is independent between the two instances.)
+    /// Note that the inner `Arc` references to the IOMMU and bitmap are cloned, i.e. both the
+    /// existing and the new `IommuMemory` object will share the IOMMU and bitmap instances.  (The
+    /// `use_iommu` flag however is copied, so is independent between the two instances.)
     pub fn inner_replaced(&self, inner: M) -> Self {
         IommuMemory {
             inner,
             iommu: Arc::clone(&self.iommu),
             use_iommu: self.use_iommu,
+            bitmap: Arc::clone(&self.bitmap),
         }
     }
 
+    /// Return a reference to the IOVA address space's dirty bitmap.
+    ///
+    /// This bitmap tracks write accesses done while the IOMMU is enabled.
+    pub fn bitmap(&self) -> &Arc<<Self as IoMemory>::Bitmap> {
+        &self.bitmap
+    }
+
     /// Enable or disable the IOMMU.
     ///
     /// Disabling the IOMMU switches to pass-through mode, where every access is done directly on
@@ -400,12 +423,42 @@ impl<M: GuestMemory + Clone, I: Iommu> Clone for IommuMemory<M, I> {
             inner: self.inner.clone(),
             iommu: Arc::clone(&self.iommu),
             use_iommu: self.use_iommu,
+            bitmap: Arc::clone(&self.bitmap),
+        }
+    }
+}
+
+impl<M: GuestMemory + Debug, I: Iommu> Debug for IommuMemory<M, I>
+where
+    <M::R as GuestMemoryRegion>::B: Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("IommuMemory")
+            .field("inner", &self.inner)
+            .field("iommu", &self.iommu)
+            .field("use_iommu", &self.use_iommu)
+            .field("bitmap", &self.bitmap)
+            .finish()
+    }
+}
+
+impl<M: GuestMemory + Default, I: Iommu + Default> Default for IommuMemory<M, I>
+where
+    <M::R as GuestMemoryRegion>::B: Default,
+{
+    fn default() -> Self {
+        IommuMemory {
+            inner: Default::default(),
+            iommu: Default::default(),
+            use_iommu: Default::default(),
+            bitmap: Default::default(),
         }
     }
 }
 
 impl<M: GuestMemory, I: Iommu> IoMemory for IommuMemory<M, I> {
     type PhysicalMemory = M;
+    type Bitmap = <M::R as GuestMemoryRegion>::B;
 
     fn check_range(&self, addr: GuestAddress, count: usize, access: Permissions) -> bool {
         if !self.use_iommu {
@@ -424,7 +477,7 @@ impl<M: GuestMemory, I: Iommu> IoMemory for IommuMemory<M, I> {
         addr: GuestAddress,
         count: usize,
         access: Permissions,
-    ) -> GuestMemoryResult<impl IoMemorySliceIterator<'a, bitmap::MS<'a, M>>> {
+    ) -> GuestMemoryResult<impl IoMemorySliceIterator<'a, bitmap::BS<'a, Self::Bitmap>>> {
         if self.use_iommu {
             IommuMemorySliceIterator::virt(self, addr, count, access)
                 .map_err(GuestMemoryError::IommuError)
@@ -445,8 +498,11 @@ impl<M: GuestMemory, I: Iommu> IoMemory for IommuMemory<M, I> {
 /// Iterates over [`VolatileSlice`]s that together form an area in an `IommuMemory`.
 ///
 /// Returned by [`IommuMemory::get_slices()`]
-#[derive(Debug)]
 pub struct IommuMemorySliceIterator<'a, M: GuestMemory, I: Iommu + 'a> {
+    /// Current IOVA (needed to access the right slice of the IOVA space dirty bitmap)
+    iova: GuestAddress,
+    /// IOVA space dirty bitmap
+    bitmap: Option<&'a <M::R as GuestMemoryRegion>::B>,
     /// Underlying physical memory (i.e. not the `IommuMemory`)
     phys_mem: &'a M,
     /// IOMMU translation result (i.e. remaining physical regions to visit)
@@ -463,6 +519,8 @@ impl<'a, M: GuestMemory, I: Iommu> IommuMemorySliceIterator<'a, M, I> {
     /// the underlying physical memory for the given address range.
     fn phys(mem: &'a IommuMemory<M, I>, addr: GuestAddress, count: usize) -> Self {
         IommuMemorySliceIterator {
+            iova: addr,
+            bitmap: None,
             phys_mem: &mem.inner,
             translation: None,
             current_translated_iter: Some(mem.inner.get_slices(addr, count)),
@@ -481,6 +539,8 @@ impl<'a, M: GuestMemory, I: Iommu> IommuMemorySliceIterator<'a, M, I> {
     ) -> Result<Self, Error> {
         let translation = mem.iommu.translate(addr, count, access)?;
         Ok(IommuMemorySliceIterator {
+            iova: addr,
+            bitmap: Some(mem.bitmap.as_ref()),
             phys_mem: &mem.inner,
             translation: Some(translation),
             current_translated_iter: None,
@@ -512,7 +572,22 @@ impl<'a, M: GuestMemory, I: Iommu> IommuMemorySliceIterator<'a, M, I> {
                 .as_mut()
                 .and_then(|iter| iter.next())
             {
-                return Some(item);
+                let mut item = match item {
+                    Ok(item) => item,
+                    Err(err) => return Some(Err(err)),
+                };
+
+                if let Some(bitmap) = self.bitmap.as_ref() {
+                    let bitmap_slice = bitmap.slice_at(self.iova.0 as usize);
+                    item = item.replace_bitmap(bitmap_slice);
+                }
+
+                self.iova = match self.iova.overflowing_add(item.len() as GuestUsize) {
+                    (x @ GuestAddress(0), _) | (x, false) => x,
+                    (_, true) => return Some(Err(GuestMemoryError::GuestAddressOverflow)),
+                };
+
+                return Some(Ok(item));
             }
 
             let next_mapping = self.translation.as_mut()?.next()?;
@@ -546,3 +621,19 @@ impl<'a, M: GuestMemory, I: Iommu> IoMemorySliceIterator<'a, bitmap::MS<'a, M>>
     for IommuMemorySliceIterator<'a, M, I>
 {
 }
+
+impl<'a, M: GuestMemory + Debug, I: Iommu> Debug for IommuMemorySliceIterator<'a, M, I>
+where
+    I::IotlbGuard<'a>: Debug,
+    <M::R as GuestMemoryRegion>::B: Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("IommuMemorySliceIterator")
+            .field("iova", &self.iova)
+            .field("bitmap", &self.bitmap)
+            .field("phys_mem", &self.phys_mem)
+            .field("translation", &self.translation)
+            .field("current_translated_iter", &self.current_translated_iter)
+            .finish()
+    }
+}

src/volatile_memory.rs

@@ -435,6 +435,17 @@ impl<'a, B: BitmapSlice> VolatileSlice<'a, B> {
         }
     }
 
+    /// Replaces the bitmap in `self` by `new_bitmap`.
+    #[cfg(feature = "iommu")]
+    pub(crate) fn replace_bitmap<NB: BitmapSlice>(self, new_bitmap: NB) -> VolatileSlice<'a, NB> {
+        VolatileSlice {
+            addr: self.addr,
+            size: self.size,
+            bitmap: new_bitmap,
+            mmap: self.mmap,
+        }
+    }
+
     /// Returns a guard for the pointer to the underlying memory.
     pub fn ptr_guard(&self) -> PtrGuard {
         PtrGuard::read(self.mmap, self.addr, self.len())