implement memory region tracking in drivers

 - Separate Vecs for initial sandbox regions and mmap regions
 - Replace unmap_regions(n) with unmap_region(region) for precise control
 - Add Hash derive for MemoryRegion and related types
 - Enable use of MemoryRegino in HashSet for efficient set operations

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -76,6 +76,7 @@ use super::{
 use super::{HyperlightExit, Hypervisor, InterruptHandle, LinuxInterruptHandle, VirtualCPU};
 #[cfg(gdb)]
 use crate::HyperlightError;
+use crate::hypervisor::get_memory_access_violation;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
 use crate::sandbox::SandboxConfiguration;
@@ -310,7 +311,8 @@ pub(crate) struct HypervLinuxDriver {
     vm_fd: VmFd,
     vcpu_fd: VcpuFd,
     entrypoint: u64,
-    mem_regions: Vec<MemoryRegion>,
+    sandbox_regions: Vec<MemoryRegion>,
+    mmap_regions: Vec<MemoryRegion>,
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
@@ -443,7 +445,8 @@ impl HypervLinuxDriver {
             page_size: 0,
             vm_fd,
             vcpu_fd,
-            mem_regions,
+            sandbox_regions: mem_regions,
+            mmap_regions: Vec::new(),
             entrypoint: entrypoint_ptr.absolute()?,
             orig_rsp: rsp_ptr,
             interrupt_handle: interrupt_handle.clone(),
@@ -534,8 +537,11 @@ impl Debug for HypervLinuxDriver {
         f.field("Entrypoint", &self.entrypoint)
             .field("Original RSP", &self.orig_rsp);
 
-        for region in &self.mem_regions {
-            f.field("Memory Region", &region);
+        for region in &self.sandbox_regions {
+            f.field("Sandbox Memory Region", &region);
+        }
+        for region in &self.mmap_regions {
+            f.field("Mapped Memory Region", &region);
         }
 
         let regs = self.vcpu_fd.get_regs();
@@ -627,20 +633,24 @@ impl Hypervisor for HypervLinuxDriver {
         }
         let mshv_region: mshv_user_mem_region = rgn.to_owned().into();
         self.vm_fd.map_user_memory(mshv_region)?;
-        self.mem_regions.push(rgn.to_owned());
+        self.mmap_regions.push(rgn.to_owned());
         Ok(())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
-        for rgn in self
-            .mem_regions
-            .split_off(self.mem_regions.len() - n as usize)
-        {
-            let mshv_region: mshv_user_mem_region = rgn.to_owned().into();
+    unsafe fn unmap_region(&mut self, region: &MemoryRegion) -> Result<()> {
+        if let Some(pos) = self.mmap_regions.iter().position(|r| r == region) {
+            let removed_region = self.mmap_regions.remove(pos);
+            let mshv_region: mshv_user_mem_region = removed_region.into();
             self.vm_fd.unmap_user_memory(mshv_region)?;
+            Ok(())
+        } else {
+            Err(new_error!("Tried to unmap region that is not mapped"))
         }
-        Ok(())
+    }
+
+    fn get_mapped_regions(&self) -> Box<dyn ExactSizeIterator<Item = &MemoryRegion> + '_> {
+        Box::new(self.mmap_regions.iter())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -842,9 +852,9 @@ impl Hypervisor for HypervLinuxDriver {
                         gpa,
                         &self
                     );
-                    match self.get_memory_access_violation(
+                    match get_memory_access_violation(
                         gpa as usize,
-                        &self.mem_regions,
+                        self.sandbox_regions.iter().chain(self.mmap_regions.iter()),
                         access_info,
                     ) {
                         Some(access_info_violation) => access_info_violation,
@@ -974,7 +984,7 @@ impl Hypervisor for HypervLinuxDriver {
             });
 
             Ok(Some(crashdump::CrashDumpContext::new(
-                &self.mem_regions,
+                &self.sandbox_regions,
                 regs,
                 xsave.buffer.to_vec(),
                 self.entrypoint,
@@ -1147,7 +1157,7 @@ impl Drop for HypervLinuxDriver {
     #[instrument(skip_all, parent = Span::current(), level = "Trace")]
     fn drop(&mut self) {
         self.interrupt_handle.dropped.store(true, Ordering::Relaxed);
-        for region in &self.mem_regions {
+        for region in self.sandbox_regions.iter().chain(self.mmap_regions.iter()) {
             let mshv_region: mshv_user_mem_region = region.to_owned().into();
             match self.vm_fd.unmap_user_memory(mshv_region) {
                 Ok(_) => (),

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -59,6 +59,7 @@ use super::{
 };
 use super::{HyperlightExit, Hypervisor, InterruptHandle, VirtualCPU};
 use crate::hypervisor::fpu::FP_CONTROL_WORD_DEFAULT;
+use crate::hypervisor::get_memory_access_violation;
 use crate::hypervisor::wrappers::WHvGeneralRegisters;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
@@ -279,7 +280,8 @@ pub(crate) struct HypervWindowsDriver {
     _surrogate_process: SurrogateProcess, // we need to keep a reference to the SurrogateProcess for the duration of the driver since otherwise it will dropped and the memory mapping will be unmapped and the surrogate process will be returned to the pool
     entrypoint: u64,
     orig_rsp: GuestPtr,
-    mem_regions: Vec<MemoryRegion>,
+    sandbox_regions: Vec<MemoryRegion>,
+    mmap_regions: Vec<MemoryRegion>,
     interrupt_handle: Arc<WindowsInterruptHandle>,
     #[cfg(gdb)]
     debug: Option<HypervDebug>,
@@ -355,7 +357,8 @@ impl HypervWindowsDriver {
             _surrogate_process: surrogate_process,
             entrypoint,
             orig_rsp: GuestPtr::try_from(RawPtr::from(rsp))?,
-            mem_regions,
+            sandbox_regions: mem_regions,
+            mmap_regions: Vec::new(),
             interrupt_handle: interrupt_handle.clone(),
             #[cfg(gdb)]
             debug,
@@ -452,8 +455,11 @@ impl Debug for HypervWindowsDriver {
         fs.field("Entrypoint", &self.entrypoint)
             .field("Original RSP", &self.orig_rsp);
 
-        for region in &self.mem_regions {
-            fs.field("Memory Region", &region);
+        for region in &self.sandbox_regions {
+            fs.field("Sandbox Memory Region", &region);
+        }
+        for region in &self.mmap_regions {
+            fs.field("Mapped Memory Region", &region);
         }
 
         // Get the registers
@@ -627,18 +633,17 @@ impl Hypervisor for HypervWindowsDriver {
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    unsafe fn map_region(&mut self, _rgn: &MemoryRegion) -> Result<()> {
+    unsafe fn map_region(&mut self, _region: &MemoryRegion) -> Result<()> {
         log_then_return!("Mapping host memory into the guest not yet supported on this platform");
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
-        if n > 0 {
-            log_then_return!(
-                "Mapping host memory into the guest not yet supported on this platform"
-            );
-        }
-        Ok(())
+    unsafe fn unmap_region(&mut self, _region: &MemoryRegion) -> Result<()> {
+        log_then_return!("Mapping host memory into the guest not yet supported on this platform");
+    }
+
+    fn get_mapped_regions(&self) -> Box<dyn ExactSizeIterator<Item = &MemoryRegion> + '_> {
+        Box::new(self.mmap_regions.iter())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -799,8 +804,11 @@ impl Hypervisor for HypervWindowsDriver {
                     gpa, access_info, &self
                 );
 
-                match self.get_memory_access_violation(gpa as usize, &self.mem_regions, access_info)
-                {
+                match get_memory_access_violation(
+                    gpa as usize,
+                    self.sandbox_regions.iter().chain(self.mmap_regions.iter()),
+                    access_info,
+                ) {
                     Some(access_info) => access_info,
                     None => HyperlightExit::Mmio(gpa),
                 }
@@ -909,7 +917,7 @@ impl Hypervisor for HypervWindowsDriver {
             });
 
             Ok(Some(crashdump::CrashDumpContext::new(
-                &self.mem_regions,
+                &self.sandbox_regions,
                 regs,
                 xsave,
                 self.entrypoint,

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -45,6 +45,7 @@ use super::{
 use super::{HyperlightExit, Hypervisor, InterruptHandle, LinuxInterruptHandle, VirtualCPU};
 #[cfg(gdb)]
 use crate::HyperlightError;
+use crate::hypervisor::get_memory_access_violation;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
 use crate::sandbox::SandboxConfiguration;
@@ -293,7 +294,9 @@ pub(crate) struct KVMDriver {
     vcpu_fd: VcpuFd,
     entrypoint: u64,
     orig_rsp: GuestPtr,
-    mem_regions: Vec<MemoryRegion>,
+    sandbox_regions: Vec<MemoryRegion>,
+    mmap_regions: Vec<(MemoryRegion, u32)>, // (region, slot_number)
+    next_slot: u32,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
     #[cfg(gdb)]
@@ -382,7 +385,9 @@ impl KVMDriver {
             vcpu_fd,
             entrypoint,
             orig_rsp: rsp_gp,
-            mem_regions,
+            next_slot: mem_regions.len() as u32,
+            sandbox_regions: mem_regions,
+            mmap_regions: Vec::new(),
             interrupt_handle: interrupt_handle.clone(),
             #[cfg(gdb)]
             debug,
@@ -430,8 +435,11 @@ impl Debug for KVMDriver {
         let mut f = f.debug_struct("KVM Driver");
         // Output each memory region
 
-        for region in &self.mem_regions {
-            f.field("Memory Region", &region);
+        for region in &self.sandbox_regions {
+            f.field("Sandbox Memory Region", &region);
+        }
+        for region in &self.mmap_regions {
+            f.field("Mapped Memory Region", &region);
         }
         let regs = self.vcpu_fd.get_regs();
         // check that regs is OK and then set field in debug struct
@@ -515,25 +523,32 @@ impl Hypervisor for KVMDriver {
         }
 
         let mut kvm_region: kvm_userspace_memory_region = region.clone().into();
-        kvm_region.slot = self.mem_regions.len() as u32;
+        kvm_region.slot = self.next_slot;
         unsafe { self.vm_fd.set_user_memory_region(kvm_region) }?;
-        self.mem_regions.push(region.to_owned());
+        self.mmap_regions.push((region.to_owned(), self.next_slot));
+        self.next_slot += 1;
         Ok(())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
-        let n_keep = self.mem_regions.len() - n as usize;
-        for (k, region) in self.mem_regions.split_off(n_keep).iter().enumerate() {
-            let mut kvm_region: kvm_userspace_memory_region = region.clone().into();
-            kvm_region.slot = (n_keep + k) as u32;
+    unsafe fn unmap_region(&mut self, region: &MemoryRegion) -> Result<()> {
+        if let Some(idx) = self.mmap_regions.iter().position(|(r, _)| r == region) {
+            let (region, slot) = self.mmap_regions.remove(idx);
+            let mut kvm_region: kvm_userspace_memory_region = region.into();
+            kvm_region.slot = slot;
             // Setting memory_size to 0 unmaps the slot's region
             // From https://docs.kernel.org/virt/kvm/api.html
             // > Deleting a slot is done by passing zero for memory_size.
             kvm_region.memory_size = 0;
             unsafe { self.vm_fd.set_user_memory_region(kvm_region) }?;
+            Ok(())
+        } else {
+            Err(new_error!("Tried to unmap region that is not mapped"))
         }
-        Ok(())
+    }
+
+    fn get_mapped_regions(&self) -> Box<dyn ExactSizeIterator<Item = &MemoryRegion> + '_> {
+        Box::new(self.mmap_regions.iter().map(|(region, _)| region))
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
@@ -694,9 +709,11 @@ impl Hypervisor for KVMDriver {
             Ok(VcpuExit::MmioRead(addr, _)) => {
                 crate::debug!("KVM MMIO Read -Details: Address: {} \n {:#?}", addr, &self);
 
-                match self.get_memory_access_violation(
+                match get_memory_access_violation(
                     addr as usize,
-                    &self.mem_regions,
+                    self.sandbox_regions
+                        .iter()
+                        .chain(self.mmap_regions.iter().map(|(r, _)| r)),
                     MemoryRegionFlags::READ,
                 ) {
                     Some(access_violation_exit) => access_violation_exit,
@@ -706,9 +723,11 @@ impl Hypervisor for KVMDriver {
             Ok(VcpuExit::MmioWrite(addr, _)) => {
                 crate::debug!("KVM MMIO Write -Details: Address: {} \n {:#?}", addr, &self);
 
-                match self.get_memory_access_violation(
+                match get_memory_access_violation(
                     addr as usize,
-                    &self.mem_regions,
+                    self.sandbox_regions
+                        .iter()
+                        .chain(self.mmap_regions.iter().map(|(r, _)| r)),
                     MemoryRegionFlags::WRITE,
                 ) {
                     Some(access_violation_exit) => access_violation_exit,
@@ -824,7 +843,7 @@ impl Hypervisor for KVMDriver {
             // The [`CrashDumpContext`] accepts xsave as a vector of u8, so we need to convert the
             // xsave region to a vector of u8
             Ok(Some(crashdump::CrashDumpContext::new(
-                &self.mem_regions,
+                &self.sandbox_regions,
                 regs,
                 xsave
                     .region

src/hyperlight_host/src/hypervisor/mod.rs

@@ -155,8 +155,13 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
     /// requirements of at least one page for base and len.
     unsafe fn map_region(&mut self, rgn: &MemoryRegion) -> Result<()>;
 
-    /// Unmap the most recent `n` regions mapped by `map_region`
-    unsafe fn unmap_regions(&mut self, n: u64) -> Result<()>;
+    /// Unmap a memory region from the sandbox
+    unsafe fn unmap_region(&mut self, rgn: &MemoryRegion) -> Result<()>;
+
+    /// Get the currently mapped dynamic memory regions (not including sandbox regions)
+    ///
+    /// Note: Box needed for trait to be object-safe :(
+    fn get_mapped_regions(&self) -> Box<dyn ExactSizeIterator<Item = &MemoryRegion> + '_>;
 
     /// Dispatch a call from the host to the guest using the given pointer
     /// to the dispatch function _in the guest's address space_.

src/hyperlight_host/src/mem/memory_region.rs

@@ -52,7 +52,7 @@ pub(crate) const DEFAULT_GUEST_BLOB_MEM_FLAGS: MemoryRegionFlags = MemoryRegionF
 
 bitflags! {
     /// flags representing memory permission for a memory region
-    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
+    #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
     pub struct MemoryRegionFlags: u32 {
         /// no permissions
         const NONE = 0;
@@ -154,7 +154,7 @@ impl TryFrom<hv_x64_memory_intercept_message> for MemoryRegionFlags {
 }
 
 // only used for debugging
-#[derive(Debug, PartialEq, Eq, Copy, Clone)]
+#[derive(Debug, PartialEq, Eq, Copy, Clone, Hash)]
 /// The type of memory region
 pub enum MemoryRegionType {
     /// The region contains the guest's page tables
@@ -181,7 +181,7 @@ pub enum MemoryRegionType {
 
 /// represents a single memory region inside the guest. All memory within a region has
 /// the same memory permissions
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct MemoryRegion {
     /// the range of guest memory addresses
     pub guest_region: Range<usize>,