Allow the guest to manage its own stack

This moves management of the guest stack into the guest, allowing the
guest to flexibly allocate pages from the scratch region to the stack.
This also makes the stack dynamically growable and essentially
unbounded, as is common on other architectures.

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

Author: syntactically

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

@@ -30,3 +30,13 @@ pub const SNAPSHOT_PT_GVA_MAX: usize = 0xffff_80ff_ffff_ffff;
 /// 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;
+
+/// On amd64, this is:
+/// - Two pages for the TSS and IDT
+/// - (up to) 4 pages for the PTEs for mapping that (including CoW'ing the root PT)
+/// - A page for the smallest possible non-exception stack
+/// - (up to) 3 pages for mapping that
+/// - Two pages for the exception stack and metadata
+pub fn min_scratch_size() -> usize {
+    12 * crate::vmem::PAGE_SIZE
+}

src/hyperlight_common/src/arch/i686/layout.rs

@@ -21,3 +21,7 @@ pub const MAX_GVA: usize = 0xffff_efff;
 pub const SNAPSHOT_PT_GVA_MIN: usize = 0xef00_0000;
 pub const SNAPSHOT_PT_GVA_MAX: usize = 0xefff_efff;
 pub const MAX_GPA: usize = 0xffff_ffff;
+
+pub fn min_scratch_size() -> usize {
+    1 * crate::vmem::PAGE_SIZE
+}

src/hyperlight_common/src/layout.rs

