Speed up full-window MSR restore audit

src/hyperlight_host/src/sandbox/initialized.rs

@@ -4135,6 +4135,42 @@ mod tests {
             assert!(sbox.status().is_poisoned());
         }
 
+        #[test]
+        #[cfg(kvm)]
+        fn batched_msr_probe_recovers_from_denied_reads() {
+            use crate::hypervisor::virtual_machine::{HypervisorType, get_available_hypervisor};
+
+            if !matches!(get_available_hypervisor(), Some(HypervisorType::Kvm)) {
+                return;
+            }
+
+            const KVM_CUSTOM_MSR: u32 = 0x4B56_4D00;
+            let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
+                .guest_msrs(&[KERNEL_GS_BASE])
+                .unwrap()
+                .build()
+                .unwrap();
+
+            assert_eq!(
+                sbox.call::<Vec<u8>>("ProbeReadableMSRs", (0u32, 4096u32))
+                    .unwrap()
+                    .len(),
+                512
+            );
+            assert_eq!(
+                sbox.call::<Vec<u8>>("ProbeReadableMSRs", (KERNEL_GS_BASE, 1u32))
+                    .unwrap(),
+                vec![1]
+            );
+            assert_eq!(
+                sbox.call::<Vec<u8>>("ProbeReadableMSRs", (KVM_CUSTOM_MSR, 1u32))
+                    .unwrap(),
+                vec![0]
+            );
+            assert!(!sbox.status().is_poisoned());
+            sbox.call::<u64>("ReadMSR", KERNEL_GS_BASE).unwrap();
+        }
+
         #[test]
         #[cfg(target_arch = "x86_64")]
         fn nested_virtualization_is_hidden_from_guest() {
@@ -4686,21 +4722,32 @@ mod tests {
             // Collect all free-running leaks for one diagnostic.
             let mut free_running_leaked: Vec<u32> = Vec::new();
 
+            const PROBE_CHUNK_SIZE: u32 = 4096;
             // Architectural and low model-specific indices.
             let low = 0x0000_0000u32..=0x0000_1FFF;
             // Hyper-V synthetic indices.
             let hyperv_synthetic = 0x4000_0000u32..=0x4000_1FFF;
             // Extended and AMD model-specific indices.
             let extended = 0xC000_0000u32..=0xC001_FFFF;
-            let windows = low.chain(hyperv_synthetic).chain(extended);
-            for msr in windows {
-                let original: u64 = match sbox.call("ReadMSR", msr) {
-                    Ok(v) => v,
-                    Err(_) => {
-                        sbox.restore(baseline.clone()).unwrap();
-                        continue;
-                    }
-                };
+
+            let mut readable_msrs = Vec::new();
+            for window in [low, hyperv_synthetic, extended] {
+                let window_end = *window.end();
+                for chunk_start in window.step_by(PROBE_CHUNK_SIZE as usize) {
+                    let chunk_len = PROBE_CHUNK_SIZE.min(window_end - chunk_start + 1);
+                    let bitmap: Vec<u8> = sbox
+                        .call("ProbeReadableMSRs", (chunk_start, chunk_len))
+                        .unwrap();
+                    assert_eq!(bitmap.len(), chunk_len.div_ceil(8) as usize);
+                    readable_msrs.extend((0..chunk_len).filter_map(|offset| {
+                        let byte = bitmap[offset as usize / 8];
+                        (byte & (1 << (offset % 8)) != 0).then_some(chunk_start + offset)
+                    }));
+                }
+            }
+
+            for msr in readable_msrs {
+                let original: u64 = sbox.call("ReadMSR", msr).unwrap();
                 readable += 1;
 
                 // A large jump distinguishes reset from normal counter progress.

src/hyperlight_host/tests/sandbox_host_tests.rs

@@ -37,7 +37,11 @@ fn pass_byte_array() {
 fn fragmented_control_round_trip_releases_buffers() {
     // The control body exceeds the four inline segment slots.
     let input = "x".repeat(5 * SandboxConfiguration::DEFAULT_H2G_BUFFER_SIZE);
-    with_all_sandboxes(|mut sbox| {
+    with_all_guests(|path| {
+        let mut sbox = SandboxBuilder::from_file(path)
+            .heap_size(256 * 1024)
+            .build()
+            .unwrap();
         let output: String = sbox.call("Echo", input.clone()).unwrap();
         assert_eq!(output, input);
 

src/tests/rust_guests/simpleguest/src/main.rs

@@ -20,7 +20,7 @@ use core::alloc::Layout;
 use core::ffi::c_char;
 use core::hint::black_box;
 #[cfg(target_arch = "x86_64")]
-use core::sync::atomic::AtomicU32;
+use core::sync::atomic::{AtomicBool, AtomicU32};
 use core::sync::atomic::{AtomicU64, Ordering};
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
@@ -60,6 +60,10 @@ const TEST_R9_MODIFIED_VALUE: u64 = 0xBADC0FFEE;
 const TEST_R10_VALUE: u64 = 0xDEADBEEF;
 #[cfg(target_arch = "x86_64")]
 const RFLAGS_DF: u64 = 1 << 10;
+#[cfg(target_arch = "x86_64")]
+static MSR_PROBE_ACTIVE: AtomicBool = AtomicBool::new(false);
+#[cfg(target_arch = "x86_64")]
+static MSR_PROBE_FAULTED: AtomicBool = AtomicBool::new(false);
 
 #[guest_function("SetStatic")]
 fn set_static() -> i32 {
@@ -1265,6 +1269,81 @@ fn read_msr(msr: u32) -> u64 {
     ((read_edx as u64) << 32) | (read_eax as u64)
 }
 
+#[cfg(target_arch = "x86_64")]
+fn msr_probe_gp_handler(
+    exception_number: u64,
+    exception_info: *mut ExceptionInfo,
+    _context: *mut Context,
+    _page_fault_address: u64,
+) -> bool {
+    if exception_number != 13 || !MSR_PROBE_ACTIVE.load(Ordering::Acquire) {
+        return false;
+    }
+
+    // SAFETY: The exception frame and faulting RIP are valid while the handler runs.
+    let (rip, opcode) = unsafe {
+        let rip = core::ptr::read_volatile(&(*exception_info).rip);
+        (rip, core::ptr::read_volatile(rip as *const [u8; 2]))
+    };
+    if opcode != [0x0F, 0x32] {
+        return false;
+    }
+
+    MSR_PROBE_FAULTED.store(true, Ordering::Release);
+    // SAFETY: RDMSR is two bytes. Resuming after it skips the expected fault.
+    unsafe {
+        core::ptr::write_volatile(&mut (*exception_info).rip, rip + 2);
+    }
+    true
+}
+
+#[guest_function("ProbeReadableMSRs")]
+#[cfg(target_arch = "x86_64")]
+fn probe_readable_msrs(start: u32, count: u32) -> Vec<u8> {
+    const MAX_MSRS_PER_PROBE: u32 = 4096;
+    const GP_VECTOR: usize = 13;
+
+    assert!(count <= MAX_MSRS_PER_PROBE);
+    let mut readable = vec![0u8; count.div_ceil(8) as usize];
+    let handler = &hyperlight_guest_bin::exception::arch::HANDLERS[GP_VECTOR];
+
+    // Touch the atomics before installing the handler so it cannot page fault
+    // while handling an MSR fault on the exception stack.
+    MSR_PROBE_FAULTED.store(false, Ordering::Relaxed);
+    MSR_PROBE_ACTIVE.store(false, Ordering::Relaxed);
+    handler
+        .compare_exchange(
+            0,
+            msr_probe_gp_handler as *const () as usize as u64,
+            Ordering::AcqRel,
+            Ordering::Acquire,
+        )
+        .expect("general protection handler already installed");
+    MSR_PROBE_ACTIVE.store(true, Ordering::Release);
+
+    for offset in 0..count {
+        MSR_PROBE_FAULTED.store(false, Ordering::Release);
+        // SAFETY: The test guest runs at CPL0. The temporary #GP handler skips
+        // unsupported RDMSR instructions and records their indices.
+        unsafe {
+            core::arch::asm!(
+                "rdmsr",
+                in("ecx") start.wrapping_add(offset),
+                lateout("eax") _,
+                lateout("edx") _,
+                options(nostack)
+            );
+        }
+        if !MSR_PROBE_FAULTED.load(Ordering::Acquire) {
+            readable[offset as usize / 8] |= 1 << (offset % 8);
+        }
+    }
+
+    MSR_PROBE_ACTIVE.store(false, Ordering::Release);
+    handler.store(0, Ordering::Release);
+    readable
+}
+
 #[guest_function("IncrementActiveSsp")]
 #[cfg(target_arch = "x86_64")]
 fn increment_active_ssp() -> u64 {