Additional changes to validate page table setup and multiple snadbox sizes (#1119)

Signed-off-by: Simon Davies <simongdavies@users.noreply.github.com>

Author: simongdavies

src/hyperlight_host/benches/benchmarks.rs

@@ -30,6 +30,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, Ret
 use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
 use hyperlight_host::GuestBinary;
 use hyperlight_host::sandbox::{MultiUseSandbox, SandboxConfiguration, UninitializedSandbox};
+use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE};
 use hyperlight_testing::{c_simple_guest_as_string, simple_guest_as_string};
 
 /// Sandbox heap size configurations for benchmarking.
@@ -54,17 +55,17 @@ impl SandboxSize {
             Self::Default => None,
             Self::Small => {
                 let mut cfg = SandboxConfiguration::default();
-                cfg.set_heap_size(8 * 1024 * 1024); // 8 MB
+                cfg.set_heap_size(SMALL_HEAP_SIZE);
                 Some(cfg)
             }
             Self::Medium => {
                 let mut cfg = SandboxConfiguration::default();
-                cfg.set_heap_size(64 * 1024 * 1024); // 64 MB
+                cfg.set_heap_size(MEDIUM_HEAP_SIZE);
                 Some(cfg)
             }
             Self::Large => {
                 let mut cfg = SandboxConfiguration::default();
-                cfg.set_heap_size(256 * 1024 * 1024); // 256 MB
+                cfg.set_heap_size(LARGE_HEAP_SIZE);
                 Some(cfg)
             }
         }

src/hyperlight_host/src/mem/mgr.rs

@@ -559,3 +559,202 @@ impl SandboxMemoryManager<HostSharedMemory> {
         }
     }
 }
