Make snapshot region aware

 - Snapshot now contains the memory region that were mapped
   at the time of snapshot
 - On restore, unmap regions that were not mapped at
   time of snapshot. Map regions that were
 - Add simple test

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/mem/mgr.rs

@@ -259,26 +259,25 @@ where
         }
     }
 
-    pub(crate) fn snapshot(&mut self) -> Result<SharedMemorySnapshot> {
-        SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)
+    /// Create a snapshot with the given mapped regions
+    pub(crate) fn snapshot(
+        &mut self,
+        mapped_regions: Vec<MemoryRegion>,
+    ) -> Result<SharedMemorySnapshot> {
+        SharedMemorySnapshot::new(&mut self.shared_mem, mapped_regions)
     }
 
     /// This function restores a memory snapshot from a given snapshot.
-    ///
-    /// Returns the number of memory regions mapped into the sandbox
-    /// that need to be unmapped in order for the restore to be
-    /// completed.
-    pub(crate) fn restore_snapshot(&mut self, snapshot: &SharedMemorySnapshot) -> Result<u64> {
+    pub(crate) fn restore_snapshot(&mut self, snapshot: &SharedMemorySnapshot) -> Result<()> {
         if self.shared_mem.mem_size() != snapshot.mem_size() {
             return Err(new_error!(
                 "Snapshot size does not match current memory size: {} != {}",
                 self.shared_mem.raw_mem_size(),
                 snapshot.mem_size()
             ));
         }
-        let old_rgns = self.mapped_rgns;
-        self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
-        Ok(old_rgns - self.mapped_rgns)
+        snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+        Ok(())
     }
 
     /// Sets `addr` to the correct offset in the memory referenced by

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -16,6 +16,7 @@ limitations under the License.
 
 use tracing::{Span, instrument};
 
+use super::memory_region::MemoryRegion;
 use super::shared_mem::SharedMemory;
 use crate::Result;
 
@@ -24,21 +25,21 @@ use crate::Result;
 #[derive(Clone)]
 pub(crate) struct SharedMemorySnapshot {
     snapshot: Vec<u8>,
-    /// How many non-main-RAM regions were mapped when this snapshot was taken?
-    mapped_rgns: u64,
+    /// The memory regions that were mapped when this snapshot was taken (excluding initial sandbox regions)
+    regions: Vec<MemoryRegion>,
 }
 
 impl SharedMemorySnapshot {
     /// Take a snapshot of the memory in `shared_mem`, then create a new
     /// instance of `Self` with the snapshot stored therein.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(super) fn new<S: SharedMemory>(shared_mem: &mut S, mapped_rgns: u64) -> Result<Self> {
+    pub(super) fn new<S: SharedMemory>(
+        shared_mem: &mut S,
+        regions: Vec<MemoryRegion>,
+    ) -> Result<Self> {
         // TODO: Track dirty pages instead of copying entire memory
         let snapshot = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
-        Ok(Self {
-            snapshot,
-            mapped_rgns,
-        })
+        Ok(Self { snapshot, regions })
     }
 
     /// Take another snapshot of the internally-stored `SharedMemory`,
@@ -51,11 +52,16 @@ impl SharedMemorySnapshot {
     }
 
     /// Copy the memory from the internally-stored memory snapshot
-    /// into the internally-stored `SharedMemory`
+    /// into the internally-stored `SharedMemory`.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(super) fn restore_from_snapshot<S: SharedMemory>(&self, shared_mem: &mut S) -> Result<u64> {
+    pub(super) fn restore_from_snapshot<S: SharedMemory>(&self, shared_mem: &mut S) -> Result<()> {
         shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))??;
-        Ok(self.mapped_rgns)
+        Ok(())
+    }
+
+    /// Get the mapped regions from this snapshot
+    pub(crate) fn regions(&self) -> &[MemoryRegion] {
+        &self.regions
     }
 
     /// Return the size of the snapshot in bytes.
