Explicitly specify the physical addresses of snapshot page tables

Previously, the guest tried to infer the addresses of the snapshot
page tables from the values of cr3 that it saw at various times, but
this was brittle and irritating to do without writing to the wrong
address at times.  This replaces that mechanism with an extra word in
the metadata portion of the scratch region that explicitly provides
this information.

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

Author: syntactically

src/hyperlight_common/src/layout.rs

@@ -22,8 +22,8 @@ pub use arch::{MAX_GPA, MAX_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN};
 
 // offsets down from the top of scratch memory for various things
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
-pub const SCRATCH_TOP_USED_OFFSET: u64 = 0x10;
-pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x18;
+pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x10;
+pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = 0x18;
 pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x20;
 
 pub fn scratch_base_gpa(size: usize) -> u64 {

src/hyperlight_guest/src/layout.rs

@@ -27,4 +27,8 @@ pub fn allocator_gva() -> *mut u64 {
     use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_ALLOCATOR_OFFSET};
     (MAX_GVA as u64 - SCRATCH_TOP_ALLOCATOR_OFFSET + 1) as *mut u64
 }
+pub fn snapshot_pt_gpa_base_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET + 1) as *mut u64
+}
 pub use arch::{scratch_base_gpa, scratch_base_gva};

src/hyperlight_guest_bin/src/paging.rs

@@ -28,11 +28,6 @@ use tracing::{Span, instrument};
 //       virtual address 0, and Rust raw pointer operations can't be
 //       used to read/write from address 0.
 
-// We get this out of CR3 the first time that we do any mapping
-// operation. In the future, if snapshot/restore changes to be able to
-// change the snapshot pt base, we will need to modify this.
-static SNAPSHOT_PT_GPA: spin::Once<u64> = spin::Once::new();
-
 #[derive(Copy, Clone)]
 struct GuestMappingOperations {
     snapshot_pt_base_gpa: u64,
@@ -43,13 +38,9 @@ struct GuestMappingOperations {
 impl GuestMappingOperations {
     fn new() -> Self {
         Self {
-            snapshot_pt_base_gpa: *SNAPSHOT_PT_GPA.call_once(|| {
-                let snapshot_pt_base_gpa: u64;
-                unsafe {
-                    asm!("mov {}, cr3", out(reg) snapshot_pt_base_gpa);
-                };
-                snapshot_pt_base_gpa
-            }),
+            snapshot_pt_base_gpa: unsafe {
+                hyperlight_guest::layout::snapshot_pt_gpa_base_gva().read_volatile()
+            },
             snapshot_pt_base_gva: hyperlight_common::layout::SNAPSHOT_PT_GVA_MIN as u64,
             scratch_base_gpa: hyperlight_guest::layout::scratch_base_gpa(),
             scratch_base_gva: hyperlight_guest::layout::scratch_base_gva(),

src/hyperlight_host/src/mem/mgr.rs

@@ -272,7 +272,9 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             stack_cookie: self.stack_cookie,
             abort_buffer: Vec::new(), // Guest doesn't need abort buffer
         };
-        host_mgr.update_scratch_bookkeeping()?;
+        host_mgr.update_scratch_bookkeeping(
+            (SandboxMemoryLayout::BASE_ADDRESS + self.layout.get_pt_offset()) as u64,
+        )?;
         Ok((host_mgr, guest_mgr))
     }
 }
@@ -431,11 +433,11 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
             Some(gscratch)
         };
-        self.update_scratch_bookkeeping()?;
+        self.update_scratch_bookkeeping(snapshot.root_pt_gpa())?;
         Ok((gsnapshot, gscratch))
     }
 
-    fn update_scratch_bookkeeping(&mut self) -> Result<()> {
+    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) -> Result<()> {
         let scratch_size = self.scratch_mem.mem_size();
 
         let size_offset =
@@ -450,6 +452,11 @@ impl SandboxMemoryManager<HostSharedMemory> {
             hyperlight_common::layout::scratch_base_gpa(scratch_size),
         )?;
 
+        let snapshot_pt_base_gpa_offset = scratch_size
+            - hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET as usize;
+        self.scratch_mem
+            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)?;
+
         Ok(())
     }
 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -1262,7 +1262,6 @@ mod tests {
 
     #[test]
     #[cfg(target_os = "linux")]
-    #[ignore] // this test will be re-enabled in the next commit
     fn snapshot_restore_handles_remapping_correctly() {
         let mut sbox: MultiUseSandbox = {
             let path = simple_guest_as_string().unwrap();