[guest/{chkstk,entrypoint},host/{drivers,sandbox_builder},common/peb,simpleguest] fixed stack setup

danbugs · danbugs · commit ae1637bbbee4 · 2025-04-24T18:23:56.000Z
(1) PE guests must have the RSP be 16-byte aligned, added stack alignment on init_rsp and re-setup to account for it.
(2) We're now properly updating the min stack address needed for chkstk via the HyperlightPEB.
(3) Added tests for PE guests.

Signed-off-by: danbugs &lt;danilochiarlone@gmail.com&gt;
diff --git a/src/hyperlight_common/src/peb.rs b/src/hyperlight_common/src/peb.rs
@@ -27,6 +27,7 @@ pub struct MemoryRegion {
 #[repr(C)]
 #[derive(Clone, Default)]
 pub struct HyperlightPEB {
+    pub min_stack_address: u64,
     // - Host configured fields
     /// Hyperlight supports two primary modes:
     /// 1. Hypervisor mode
@@ -86,6 +87,7 @@ impl HyperlightPEB {
     /// Creates a new HyperlightPEB with the basic configuration based on the provided guest memory
     /// layout and default guest heap/stack sizes. The guest can later fill additional fields.
     pub fn new(
+        min_stack_address: u64,
         run_mode: RunMode,
         guest_heap_size: u64,
         guest_stack_size: u64,
@@ -94,6 +96,7 @@ impl HyperlightPEB {
         guest_memory_size: u64,
     ) -> Self {
         Self {
+            min_stack_address,
             run_mode,
             outb_ptr: 0,
             outb_ptr_ctx: 0,
@@ -222,6 +225,16 @@ impl HyperlightPEB {
         region.offset.unwrap() + self.guest_memory_base_address + region.size
     }
 
+    /// Calculate the minimum guest stack address (start of guest stack data region in the guest
+    /// address space).
+    pub fn calculate_min_stack_address(&self) -> u64 {
+        let region = self
+            .guest_stack_data
+            .as_ref()
+            .expect("Guest stack data region not set");
+        region.offset.unwrap() + self.guest_memory_base_address
+    }
+
     /// Sets the guest heap data region.
     /// - HyperlightPEB is always set with a default size for heap from the guest binary, there's an
     /// option to override this size with the `size_override` parameter.
diff --git a/src/hyperlight_guest/src/chkstk.rs b/src/hyperlight_guest/src/chkstk.rs
@@ -15,13 +15,12 @@ limitations under the License.
 */
 
 use core::arch::global_asm;
-use core::mem::size_of;
+use core::mem::{offset_of, size_of};
 
-use hyperlight_common::peb::RunMode;
-use hyperlight_common::RUNNING_MODE;
+use hyperlight_common::peb::{HyperlightPEB, RunMode};
+use hyperlight_common::{PEB, RUNNING_MODE};
 
 use crate::guest_error::{set_invalid_runmode_error, set_stack_allocate_error};
-use crate::MIN_STACK_ADDRESS;
 
 extern "win64" {
     fn __chkstk();
@@ -52,20 +51,22 @@ global_asm!(
 
     handle_hypervisor:
         /* Load the minimum stack address from the PEB */
-        mov r11, [rip+{min_stack_addr}]  
+        /* min_stack_address is offset 0x0 in the PEB struct */
+        mov r11, [rip+{peb_ptr}]
+        mov r11, qword ptr [r11]
 
         /* Get the current stack pointer */
-        lea r10, [rsp+0x18]  
+        lea r10, [rsp+0x18]
 
         /* Calculate what the new stack pointer will be */
         sub r10, rax
-        
+
         /* If result is negative, cause StackOverflow */
         js call_set_error
-        
+
         /* Compare the new stack pointer with the minimum stack address */
-        cmp r10, r11   
-        /* If the new stack pointer is greater or equal to the minimum stack address,  
+        cmp r10, r11
+        /* If the new stack pointer is greater or equal to the minimum stack address,
             then we are good. Otherwise set the error code to 9 (stack overflow) call set_error and halt */
         jae cs_ret
 
@@ -90,7 +91,7 @@ global_asm!(
         cmp r10, r11
         jne csip_stackprobe
     cs_ret:
-        /* Restore RAX, R11 */
+        /* Restore R10, R11 */
         pop r11
         pop r10
         ret
@@ -104,7 +105,7 @@ global_asm!(
     handle_invalid:
         call {invalid_runmode}",
     run_mode = sym RUNNING_MODE,
-    min_stack_addr = sym MIN_STACK_ADDRESS,
+    peb_ptr = sym PEB,
     set_error = sym set_stack_allocate_error,
     invalid_runmode = sym set_invalid_runmode_error
 );
@@ -118,4 +119,5 @@ const _: () = {
     assert!(RunMode::InProcessWindows as u64 == 2);
     assert!(RunMode::InProcessLinux as u64 == 3);
     assert!(RunMode::Invalid as u64 == 4);
+    assert!(offset_of!(HyperlightPEB, min_stack_address) == 0x0);
 };
diff --git a/src/hyperlight_guest/src/entrypoint.rs b/src/hyperlight_guest/src/entrypoint.rs
@@ -27,7 +27,7 @@ use crate::gdt::load_gdt;
 use crate::guest_function_call::dispatch_function;
 use crate::guest_logger::init_logger;
 use crate::idtr::load_idt;
-use crate::{__security_cookie, HEAP_ALLOCATOR, MIN_STACK_ADDRESS};
+use crate::{__security_cookie, HEAP_ALLOCATOR};
 
 #[inline(never)]
 pub fn halt() {
@@ -111,12 +111,6 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, max_log_level: u64
 
         match RUNNING_MODE {
             RunMode::Hypervisor => {
-                // This static is to make it easier to implement the __chkstk function in assembly.
-                // It also means that, should we change the layout of the struct in the future, we
-                // don't have to change the assembly code. Plus, while this could be accessible via
-                // the PEB, we don't want to expose it entirely to user code.
-                MIN_STACK_ADDRESS = (*PEB).get_stack_data_address();
-
                 // Setup GDT and IDT
                 load_gdt();
                 load_idt();
diff --git a/src/hyperlight_guest/src/lib.rs b/src/hyperlight_guest/src/lib.rs
@@ -85,7 +85,5 @@ pub(crate) static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::empty();
 #[no_mangle]
 pub(crate) static mut __security_cookie: u64 = 0;
 
-pub static mut MIN_STACK_ADDRESS: u64 = 0;
-
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
     GuestFunctionRegister::new();
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_linux.rs b/src/hyperlight_host/src/hypervisor/hyperv_linux.rs
@@ -62,7 +62,7 @@ use super::{
 use crate::hypervisor::hypervisor_handler::HypervisorHandler;
 use crate::hypervisor::HyperlightExit;
 use crate::mem::ptr::{GuestPtr, RawPtr};
-use crate::sandbox::sandbox_builder::{MemoryRegionFlags, SandboxMemorySections};
+use crate::sandbox::sandbox_builder::{MemoryRegionFlags, SandboxMemorySections, STACK_ALIGNMENT};
 #[cfg(gdb)]
 use crate::HyperlightError;
 use crate::{log_then_return, new_error, Result};
@@ -489,13 +489,17 @@ impl Hypervisor for HypervLinuxDriver {
         )?;
 
         // 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()?;
+        let mut 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))?;
+                self.orig_rsp = GuestPtr::try_from(RawPtr::from(rsp - STACK_ALIGNMENT))?;
+
+                // Need to update the min stack address from tmp_stack address to the new stack
+                hyperlight_peb.min_stack_address = hyperlight_peb.calculate_min_stack_address();
+                self.mem_sections.write_hyperlight_peb(hyperlight_peb)?;
             }
         }
 
@@ -730,67 +734,3 @@ impl Drop for HypervLinuxDriver {
         }
     }
 }
-
-// TODO(danbugs:297): bring back
-// #[cfg(test)]
-// mod tests {
-//     use super::*;
-//     use crate::mem::memory_region::MemoryRegionVecBuilder;
-//     use crate::mem::shared_mem::{ExclusiveSharedMemory, SharedMemory};
-//
-//     #[rustfmt::skip]
-//     const CODE: [u8; 12] = [
-//         0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */
-//         0x00, 0xd8, /* add %bl, %al */
-//         0x04, b'0', /* add $'0', %al */
-//         0xee, /* out %al, (%dx) */
-//         /* send a 0 to indicate we're done */
-//         0xb0, b'\0', /* mov $'\0', %al */
-//         0xee, /* out %al, (%dx) */
-//         0xf4, /* HLT */
-//     ];
-//
-//     fn shared_mem_with_code(
-//         code: &[u8],
-//         mem_size: usize,
-//         load_offset: usize,
-//     ) -> Result<Box<ExclusiveSharedMemory>> {
-//         if load_offset > mem_size {
-//             log_then_return!(
-//                 "code load offset ({}) > memory size ({})",
-//                 load_offset,
-//                 mem_size
-//             );
-//         }
-//         let mut shared_mem = ExclusiveSharedMemory::new(mem_size)?;
-//         shared_mem.copy_from_slice(code, load_offset)?;
-//         Ok(Box::new(shared_mem))
-//     }
-//
-//     #[test]
-//     fn create_driver() {
-//         if !super::is_hypervisor_present() {
-//             return;
-//         }
-//         const MEM_SIZE: usize = 0x3000;
-//         let gm = shared_mem_with_code(CODE.as_slice(), MEM_SIZE, 0).unwrap();
-//         let rsp_ptr = GuestPtr::try_from(0).unwrap();
-//         let pml4_ptr = GuestPtr::try_from(0).unwrap();
-//         let entrypoint_ptr = GuestPtr::try_from(0).unwrap();
-//         let mut regions = MemoryRegionVecBuilder::new(0, gm.base_addr());
-//         regions.push_page_aligned(
-//             MEM_SIZE,
-//             MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
-//             crate::mem::memory_region::MemoryRegionType::Code,
-//         );
-//         super::HypervLinuxDriver::new(
-//             regions.build(),
-//             entrypoint_ptr,
-//             rsp_ptr,
-//             pml4_ptr,
-//             #[cfg(gdb)]
-//             None,
-//         )
-//         .unwrap();
-//     }
-// }
diff --git a/src/hyperlight_host/src/hypervisor/hyperv_windows.rs b/src/hyperlight_host/src/hypervisor/hyperv_windows.rs
@@ -44,7 +44,7 @@ use crate::hypervisor::fpu::FP_CONTROL_WORD_DEFAULT;
 use crate::hypervisor::hypervisor_handler::HypervisorHandler;
 use crate::hypervisor::wrappers::WHvGeneralRegisters;
 use crate::mem::ptr::{GuestPtr, RawPtr};
-use crate::sandbox::sandbox_builder::{MemoryRegionFlags, SandboxMemorySections};
+use crate::sandbox::sandbox_builder::{MemoryRegionFlags, SandboxMemorySections, STACK_ALIGNMENT};
 use crate::{debug, new_error, Result};
 
 /// A Hypervisor driver for HyperV-on-Windows.
@@ -339,13 +339,17 @@ impl Hypervisor for HypervWindowsDriver {
         )?;
 
         // 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()?;
+        let mut 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))?;
+                self.orig_rsp = GuestPtr::try_from(RawPtr::from(rsp - STACK_ALIGNMENT))?;
+
+                // Need to update the min stack address from tmp_stack address to the new stack
+                hyperlight_peb.min_stack_address = hyperlight_peb.calculate_min_stack_address();
+                self.mem_sections.write_hyperlight_peb(hyperlight_peb)?;
             }
         }
 
@@ -504,29 +508,28 @@ impl Hypervisor for HypervWindowsDriver {
     }
 }
 
-// TODO(danbugs:297): bring back
-// #[cfg(test)]
-// pub mod tests {
-//     use std::sync::{Arc, Mutex};
-//
-//     use serial_test::serial;
-//
-//     use crate::hypervisor::handlers::{MemAccessHandler, OutBHandler};
-//     use crate::hypervisor::tests::test_initialise;
-//     use crate::Result;
-//
-//     #[test]
-//     #[serial]
-//     fn test_init() {
-//         let outb_handler = {
-//             let func: Box<dyn FnMut(u16, u64) -> Result<()> + Send> =
-//                 Box::new(|_, _| -> Result<()> { Ok(()) });
-//             Arc::new(Mutex::new(OutBHandler::from(func)))
-//         };
-//         let mem_access_handler = {
-//             let func: Box<dyn FnMut() -> Result<()> + Send> = Box::new(|| -> Result<()> { Ok(()) });
-//             Arc::new(Mutex::new(MemAccessHandler::from(func)))
-//         };
-//         test_initialise(outb_handler, mem_access_handler).unwrap();
-//     }
-// }
+#[cfg(test)]
+pub mod tests {
+    use std::sync::{Arc, Mutex};
+
+    use serial_test::serial;
+
+    use crate::hypervisor::handlers::{MemAccessHandler, OutBHandler};
+    use crate::hypervisor::tests::test_initialise;
+    use crate::Result;
+
+    #[test]
+    #[serial]
+    fn test_init() {
+        let outb_handler = {
+            let func: Box<dyn FnMut(u16, u64) -> Result<()> + Send> =
+                Box::new(|_, _| -> Result<()> { Ok(()) });
+            Arc::new(Mutex::new(OutBHandler::from(func)))
+        };
+        let mem_access_handler = {
+            let func: Box<dyn FnMut() -> Result<()> + Send> = Box::new(|| -> Result<()> { Ok(()) });
+            Arc::new(Mutex::new(MemAccessHandler::from(func)))
+        };
+        test_initialise(outb_handler, mem_access_handler).unwrap();
+    }
+}
diff --git a/src/hyperlight_host/src/hypervisor/kvm.rs b/src/hyperlight_host/src/hypervisor/kvm.rs
diff --git a/src/hyperlight_host/src/sandbox/sandbox_builder.rs b/src/hyperlight_host/src/sandbox/sandbox_builder.rs
diff --git a/src/tests/rust_guests/simpleguest/src/main.rs b/src/tests/rust_guests/simpleguest/src/main.rs