@@ -32,3 +32,6 @@ pub fn scratch_base_gpa(size: usize) -> u64 {
 pub fn scratch_base_gva(size: usize) -> u64 {
     (MAX_GVA - size + 1) as u64
 }
+
+/// Compute the minimum scratch region size needed for a sandbox.
+pub use arch::min_scratch_size;

src/hyperlight_common/src/mem.rs

@@ -46,5 +46,4 @@ pub struct HyperlightPEB {
     pub output_stack: GuestMemoryRegion,
     pub init_data: GuestMemoryRegion,
     pub guest_heap: GuestMemoryRegion,
-    pub guest_stack: GuestStack,
 }

src/hyperlight_guest/src/arch/amd64/layout.rs

@@ -18,7 +18,13 @@ limitations under the License.
 // src/hyperlight_common/src/arch/amd64/layout.rs and
 // src/hyperlight_guest_bin/src/arch/amd64/layout.rs
 
-pub const MAIN_STACK_TOP_GVA: u64 = 0xffff_feff_ffff_f000;
+/// Note that the x86-64 ELF psABI requires that the stack be 16-byte
+/// aligned before a call instruction; we use the aligned version
+/// here, even though this requires adjusting the pointer by 8 bytes
+/// when entering the guest without a call instruction to push a
+/// return address.
+pub const MAIN_STACK_TOP_GVA: u64 = 0xffff_ff00_0000_0000;
+pub const MAIN_STACK_LIMIT_GVA: u64 = 0xffff_fe00_0000_0000;
 
 pub fn scratch_size() -> u64 {
     let addr = crate::layout::scratch_size_gva();

src/hyperlight_guest/src/arch/i686/layout.rs

@@ -18,6 +18,7 @@ limitations under the License.
 // allow compiling the guest for real mode boot scenarios.
 
 pub const MAIN_STACK_TOP_GVA: usize = 0xdfff_efff;
+pub const MAIN_STACK_LIMIT_GVA: usize = 0xdf00_0000;
 
 pub fn scratch_size() -> u64 {
     hyperlight_common::vmem::PAGE_SIZE as u64

src/hyperlight_guest/src/layout.rs

@@ -18,7 +18,7 @@ limitations under the License.
 #[cfg_attr(target_arch = "x86", path = "arch/i686/layout.rs")]
 mod arch;
 
-pub use arch::MAIN_STACK_TOP_GVA;
+pub use arch::{MAIN_STACK_LIMIT_GVA, MAIN_STACK_TOP_GVA};
 pub fn scratch_size_gva() -> *mut u64 {
     use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SIZE_OFFSET};
     (MAX_GVA as u64 - SCRATCH_TOP_SIZE_OFFSET + 1) as *mut u64

src/hyperlight_guest_bin/src/arch/amd64/exn/handle.rs

@@ -18,6 +18,7 @@ use core::fmt::Write;
 
 use hyperlight_common::outb::Exception;
 use hyperlight_guest::exit::write_abort;
+use hyperlight_guest::layout::{MAIN_STACK_LIMIT_GVA, MAIN_STACK_TOP_GVA};
 
 use super::super::context::Context;
 use super::super::machine::ExceptionInfo;
@@ -77,6 +78,23 @@ pub(crate) extern "C" fn hl_exception_handler(
     let saved_rip = unsafe { (&raw const (*exn_info).rip).read_volatile() };
     let error_code = unsafe { (&raw const (*exn_info).error_code).read_volatile() };
 
+    if exception_number == 14
+        && (MAIN_STACK_LIMIT_GVA..MAIN_STACK_TOP_GVA).contains(&page_fault_address)
+    {
+        // TODO: perhaps we should have a sanity check that the
+        // stack grows only one page at a time, which should be
+        // ensured by our stack probing discipline?
+        unsafe {
+            let new_page = hyperlight_guest::prim_alloc::alloc_phys_pages(1);
+            crate::paging::map_region(
+                new_page,
+                (page_fault_address & !0xfff) as *mut u8,
+                hyperlight_common::vmem::PAGE_SIZE as u64,
+            );
+            return;
+        }
+    }
+
     // Check for registered user handlers (only for architecture-defined vectors 0-30)
     if exception_number < 31 {
         let handler =

src/hyperlight_guest_bin/src/arch/amd64/init.rs

@@ -86,6 +86,19 @@ unsafe fn init_tss(pc: *mut ProcCtrl) {
     }
 }
 
+unsafe fn init_stack() -> u64 {
+    use hyperlight_guest::layout::MAIN_STACK_TOP_GVA;
+    let stack_top_page_base = (MAIN_STACK_TOP_GVA - 1) & !0xfff;
+    unsafe {
+        crate::paging::map_region(
+            hyperlight_guest::prim_alloc::alloc_phys_pages(1),
+            stack_top_page_base as *mut u8,
+            hyperlight_common::vmem::PAGE_SIZE as u64,
+        );
+    }
+    MAIN_STACK_TOP_GVA
+}
+
 /// Machine-specific initialisation; calls [`crate::generic_init`]
 /// once stack, CoW, etc have been set up.
 #[unsafe(no_mangle)]
@@ -95,18 +108,25 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
         init_gdt(pc);
         init_tss(pc);
         init_idt(pc);
-        call_generic_init(peb_address, seed, ops, max_log_level);
+        let stack_top = init_stack();
+        pivot_stack(peb_address, seed, ops, max_log_level, stack_top);
     }
 }
 
 unsafe extern "C" {
-    unsafe fn call_generic_init(peb_address: u64, seed: u64, ops: u64, max_log_level: u64) -> !;
+    unsafe fn pivot_stack(
+        peb_address: u64,
+        seed: u64,
+        ops: u64,
+        max_log_level: u64,
+        stack_top: u64,
+    ) -> !;
 }
 
 core::arch::global_asm!("
-    .global call_generic_init\n
-    call_generic_init:\n
-    sub rsp, 0x8\n
+    .global pivot_stack\n
+    pivot_stack:\n
+    mov rsp, r8\n
     call {generic_init}\n
     hlt\n
 ", generic_init = sym crate::generic_init);

src/hyperlight_guest_bin/src/lib.rs

@@ -119,8 +119,10 @@ pub static mut GUEST_HANDLE: GuestHandle = GuestHandle::new();
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
     GuestFunctionRegister::new();
 
-pub static mut MIN_STACK_ADDRESS: u64 = 0;
-
+/// The size of one page in the host OS, which may have some impacts
+/// on how buffers for host consumption should be aligned. Code only
+/// working with the guest page tables should use
+/// [`hyperlight_common::vm::PAGE_SIZE`] instead.
 pub static mut OS_PAGE_SIZE: u32 = 0;
 
 // === Panic Handler ===
@@ -202,12 +204,6 @@ pub(crate) extern "C" fn generic_init(peb_address: u64, seed: u64, ops: u64, max
     unsafe {
         (*peb_ptr).guest_function_dispatch_ptr = dispatch_function as usize as u64;
 
-        // 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.
-        MIN_STACK_ADDRESS = (*peb_ptr).guest_stack.min_user_stack_address;
-
         let srand_seed = (((peb_address << 8) ^ (seed >> 4)) >> 32) as u32;
         // Set the seed for the random number generator for C code using rand;
         srand(srand_seed);