feat: add snapshot generation counter to scratch bookkeeping

Signed-off-by: Tomasz Andrzejak <andreiltd@gmail.com>

Author: andreiltd

src/hyperlight_common/src/layout.rs

@@ -37,7 +37,8 @@ pub use arch::{SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN};
 pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08;
 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 const SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET: u64 = 0x20;
+pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x30;
 
 /// Offset from the top of scratch memory for a shared host-guest u64 counter.
 ///

src/hyperlight_guest/src/layout.rs

@@ -32,6 +32,10 @@ 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 fn snapshot_generation_gva() -> *mut u64 {
+    use hyperlight_common::layout::{MAX_GVA, SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET};
+    (MAX_GVA as u64 - SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET + 1) as *mut u64
+}
 pub use arch::{scratch_base_gpa, scratch_base_gva};
 
 /// Returns a pointer to the guest counter u64 in scratch memory.

src/hyperlight_host/src/mem/mgr.rs

@@ -148,6 +148,9 @@ pub(crate) struct SandboxMemoryManager<S: SharedMemory> {
     pub(crate) mapped_rgns: u64,
     /// Buffer for accumulating guest abort messages
     pub(crate) abort_buffer: Vec<u8>,
+    /// Snapshot restore generation counter. 0 means no restore
+    /// and is incremented on each `restore_snapshot` call.
+    pub(crate) restore_count: u64,
 }
 
 /// Buffer for building guest page tables during snapshot creation.
@@ -274,6 +277,7 @@ where
             entrypoint,
             mapped_rgns: 0,
             abort_buffer: Vec::new(),
+            restore_count: 0,
         }
     }
 
@@ -341,6 +345,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             entrypoint: self.entrypoint,
             mapped_rgns: self.mapped_rgns,
             abort_buffer: self.abort_buffer,
+            restore_count: self.restore_count,
         };
         let guest_mgr = SandboxMemoryManager {
             shared_mem: gshm,
@@ -349,6 +354,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             entrypoint: self.entrypoint,
             mapped_rgns: self.mapped_rgns,
             abort_buffer: Vec::new(), // Guest doesn't need abort buffer
+            restore_count: self.restore_count,
         };
         host_mgr.update_scratch_bookkeeping()?;
         host_mgr.copy_pt_to_scratch()?;
@@ -541,6 +547,8 @@ impl SandboxMemoryManager<HostSharedMemory> {
             Some(gscratch)
         };
         self.layout = *snapshot.layout();
+        self.restore_count += 1;
+
         self.update_scratch_bookkeeping()?;
         self.copy_pt_to_scratch()?;
         Ok((gsnapshot, gscratch))
@@ -565,6 +573,10 @@ impl SandboxMemoryManager<HostSharedMemory> {
             SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET,
             self.layout.get_pt_base_gpa(),
         )?;
+        self.update_scratch_bookkeeping_item(
+            SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET,
+            self.restore_count,
+        )?;
 
         // Initialise the guest input and output data buffers in
         // scratch memory. TODO: remove the need for this.