Expose how much scratch a sandbox has allocated

Signed-off-by: danbugs <danilochiarlone@gmail.com>

Author: danbugs

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 * `Sandbox::snapshot_with_options` takes a snapshot with `SnapshotOptions`. `SnapshotOptions::scratch_size` gives the sandbox, and later restores of the snapshot, a different scratch size. A guest can boot with the scratch its initialisation needs and run guest function calls from a snapshot with only the scratch they need, so each restore costs less. It can also give a guest more scratch than it booted with.
+* `Sandbox::scratch_allocated_bytes` returns how much scratch the guest has allocated since the sandbox was last restored or snapshotted. The guest controls the value it is read from.
 
 ### Changed
 * Core dumps no longer record the guest binary path. The process name and command line are always `<unknown>`, as they already were for sandboxes created from snapshots by @ludfjig in https://github.com/hyperlight-dev/hyperlight/pull/1975

src/hyperlight_host/src/mem/mgr.rs

@@ -1005,6 +1005,14 @@ impl SandboxMemoryManager<HostSharedMemory> {
         Ok(end.min(self.layout.get_transport_arena().base_addr()))
     }
 
+    /// Bytes of scratch the guest has allocated since the allocator was
+    /// last reset, at build, restore or snapshot.
+    pub(crate) fn scratch_allocated_bytes(&self) -> Result<u64> {
+        Ok(self
+            .used_scratch_end_gpa()?
+            .saturating_sub(self.first_free_scratch_gpa()))
+    }
+
     /// Zeroes the scratch pages freed by resetting the allocator, up to
     /// `end`, the allocator's end as `used_scratch_end_gpa` read it.
     fn zero_freed_scratch(&mut self, end: u64) -> Result<()> {

src/hyperlight_host/src/sandbox/initialized.rs

@@ -1279,6 +1279,35 @@ impl Sandbox {
     pub fn status(&self) -> SandboxStatus {
         self.status
     }
+
+    /// Returns how many bytes of scratch memory the guest has allocated
+    /// since the sandbox was last restored or snapshotted.
+    /// Before the first restore or snapshot, this includes allocations
+    /// made during guest initialisation.
+    ///
+    /// Host-reserved scratch, such as the transport and exception stack,
+    /// is excluded.
+    ///
+    /// The guest controls this value. A malicious guest can make it
+    /// report less or more scratch than it actually allocated.
+    ///
+    /// ```no_run
+    /// # use hyperlight_host::SandboxBuilder;
+    /// # fn example() -> Result<(), Box<dyn std::error::Error>> {
+    /// let mut sandbox = SandboxBuilder::from_file("guest.bin").build()?;
+    /// let snapshot = sandbox.snapshot()?;
+    /// sandbox.call::<()>("WriteSomething", ())?;
+    /// println!(
+    ///     "guest-reported scratch allocation: {} bytes",
+    ///     sandbox.scratch_allocated_bytes()?
+    /// );
+    /// sandbox.restore(snapshot)?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn scratch_allocated_bytes(&self) -> Result<u64> {
+        self.mem_mgr.scratch_allocated_bytes()
+    }
 }
 
 impl Callable for Sandbox {
@@ -1673,6 +1702,35 @@ mod tests {
         assert_eq!(echo, "recovered");
     }
 
+    /// The scratch a guest has allocated covers every page it wrote, grows
+    /// as it writes more, and starts again from zero at a snapshot and at a
+    /// restore.
+    #[test]
+    fn scratch_allocated_bytes_counts_written_pages() {
+        let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
+            .heap_size(8 << 20)
+            .scratch_size(16 << 20)
+            .build()
+            .unwrap();
+        let snapshot = sbox.snapshot().unwrap();
+        assert_eq!(sbox.scratch_allocated_bytes().unwrap(), 0);
+
+        sbox.call::<()>("AllocAndWritePattern", 1u64 << 20).unwrap();
+        let one = sbox.scratch_allocated_bytes().unwrap();
+        assert!(one >= 1 << 20, "{one}");
+        // The new pattern is built before the old one is dropped.
+        sbox.call::<()>("AllocAndWritePattern", 2u64 << 20).unwrap();
+        let three = sbox.scratch_allocated_bytes().unwrap();
+        assert!(three >= 3 << 20, "{three}");
+
+        sbox.restore(snapshot).unwrap();
+        assert_eq!(sbox.scratch_allocated_bytes().unwrap(), 0);
+        sbox.call::<()>("AllocAndWritePattern", 1u64 << 20).unwrap();
+        assert_eq!(sbox.scratch_allocated_bytes().unwrap(), one);
+        sbox.snapshot().unwrap();
+        assert_eq!(sbox.scratch_allocated_bytes().unwrap(), 0);
+    }
+
     /// A resized capture installs like any capture. The sandbox keeps its
     /// state, folds mapped regions into the snapshot, and caches it.
     #[test]
@@ -3121,6 +3179,15 @@ mod tests {
                     allocator,
                 )
                 .unwrap();
+            // The getter holds a corrupt allocator to the scratch the guest
+            // can allocate.
+            let arena = sbox.mem_mgr.layout.get_transport_arena().base_addr();
+            let allocated = if allocator == 0 {
+                0
+            } else {
+                arena - first_free
+            };
+            assert_eq!(sbox.scratch_allocated_bytes().unwrap(), allocated);
 
             // `map_region` drops the cached snapshot without running the guest.
             let map_mem = allocate_guest_memory();