[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

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -14,6 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::mem::MaybeUninit;
+use std::os::unix::ffi::OsStrExt;
+use std::path::Path;
 use std::sync::{Arc, Mutex};
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
@@ -31,12 +34,13 @@ use crate::func::{ParameterTuple, SupportedReturnType};
 use crate::hypervisor::handlers::DbgMemAccessHandlerWrapper;
 use crate::hypervisor::handlers::{MemAccessHandlerCaller, OutBHandlerCaller};
 use crate::hypervisor::{Hypervisor, InterruptHandle};
+use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
 use crate::mem::ptr::RawPtr;
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::metrics::maybe_time_and_emit_guest_call;
 use crate::sandbox_state::sandbox::{DevolvableSandbox, EvolvableSandbox, Sandbox};
 use crate::sandbox_state::transition::{MultiUseContextCallback, Noop};
-use crate::{HyperlightError, Result};
+use crate::{HyperlightError, Result, log_then_return};
 
 /// A sandbox that supports being used Multiple times.
 /// The implication of being used multiple times is two-fold:
@@ -173,6 +177,80 @@ impl MultiUseSandbox {
         })
     }
 
+    /// 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.
+    ///
+    /// 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
+    #[instrument(err(Debug), skip(self, rgn), parent = Span::current())]
+    pub unsafe fn map_region(&mut self, rgn: &MemoryRegion) -> Result<()> {
+        if rgn.flags.contains(MemoryRegionFlags::STACK_GUARD) {
+            // Stack guard pages are an internal implementation detail
+            // (which really should be moved into the guest)
+            log_then_return!("Cannot map host memory as a stack guard page");
+        }
+        if rgn.flags.contains(MemoryRegionFlags::WRITE) {
+            // TODO: Implement support for writable mappings, which
+            // need to be registered with the memory manager so that
+            // writes can be rolled back when necessary.
+            log_then_return!("TODO: Writable mappings not yet supported");
+        }
+        unsafe { self.vm.map_region(rgn) }?;
+        self.mem_mgr.unwrap_mgr_mut().mapped_rgns += 1;
+        Ok(())
+    }
+
+    /// Map the contents of a file into the guest at a particular address
+    ///
+    /// Returns the length of the mapping
+    #[instrument(err(Debug), skip(self, fp, guest_base), parent = Span::current())]
+    pub(crate) fn map_file_cow(&mut self, fp: &Path, guest_base: u64) -> Result<u64> {
+        #[cfg(windows)]
+        log_then_return!("mmap'ing a file into the guest is not yet supported on Windows");
+        unsafe {
+            let fd = libc::open(
+                std::ffi::CString::new(fp.as_os_str().as_bytes())?.as_ptr(),
+                libc::O_RDWR,
+            );
+            if fd < 0 {
+                log_then_return!("open error: {:?}", std::io::Error::last_os_error());
+            }
+            let mut st = MaybeUninit::<libc::stat>::uninit();
+            libc::fstat(fd, st.as_mut_ptr());
+            let st = st.assume_init();
+            let page_size = page_size::get();
+            let size = (st.st_size as usize).div_ceil(page_size) * page_size;
+            let base = libc::mmap(
+                std::ptr::null_mut(),
+                size,
+                libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
+                libc::MAP_PRIVATE,
+                fd,
+                0,
+            );
+            if base == libc::MAP_FAILED {
+                log_then_return!("mmap error: {:?}", std::io::Error::last_os_error());
+            }
+            libc::close(fd);
+
+            self.map_region(&MemoryRegion {
+                host_region: base as usize..base.wrapping_add(size) as usize,
+                guest_region: guest_base as usize..guest_base as usize + size,
+                flags: MemoryRegionFlags::READ | MemoryRegionFlags::EXECUTE,
+                region_type: MemoryRegionType::Heap,
+            })?;
+
+            Ok(size as u64)
+        }
+    }
+
     /// This function is kept here for fuzz testing the parameter and return types
     #[cfg(feature = "fuzzing")]
     #[instrument(err(Debug), skip(self, args), parent = Span::current())]
@@ -193,7 +271,9 @@ impl MultiUseSandbox {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn restore_state(&mut self) -> Result<()> {
         let mem_mgr = self.mem_mgr.unwrap_mgr_mut();
-        mem_mgr.restore_state_from_last_snapshot()
+        let rgns_to_unmap = mem_mgr.restore_state_from_last_snapshot()?;
+        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
+        Ok(())
     }
 
     pub(crate) fn call_guest_function_by_name_no_reset(
@@ -275,9 +355,11 @@ impl DevolvableSandbox<MultiUseSandbox, MultiUseSandbox, Noop<MultiUseSandbox, M
     /// The devolve can be used to return the MultiUseSandbox to the state before the code was loaded. Thus avoiding initialisation overhead
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn devolve(mut self, _tsn: Noop<MultiUseSandbox, MultiUseSandbox>) -> Result<MultiUseSandbox> {
-        self.mem_mgr
+        let rgns_to_unmap = self
+            .mem_mgr
             .unwrap_mgr_mut()
             .pop_and_restore_state_from_snapshot()?;
+        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
         Ok(self)
     }
 }