Add scratch region

This adds a second region of guest physical memory, which is intended
to become the only mutable region of memory, but which is currently
unused.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_common/src/arch/amd64/layout.rs

@@ -14,5 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
  */
 
-// Keep in mind that the minimum upper half GVA is 0xffff_8000_0000_0000
-pub const SNAPSHOT_PT_GVA: usize = 0xffff_ff00_0000_0000;
+/// We have this the top of the page below the top of memory in order
+/// to make working with start/end ptrs in a few places more
+/// convenient (not needing to worry about overflow)
+pub const MAX_GVA: usize = 0xffff_ffff_ffff_efff;
+pub const SNAPSHOT_PT_GVA: usize = 0xffff_8000_0000_0000;
+
+/// We assume 36-bit IPAs for now, since every amd64 processor
+/// supports at least 36 bits.  Almost all of them support at least 40
+/// bits, so we could consider bumping this in the future if we were
+/// ever memory-constrained.
+pub const MAX_GPA: usize = 0x0000_000f_ffff_ffff;

src/hyperlight_common/src/arch/amd64/vmem.rs

@@ -145,8 +145,15 @@ impl<const HIGH_BIT: u8, const LOW_BIT: u8, Op: TableOps> Iterator
         let next_vmin = if self.n == 0 {
             self.request.vmin
         } else {
-            // Align to the next boundary by adding one entry's worth and masking off lower bits
-            (self.request.vmin + (self.n << LOW_BIT)) & !lower_bits_mask
+            // Align to the next boundary by adding one entry's worth
+            // and masking off lower bits. Masking off before adding
+            // is safe, since n << LOW_BIT must always have zeros in
+            // these positions.
+            let aligned_min = self.request.vmin & !lower_bits_mask;
+            // Use checked_add here because going past the end of the
+            // address space counts as "the next one would be out of
+            // range"
+            aligned_min.checked_add(self.n << LOW_BIT)?
         };
 
         // Check if we've processed the entire requested range

src/hyperlight_common/src/layout.rs

@@ -21,4 +21,21 @@ mod arch;
 
 // The constraint on the feature is temporary and will be removed when other arch i686 is added
 #[cfg(feature = "init-paging")]
-pub use arch::SNAPSHOT_PT_GVA;
+pub use arch::MAX_GPA;
+#[cfg(feature = "init-paging")]
+pub use arch::{MAX_GVA, SNAPSHOT_PT_GVA};
+
+// offsets down from the top of scratch memory for various things
+pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
+pub const SCRATCH_TOP_USED_OFFSET: u64 = 0x10;
+pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x18;
+pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x20;
+
+#[cfg(feature = "init-paging")]
+pub fn scratch_base_gpa(size: usize) -> u64 {
+    (MAX_GPA - size + 1) as u64
+}
+#[cfg(feature = "init-paging")]
+pub fn scratch_base_gva(size: usize) -> u64 {
+    (MAX_GVA - size + 1) as u64
+}

src/hyperlight_host/src/error.rs

