[host/mem] added 'user memory' region to store additional guest blob

+ added some extra configuration to set the maximum memory size (needed for extra blob)
+ if guest blob is provided, we now write it to shared mem when creating a sandbox

Signed-off-by: danbugs <[email protected]>

Author: danbugs

src/hyperlight_host/src/mem/layout.rs

@@ -31,6 +31,8 @@ use crate::sandbox::SandboxConfiguration;
 use crate::{Result, new_error};
 
 // +-------------------------------------------+
+// |              User Memory                  | (get_total_memory - size_a)
+// +-------------------------------------------+ (size_a)
 // |             Guest (User) Stack            |
 // +-------------------------------------------+
 // |             Guard Page (4KiB)             |
@@ -103,6 +105,7 @@ pub(crate) struct SandboxMemoryLayout {
     guest_heap_buffer_offset: usize,
     guard_page_offset: usize,
     guest_user_stack_buffer_offset: usize, // the lowest address of the user stack
+    user_memory_offset: usize,
 
     // other
     pub(crate) peb_address: usize,
@@ -181,6 +184,10 @@ impl Debug for SandboxMemoryLayout {
                 "Guest User Stack Buffer Offset",
                 &format_args!("{:#x}", self.guest_user_stack_buffer_offset),
             )
+            .field(
+                "User Memory Offset",
+                &format_args!("{:#x}", self.user_memory_offset),
+            )
             .field(
                 "Page Table Size",
                 &format_args!("{:#x}", self.total_page_table_size),
@@ -275,6 +282,7 @@ impl SandboxMemoryLayout {
         let guest_user_stack_buffer_offset = guard_page_offset + PAGE_SIZE_USIZE;
         // round up stack size to page size. This is needed for MemoryRegion
         let stack_size_rounded = round_up_to(stack_size, PAGE_SIZE_USIZE);
+        let user_memory_offset = guest_user_stack_buffer_offset + stack_size_rounded;
 
         Ok(Self {
             peb_offset,
@@ -299,6 +307,7 @@ impl SandboxMemoryLayout {
             guard_page_offset,
             total_page_table_size,
             guest_code_offset,
+            user_memory_offset,
         })
     }
 
@@ -444,7 +453,11 @@ impl SandboxMemoryLayout {
     /// layout.
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     fn get_unaligned_memory_size(&self) -> usize {
-        self.get_top_of_user_stack_offset() + self.stack_size
+        if self.sandbox_memory_config.get_guest_memory_size() == 0 {
+            self.user_memory_offset
+        } else {
+            self.sandbox_memory_config.get_guest_memory_size()
+        }
     }
 
     /// get the code offset
@@ -682,12 +695,34 @@ impl SandboxMemoryLayout {
         }
 
         // stack
-        let final_offset = builder.push_page_aligned(
+        let user_memory_offset = builder.push_page_aligned(
             self.get_guest_stack_size(),
             MemoryRegionFlags::READ | MemoryRegionFlags::WRITE,
             Stack,
         );
 
+        let expected_user_memory_offset = TryInto::<usize>::try_into(self.user_memory_offset)?;
+
+        if user_memory_offset != expected_user_memory_offset {
+            return Err(new_error!(
+                "User Memory offset does not match expected User Memory offset expected:  {}, actual:  {}",
+                expected_user_memory_offset,
+                user_memory_offset
+            ));
+        }
+
+        let user_region_size = self.get_memory_size()? - user_memory_offset;
+
+        let final_offset = if user_region_size > 0 {
+            builder.push_page_aligned(
+                user_region_size,
+                MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
+                Code,
+            )
+        } else {
+            user_memory_offset
+        };
+
         let expected_final_offset = TryInto::<usize>::try_into(self.get_memory_size()?)?;
 
         if final_offset != expected_final_offset {
@@ -701,6 +736,21 @@ impl SandboxMemoryLayout {
         Ok(builder.build())
     }
 
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub(crate) fn write_user_memory(
+        &self,
+        shared_mem: &mut ExclusiveSharedMemory,
+        bytes: &[u8],
+    ) -> Result<(u64, usize)> {
+        if self.user_memory_offset + bytes.len() > self.get_memory_size()? {
+            return Err(new_error!(
+                "Not enough memory in shared memory to write user memory, increase guest memory size with set_guest_memory_size"
+            ));
+        }
+        shared_mem.copy_from_slice(bytes, self.user_memory_offset)?;
+        Ok((self.user_memory_offset as u64, bytes.len()))
+    }
+
     /// Write the finished memory layout to `shared_mem` and return
     /// `Ok` if successful.
     ///

src/hyperlight_host/src/sandbox/config.rs

@@ -55,6 +55,8 @@ pub struct SandboxConfiguration {
     /// The size of the memory buffer that is made available for input to the
     /// Guest Binary
     output_data_size: usize,
+    /// Total guest memory size
+    guest_memory_size: usize,
     /// The stack size to use in the guest sandbox. If set to 0, the stack
     /// size will be determined from the PE file header.
     ///
@@ -133,6 +135,7 @@ impl SandboxConfiguration {
             interrupt_vcpu_sigrtmin_offset,
             #[cfg(gdb)]
             guest_debug_info,
+            guest_memory_size: 0,
             #[cfg(crashdump)]
             guest_core_dump,
         }
@@ -262,6 +265,16 @@ impl SandboxConfiguration {
         (self.heap_size_override > 0).then_some(self.heap_size_override)
     }
 
+    /// Set the total memory size.
+    pub fn set_guest_memory_size(&mut self, guest_memory_size: usize) {
+        self.guest_memory_size = guest_memory_size;
+    }
+
+    /// Get the total memory size.
+    pub(crate) fn get_guest_memory_size(&self) -> usize {
+        self.guest_memory_size
+    }
+
     /// If self.stack_size is non-zero, return it. Otherwise,
     /// return exe_info.stack_reserve()
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]

src/hyperlight_host/src/sandbox/mem_mgr.rs

@@ -98,6 +98,15 @@ impl MemMgrWrapper<ExclusiveSharedMemory> {
         let mem_size = shared_mem.mem_size();
         layout.write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
     }
+
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub(super) fn write_user_memory(&mut self, user_memory: &[u8]) -> Result<()> {
+        let mgr = self.unwrap_mgr_mut();
+        let layout = mgr.layout;
+        let shared_mem = mgr.get_shared_mem_mut();
+        layout.write_user_memory(shared_mem, user_memory)?;
+        Ok(())
+    }
 }
 
 impl MemMgrWrapper<HostSharedMemory> {

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -240,13 +240,19 @@ impl UninitializedSandbox {
 
         let mut mem_mgr_wrapper = {
             let mut mgr = UninitializedSandbox::load_guest_binary(sandbox_cfg, &guest_binary)?;
+
             let stack_guard = Self::create_stack_guard();
             mgr.set_stack_guard(&stack_guard)?;
             MemMgrWrapper::new(mgr, stack_guard)
         };
 
         mem_mgr_wrapper.write_memory_layout()?;
 
+        // if env has a guest blob, load it into shared mem
+        if let Some(bytes) = env.guest_blob {
+            mem_mgr_wrapper.write_user_memory(bytes)?;
+        }
+
         let host_funcs = Arc::new(Mutex::new(FunctionRegistry::default()));
 
         let mut sandbox = Self {