[common,guest,host] modified HyperlightPEB API + added rsp mod in hosts

- cleaned up the HyperlightPEB API in the common library.
-- added MemoryRegion struct to better group related offsets and sizes.
-- removed pub fields from HyperlightPEB struct making fields accessible only via getters/setters.
-- cleaned up, commented, and re-organized existing fxns for HyperlightPEB struct.
- now we modify the rsp in the host if the guest sets up a new stack region (i.e., essentially dropping
the tmp stack).

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

Author: danbugs

Diff for: src/hyperlight_common/src/outb.rs

@@ -45,7 +45,7 @@ pub fn outb(port: u16, value: u8) {
             }
             RunMode::InProcessLinux | RunMode::InProcessWindows => {
                 if let Some(outb_func) = OUTB_HANDLER_CTX {
-                    outb_func((*PEB).outb_ptr_ctx as *mut core::ffi::c_void, port, value);
+                    outb_func((*PEB).get_outb_ptr_ctx() as *mut core::ffi::c_void, port, value);
                 } else if let Some(outb_func) = OUTB_HANDLER {
                     outb_func(port, value);
                 } else {

Diff for: src/hyperlight_common/src/peb.rs

@@ -76,7 +76,7 @@ static INIT: Once = Once::new();
 pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, max_log_level: u64) {
     INIT.call_once(|| unsafe {
         PEB = peb_address as *mut HyperlightPEB;
-        RUNNING_MODE = (*PEB).clone().run_mode;
+        RUNNING_MODE = (*PEB).clone().get_run_mode();
 
         // The guest receives an undifferentiated block of memory that it can address as it sees fit.
         // This 'addressing' is done by writing to the PEB the guest's memory layout via this function,
@@ -87,15 +87,17 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, max_log_level: u64
 
         // The guest sets the address to a "guest function dispatch" function, which is a function
         // that is called by the host to dispatch calls to guest functions.
-        (*PEB).guest_function_dispatch_ptr = dispatch_function as usize as u64;
+        (*PEB).set_guest_function_dispatch_ptr(
+            dispatch_function as u64,
+        );
 
         // Set up the guest heap
         HEAP_ALLOCATOR
             .try_lock()
             .expect("Failed to access HEAP_ALLOCATOR")
             .init(
                 (*PEB).get_heap_data_address() as usize,
-                (*PEB).guest_heap_data_size as usize,
+                (*PEB).get_guest_heap_data_size() as usize,
             );
 
         __security_cookie = peb_address ^ seed;
@@ -125,17 +127,17 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, max_log_level: u64
             RunMode::InProcessLinux | RunMode::InProcessWindows => {
                 OUTB_HANDLER = {
                     let outb_handler: extern "C" fn(u16, u8) =
-                        core::mem::transmute((*PEB).outb_ptr);
+                        core::mem::transmute((*PEB).get_outb_ptr());
                     Some(outb_handler)
                 };
 
-                if (*PEB).outb_ptr_ctx == 0 {
+                if (*PEB).get_outb_ptr_ctx() == 0 {
                     panic!("outb_ptr_ctx is null");
                 }
 
                 OUTB_HANDLER_CTX = {
                     let outb_handler_ctx: extern "C" fn(*mut core::ffi::c_void, u16, u8) =
-                        core::mem::transmute((*PEB).outb_ptr);
+                        core::mem::transmute((*PEB).get_outb_ptr());
                     Some(outb_handler_ctx)
                 };
             }

Diff for: src/hyperlight_guest/src/entrypoint.rs

@@ -38,17 +38,17 @@ pub(crate) fn write_error(error_code: ErrorCode, message: Option<&str>) {
     unsafe {
         assert_ne!(!peb.get_guest_error_data_address(), 0);
         let len = guest_error_buffer.len();
-        if guest_error_buffer.len() > peb.guest_error_data_size as usize {
+        if guest_error_buffer.len() > peb.get_guest_error_data_size() as usize {
             error!(
                 "Guest error buffer is too small to hold the error message: size {} buffer size {} message may be truncated",
                 guest_error_buffer.len(),
-                peb.guest_error_data_size as usize
+                peb.get_guest_error_data_size() as usize
             );
             // get the length of the message
             let message_len = message.map_or("".to_string(), |m| m.to_string()).len();
             // message is too long, truncate it
             let truncate_len =
-                message_len - (guest_error_buffer.len() - peb.guest_error_data_size as usize);
+                message_len - (guest_error_buffer.len() - peb.get_guest_error_data_size() as usize);
             let truncated_message = message
                 .map_or("".to_string(), |m| m.to_string())
                 .chars()
@@ -77,7 +77,7 @@ pub(crate) fn reset_error() {
         core::ptr::write_bytes(
             (*PEB).get_guest_error_data_address() as *mut u8,
             0,
-            (*PEB).guest_error_data_size as usize,
+            (*PEB).get_guest_error_data_size() as usize,
         );
     }
 }

Diff for: src/hyperlight_guest/src/guest_error.rs

@@ -72,7 +72,7 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
         copy_nonoverlapping(
             info.to_string().as_ptr(),
             (*PEB).get_guest_panic_context_address() as *mut u8,
-            (*PEB).guest_panic_context_size as usize,
+            (*PEB).get_guest_panic_context_size() as usize,
         );
     }
     outb(OutBAction::Abort as u16, ErrorCode::UnknownError as u8);

Diff for: src/hyperlight_guest/src/lib.rs

@@ -54,7 +54,10 @@ pub(crate) fn call_function_on_guest(
         .try_into()
         .map_err(|_| HyperlightError::Error("Failed to serialize FunctionCall".to_string()))?;
 
-    let input_data_region = mem_mgr.read_hyperlight_peb()?.get_input_data_guest_region();
+    let input_data_region = mem_mgr
+        .memory_sections
+        .read_hyperlight_peb()?
+        .get_input_data_guest_region();
 
     mem_mgr.write_guest_function_call(input_data_region, &buffer)?;
 
@@ -81,6 +84,7 @@ pub(crate) fn call_function_on_guest(
     check_for_guest_error(mem_mgr)?;
 
     let output_data_region = mem_mgr
+        .memory_sections
         .read_hyperlight_peb()?
         .get_output_data_guest_region();
 

Diff for: src/hyperlight_host/src/func/guest_dispatch.rs

@@ -493,7 +493,16 @@ impl Hypervisor for HypervLinuxDriver {
             dbg_mem_access_fn,
         )?;
 
-        // TODO(danbugs:297): here, we should update the rsp to what the guest configured.
+        // The guest may have chosen a different stack region. If so, we drop usage of our tmp stack.
+        let hyperlight_peb = self.mem_sections.read_hyperlight_peb()?;
+
+        if let Some(guest_stack_data) = &hyperlight_peb.get_guest_stack_data_region() {
+            if guest_stack_data.offset.is_some() {
+                // If we got here, it means the guest has set up a new stack
+                let rsp = hyperlight_peb.get_top_of_guest_stack_data();
+                self.orig_rsp = GuestPtr::try_from(RawPtr::from(rsp))?;
+            }
+        }
 
         Ok(())
     }

Diff for: src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -338,6 +338,17 @@ impl Hypervisor for HypervWindowsDriver {
             dbg_mem_access_hdl,
         )?;
 
+        // The guest may have chosen a different stack region. If so, we drop usage of our tmp stack.
+        let hyperlight_peb = self.mem_sections.read_hyperlight_peb()?;
+
+        if let Some(guest_stack_data) = &hyperlight_peb.get_guest_stack_data_region() {
+            if guest_stack_data.offset.is_some() {
+                // If we got here, it means the guest has set up a new stack
+                let rsp = hyperlight_peb.get_top_of_guest_stack_data();
+                self.orig_rsp = GuestPtr::try_from(RawPtr::from(rsp))?;
+            }
+        }
+
         Ok(())
     }
 

Diff for: src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -137,7 +137,7 @@ impl HvHandlerExecVars {
             .thread_id
             .try_lock()
             .map_err(|_| new_error!("Failed to get_thread_id"))?)
-        .ok_or_else(|| new_error!("thread_id not set"))
+            .ok_or_else(|| new_error!("thread_id not set"))
     }
 
     #[cfg(target_os = "windows")]
@@ -405,7 +405,7 @@ impl HypervisorHandler {
                                     .as_mut()
                                     .ok_or_else(|| {
                                         new_error!("guest shm lock: {}:{}", file!(), line!())
-                                    })?.read_hyperlight_peb()?.guest_function_dispatch_ptr);
+                                    })?.memory_sections.read_hyperlight_peb()?.get_guest_function_dispatch_ptr());
 
                                 if dispatch_function_addr == RawPtr(0) {
                                     log_then_return!(
@@ -795,7 +795,7 @@ impl HypervisorHandler {
                         0,
                         0,
                     )
-                    .map_err(|e| new_error!("Failed to cancel guest execution {:?}", e))?;
+                        .map_err(|e| new_error!("Failed to cancel guest execution {:?}", e))?;
                 }
             }
             // if running in-process on windows, we currently have no way of cancelling the execution

Diff for: src/hyperlight_host/src/hypervisor/hypervisor_handler.rs

@@ -438,7 +438,16 @@ impl Hypervisor for KVMDriver {
             dbg_mem_access_fn,
         )?;
 
-        // TODO(danbugs:297): here, we should update the rsp to what the guest configured.
+        // The guest may have chosen a different stack region. If so, we drop usage of our tmp stack.
+        let hyperlight_peb = self.mem_sections.read_hyperlight_peb()?;
+
+        if let Some(guest_stack_data) = &hyperlight_peb.get_guest_stack_data_region() {
+            if guest_stack_data.offset.is_some() {
+                // If we got here, it means the guest has set up a new stack
+                let rsp = hyperlight_peb.get_top_of_guest_stack_data();
+                self.orig_rsp = GuestPtr::try_from(RawPtr::from(rsp))?;
+            }
+        }
 
         Ok(())
     }

Diff for: src/hyperlight_host/src/hypervisor/kvm.rs

@@ -25,7 +25,6 @@ use hyperlight_common::flatbuffer_wrappers::function_call::{
 use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
-use hyperlight_common::peb::HyperlightPEB;
 use serde_json::from_str;
 use tracing::{instrument, Span};
 
@@ -236,28 +235,6 @@ impl<S> SandboxMemoryManager<S>
 where
     S: SharedMemory,
 {
-    pub(crate) fn write_hyperlight_peb(&mut self, hyperlight_peb: HyperlightPEB) -> Result<()> {
-        let peb_offset = self
-            .memory_sections
-            .get_hyperlight_peb_section_host_address()
-            .unwrap() as *mut HyperlightPEB;
-
-        unsafe {
-            peb_offset.copy_from(&hyperlight_peb, 1);
-        }
-
-        Ok(())
-    }
-
-    pub(crate) fn read_hyperlight_peb(&self) -> Result<HyperlightPEB> {
-        let peb_offset = self
-            .memory_sections
-            .get_hyperlight_peb_section_host_address()
-            .unwrap() as *mut HyperlightPEB;
-
-        Ok(unsafe { peb_offset.read() })
-    }
-
     /// 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
@@ -376,7 +353,10 @@ impl SandboxMemoryManager<HostSharedMemory> {
     /// Reads a host function call from memory
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn get_host_function_call(&mut self) -> Result<FunctionCall> {
-        let (ptr, size) = self.read_hyperlight_peb()?.get_output_data_guest_region();
+        let (ptr, size) = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_output_data_guest_region();
         self.shared_mem
             .try_pop_buffer_into::<FunctionCall>(ptr as usize, size as usize)
     }
@@ -390,7 +370,10 @@ impl SandboxMemoryManager<HostSharedMemory> {
             )
         })?;
 
-        let (ptr, size) = self.read_hyperlight_peb()?.get_input_data_guest_region();
+        let (ptr, size) = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_input_data_guest_region();
 
         self.shared_mem.push_buffer(
             ptr as usize,
@@ -434,23 +417,32 @@ impl SandboxMemoryManager<HostSharedMemory> {
     /// Read guest log data from the `SharedMemory` contained within `self`
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn read_guest_log_data(&mut self) -> Result<GuestLogData> {
-        let (ptr, size) = self.read_hyperlight_peb()?.get_output_data_guest_region();
+        let (ptr, size) = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_output_data_guest_region();
         self.shared_mem
             .try_pop_buffer_into::<GuestLogData>(ptr as usize, size as usize)
     }
 
     /// Get the length of the host exception
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     fn get_host_error_length(&self) -> Result<i32> {
-        let offset = self.read_hyperlight_peb()?.get_host_error_guest_offset() as usize;
+        let offset = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_host_error_guest_address() as usize;
         // The host exception field is expected to contain a 32-bit length followed by the exception data.
         self.shared_mem.read::<i32>(offset)
     }
 
     /// Get a bool indicating if there is a host error
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     fn has_host_error(&self) -> Result<bool> {
-        let offset = self.read_hyperlight_peb()?.get_host_error_guest_offset() as usize;
+        let offset = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_host_error_guest_address() as usize;
         // The host exception field is expected to contain a 32-bit length followed by the exception data.
         let len = self.shared_mem.read::<i32>(offset)?;
         Ok(len != 0)
@@ -465,7 +457,10 @@ impl SandboxMemoryManager<HostSharedMemory> {
     /// self.get_host_error_length()
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     fn get_host_error_data(&self, exception_data_slc: &mut [u8]) -> Result<()> {
-        let offset = self.read_hyperlight_peb()?.get_host_error_guest_offset() as usize;
+        let offset = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_host_error_guest_address() as usize;
         let len = self.get_host_error_length()?;
 
         let exception_data_slc_len = exception_data_slc.len();
@@ -510,9 +505,14 @@ impl SandboxMemoryManager<HostSharedMemory> {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn get_guest_error(&self) -> Result<GuestError> {
         // get memory buffer max size
-        let guest_error_data_ptr =
-            self.read_hyperlight_peb()?.get_guest_error_guest_offset() as usize;
-        let guest_error_data_size = self.read_hyperlight_peb()?.guest_error_data_size as usize;
+        let guest_error_data_ptr = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_error_guest_address() as usize;
+        let guest_error_data_size = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_error_data_size() as usize;
 
         // get guest error from layout and shared mem
         let mut guest_error_buffer = vec![b'0'; guest_error_data_size];
@@ -541,18 +541,29 @@ impl SandboxMemoryManager<HostSharedMemory> {
             .try_into()
             .map_err(|_| new_error!("write_outb_error: failed to convert GuestError to Vec<u8>"))?;
 
-        let guest_error_data_ptr = self.read_hyperlight_peb()?.get_guest_error_guest_offset();
-        let guest_error_data_size = self.read_hyperlight_peb()?.guest_error_data_size;
+        let guest_error_data_ptr = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_error_guest_address();
+        let guest_error_data_size = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_error_data_size();
 
         if guest_error_buffer.len() as u64 > guest_error_data_size {
             log_then_return!("The guest error message is too large to fit in the shared memory");
         }
         self.shared_mem
             .copy_from_slice(guest_error_buffer.as_slice(), guest_error_data_ptr as usize)?;
 
-        let host_error_data_ptr =
-            self.read_hyperlight_peb()?.get_host_error_guest_offset() as usize;
-        let host_error_data_size = self.read_hyperlight_peb()?.host_error_data_size as usize;
+        let host_error_data_ptr = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_host_error_guest_address() as usize;
+        let host_error_data_size = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_host_error_data_size() as usize;
 
         // First four bytes of host exception are length
 
@@ -574,8 +585,14 @@ impl SandboxMemoryManager<HostSharedMemory> {
     /// Read guest panic data from the `SharedMemory` contained within `self`
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub fn read_guest_panic_context_data(&self) -> Result<Vec<u8>> {
-        let offset = self.read_hyperlight_peb()?.get_guest_panic_context_offset() as usize;
-        let size = self.read_hyperlight_peb()?.guest_panic_context_size as usize;
+        let offset = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_panic_context_guest_address() as usize;
+        let size = self
+            .memory_sections
+            .read_hyperlight_peb()?
+            .get_guest_panic_context_size() as usize;
         let mut vec_out = vec![0; size];
         self.shared_mem
             .copy_to_slice(vec_out.as_mut_slice(), offset)?;

Diff for: src/hyperlight_host/src/mem/mgr.rs

@@ -99,10 +99,10 @@ impl<'a> LeakedOutBWrapper<'a> {
         };
 
         let addr: u64 = res.hdl_wrapper_addr()?;
-        let mut peb = mgr.read_hyperlight_peb()?;
+        let mut peb = mgr.memory_sections.read_hyperlight_peb()?;
         peb.set_outb_ptr(Self::outb_addr());
         peb.set_outb_ptr_ctx(addr);
-        mgr.write_hyperlight_peb(peb)?;
+        mgr.memory_sections.write_hyperlight_peb(peb)?;
 
         Ok(res)
     }