@@ -351,6 +351,7 @@ impl HyperlightError {
             | HyperlightError::HyperlightVmError(HyperlightVmError::Initialize(_))
             | HyperlightError::HyperlightVmError(HyperlightVmError::MapRegion(_))
             | HyperlightError::HyperlightVmError(HyperlightVmError::UnmapRegion(_))
+            | HyperlightError::HyperlightVmError(HyperlightVmError::UpdateScratch(_))
             | HyperlightError::IOError(_)
             | HyperlightError::IntConversionFailure(_)
             | HyperlightError::InvalidFlatBuffer(_)

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -65,7 +65,7 @@ use crate::hypervisor::{InterruptHandle, InterruptHandleImpl, get_max_log_level}
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::ptr::{GuestPtr, RawPtr};
-use crate::mem::shared_mem::HostSharedMemory;
+use crate::mem::shared_mem::{GuestSharedMemory, HostSharedMemory, SharedMemory};
 use crate::metrics::{METRIC_ERRONEOUS_VCPU_KICKS, METRIC_GUEST_CANCELLATION};
 use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::host_funcs::FunctionRegistry;
@@ -95,6 +95,10 @@ pub(crate) struct HyperlightVm {
     mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region)
     next_slot: u32,                     // Monotonically increasing slot number
     freed_slots: Vec<u32>,              // Reusable slots from unmapped regions
+    scratch_slot: u32,                  // The slot number used for the scratch region
+    // The current scratch region, used to keep it alive as long as it
+    // is used & when unmapping
+    scratch_memory: Option<GuestSharedMemory>,
 
     #[cfg(gdb)]
     gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
@@ -247,6 +251,15 @@ pub enum UnmapRegionError {
     UnmapMemory(#[from] UnmapMemoryError),
 }
 
+/// Errors that can occur when updating the scratch mapping
+#[derive(Debug, thiserror::Error)]
+pub enum UpdateScratchError {
+    #[error("VM map memory error: {0}")]
+    MapMemory(#[from] MapMemoryError),
+    #[error("VM unmap memory error: {0}")]
+    UnmapMemory(#[from] UnmapMemoryError),
+}
+
 /// Errors that can occur during HyperlightVm creation
 #[derive(Debug, thiserror::Error)]
 pub enum CreateHyperlightVmError {
@@ -262,6 +275,8 @@ pub enum CreateHyperlightVmError {
     SendDbgMsg(#[from] SendDbgMsgError),
     #[error("VM operation error: {0}")]
     Vm(#[from] VmError),
+    #[error("Set scratch error: {0}")]
+    UpdateScratch(#[from] UpdateScratchError),
 }
 
 /// Errors that can occur during debug exit handling
@@ -311,6 +326,8 @@ pub enum HyperlightVmError {
     MapRegion(#[from] MapRegionError),
     #[error("Unmap region error: {0}")]
     UnmapRegion(#[from] UnmapRegionError),
+    #[error("Update scratch error: {0}")]
+    UpdateScratch(#[from] UpdateScratchError),
 }
 
 impl HyperlightVm {
@@ -319,6 +336,7 @@ impl HyperlightVm {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         mem_regions: Vec<MemoryRegion>,
+        scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
         entrypoint: u64,
         rsp: u64,
@@ -395,6 +413,7 @@ impl HyperlightVm {
             }),
         });
 
+        let scratch_slot = mem_regions.len() as u32;
         #[cfg_attr(not(gdb), allow(unused_mut))]
         let mut ret = Self {
             vm,
@@ -403,10 +422,12 @@ impl HyperlightVm {
             interrupt_handle,
             page_size: 0, // Will be set in `initialise`
 
-            next_slot: mem_regions.len() as u32,
+            next_slot: scratch_slot + 1,
             sandbox_regions: mem_regions,
             mmap_regions: Vec::new(),
             freed_slots: Vec::new(),
+            scratch_slot,
+            scratch_memory: None,
 
             #[cfg(gdb)]
             gdb_conn,
@@ -418,6 +439,8 @@ impl HyperlightVm {
             rt_cfg,
         };
 
+        ret.update_scratch_mapping(scratch_mem)?;
+
         // Send the interrupt handle to the GDB thread if debugging is enabled
         // This is used to allow the GDB thread to stop the vCPU
         #[cfg(gdb)]
@@ -542,6 +565,24 @@ impl HyperlightVm {
         self.mmap_regions.iter().map(|(_, region)| region)
     }
 
+    /// Update the scratch mapping to point to a new GuestSharedMemory
+    pub(crate) fn update_scratch_mapping(
+        &mut self,
+        scratch: GuestSharedMemory,
+    ) -> Result<(), UpdateScratchError> {
+        let guest_base = hyperlight_common::layout::scratch_base_gpa(scratch.mem_size());
+        let rgn = scratch.mapping_at(guest_base, MemoryRegionType::Scratch);
+
+        if let Some(old_scratch) = self.scratch_memory.replace(scratch) {
+            let old_base = hyperlight_common::layout::scratch_base_gpa(old_scratch.mem_size());
+            let old_rgn = old_scratch.mapping_at(old_base, MemoryRegionType::Scratch);
+            self.vm.unmap_memory((self.scratch_slot, &old_rgn))?;
+        }
+        unsafe { self.vm.map_memory((self.scratch_slot, &rgn))? };
+
+        Ok(())
+    }
+
     /// 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/hypervisor/mod.rs

@@ -525,9 +525,9 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mut mem_mgr, mut gshm) = sandbox.mgr.build();
+        let (mut mem_mgr, gshm) = sandbox.mgr.build();
         let mut vm = set_up_hypervisor_partition(
-            &mut gshm,
+            gshm,
             &config,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,

src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs

@@ -83,6 +83,23 @@ impl KvmVm {
             .create_vcpu(0)
             .map_err(|e| CreateVmError::CreateVcpuFd(e.into()))?;
 
+        // Set the CPUID leaf for MaxPhysAddr. KVM allows this to
+        // easily be overridden by the hypervisor and defaults it very
+        // low, while mshv passes it through from hardware unless an
+        // intercept is installed.
+        let mut kvm_cpuid = hv
+            .get_supported_cpuid(kvm_bindings::KVM_MAX_CPUID_ENTRIES)
+            .map_err(|e| CreateVmError::InitializeVm(e.into()))?;
+        for entry in kvm_cpuid.as_mut_slice().iter_mut() {
+            if entry.function == 0x8000_0008 {
+                entry.eax &= !0xff;
+                entry.eax |= hyperlight_common::layout::MAX_GPA.ilog2() + 1;
+            }
+        }
+        vcpu_fd
+            .set_cpuid2(&kvm_cpuid)
+            .map_err(|e| CreateVmError::InitializeVm(e.into()))?;
+
         Ok(Self {
             vm_fd,
             vcpu_fd,

src/hyperlight_host/src/mem/layout.rs

@@ -126,6 +126,10 @@ pub(crate) struct SandboxMemoryLayout {
     // The offset in the sandbox memory where the code starts
     guest_code_offset: usize,
     pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
+
+    // The size of the scratch region in physical memory; note that
+    // this will appear under the top of physical memory.
+    scratch_size: usize,
 }
 
 impl Debug for SandboxMemoryLayout {
@@ -202,6 +206,10 @@ impl Debug for SandboxMemoryLayout {
                 "Guest Code Offset",
                 &format_args!("{:#x}", self.guest_code_offset),
             )
+            .field(
+                "Scratch region size",
+                &format_args!("{:#x}", self.scratch_size),
+            )
             .finish()
     }
 }
@@ -229,6 +237,7 @@ impl SandboxMemoryLayout {
         code_size: usize,
         stack_size: usize,
         heap_size: usize,
+        scratch_size: usize,
         init_data_size: usize,
         init_data_permissions: Option<MemoryRegionFlags>,
     ) -> Result<Self> {
@@ -295,6 +304,7 @@ impl SandboxMemoryLayout {
             init_data_permissions,
             pt_offset,
             pt_size: None,
+            scratch_size,
         })
     }
 
@@ -324,6 +334,11 @@ impl SandboxMemoryLayout {
         self.stack_size
     }
 
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
+    pub(super) fn get_scratch_size(&self) -> usize {
+        self.scratch_size
+    }
+
     /// Get the offset in guest memory to the output data pointer.
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     fn get_output_data_pointer_offset(&self) -> usize {
@@ -818,7 +833,7 @@ mod tests {
     fn test_get_memory_size() {
         let sbox_cfg = SandboxConfiguration::default();
         let sbox_mem_layout =
-            SandboxMemoryLayout::new(sbox_cfg, 4096, 2048, 4096, 0, None).unwrap();
+            SandboxMemoryLayout::new(sbox_cfg, 4096, 2048, 4096, 0x3000, 0, None).unwrap();
         assert_eq!(
             sbox_mem_layout.get_memory_size().unwrap(),
             get_expected_memory_size(&sbox_mem_layout)

src/hyperlight_host/src/mem/memory_region.rs

@@ -138,6 +138,8 @@ pub enum MemoryRegionType {
     GuardPage,
     /// The region contains the Stack
     Stack,
+    /// The scratch region
+    Scratch,
 }
 
 /// A trait that distinguishes between different kinds of memory region representations.

src/hyperlight_host/src/mem/mgr.rs

@@ -45,6 +45,8 @@ pub(crate) const STACK_COOKIE_LEN: usize = 16;
 pub(crate) struct SandboxMemoryManager<S> {
     /// Shared memory for the Sandbox
     pub(crate) shared_mem: S,
+    /// Scratch memory for the Sandbox
+    pub(crate) scratch_mem: S,
     /// The memory layout of the underlying shared memory
     pub(crate) layout: SandboxMemoryLayout,
     /// Pointer to where to load memory from
@@ -145,13 +147,15 @@ where
     pub(crate) fn new(
         layout: SandboxMemoryLayout,
         shared_mem: S,
+        scratch_mem: S,
         load_addr: RawPtr,
         entrypoint_offset: Option<Offset>,
         stack_cookie: [u8; STACK_COOKIE_LEN],
     ) -> Self {
         Self {
             layout,
             shared_mem,
+            scratch_mem,
             load_addr,
             entrypoint_offset,
             mapped_rgns: 0,
@@ -191,26 +195,22 @@ where
             mapped_regions,
         )
     }
-
-    /// This function restores a memory snapshot from a given snapshot.
-    pub(crate) fn restore_snapshot(&mut self, snapshot: &Snapshot) -> Result<()> {
-        self.shared_mem.restore_from_snapshot(snapshot)?;
-        Ok(())
-    }
 }
 
 impl SandboxMemoryManager<ExclusiveSharedMemory> {
     pub(crate) fn from_snapshot(s: &Snapshot) -> Result<Self> {
         let layout = *s.layout();
         let mut shared_mem = ExclusiveSharedMemory::new(s.mem_size())?;
         shared_mem.copy_from_slice(s.memory(), 0)?;
+        let scratch_mem = ExclusiveSharedMemory::new(s.layout().get_scratch_size())?;
         let load_addr: RawPtr = RawPtr::try_from(layout.get_guest_code_address())?;
         let stack_cookie = rand::random::<[u8; STACK_COOKIE_LEN]>();
         let entrypoint_gva = s.preinitialise();
         let entrypoint_offset = entrypoint_gva.map(|x| (x - u64::from(&load_addr)).into());
         Ok(Self::new(
             layout,
             shared_mem,
+            scratch_mem,
             load_addr,
             entrypoint_offset,
             stack_cookie,
@@ -236,9 +236,11 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         SandboxMemoryManager<GuestSharedMemory>,
     ) {
         let (hshm, gshm) = self.shared_mem.build();
+        let (hscratch, gscratch) = self.scratch_mem.build();
         (
             SandboxMemoryManager {
                 shared_mem: hshm,
+                scratch_mem: hscratch,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
@@ -248,6 +250,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             },
             SandboxMemoryManager {
                 shared_mem: gshm,
+                scratch_mem: gscratch,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
@@ -382,6 +385,37 @@ impl SandboxMemoryManager<HostSharedMemory> {
             };
         }
     }
+
+    /// This function restores a memory snapshot from a given snapshot.
+    pub(crate) fn restore_snapshot(
+        &mut self,
+        snapshot: &Snapshot,
+    ) -> Result<Option<GuestSharedMemory>> {
+        if self.shared_mem.mem_size() != snapshot.mem_size() {
+            return Err(new_error!(
+                "Snapshot size does not match current memory size: {} != {}",
+                self.shared_mem.raw_mem_size(),
+                snapshot.mem_size()
+            ));
+        }
+        self.shared_mem.restore_from_snapshot(snapshot)?;
+        let new_scratch_size = snapshot.layout().get_scratch_size();
+        if new_scratch_size == self.scratch_mem.mem_size() {
+            self.scratch_mem.zero()?;
+            Ok(None)
+        } else {
+            let new_scratch_mem = ExclusiveSharedMemory::new(new_scratch_size)?;
+            let (hscratch, gscratch) = new_scratch_mem.build();
+            // Even though this destroys the reference to the host
+            // side of the old scratch mapping, the VM should still
+            // own the reference to the guest side of the old scratch
+            // mapping, so it won't actually be deallocated until it
+            // has been unmapped from the VM.
+            self.scratch_mem = hscratch;
+
+            Ok(Some(gscratch))
+        }
+    }
 }
 
 #[cfg(test)]