Add one more test

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -2459,5 +2459,303 @@ mod tests {
             normalize_sregs_no_granularity(&mut expected_reset, &sregs);
             assert_eq!(sregs, expected_reset);
         }
+
+        /// Verifies guest-visible FPU state (via FXSAVE) is properly reset.
+        /// Unlike tests using hypervisor API, this runs actual guest code with FXSAVE.
+        #[test]
+        fn reset_vcpu_fpu_guest_visible_state() {
+            let mut ctx = hyperlight_vm_with_mem_mgr_fxsave();
+
+            // Verify FPU was dirtied after first run
+            let fpu_before_reset = ctx.vm.vm.fpu().unwrap();
+            assert_eq!(
+                fpu_before_reset.fcw, 0x0F7F,
+                "FCW should be dirty after first run"
+            );
+            assert_ne!(
+                fpu_before_reset.ftwx, 0,
+                "FTW should indicate valid registers after first run"
+            );
+
+            let fxsave_before = ctx.read_fxsave();
+            let fcw_before = u16::from_le_bytes(fxsave_before[0..2].try_into().unwrap());
+            assert_eq!(fcw_before, 0x0F7F, "Guest FXSAVE FCW should be dirty");
+            let mxcsr_before = u32::from_le_bytes(fxsave_before[24..28].try_into().unwrap());
+            assert_eq!(mxcsr_before, 0x3F80, "Guest FXSAVE MXCSR should be dirty");
+            let xmm0_before = u32::from_le_bytes(fxsave_before[160..164].try_into().unwrap());
+            assert_eq!(xmm0_before, 0x11111111, "Guest FXSAVE XMM0 should be dirty");
+
+            ctx.vm.reset_vcpu().unwrap();
+
+            // Re-run from entrypoint (flag=1 means guest skips dirty phase, just does FXSAVE)
+            let regs = CommonRegisters {
+                rip: ctx.vm.entrypoint,
+                rsp: ctx.vm.orig_rsp.absolute().unwrap(),
+                rflags: 1 << 1,
+                ..Default::default()
+            };
+            ctx.vm.vm.set_regs(&regs).unwrap();
+            ctx.run();
+
+            // Verify guest-visible state is reset
+            let fxsave_after = ctx.read_fxsave();
+            let fcw_after = u16::from_le_bytes(fxsave_after[0..2].try_into().unwrap());
+            assert_eq!(
+                fcw_after, 0x037F,
+                "Guest FXSAVE FCW should be reset to default 0x037F, got 0x{:04X}",
+                fcw_after
+            );
+
+            let fsw_after = u16::from_le_bytes(fxsave_after[2..4].try_into().unwrap());
+            assert_eq!(fsw_after, 0, "FSW should be reset");
+
+            let ftw_after = fxsave_after[4];
+            assert_eq!(ftw_after, 0, "FTW should be 0 (all empty)");
+
+            let fop_after = u16::from_le_bytes(fxsave_after[6..8].try_into().unwrap());
+            assert_eq!(fop_after, 0, "FOP should be 0");
+
+            let fip_after = u64::from_le_bytes(fxsave_after[8..16].try_into().unwrap());
+            assert_eq!(fip_after, 0, "FIP should be 0");
+
+            let fdp_after = u64::from_le_bytes(fxsave_after[16..24].try_into().unwrap());
+            assert_eq!(fdp_after, 0, "FDP should be 0");
+
+            let mxcsr_after = u32::from_le_bytes(fxsave_after[24..28].try_into().unwrap());
+            assert_eq!(
+                mxcsr_after, MXCSR_DEFAULT,
+                "Guest FXSAVE MXCSR should be reset to 0x{:08X}, got 0x{:08X}",
+                MXCSR_DEFAULT, mxcsr_after
+            );
+
+            // ST0-ST7 should be zeroed
+            for i in 0..8 {
+                let offset = 32 + i * 16;
+                let st_bytes = &fxsave_after[offset..offset + 10];
+                assert!(st_bytes.iter().all(|&b| b == 0), "ST{} should be zeroed", i);
+            }
+
+            // XMM0-XMM15 should be zeroed
+            for i in 0..16 {
+                let offset = 160 + i * 16;
+                let xmm_bytes = &fxsave_after[offset..offset + 16];
+                assert!(
+                    xmm_bytes.iter().all(|&b| b == 0),
+                    "XMM{} should be zeroed",
+                    i
+                );
+            }
+        }
+
+        struct FxsaveTestContext {
+            vm: HyperlightVm,
+            hshm: SandboxMemoryManager<HostSharedMemory>,
+            host_funcs: Arc<Mutex<FunctionRegistry>>,
+            fxsave_offset: usize,
+            #[cfg(gdb)]
+            dbg_mem_access_hdl: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
+        }
+
+        impl FxsaveTestContext {
+            fn run(&mut self) {
+                self.vm
+                    .run(
+                        &mut self.hshm,
+                        &self.host_funcs,
+                        #[cfg(gdb)]
+                        self.dbg_mem_access_hdl.clone(),
+                    )
+                    .unwrap();
+            }
+
+            fn read_fxsave(&self) -> [u8; 512] {
+                let mut fxsave = [0u8; 512];
+                self.hshm
+                    .shared_mem
+                    .copy_to_slice(&mut fxsave, self.fxsave_offset)
+                    .unwrap();
+                fxsave
+            }
+        }
+
+        /// Creates VM with guest code that: dirtys FPU (if flag==0), does FXSAVE, sets flag=1.
+        fn hyperlight_vm_with_mem_mgr_fxsave() -> FxsaveTestContext {
+            use iced_x86::code_asm::*;
+
+            // Stack: RSP+0 = FXSAVE (512 bytes), RSP+512 = flag
+            const FXSAVE_STACK_OFFSET: i32 = 0;
+            const FLAG_STACK_OFFSET: i32 = 512;
+            const STACK_ALLOC_SIZE: i32 = 528;
+
+            let mut a = CodeAssembler::new(64).unwrap();
+
+            // Align stack and allocate space
+            a.and(rsp, -16i32).unwrap();
+            a.sub(rsp, STACK_ALLOC_SIZE).unwrap();
+
+            // Check flag: if flag != 0, skip dirty phase
+            a.mov(al, byte_ptr(rsp + FLAG_STACK_OFFSET)).unwrap();
+            a.test(al, al).unwrap();
+            let mut skip_dirty = a.create_label();
+            a.jnz(skip_dirty).unwrap();
+
+            // Dirty x87 FPU (7 pushes so TOP=1)
+            a.fldz().unwrap();
+            a.fldl2e().unwrap();
+            a.fldl2t().unwrap();
+            a.fldlg2().unwrap();
+            a.fldln2().unwrap();
+            a.fldpi().unwrap();
+            a.fld1().unwrap();
+
+            // Dirty FCW (0x0F7F vs default 0x037F)
+            a.sub(rsp, 16i32).unwrap();
+            a.mov(dword_ptr(rsp), 0x0F7Fu32).unwrap();
+            a.fldcw(word_ptr(rsp)).unwrap();
+            a.add(rsp, 16i32).unwrap();
+
+            // Dirty MXCSR (0x3F80 vs default 0x1F80)
+            a.sub(rsp, 16i32).unwrap();
+            a.mov(dword_ptr(rsp), 0x3F80u32).unwrap();
+            a.ldmxcsr(dword_ptr(rsp)).unwrap();
+            a.add(rsp, 16i32).unwrap();
+
+            // Dirty XMM0-7
+            let xmm_regs = [xmm0, xmm1, xmm2, xmm3, xmm4, xmm5, xmm6, xmm7];
+            for (i, xmm) in xmm_regs.iter().enumerate() {
+                a.mov(eax, 0x11111111u32 * (i as u32 + 1)).unwrap();
+                a.movd(*xmm, eax).unwrap();
+                a.pshufd(*xmm, *xmm, 0).unwrap();
+            }
+
+            a.mov(byte_ptr(rsp + FLAG_STACK_OFFSET), 1u32).unwrap();
+
+            // FXSAVE (runs on both executions)
+            a.set_label(&mut skip_dirty).unwrap();
+            a.lea(rax, ptr(rsp + FXSAVE_STACK_OFFSET)).unwrap();
+            a.fxsave(ptr(rax)).unwrap();
+            a.mov(r15, rsp).unwrap(); // Store RSP for host to read
+
+            a.hlt().unwrap();
+
+            let code = a.assemble(0).unwrap();
+
+            // Create VM with the code
+            let config: SandboxConfiguration = Default::default();
+            #[cfg(any(crashdump, gdb))]
+            let rt_cfg: SandboxRuntimeConfig = Default::default();
+
+            let mut layout =
+                SandboxMemoryLayout::new(config, code.len(), 4096, 4096, 0x3000, 0, None).unwrap();
+
+            let pt_base_gpa = SandboxMemoryLayout::BASE_ADDRESS + layout.get_pt_offset();
+            let pt_buf = GuestPageTableBuffer::new(pt_base_gpa);
+
+            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 scratch_mem = ExclusiveSharedMemory::new(config.get_scratch_size()).unwrap();
+            let mut mem_mgr = SandboxMemoryManager::new(
+                layout,
+                eshm,
+                scratch_mem,
+                load_addr,
+                Some(Offset::from(0)),
+                [0u8; 16],
+            );
+            mem_mgr.write_memory_layout().unwrap();
+
+            let (mut hshm, gshm) = mem_mgr.build();
+
+            let peb_address = gshm.layout.peb_address;
+            let mut vm = set_up_hypervisor_partition(
+                gshm,
+                &config,
+                #[cfg(any(crashdump, gdb))]
+                &rt_cfg,
+                crate::mem::exe::LoadInfo::dummy(),
+            )
+            .unwrap();
+
+            let seed = rand::rng().random::<u64>();
+            let peb_addr = RawPtr::from(u64::try_from(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()));
+
+            vm.initialise(
+                peb_addr,
+                seed,
+                page_size,
+                &mut hshm,
+                &host_funcs,
+                None,
+                #[cfg(gdb)]
+                dbg_mem_access_hdl.clone(),
+            )
+            .unwrap();
+
+            // R15 = RSP after stack allocation, convert to host offset
+            let regs = vm.vm.regs().unwrap();
+            let fxsave_offset = regs.r15 as usize - SandboxMemoryLayout::BASE_ADDRESS;
+
+            FxsaveTestContext {
+                vm,
+                hshm,
+                host_funcs,
+                fxsave_offset,
+                #[cfg(gdb)]
+                dbg_mem_access_hdl,
+            }
+        }
     }
 }