fix: size foundation memory tests

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -1406,15 +1406,12 @@ mod tests {
         assert_eq!(res, 0);
     }
 
-    // Tests to ensure that many (1000) function calls can be made in a call context with a small stack (24K) and heap(32K).
-    // This test effectively ensures that the stack is being properly reset after each call and we are not leaking memory in the Guest.
+    // Checks that 1,000 calls work with constrained guest memory.
+    // This catches guest stack reset and heap leaks.
     #[test]
     fn test_with_small_stack_and_heap() {
-        const HEAP_SIZE: u64 = 32 * 1024;
-        // min_scratch_size already includes 1 page (4k on most
-        // platforms) of guest stack, so add 20k more to get 24k
-        // total, and then add some more for the eagerly-copied page
-        // tables on amd64
+        const HEAP_SIZE: u64 = 128 * 1024;
+        // Leave headroom for legacy transport and eagerly copied page tables.
         let scratch_size = {
             let defaults = SandboxConfiguration::default();
             hyperlight_common::layout::min_scratch_size(
@@ -1425,8 +1422,7 @@ mod tests {
                 defaults.get_g2h_pool_pages(),
                 defaults.get_h2g_pool_pages(),
             )
-        } + 0x10000
-            + 0x10000;
+        } + 0x40000;
 
         let mut sbox1 = SandboxBuilder::from_file(simple_guest_as_pathbuf())
             .heap_size(HEAP_SIZE)
@@ -2147,7 +2143,7 @@ mod tests {
     #[test]
     fn snapshot_restore_recovers_oom_with_larger_heap() {
         let mut source_cfg = SandboxConfiguration::default();
-        source_cfg.set_heap_size(0x20_000);
+        source_cfg.set_heap_size(0x40_000);
         let path = simple_guest_as_pathbuf();
         let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(source_cfg))
             .unwrap()
@@ -2156,7 +2152,7 @@ mod tests {
         let snapshot = source.snapshot().unwrap();
 
         let mut target_cfg = SandboxConfiguration::default();
-        target_cfg.set_heap_size(0x8000);
+        target_cfg.set_heap_size(0x20_000);
         let path = simple_guest_as_pathbuf();
         let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(target_cfg))
             .unwrap()
@@ -2177,7 +2173,7 @@ mod tests {
     #[test]
     fn snapshot_restore_applies_smaller_heap_limit() {
         let mut source_cfg = SandboxConfiguration::default();
-        source_cfg.set_heap_size(0x8000);
+        source_cfg.set_heap_size(0x20_000);
         let path = simple_guest_as_pathbuf();
         let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(source_cfg))
             .unwrap()
@@ -2186,26 +2182,28 @@ mod tests {
         let snapshot = source.snapshot().unwrap();
 
         let mut target_cfg = SandboxConfiguration::default();
-        target_cfg.set_heap_size(0x20_000);
+        target_cfg.set_heap_size(0x80_000);
         let path = simple_guest_as_pathbuf();
         let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(target_cfg))
             .unwrap()
             .evolve()
             .unwrap();
 
         assert_eq!(
-            target.call::<i32>("CallMalloc", 0x10_000i32).unwrap(),
-            0x10_000
+            target.call::<i32>("CallMalloc", 0x30_000i32).unwrap(),
+            0x30_000
         );
         target.restore(snapshot).unwrap();
-        assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000);
-        assert!(target.call::<i32>("CallMalloc", 0x10_000i32).is_err());
+        assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000);
+        assert!(target.call::<i32>("CallMalloc", 0x30_000i32).is_err());
         assert!(target.status().is_poisoned());
     }
 
     #[test]
     fn snapshot_restore_applies_smaller_io_limits() {
         let mut source_cfg = SandboxConfiguration::default();
+        source_cfg.set_heap_size(0x40_000);
+        source_cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024);
         source_cfg.set_input_data_size(0x2000);
         source_cfg.set_output_data_size(0x2000);
         let path = simple_guest_as_pathbuf();
@@ -2216,6 +2214,8 @@ mod tests {
         let snapshot = source.snapshot().unwrap();
 
         let mut target_cfg = SandboxConfiguration::default();
+        target_cfg.set_heap_size(0x40_000);
+        target_cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024);
         target_cfg.set_input_data_size(0x8000);
         target_cfg.set_output_data_size(0x8000);
         let path = simple_guest_as_pathbuf();
@@ -2242,7 +2242,7 @@ mod tests {
         let mut small_cfg = SandboxConfiguration::default();
         small_cfg.set_input_data_size(0x2000);
         small_cfg.set_output_data_size(0x2000);
-        small_cfg.set_heap_size(0x8000);
+        small_cfg.set_heap_size(0x20_000);
         let path = simple_guest_as_pathbuf();
         let mut small = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(small_cfg))
             .unwrap()
