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

@@ -13,8 +13,6 @@ use std::sync::{Arc, Mutex};
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use libc::EFD_NONBLOCK;
 use linux_loader::cmdline::Cmdline as LoaderKernelCmdline;
-#[cfg(target_arch = "aarch64")]
-use linux_loader::loader::pe::PE as Loader;
 use utils::time::TimestampUs;
 #[cfg(target_arch = "aarch64")]
 use vm_superio::Rtc;
@@ -64,7 +62,6 @@ use crate::vmm_config::machine_config::MachineConfigError;
 use crate::vmm_config::snapshot::{MemBackendConfig, MemBackendType};
 use crate::vstate::kvm::Kvm;
 use crate::vstate::memory;
-use crate::vstate::memory::GuestMemory;
 use crate::vstate::vcpu::{Vcpu, VcpuError};
 use crate::vstate::vm::Vm;
 use crate::{EventManager, Vmm, VmmError, device_manager};
@@ -648,7 +645,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/lib.rs

@@ -700,7 +700,7 @@ impl Vmm {
     pub fn update_balloon_config(&mut self, amount_mib: u32) -> Result<(), BalloonError> {
         // The balloon cannot have a target size greater than the size of
         // the guest memory.
-        if u64::from(amount_mib) > mem_size_mib(&self.vm.guest_memory()) {
+        if u64::from(amount_mib) > mem_size_mib(self.vm.guest_memory()) {
             return Err(BalloonError::TooManyPagesRequested);
         }
 

src/vmm/src/persist.rs

@@ -31,10 +31,10 @@ use crate::snapshot::Snapshot;
 use crate::utils::u64_to_usize;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
-use crate::vmm_config::snapshot::{CreateSnapshotParams, LoadSnapshotParams};
 use crate::vmm_config::machine_config::{
     HugePageConfig, MachineConfigError, MachineConfigUpdate, MemoryConfig,
 };
+use crate::vmm_config::snapshot::{CreateSnapshotParams, LoadSnapshotParams};
 use crate::vstate::kvm::KvmState;
 use crate::vstate::memory::MemoryError;
 use crate::vstate::vcpu::{VcpuSendEventError, VcpuState};
@@ -138,7 +138,7 @@ pub enum CreateSnapshotError {
     /// Cannot get dirty bitmap: {0}
     DirtyBitmap(kvm_ioctls::Error),
     /// Cannot write memory file: {0}
-    Memory(MemoryError),
+    Memory(#[from] MemoryError),
     /// Cannot perform {0} on the memory backing file: {1}
     MemoryBackingFile(&'static str, io::Error),
     /// Cannot save the microVM state: {0}

src/vmm/src/vstate/vm.rs

@@ -17,10 +17,9 @@ use vmm_sys_util::eventfd::EventFd;
 
 pub use crate::arch::{ArchVm as Vm, ArchVmError, VmState};
 use crate::logger::info;
-use crate::persist::CreateSnapshotError;
+use crate::persist::{CreateSnapshotError, GuestRegionUffdMapping};
 use crate::utils::u64_to_usize;
 use crate::vmm_config::snapshot::SnapshotType;
-use crate::persist::GuestRegionUffdMapping;
 use crate::vstate::memory::{
     GuestMemory, GuestMemoryExtension, GuestMemoryMmap, GuestMemoryRegion, GuestRegionMmap,
     KvmRegion,
@@ -36,6 +35,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.
@@ -103,6 +104,7 @@ impl Vm {
             fd,
             max_memslots: kvm.max_nr_memslots(),
             guest_memory: GuestMemoryMmap::default(),
+            swiotlb_regions: GuestMemoryMmap::default(),
         })
     }
 
@@ -142,6 +144,8 @@ impl Vm {
         let next_slot = self
             .guest_memory()
             .num_regions()
+            .checked_add(self.io_memory().num_regions())
+            .ok_or(VmError::NotEnoughMemorySlots)?
             .try_into()
             .map_err(|_| VmError::NotEnoughMemorySlots)?;
 
@@ -177,16 +181,56 @@ 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())
+    }
+
     pub(crate) fn create_uffd(
         &self,
     ) -> Result<(Uffd, Vec<GuestRegionUffdMapping>), userfaultfd::Error> {
@@ -203,24 +247,27 @@ impl Vm {
 
         let mut offset = 0;
 
-        for mem_region in self.common.guest_memory.iter() {
-            uffd.register(mem_region.as_ptr().cast(), mem_region.size() as _)?;
+        for mem_region in self.all_regions() {
+            uffd.register(
+                mem_region.inner().userspace_addr as *mut _,
+                u64_to_usize(mem_region.len()),
+            )?;
             mappings.push(GuestRegionUffdMapping {
-                base_host_virt_addr: mem_region.as_ptr() as u64,
-                size: mem_region.size(),
+                base_host_virt_addr: mem_region.inner().userspace_addr,
+                size: u64_to_usize(mem_region.len()),
                 offset,
                 ..Default::default()
             });
 
-            offset += mem_region.size() as u64;
+            offset += mem_region.len();
         }
 
         Ok((uffd, mappings))
     }
 
     /// 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()));
@@ -230,7 +277,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))
@@ -293,18 +340,22 @@ impl Vm {
                 let dirty_bitmap = self.get_dirty_bitmap().map_err(DirtyBitmap)?;
                 self.guest_memory()
                     .dump_dirty(&mut file, &dirty_bitmap)
-                    .map_err(Memory)
+                    .and_then(|_| self.swiotlb_regions().dump_dirty(&mut file, &dirty_bitmap))
             }
             SnapshotType::Full => {
-                let dump_res = self.guest_memory().dump(&mut file).map_err(Memory);
+                let dump_res = self
+                    .guest_memory()
+                    .dump(&mut file)
+                    .and_then(|_| self.swiotlb_regions().dump(&mut file));
                 if dump_res.is_ok() {
                     self.reset_dirty_bitmap();
                     self.guest_memory().reset_dirty();
                 }
 
                 dump_res
             }
-        }?;
+        }
+        .map_err(Memory)?;
 
         file.flush()
             .map_err(|err| MemoryBackingFile("flush", err))?;