SandboxMemoryLayout: do not pass heap and scratch size explicitly

These values could conflict with those in the SandboxConfiguration
also passed in.

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

Author: syntactically

src/hyperlight_host/src/mem/layout.rs

@@ -200,11 +200,22 @@ impl SandboxMemoryLayout {
     pub(crate) fn new(
         cfg: SandboxConfiguration,
         code_size: usize,
-        heap_size: usize,
-        scratch_size: usize,
         init_data_size: usize,
         init_data_permissions: Option<MemoryRegionFlags>,
     ) -> Result<Self> {
+        let heap_size = usize::try_from(cfg.get_heap_size())?;
+        let scratch_size = cfg.get_scratch_size();
+        if scratch_size > Self::MAX_MEMORY_SIZE {
+            return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE));
+        }
+        let min_scratch_size = hyperlight_common::layout::min_scratch_size(
+            cfg.get_input_data_size(),
+            cfg.get_output_data_size(),
+        );
+        if scratch_size < min_scratch_size {
+            return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size));
+        }
+
         let guest_code_offset = 0;
         // The following offsets are to the fields of the PEB struct itself!
         let peb_offset = code_size.next_multiple_of(PAGE_SIZE_USIZE);
@@ -655,8 +666,7 @@ mod tests {
     #[test]
     fn test_get_memory_size() {
         let sbox_cfg = SandboxConfiguration::default();
-        let sbox_mem_layout =
-            SandboxMemoryLayout::new(sbox_cfg, 4096, 4096, 0x3000, 0, None).unwrap();
+        let sbox_mem_layout = SandboxMemoryLayout::new(sbox_cfg, 4096, 0, None).unwrap();
         assert_eq!(
             sbox_mem_layout.get_memory_size().unwrap(),
             get_expected_memory_size(&sbox_mem_layout)
@@ -667,8 +677,7 @@ mod tests {
     fn test_max_memory_sandbox() {
         let mut cfg = SandboxConfiguration::default();
         cfg.set_input_data_size(0x40000000);
-        let layout = SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 4096, None).unwrap();
-        let result = layout.get_memory_size();
-        assert!(matches!(result.unwrap_err(), MemoryRequestTooBig(..)));
+        let layout = SandboxMemoryLayout::new(cfg, 4096, 4096, None);
+        assert!(matches!(layout.unwrap_err(), MemoryRequestTooBig(..)));
     }
 }

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -359,8 +359,6 @@ impl Snapshot {
         let mut layout = crate::mem::layout::SandboxMemoryLayout::new(
             cfg,
             exe_info.loaded_size(),
-            usize::try_from(cfg.get_heap_size())?,
-            cfg.get_scratch_size(),
             guest_blob_size,
             guest_blob_mem_flags,
         )?;
@@ -612,9 +610,8 @@ mod tests {
         let mut snapshot_mem = ExclusiveSharedMemory::new(PAGE_SIZE + pt_bytes.len()).unwrap();
 
         snapshot_mem.copy_from_slice(&pt_bytes, PAGE_SIZE).unwrap();
-        let cfg = crate::sandbox::SandboxConfiguration::default();
         let mgr = SandboxMemoryManager::new(
-            SandboxMemoryLayout::new(cfg, 4096, 2048, 4096, 0x3000, None).unwrap(),
+            SandboxMemoryLayout::new(cfg, 4096, 0x3000, None).unwrap(),
             snapshot_mem,
             scratch_mem,
             0.into(),