@@ -2272,15 +2272,15 @@ mod tests {
 
         target.restore(small_snapshot.clone()).unwrap();
         assert_eq!(target.call::<i32>("GetStatic", ()).unwrap(), 11);
-        assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000);
+        assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000);
 
         target.restore(large_snapshot).unwrap();
         assert_eq!(target.call::<i32>("GetStatic", ()).unwrap(), 22);
         assert_eq!(target.mem_mgr.layout.heap_size(), 0x40_000);
 
         target.restore(small_snapshot).unwrap();
         assert_eq!(target.call::<i32>("GetStatic", ()).unwrap(), 11);
-        assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000);
+        assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000);
     }
 
     #[test]

src/hyperlight_host/src/sandbox/snapshot/file_tests.rs

@@ -2781,7 +2781,9 @@ fn round_trip_preserves_stack_top_gva() {
 
 #[test]
 fn round_trip_preserves_non_default_scratch_size() {
-    let custom_scratch: usize = 256 * 1024;
+    use crate::sandbox::SandboxConfiguration;
+
+    let custom_scratch = SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 64 * 1024;
     let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
         .scratch_size(custom_scratch)
         .build()

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -1163,6 +1163,7 @@ mod tests {
         {
             let mut cfg = SandboxConfiguration::default();
             cfg.set_heap_size(16 * 1024 * 1024); // 16MB heap
+            cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024);
 
             let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None);
 
@@ -1185,7 +1186,7 @@ mod tests {
         // Test 3: Create snapshot with custom scratch size
         {
             let mut cfg = SandboxConfiguration::default();
-            cfg.set_scratch_size(256 * 1024); // 256KB scratch
+            cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 64 * 1024);
 
             let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None);
 
@@ -1208,6 +1209,7 @@ mod tests {
         // Test 4: Create snapshot with custom input/output buffer sizes
         {
             let mut cfg = SandboxConfiguration::default();
+            cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 128 * 1024);
             cfg.set_input_data_size(64 * 1024); // 64KB input
             cfg.set_output_data_size(64 * 1024); // 64KB output
 
@@ -1233,7 +1235,7 @@ mod tests {
         {
             let mut cfg = SandboxConfiguration::default();
             cfg.set_heap_size(32 * 1024 * 1024); // 32MB heap
-            cfg.set_scratch_size(256 * 1024 * 2); // 512KB scratch (256KB will be input/output)
+            cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1024 * 1024);
             cfg.set_input_data_size(128 * 1024); // 128KB input
             cfg.set_output_data_size(128 * 1024); // 128KB output
 

src/hyperlight_host/tests/integration_test.rs

@@ -523,8 +523,8 @@ fn guest_malloc_abort() {
     });
 
     // allocate a vector (on heap) that is bigger than the heap
-    let heap_size = 0x8000;
-    let size_to_allocate = 0x10000;
+    let heap_size = 128 * 1024;
+    let size_to_allocate = 256 * 1024;
     assert!(
         size_to_allocate > heap_size,
         "precondition: size_to_allocate ({size_to_allocate}) must be > heap_size ({heap_size})"
@@ -601,7 +601,7 @@ fn corrupt_output_back_pointer_rejected() {
 
 #[test]
 fn guest_panic_no_alloc() {
-    let heap_size = 0x8000;
+    let heap_size = 128 * 1024;
 
     let configure = |builder: SandboxBuilder| builder.heap_size(heap_size);
     with_rust_sandbox_from(configure, |mut sbox| {
@@ -612,10 +612,15 @@ fn guest_panic_no_alloc() {
             )
             .unwrap_err();
 
+        // Legacy transport may report its own allocation failure.
         assert!(
             matches!(
                 &res,
-                HyperlightError::GuestAborted(code, msg) if *code == ErrorCode::UnknownError as u8 && msg.contains("memory allocation of ") && msg.contains("bytes failed")
+                HyperlightError::GuestAborted(code, msg)
+                    if (*code == ErrorCode::UnknownError as u8
+                        && msg.contains("memory allocation of ")
+                        && msg.contains("bytes failed"))
+                        || *code == ErrorCode::MallocFailed as u8
             ),
             "unexpected error: {res:?}"
         );
@@ -1664,6 +1669,8 @@ fn fill_heap_and_cause_exception() {
 
         let err = result.unwrap_err();
         match &err {
+            // Legacy transport may report its own allocation failure.
+            HyperlightError::GuestAborted(code, _) if *code == ErrorCode::MallocFailed as u8 => {}
             HyperlightError::GuestAborted(code, message) => {
                 assert_eq!(*code, ErrorCode::GuestError as u8, "Full error: {:?}", err);