fixup! Reset vcpu state

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -1120,6 +1120,7 @@ mod debug {
 mod tests {
     use std::sync::{Arc, Mutex};
 
+    use hyperlight_common::vmem::{self, BasicMapping, Mapping, MappingKind};
     use rand::Rng;
 
     use super::*;
@@ -1128,9 +1129,9 @@ mod tests {
     use crate::hypervisor::regs::{CommonSegmentRegister, CommonTableRegister, MXCSR_DEFAULT};
     #[cfg(target_os = "windows")]
     use crate::hypervisor::wrappers::HandleWrapper;
-    use crate::mem::exe::LoadInfo;
     use crate::mem::layout::SandboxMemoryLayout;
-    use crate::mem::mgr::SandboxMemoryManager;
+    use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegionFlags};
+    use crate::mem::mgr::{GuestPageTableBuffer, SandboxMemoryManager};
     use crate::mem::ptr::RawPtr;
     use crate::mem::ptr_offset::Offset;
     use crate::mem::shared_mem::ExclusiveSharedMemory;
@@ -1473,24 +1474,68 @@ mod tests {
         xsave
     }
 
-    fn hyperlight_vm(code: &[u8]) -> Result<HyperlightVm> {
+    fn hyperlight_vm(code: &[u8]) -> HyperlightVm {
         let config: SandboxConfiguration = Default::default();
         #[cfg(crashdump)]
         let rt_cfg: SandboxRuntimeConfig = Default::default();
 
-        let layout = SandboxMemoryLayout::new(config, code.len(), 4096, 0, 0, None)?;
+        let mut layout = SandboxMemoryLayout::new(config, code.len(), 4096, 4096, 0, None).unwrap();
 
-        let mem_size = layout.get_memory_size()?;
-        let eshm = ExclusiveSharedMemory::new(mem_size)?;
+        let pt_base_gpa = SandboxMemoryLayout::BASE_ADDRESS + layout.get_pt_offset();
+        let pt_buf = GuestPageTableBuffer::new(pt_base_gpa);
 
-        let stack_cookie = [0u8; 16];
-        let mem_mgr = SandboxMemoryManager::new(
-            layout,
-            eshm,
-            RawPtr::from(0),
-            Some(Offset::from(0)),
-            stack_cookie,
-        );
+        for rgn in layout
+            .get_memory_regions_::<GuestMemoryRegion>(())
+            .unwrap()
+            .iter()
+        {
+            let readable = rgn.flags.contains(MemoryRegionFlags::READ);
+            let writable = rgn.flags.contains(MemoryRegionFlags::WRITE)
+                || rgn.flags.contains(MemoryRegionFlags::STACK_GUARD);
+            let executable = rgn.flags.contains(MemoryRegionFlags::EXECUTE);
+            let mapping = Mapping {
+                phys_base: rgn.guest_region.start as u64,
+                virt_base: rgn.guest_region.start as u64,
+                len: rgn.guest_region.len() as u64,
+                kind: MappingKind::BasicMapping(BasicMapping {
+                    readable,
+                    writable,
+                    executable,
+                }),
+            };
+            unsafe { vmem::map(&pt_buf, mapping) };
+        }
+
+        let mut pt_size_mapped = 0;
+        while pt_buf.size() > pt_size_mapped {
+            let mapping = Mapping {
+                phys_base: (pt_base_gpa + pt_size_mapped) as u64,
+                virt_base: (hyperlight_common::layout::SNAPSHOT_PT_GVA + pt_size_mapped) as u64,
+                len: (pt_buf.size() - pt_size_mapped) as u64,
+                kind: MappingKind::BasicMapping(BasicMapping {
+                    readable: true,
+                    writable: true,
+                    executable: false,
+                }),
+            };
+            unsafe { vmem::map(&pt_buf, mapping) };
+            pt_size_mapped = pt_buf.size();
+        }
+
+        let pt_bytes = pt_buf.into_bytes();
+        layout.set_pt_size(pt_bytes.len());
+
+        let mem_size = layout.get_memory_size().unwrap();
+        let mut eshm = ExclusiveSharedMemory::new(mem_size).unwrap();
+        eshm.copy_from_slice(&pt_bytes, layout.get_pt_offset())
+            .unwrap();
+        eshm.copy_from_slice(code, layout.get_guest_code_offset())
+            .unwrap();
+
+        let load_addr = RawPtr::from(layout.get_guest_code_address() as u64);
+        let mut mem_mgr =
+            SandboxMemoryManager::new(layout, eshm, load_addr, Some(Offset::from(0)), [0u8; 16]);
+        mem_mgr.write_memory_layout().unwrap();
 
         let (mut hshm, mut gshm) = mem_mgr.build();
 
@@ -1499,30 +1544,19 @@ mod tests {
             &config,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,
-            LoadInfo::dummy(),
-        )?;
-
-        // Write code
-        let code_offset = layout.get_guest_code_offset();
-        hshm.shared_mem.copy_from_slice(code, code_offset)?;
-
-        let seed = {
-            let mut rng = rand::rng();
-            rng.random::<u64>()
-        };
-        let peb_addr = {
-            let peb_u64 = u64::try_from(gshm.layout.peb_address)?;
-            RawPtr::from(peb_u64)
-        };
+            crate::mem::exe::LoadInfo::dummy(),
+        )
+        .unwrap();
 
-        let page_size = u32::try_from(page_size::get())?;
+        let seed = rand::rng().random::<u64>();
+        let peb_addr = RawPtr::from(u64::try_from(gshm.layout.peb_address).unwrap());
+        let page_size = u32::try_from(page_size::get()).unwrap();
 
         #[cfg(gdb)]
         let dbg_mem_access_hdl = Arc::new(Mutex::new(hshm.clone()));
 
         let host_funcs = Arc::new(Mutex::new(FunctionRegistry::default()));
 
-        // Run the VM
         vm.initialise(
             peb_addr,
             seed,
@@ -1532,14 +1566,15 @@ mod tests {
             None,
             #[cfg(gdb)]
             dbg_mem_access_hdl.clone(),
-        )?;
-        Ok(vm)
+        )
+        .unwrap();
+        vm
     }
 
     #[test]
     fn reset_vcpu_simple() {
         const CODE: [u8; 1] = [0xf4]; // hlt
-        let hyperlight_vm = hyperlight_vm(&CODE).unwrap();
+        let hyperlight_vm = hyperlight_vm(&CODE);
         let available_hv = *get_available_hypervisor().as_ref().unwrap();
 
         // Set all vCPU state to dirty values
@@ -1757,7 +1792,7 @@ mod tests {
                 0xf4, // hlt
             ];
 
-            let hyperlight_vm = hyperlight_vm(&CODE).unwrap();
+            let hyperlight_vm = hyperlight_vm(&CODE);
 
             // After run, check registers match expected dirty state
             let regs = hyperlight_vm.vm.regs().unwrap();
@@ -1905,7 +1940,7 @@ mod tests {
                 0xf4, // hlt
             ];
 
-            let hyperlight_vm = hyperlight_vm(&CODE).unwrap();
+            let hyperlight_vm = hyperlight_vm(&CODE);
 
             // After run, check FPU state matches expected dirty values
             let fpu = hyperlight_vm.vm.fpu().unwrap();
@@ -2001,7 +2036,7 @@ mod tests {
                 0xf4,                                                       // hlt
             ];
 
-            let hyperlight_vm = hyperlight_vm(code).unwrap();
+            let hyperlight_vm = hyperlight_vm(code);
 
             // Verify debug registers are dirty
             let debug_regs = hyperlight_vm.vm.debug_regs().unwrap();
@@ -2058,7 +2093,7 @@ mod tests {
                 0xf4,                                                       // hlt
             ];
 
-            let hyperlight_vm = hyperlight_vm(code).unwrap();
+            let hyperlight_vm = hyperlight_vm(code);
 
             // Get the expected defaults
             let defaults = CommonSpecialRegisters::standard_64bit_defaults(hyperlight_vm.pml4_addr);