fixup! snapshot: handle the scratch region correctly

Be slightly more careful about handling offsets from the guest in
snapshot.rs: previously, a maliciously-constructed sandbox could
possibly cause underflow panics, and another function could be misused
in a way that would result in a slice out-of-bounds panic (although
it was not currently being misused in such a way).

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

Author: syntactically

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -177,12 +177,14 @@ fn access_gpa<'a>(
     scratch: &'a ExclusiveSharedMemory,
     scratch_size: usize,
     gpa: u64,
-) -> (&'a ExclusiveSharedMemory, usize) {
+) -> Option<(&'a ExclusiveSharedMemory, usize)> {
     let scratch_base = scratch_base_gpa(scratch_size);
     if gpa >= scratch_base {
-        (scratch, (gpa - scratch_base) as usize)
+        Some((scratch, (gpa - scratch_base) as usize))
+    } else if gpa >= SandboxMemoryLayout::BASE_ADDRESS as u64 {
+        Some((snap, gpa as usize - SandboxMemoryLayout::BASE_ADDRESS))
     } else {
-        (snap, gpa as usize - SandboxMemoryLayout::BASE_ADDRESS)
+        None
     }
 }
 
@@ -222,8 +224,8 @@ impl<'a> hyperlight_common::vmem::TableOps for SharedMemoryPageTableBuffer<'a> {
         addr + offset
     }
     unsafe fn read_entry(&self, addr: u64) -> u64 {
-        let (mem, off) = access_gpa(self.snap, self.scratch, self.scratch_size, addr);
-        let Some(pte_bytes) = mem.as_slice().get(off..off + 8) else {
+        let memoff = access_gpa(self.snap, self.scratch, self.scratch_size, addr);
+        let Some(pte_bytes) = memoff.and_then(|(mem, off)| mem.as_slice().get(off..off + 8)) else {
             // Attacker-controlled data pointed out-of-bounds. We'll
             // default to returning 0 in this case, which, for most
             // architectures (including x86-64 and arm64, the ones we
@@ -303,9 +305,12 @@ unsafe fn guest_page<'a>(
     scratch_size: usize,
     gpa: u64,
 ) -> Option<&'a [u8]> {
+    if !gpa.is_multiple_of(PAGE_SIZE as u64) {
+        return None;
+    }
     let gpa_u = gpa as usize;
     for rgn in regions {
-        if gpa_u >= rgn.guest_region.start && gpa_u < rgn.guest_region.end {
+        if gpa_u >= rgn.guest_region.start && gpa_u + PAGE_SIZE <= rgn.guest_region.end {
             let off = gpa_u - rgn.guest_region.start;
             #[allow(clippy::useless_conversion)]
             let host_region_base: usize = rgn.host_region.start.into();
@@ -314,7 +319,7 @@ unsafe fn guest_page<'a>(
             });
         }
     }
-    let (mem, off) = access_gpa(snap, scratch, scratch_size, gpa);
+    let (mem, off) = access_gpa(snap, scratch, scratch_size, gpa)?;
     if off + PAGE_SIZE <= mem.as_slice().len() {
         Some(&mem.as_slice()[off..off + PAGE_SIZE])
     } else {