Review updates

Signed-off-by: Simon Davies <simongdavies@users.noreply.github.com>

Author: simongdavies

src/hyperlight_host/src/mem/layout.rs

@@ -478,7 +478,7 @@ impl SandboxMemoryLayout {
         self.pt_offset
     }
 
-    /// Get the offset into the snapshot region of the page tables
+    /// Sets the size of the memory region used for page tables
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     #[cfg(feature = "init-paging")]
     pub(crate) fn set_pt_size(&mut self, size: usize) {
@@ -668,8 +668,8 @@ impl SandboxMemoryLayout {
             if after_init_offset != expected_pt_offset {
                 return Err(new_error!(
                     "Page table offset does not match expected:  {}, actual:  {}",
-                    expected_init_data_offset,
-                    init_data_offset
+                    expected_pt_offset,
+                    after_init_offset
                 ));
             }
 

src/hyperlight_host/src/mem/mgr.rs

@@ -127,7 +127,7 @@ impl GuestPageTableBuffer {
     }
 
     pub(crate) fn size(&self) -> usize {
-        self.buffer.borrow().len() * PAGE_TABLE_SIZE
+        self.buffer.borrow().len()
     }
 
     pub(crate) fn into_bytes(self) -> Box<[u8]> {
@@ -409,8 +409,6 @@ impl SandboxMemoryManager<HostSharedMemory> {
     }
 }
 
-// ...existing code...
-
 #[cfg(test)]
 #[cfg(all(feature = "init-paging", target_arch = "x86_64"))]
 mod tests {

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -62,14 +62,19 @@ pub struct Snapshot {
     /// require constant-time equality checking
     hash: [u8; 32],
 
-    /// TODO: this should not necessarily be around in the long term...
+    /// Preinitialisation entry point for snapshots created directly from a
+    /// guest binary.
     ///
-    /// When creating a snapshot directly from a guest binary, this
-    /// tracks the address that we need to call into before actually
-    /// using a sandbox from this snapshot in order to do
-    /// preinitialisation. Ideally we would either not need to do this
-    /// at all, or do it as part of the snapshot creation process and
-    /// never need this.
+    /// When creating a snapshot directly from a guest binary, this tracks
+    /// the address that we need to call into before actually using a
+    /// sandbox from this snapshot in order to perform guest-side
+    /// preinitialisation.
+    ///
+    /// Long-term, the intention is to run this preinitialisation eagerly as
+    /// part of the snapshot creation process so that restored sandboxes can
+    /// begin executing from their normal entry point without requiring this
+    /// field. Until that refactoring happens, this remains part of the
+    /// snapshot format and must be preserved.
     preinitialise: Option<u64>,
 }
 
@@ -105,7 +110,8 @@ fn hash(memory: &[u8], regions: &[MemoryRegion]) -> Result<[u8; 32]> {
 }
 
 impl Snapshot {
-    /// Create a new snapshot that runs the guest binary identified by env
+    /// Create a new snapshot from the guest binary identified by `env`. With the configuration 
+    /// specified in `cfg`.
     pub(crate) fn from_env<'a, 'b>(
         env: impl Into<GuestEnvironment<'a, 'b>>,
         cfg: SandboxConfiguration,

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -88,6 +88,8 @@ pub enum GuestBinary<'a> {
 impl<'a> GuestBinary<'a> {
     /// If the guest binary is identified by a file, canonicalise the path
     ///
+    /// For [`GuestBinary::FilePath`], this resolves the path to its canonical
+    /// form. For [`GuestBinary::Buffer`], this method is a no-op.
     /// TODO: Maybe we should make the GuestEnvironment or
     ///       GuestBinary constructors crate-private and turn this
     ///       into an invariant on one of those types.
@@ -161,7 +163,7 @@ impl UninitializedSandbox {
     // that can be changed (from the original snapshot) is the configuration defines the behaviour of
     // `InterruptHandler` on Linux.
     //
-    // This is ok for now as this is not a public
+    // This is ok for now as this is not a public function
     fn from_snapshot(
         snapshot: Arc<Snapshot>,
         cfg: Option<SandboxConfiguration>,

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -120,11 +120,11 @@ pub(crate) fn set_up_hypervisor_partition(
     };
     let entrypoint_ptr = mgr
         .entrypoint_offset
-        .map(|x| {
+        .ok_or_else(|| new_error!("Entrypoint offset is None")) 
+        .and_then(|x| {
             let entrypoint_total_offset = mgr.load_addr.clone() + x;
             GuestPtr::try_from(entrypoint_total_offset)
-        })
-        .ok_or(new_error!("Failed to create entrypoint pointer"))??;
+        })?;
 
     // Create gdb thread if gdb is enabled and the configuration is provided
     #[cfg(gdb)]