[hyperlight_host] Allow mapping a host memory region into a guest

Signed-off-by: Lucy Menon <[email protected]>

Author: syntactically

src/hyperlight_host/src/func/call_ctx.rs

@@ -17,6 +17,7 @@ limitations under the License.
 use tracing::{Span, instrument};
 
 use super::{ParameterTuple, SupportedReturnType};
+use crate::mem::memory_region::MemoryRegion;
 use crate::sandbox::Callable;
 use crate::{MultiUseSandbox, Result};
 /// A context for calling guest functions.
@@ -70,6 +71,30 @@ impl MultiUseGuestCallContext {
     pub(crate) fn finish_no_reset(self) -> MultiUseSandbox {
         self.sbox
     }
+
+    /// Map a region of host memory into the sandbox.
+    ///
+    /// Depending on the host platform, there are likely alignment
+    /// requirements of at least one page for base and len.
+    ///
+    /// `rgn.region_type` is ignored, since guest PTEs are not created
+    /// for the new memory.
+    ///
+    /// # Safety
+    /// It is the caller's responsibility to ensure that the host side
+    /// of the region remains intact and is not written to until this
+    /// mapping is removed, either due to the destruction of the
+    /// sandbox or due to a state rollback
+    pub unsafe fn map_region(&mut self, rgn: &MemoryRegion) -> Result<()> {
+        unsafe { self.sbox.map_region(rgn) }
+    }
+
+    /// Map the contents of a file into the guest at a particular address
+    ///
+    /// Returns the length of the mapping
+    pub fn map_file_cow(&mut self, fp: &std::path::Path, guest_base: u64) -> Result<u64> {
+        self.sbox.map_file_cow(fp, guest_base)
+    }
 }
 
 impl Callable for MultiUseGuestCallContext {

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -297,6 +297,7 @@ pub(crate) fn is_hypervisor_present() -> bool {
 /// called the Microsoft Hypervisor (MSHV)
 pub(crate) struct HypervLinuxDriver {
     _mshv: Mshv,
+    page_size: usize,
     vm_fd: VmFd,
     vcpu_fd: VcpuFd,
     entrypoint: u64,
@@ -424,6 +425,7 @@ impl HypervLinuxDriver {
         #[allow(unused_mut)]
         let mut hv = Self {
             _mshv: mshv,
+            page_size: 0,
             vm_fd,
             vcpu_fd,
             mem_regions,
@@ -525,6 +527,8 @@ impl Hypervisor for HypervLinuxDriver {
         max_guest_log_level: Option<LevelFilter>,
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()> {
+        self.page_size = page_size as usize;
+
         let max_guest_log_level: u64 = match max_guest_log_level {
             Some(level) => level as u64,
             None => self.get_max_log_level().into(),
@@ -556,6 +560,37 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(())
     }
 
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
+    unsafe fn map_region(&mut self, rgn: &MemoryRegion) -> Result<()> {
+        if [
+            rgn.guest_region.start,
+            rgn.guest_region.end,
+            rgn.host_region.start,
+            rgn.host_region.end,
+        ]
+        .iter()
+        .any(|x| x % self.page_size != 0)
+        {
+            log_then_return!("region is not page-aligned");
+        }
+        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());
+        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();
+            self.vm_fd.unmap_user_memory(mshv_region)?;
+        }
+        Ok(())
+    }
+
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn dispatch_call_from_host(
         &mut self,

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -606,6 +606,16 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(())
     }
 
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
+    unsafe fn map_region(&mut self, _rgn: &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<()> {
+        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")]
     fn dispatch_call_from_host(
         &mut self,

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -493,6 +493,16 @@ impl Hypervisor for KVMDriver {
         Ok(())
     }
 
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
+    unsafe fn map_region(&mut self, _rgn: &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<()> {
+        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")]
     fn dispatch_call_from_host(
         &mut self,

src/hyperlight_host/src/hypervisor/mod.rs

@@ -132,6 +132,15 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
         #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper,
     ) -> Result<()>;
 
+    /// Map a region of host memory into the sandbox.
+    ///
+    /// Depending on the host platform, there are likely alignment
+    /// 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<()>;
+
     /// 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/mgr.rs

@@ -73,6 +73,8 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) load_addr: RawPtr,
     /// Offset for the execution entrypoint from `load_addr`
     pub(crate) entrypoint_offset: Offset,
+    /// How many memory regions were mapped after sandbox creation
+    pub(crate) mapped_rgns: u64,
     /// A vector of memory snapshots that can be used to save and  restore the state of the memory
     /// This is used by the Rust Sandbox implementation (rather than the mem_snapshot field above which only exists to support current C API)
     snapshots: Arc<Mutex<Vec<SharedMemorySnapshot>>>,
@@ -95,6 +97,7 @@ where
             shared_mem,
             load_addr,
             entrypoint_offset,
+            mapped_rgns: 0,
             snapshots: Arc::new(Mutex::new(Vec::new())),
         }
     }
@@ -265,7 +268,7 @@ where
     /// this function will create a memory snapshot and push it onto the stack of snapshots
     /// It should be used when you want to save the state of the memory, for example, when evolving a sandbox to a new state
     pub(crate) fn push_state(&mut self) -> Result<()> {
-        let snapshot = SharedMemorySnapshot::new(&mut self.shared_mem)?;
+        let snapshot = SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)?;
         self.snapshots
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
@@ -277,7 +280,11 @@ where
     /// off the stack
     /// It should be used when you want to restore the state of the memory to a previous state but still want to
     /// retain that state, for example after calling a function in the guest
-    pub(crate) fn restore_state_from_last_snapshot(&mut self) -> Result<()> {
+    ///
+    /// Returns the number of memory regions mapped into the sandbox
+    /// that need to be unmapped in order for the restore to be
+    /// completed.
+    pub(crate) fn restore_state_from_last_snapshot(&mut self) -> Result<u64> {
         let mut snapshots = self
             .snapshots
             .try_lock()
@@ -288,13 +295,15 @@ where
         }
         #[allow(clippy::unwrap_used)] // We know that last is not None because we checked it above
         let snapshot = last.unwrap();
-        snapshot.restore_from_snapshot(&mut self.shared_mem)
+        let old_rgns = self.mapped_rgns;
+        self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+        Ok(old_rgns - self.mapped_rgns)
     }
 
     /// this function pops the last snapshot off the stack and restores the memory to the previous state
     /// It should be used when you want to restore the state of the memory to a previous state and do not need to retain that state
     /// for example when devolving a sandbox to a previous state.
-    pub(crate) fn pop_and_restore_state_from_snapshot(&mut self) -> Result<()> {
+    pub(crate) fn pop_and_restore_state_from_snapshot(&mut self) -> Result<u64> {
         let last = self
             .snapshots
             .try_lock()
@@ -430,13 +439,15 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
+                mapped_rgns: 0,
                 snapshots: Arc::new(Mutex::new(Vec::new())),
             },
             SandboxMemoryManager {
                 shared_mem: gshm,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
+                mapped_rgns: 0,
                 snapshots: Arc::new(Mutex::new(Vec::new())),
             },
         )

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -24,16 +24,21 @@ use crate::Result;
 #[derive(Clone)]
 pub(super) struct SharedMemorySnapshot {
     snapshot: Vec<u8>,
+    /// How many non-main-RAM regions were mapped when this snapshot was taken?
+    mapped_rgns: u64,
 }
 
 impl SharedMemorySnapshot {
     /// Take a snapshot of the memory in `shared_mem`, then create a new
     /// instance of `Self` with the snapshot stored therein.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(super) fn new<S: SharedMemory>(shared_mem: &mut S) -> Result<Self> {
+    pub(super) fn new<S: SharedMemory>(shared_mem: &mut S, mapped_rgns: u64) -> Result<Self> {
         // TODO: Track dirty pages instead of copying entire memory
         let snapshot = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
-        Ok(Self { snapshot })
+        Ok(Self {
+            snapshot,
+            mapped_rgns,
+        })
     }
 
     /// Take another snapshot of the internally-stored `SharedMemory`,
@@ -51,8 +56,9 @@ impl SharedMemorySnapshot {
     pub(super) fn restore_from_snapshot<S: SharedMemory>(
         &mut self,
         shared_mem: &mut S,
-    ) -> Result<()> {
-        shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))?
+    ) -> Result<u64> {
+        shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))??;
+        Ok(self.mapped_rgns)
     }
 }
 
@@ -69,7 +75,7 @@ mod tests {
         let data2 = data1.iter().map(|b| b + 1).collect::<Vec<u8>>();
         let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
         gm.copy_from_slice(data1.as_slice(), 0).unwrap();
-        let mut snap = super::SharedMemorySnapshot::new(&mut gm).unwrap();
+        let mut snap = super::SharedMemorySnapshot::new(&mut gm, 0).unwrap();
         {
             // after the first snapshot is taken, make sure gm has the equivalent
             // of data1