@@ -78,7 +84,7 @@ mod tests {
         let data2 = data1.iter().map(|b| b + 1).collect::<Vec<u8>>();
         let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
         gm.copy_from_slice(data1.as_slice(), 0).unwrap();
-        let mut snap = super::SharedMemorySnapshot::new(&mut gm, 0).unwrap();
+        let mut snap = super::SharedMemorySnapshot::new(&mut gm, Vec::new()).unwrap();
         {
             // after the first snapshot is taken, make sure gm has the equivalent
             // of data1

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::collections::HashSet;
 #[cfg(unix)]
 use std::os::fd::AsRawFd;
 #[cfg(unix)]
@@ -95,18 +96,35 @@ impl MultiUseSandbox {
     /// Create a snapshot of the current state of the sandbox's memory.
     #[instrument(err(Debug), skip_all, parent = Span::current())]
     pub fn snapshot(&mut self) -> Result<Snapshot> {
-        let snapshot = self.mem_mgr.unwrap_mgr_mut().snapshot()?;
-        Ok(Snapshot { inner: snapshot })
+        let mapped_regions_iter = self.vm.get_mapped_regions();
+        let mapped_regions_vec: Vec<MemoryRegion> = mapped_regions_iter.cloned().collect();
+        let memory_snapshot = self.mem_mgr.unwrap_mgr_mut().snapshot(mapped_regions_vec)?;
+        Ok(Snapshot {
+            inner: memory_snapshot,
+        })
     }
 
     /// Restore the sandbox's memory to the state captured in the given snapshot.
     #[instrument(err(Debug), skip_all, parent = Span::current())]
     pub fn restore(&mut self, snapshot: &Snapshot) -> Result<()> {
-        let rgns_to_unmap = self
-            .mem_mgr
+        self.mem_mgr
             .unwrap_mgr_mut()
             .restore_snapshot(&snapshot.inner)?;
-        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
+
+        let current_regions: HashSet<_> = self.vm.get_mapped_regions().cloned().collect();
+        let snapshot_regions: HashSet<_> = snapshot.inner.regions().iter().cloned().collect();
+
+        let regions_to_unmap = current_regions.difference(&snapshot_regions);
+        let regions_to_map = snapshot_regions.difference(&current_regions);
+
+        for region in regions_to_unmap {
+            unsafe { self.vm.unmap_region(region)? };
+        }
+
+        for region in regions_to_map {
+            unsafe { self.vm.map_region(region)? };
+        }
+
         Ok(())
     }
 
@@ -645,4 +663,57 @@ mod tests {
             region_type: MemoryRegionType::Heap,
         }
     }
+
+    #[cfg(target_os = "linux")]
+    fn allocate_guest_memory() -> GuestSharedMemory {
+        page_aligned_memory(b"test data for snapshot")
+    }
+
+    #[test]
+    #[cfg(target_os = "linux")]
+    fn snapshot_restore_handles_remapping_correctly() {
+        let mut sbox: MultiUseSandbox = {
+            let path = simple_guest_as_string().unwrap();
+            let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap();
+            u_sbox.evolve().unwrap()
+        };
+
+        // 1. Take snapshot 1 with no additional regions mapped
+        let snapshot1 = sbox.snapshot().unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 0);
+
+        // 2. Map a memory region
+        let map_mem = allocate_guest_memory();
+        let guest_base = 0x200000000_usize;
+        let region = region_for_memory(&map_mem, guest_base);
+
+        unsafe { sbox.map_region(&region).unwrap() };
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 1);
+
+        // 3. Take snapshot 2 with 1 region mapped
+        let snapshot2 = sbox.snapshot().unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 1);
+
+        // 4. Restore to snapshot 1 (should unmap the region)
+        sbox.restore(&snapshot1).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 0);
+
+        // 5. Restore forward to snapshot 2 (should remap the region)
+        sbox.restore(&snapshot2).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 1);
+
+        // Verify the region is the same
+        let mut restored_regions = sbox.vm.get_mapped_regions();
+        assert_eq!(*restored_regions.next().unwrap(), region);
+        assert!(restored_regions.next().is_none());
+        drop(restored_regions);
+
+        // 6. Try map the region again (should fail since already mapped)
+        let err = unsafe { sbox.map_region(&region) };
+        assert!(
+            err.is_err(),
+            "Expected error when remapping existing region: {:?}",
+            err
+        );
+    }
 }