+
+#[cfg(test)]
+#[cfg(feature = "init-paging")]
+mod tests {
+    use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE};
+    use hyperlight_testing::simple_guest_as_string;
+
+    use super::*;
+    use crate::GuestBinary;
+    use crate::mem::shared_mem::GuestSharedMemory;
+    use crate::sandbox::SandboxConfiguration;
+    use crate::sandbox::uninitialized::UninitializedSandbox;
+
+    /// Helper to create a sandbox with page tables set up and return the manager
+    fn create_sandbox_with_page_tables(
+        config: Option<SandboxConfiguration>,
+    ) -> Result<SandboxMemoryManager<GuestSharedMemory>> {
+        let path = simple_guest_as_string().expect("failed to get simple guest path");
+        let sandbox = UninitializedSandbox::new(GuestBinary::FilePath(path), config)
+            .expect("failed to create sandbox");
+
+        let mem_size = sandbox.mgr.layout.get_memory_size()?;
+
+        // Build the shared memory to get GuestSharedMemory
+        let (_host_mem, guest_mem) = sandbox.mgr.shared_mem.build();
+        let mut mgr = SandboxMemoryManager {
+            shared_mem: guest_mem,
+            layout: sandbox.mgr.layout,
+            load_addr: sandbox.mgr.load_addr,
+            entrypoint_offset: sandbox.mgr.entrypoint_offset,
+            mapped_rgns: sandbox.mgr.mapped_rgns,
+            stack_cookie: sandbox.mgr.stack_cookie,
+            abort_buffer: sandbox.mgr.abort_buffer,
+        };
+
+        // Get regions and set up page tables
+        let mut regions = mgr.layout.get_memory_regions(&mgr.shared_mem)?;
+        let mem_size_u64 = u64::try_from(mem_size)?;
+        // set_up_shared_memory builds the page tables in shared memory
+        mgr.set_up_shared_memory(mem_size_u64, &mut regions)?;
+
+        Ok(mgr)
+    }
+
+    /// Verify a range of pages all have the same expected flags
+    fn verify_page_range(
+        excl_mem: &mut ExclusiveSharedMemory,
+        start_addr: usize,
+        end_addr: usize,
+        expected_flags: u64,
+        region_name: &str,
+    ) -> Result<()> {
+        let mut addr = start_addr;
+
+        while addr < end_addr {
+            let p = addr >> 21;
+            let i = (addr >> 12) & 0x1ff;
+            let pte_idx = p * 512 + i;
+            let offset = SandboxMemoryLayout::PT_OFFSET + (pte_idx * 8);
+
+            let pte_val = excl_mem.read_u64(offset)?;
+            let expected_pte = (addr as u64) | expected_flags;
+
+            if pte_val != expected_pte {
+                return Err(new_error!(
+                    "{} region: addr 0x{:x}: expected PTE 0x{:x}, got 0x{:x}",
+                    region_name,
+                    addr,
+                    expected_pte,
+                    pte_val
+                ));
+            }
+
+            addr += 0x1000;
+        }
+
+        Ok(())
+    }
+
+    /// Get expected flags for a memory region type
+    fn get_expected_flags(region: &MemoryRegion) -> u64 {
+        match region.region_type {
+            MemoryRegionType::Code => PAGE_PRESENT | PAGE_RW | PAGE_USER,
+            MemoryRegionType::Stack => PAGE_PRESENT | PAGE_RW | PAGE_USER | PAGE_NX,
+            #[cfg(feature = "executable_heap")]
+            MemoryRegionType::Heap => PAGE_PRESENT | PAGE_RW | PAGE_USER,
+            #[cfg(not(feature = "executable_heap"))]
+            MemoryRegionType::Heap => PAGE_PRESENT | PAGE_RW | PAGE_USER | PAGE_NX,
+            MemoryRegionType::GuardPage => PAGE_PRESENT | PAGE_RW | PAGE_USER | PAGE_NX,
+            MemoryRegionType::InputData => PAGE_PRESENT | PAGE_RW | PAGE_NX,
+            MemoryRegionType::OutputData => PAGE_PRESENT | PAGE_RW | PAGE_NX,
+            MemoryRegionType::Peb => PAGE_PRESENT | PAGE_RW | PAGE_NX,
+            MemoryRegionType::HostFunctionDefinitions => PAGE_PRESENT | PAGE_NX,
+            MemoryRegionType::PageTables => PAGE_PRESENT | PAGE_RW | PAGE_NX,
+            MemoryRegionType::InitData => region.flags.translate_flags(),
+        }
+    }
+
+    /// Verify the complete paging structure for a sandbox configuration
+    fn verify_paging_structure(name: &str, config: Option<SandboxConfiguration>) -> Result<()> {
+        let mut mgr = create_sandbox_with_page_tables(config)?;
+
+        let regions = mgr.layout.get_memory_regions(&mgr.shared_mem)?;
+
+        mgr.shared_mem.with_exclusivity(|excl_mem| {
+            // Verify PML4 entry (single entry pointing to PDPT)
+            let pml4_val = excl_mem.read_u64(SandboxMemoryLayout::PML4_OFFSET)?;
+            let expected_pml4 =
+                SandboxMemoryLayout::PDPT_GUEST_ADDRESS as u64 | PAGE_PRESENT | PAGE_RW;
+            if pml4_val != expected_pml4 {
+                return Err(new_error!(
+                    "{}: PML4[0] incorrect: expected 0x{:x}, got 0x{:x}",
+                    name,
+                    expected_pml4,
+                    pml4_val
+                ));
+            }
+
+            // Verify PDPT entry (single entry pointing to PD)
+            let pdpt_val = excl_mem.read_u64(SandboxMemoryLayout::PDPT_OFFSET)?;
+            let expected_pdpt =
+                SandboxMemoryLayout::PD_GUEST_ADDRESS as u64 | PAGE_PRESENT | PAGE_RW;
+            if pdpt_val != expected_pdpt {
+                return Err(new_error!(
+                    "{}: PDPT[0] incorrect: expected 0x{:x}, got 0x{:x}",
+                    name,
+                    expected_pdpt,
+                    pdpt_val
+                ));
+            }
+
+            // Verify all 512 PD entries (each pointing to a PT)
+            for i in 0..512 {
+                let offset = SandboxMemoryLayout::PD_OFFSET + (i * 8);
+                let pd_val = excl_mem.read_u64(offset)?;
+                let expected_pt_addr =
+                    SandboxMemoryLayout::PT_GUEST_ADDRESS as u64 + (i as u64 * 4096);
+                let expected_pd = expected_pt_addr | PAGE_PRESENT | PAGE_RW;
+                if pd_val != expected_pd {
+                    return Err(new_error!(
+                        "{}: PD[{}] incorrect: expected 0x{:x}, got 0x{:x}",
+                        name,
+                        i,
+                        expected_pd,
+                        pd_val
+                    ));
+                }
+            }
+
+            // Verify PTEs for each memory region
+            for region in &regions {
+                let start = region.guest_region.start;
+                let end = region.guest_region.end;
+                let expected_flags = get_expected_flags(region);
+
+                verify_page_range(
+                    excl_mem,
+                    start,
+                    end,
+                    expected_flags,
+                    &format!("{} {:?}", name, region.region_type),
+                )?;
+            }
+
+            Ok(())
+        })??;
+
+        Ok(())
+    }
+
+    /// Test the complete paging structure (PML4, PDPT, PD, and all PTEs) for
+    /// sandboxes of different sizes: default, small (8MB), medium (64MB), and large (256MB)
+    #[test]
+    fn test_page_table_contents() {
+        let test_cases: [(&str, Option<SandboxConfiguration>); 4] = [
+            ("default", None),
+            ("small (8MB heap)", {
+                let mut cfg = SandboxConfiguration::default();
+                cfg.set_heap_size(SMALL_HEAP_SIZE);
+                Some(cfg)
+            }),
+            ("medium (64MB heap)", {
+                let mut cfg = SandboxConfiguration::default();
+                cfg.set_heap_size(MEDIUM_HEAP_SIZE);
+                Some(cfg)
+            }),
+            ("large (256MB heap)", {
+                let mut cfg = SandboxConfiguration::default();
+                cfg.set_heap_size(LARGE_HEAP_SIZE);
+                Some(cfg)
+            }),
+        ];
+
+        for (name, config) in test_cases {
+            verify_paging_structure(name, config)
+                .unwrap_or_else(|e| panic!("Page table verification failed for {}: {}", name, e));
+        }
+    }
+}

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -787,6 +787,7 @@ mod tests {
     use std::thread;
 
     use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+    use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE};
     use hyperlight_testing::simple_guest_as_string;
 
     #[cfg(target_os = "linux")]
