add concept of "swiotlb regions" to struct Vm

Add a special `GuestMemoryMmap` to `struct Vm` that has the semantics
that if its not empty, it represents the swiotlb area in the guest.

Have all virtio devices only get access to swiotlb regions if they
exist, only falling back onto having access to all of guest memory
otherwise. This means that in case of swiotlb regions being set up,
virtio devices cannot access generic memory, so if the guest decides to
(incorrectly) place I/O buffers outside of swiotlb we'll get a somewhat
meaningful error of "invalid guest address" (or similar), instead of
random things failing further down the line in case generic memory isn't
accessible by firecracker/the host kernel for some reason (e.g. if its
guest_memfd backed in the future).

When a VM with swiotlb regions is being snapshotted, the regions are
simply concatenated to the end of the snapshot memory file. This changes
the snapshot memory file format in the sense that regions in the file
are not necessarily sorted by guest address anymore.

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

Author: roypat

src/vmm/src/arch/aarch64/vm.rs

@@ -70,6 +70,7 @@ impl ArchVm {
     pub fn save_state(&self, mpidrs: &[u64]) -> Result<VmState, ArchVmError> {
         Ok(VmState {
             memory: self.common.guest_memory.describe(),
+            io_memory: self.common.swiotlb_regions.describe(),
             gic: self
                 .get_irqchip()
                 .save_device(mpidrs)
@@ -96,6 +97,8 @@ impl ArchVm {
 pub struct VmState {
     /// Guest memory state
     pub memory: GuestMemoryState,
+    /// io memory state
+    pub io_memory: GuestMemoryState,
     /// GIC state.
     pub gic: GicState,
 }

src/vmm/src/arch/x86_64/vm.rs

@@ -187,6 +187,7 @@ impl ArchVm {
 
         Ok(VmState {
             memory: self.common.guest_memory.describe(),
+            io_memory: self.common.swiotlb_regions.describe(),
             pitstate,
             clock,
             pic_master,
@@ -211,6 +212,8 @@ impl ArchVm {
 pub struct VmState {
     /// guest memory state
     pub memory: GuestMemoryState,
+    /// io memory state
+    pub io_memory: GuestMemoryState,
     pitstate: kvm_pit_state2,
     clock: kvm_clock_data,
     // TODO: rename this field to adopt inclusive language once Linux updates it, too.

src/vmm/src/builder.rs

@@ -471,7 +471,7 @@ pub fn build_microvm_from_snapshot(
 
     // Restore devices states.
     let mmio_ctor_args = MMIODevManagerConstructorArgs {
-        mem: vmm.vm.guest_memory(),
+        mem: vmm.vm.io_memory(),
         vm: vmm.vm.fd(),
         event_manager,
         resource_allocator: &mut vmm.resource_allocator,
@@ -596,7 +596,7 @@ fn attach_virtio_device<T: 'static + VirtioDevice + MutEventSubscriber + Debug>(
     event_manager.add_subscriber(device.clone());
 
     // The device mutex mustn't be locked here otherwise it will deadlock.
-    let device = MmioTransport::new(vmm.vm.guest_memory().clone(), device, is_vhost_user);
+    let device = MmioTransport::new(vmm.vm.io_memory().clone(), device, is_vhost_user);
     vmm.mmio_device_manager
         .register_mmio_virtio_for_boot(
             vmm.vm.fd(),

src/vmm/src/persist.rs

@@ -180,7 +180,7 @@ pub fn create_snapshot(
         .for_each_virtio_device(|_, _, _, dev| {
             let d = dev.lock().unwrap();
             if d.is_activated() {
-                d.mark_queue_memory_dirty(vmm.vm.guest_memory())
+                d.mark_queue_memory_dirty(vmm.vm.io_memory())
             } else {
                 Ok(())
             }

src/vmm/src/vstate/vm.rs

@@ -34,6 +34,8 @@ pub struct VmCommon {
     max_memslots: usize,
     /// The guest memory of this Vm.
     pub guest_memory: GuestMemoryMmap,
+    /// The swiotlb regions of this Vm.
+    pub swiotlb_regions: GuestMemoryMmap,
 }
 
 /// Errors associated with the wrappers over KVM ioctls.
@@ -101,6 +103,7 @@ impl Vm {
             fd,
             max_memslots: kvm.max_nr_memslots(),
             guest_memory: GuestMemoryMmap::default(),
+            swiotlb_regions: GuestMemoryMmap::default(),
         })
     }
 
@@ -140,6 +143,8 @@ impl Vm {
         let next_slot = self
             .guest_memory()
             .num_regions()
+            .checked_add(self.swiotlb_regions().num_regions())
+            .ok_or(VmError::NotEnoughMemorySlots)?
             .try_into()
             .map_err(|_| VmError::NotEnoughMemorySlots)?;
 
@@ -175,19 +180,59 @@ impl Vm {
         Ok(())
     }
 
+    /// Registers a new io memory region to this [`Vm`].
+    pub fn register_swiotlb_region(&mut self, region: GuestRegionMmap) -> Result<(), VmError> {
+        let arcd_region = Arc::new(self.kvmify_region(region)?);
+        let new_collection = self
+            .common
+            .swiotlb_regions
+            .insert_region(Arc::clone(&arcd_region))?;
+
+        self.register_kvm_region(arcd_region.as_ref())?;
+
+        self.common.swiotlb_regions = new_collection;
+        Ok(())
+    }
+
     /// Gets a reference to the kvm file descriptor owned by this VM.
     pub fn fd(&self) -> &VmFd {
         &self.common.fd
     }
 
-    /// Gets a reference to this [`Vm`]'s [`GuestMemoryMmap`] object
+    /// Gets a reference to this [`Vm`]'s [`GuestMemoryMmap`] object, which
+    /// contains all non-swiotlb guest memory regions.
     pub fn guest_memory(&self) -> &GuestMemoryMmap {
         &self.common.guest_memory
     }
 
+    /// Returns a reference to the [`GuestMemoryMmap`] that I/O devices attached to this [`Vm`]
+    /// have access to. If no I/O regions were registered, return the same as [`Vm::guest_memory`],
+    /// otherwise returns the [`GuestMemoryMmap`] describing a specific swiotlb region.
+    pub fn io_memory(&self) -> &GuestMemoryMmap {
+        if self.common.swiotlb_regions.num_regions() > 0 {
+            &self.common.swiotlb_regions
+        } else {
+            &self.common.guest_memory
+        }
+    }
+
+    /// Gets a reference to the [`GuestMemoryMmap`] holding the swiotlb regions registered to
+    /// this [`Vm`]. Unlike [`Vm::io_memory`], does not fall back to returning the
+    /// [`GuestMemoryMmap`] of normal memory when no swiotlb regions were registered.
+    pub fn swiotlb_regions(&self) -> &GuestMemoryMmap {
+        &self.common.swiotlb_regions
+    }
+
+    /// Returns an iterator over all regions, normal and swiotlb.
+    fn all_regions(&self) -> impl Iterator<Item = &KvmRegion> {
+        self.guest_memory()
+            .iter()
+            .chain(self.common.swiotlb_regions.iter())
+    }
+
     /// Resets the KVM dirty bitmap for each of the guest's memory regions.
     pub fn reset_dirty_bitmap(&self) {
-        self.guest_memory().iter().for_each(|region| {
+        self.all_regions().for_each(|region| {
             let _ = self
                 .fd()
                 .get_dirty_log(region.inner().slot, u64_to_usize(region.len()));
@@ -197,7 +242,7 @@ impl Vm {
     /// Retrieves the KVM dirty bitmap for each of the guest's memory regions.
     pub fn get_dirty_bitmap(&self) -> Result<DirtyBitmap, vmm_sys_util::errno::Error> {
         let mut bitmap: DirtyBitmap = HashMap::new();
-        self.guest_memory().iter().try_for_each(|region| {
+        self.all_regions().try_for_each(|region| {
             self.fd()
                 .get_dirty_log(region.inner().slot, u64_to_usize(region.len()))
                 .map(|bitmap_region| _ = bitmap.insert(region.inner().slot, bitmap_region))
@@ -258,10 +303,14 @@ impl Vm {
         match snapshot_type {
             SnapshotType::Diff => {
                 let dirty_bitmap = self.get_dirty_bitmap()?;
-                self.guest_memory().dump_dirty(&mut file, &dirty_bitmap)?;
+                self.guest_memory()
+                    .dump_dirty(&mut file, &dirty_bitmap)
+                    .and_then(|_| self.swiotlb_regions().dump_dirty(&mut file, &dirty_bitmap))?;
             }
             SnapshotType::Full => {
-                self.guest_memory().dump(&mut file)?;
+                self.guest_memory()
+                    .dump(&mut file)
+                    .and_then(|_| self.swiotlb_regions().dump(&mut file))?;
                 self.reset_dirty_bitmap();
                 self.guest_memory().reset_dirty();
             }