@@ -1296,4 +1297,27 @@ mod tests {
         };
         assert_ne!(sandbox3.id, sandbox_id);
     }
+
+    /// Test that sandboxes can be created and evolved with different heap sizes
+    #[test]
+    fn test_sandbox_creation_various_sizes() {
+        let test_cases: [(&str, u64); 3] = [
+            ("small (8MB heap)", SMALL_HEAP_SIZE),
+            ("medium (64MB heap)", MEDIUM_HEAP_SIZE),
+            ("large (256MB heap)", LARGE_HEAP_SIZE),
+        ];
+
+        for (name, heap_size) in test_cases {
+            let mut cfg = SandboxConfiguration::default();
+            cfg.set_heap_size(heap_size);
+
+            let path = simple_guest_as_string().unwrap();
+            let sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(cfg))
+                .unwrap_or_else(|e| panic!("Failed to create {} sandbox: {}", name, e))
+                .evolve()
+                .unwrap_or_else(|e| panic!("Failed to evolve {} sandbox: {}", name, e));
+
+            drop(sbox);
+        }
+    }
 }

src/hyperlight_testing/src/lib.rs

@@ -130,3 +130,15 @@ pub fn simple_guest_for_fuzzing_as_string() -> Result<String> {
 
     simple_guest_as_string()
 }
+
+/// Standard sandbox heap sizes for benchmarking and testing.
+/// These constants define common heap sizes used across benchmarks and tests
+/// to ensure consistency.
+pub mod sandbox_sizes {
+    /// Small heap size: 8 MB
+    pub const SMALL_HEAP_SIZE: u64 = 8 * 1024 * 1024;
+    /// Medium heap size: 64 MB
+    pub const MEDIUM_HEAP_SIZE: u64 = 64 * 1024 * 1024;
+    /// Large heap size: 256 MB
+    pub const LARGE_HEAP_SIZE: u64 = 256 * 1024 * 1024;
+}