refactor: move deferred_hshm and counter_taken to UninitializedSandbox

ExclusiveSharedMemory's purpose is lock-free bulk memory operations.
Storing deferred_hshm (Arc<Mutex<Option<HostSharedMemory>>>) and
counter_taken (AtomicBool) on it added state that conceptually belongs
to the sandbox lifetime, not the memory region lifecycle.

This commit moves both fields to UninitializedSandbox, populates
deferred_hshm in evolve_impl_multi_use() after build() returns the
HostSharedMemory, and adds as_host_shared_memory() on
ExclusiveSharedMemory (test-only) so UninitializedSandbox::simulate_build()
can construct the HostSharedMemory without accessing private fields
across modules.

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

Author: danbugs

src/hyperlight_host/src/mem/shared_mem.rs

@@ -20,9 +20,7 @@ use std::io::Error;
 use std::mem::{align_of, size_of};
 #[cfg(target_os = "linux")]
 use std::ptr::null_mut;
-#[cfg(feature = "nanvix-unstable")]
-use std::sync::atomic::AtomicBool;
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::{Arc, RwLock};
 
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use tracing::{Span, instrument};
@@ -138,17 +136,7 @@ impl Drop for HostMapping {
 /// and taking snapshots.
 #[derive(Debug)]
 pub struct ExclusiveSharedMemory {
-    pub(crate) region: Arc<HostMapping>,
-    /// Populated by [`build()`](Self::build) with a [`HostSharedMemory`]
-    /// view of this region. Code that needs host-style volatile access
-    /// before `build()` (e.g. `GuestCounter`) can clone this `Arc` and
-    /// will see `Some` once `build()` completes.
-    pub(crate) deferred_hshm: Arc<Mutex<Option<HostSharedMemory>>>,
-    /// Set to `true` once a [`GuestCounter`] has been handed out via
-    /// [`UninitializedSandbox::guest_counter()`]. Prevents creating
-    /// multiple counters that would have divergent cached values.
-    #[cfg(feature = "nanvix-unstable")]
-    pub(crate) counter_taken: AtomicBool,
+    region: Arc<HostMapping>,
 }
 unsafe impl Send for ExclusiveSharedMemory {}
 
@@ -436,9 +424,6 @@ impl ExclusiveSharedMemory {
                 ptr: addr as *mut u8,
                 size: total_size,
             }),
-            deferred_hshm: Arc::new(Mutex::new(None)),
-            #[cfg(feature = "nanvix-unstable")]
-            counter_taken: AtomicBool::new(false),
         })
     }
 
@@ -557,9 +542,6 @@ impl ExclusiveSharedMemory {
                 size: total_size,
                 handle,
             }),
-            deferred_hshm: Arc::new(Mutex::new(None)),
-            #[cfg(feature = "nanvix-unstable")]
-            counter_taken: AtomicBool::new(false),
         })
     }
 
@@ -670,15 +652,6 @@ impl ExclusiveSharedMemory {
             region: self.region.clone(),
             lock: lock.clone(),
         };
-        // Publish the HostSharedMemory so any pre-existing GuestCounter
-        // can begin issuing volatile writes via the proper protocol.
-        // The mutex is only locked here and during GuestCounter
-        // operations; since build() consumes self, there is no
-        // concurrent access that could poison it.
-        #[allow(clippy::unwrap_used)]
-        {
-            *self.deferred_hshm.lock().unwrap() = Some(hshm.clone());
-        }
         (
             hshm,
             GuestSharedMemory {
@@ -688,24 +661,23 @@ impl ExclusiveSharedMemory {
         )
     }
 
-    /// Populate the deferred `HostSharedMemory` slot without consuming
-    /// `self`. Used in tests where `evolve()` / full `build()` is not
-    /// available.
-    #[cfg(all(test, feature = "nanvix-unstable"))]
-    pub(crate) fn simulate_build(&self) {
-        let lock = Arc::new(RwLock::new(()));
-        let hshm = HostSharedMemory {
-            region: self.region.clone(),
-            lock,
-        };
-        *self.deferred_hshm.lock().unwrap() = Some(hshm);
-    }
-
     /// Gets the file handle of the shared memory region for this Sandbox
     #[cfg(target_os = "windows")]
     pub fn get_mmap_file_handle(&self) -> HANDLE {
         self.region.handle
     }
+
+    /// Create a [`HostSharedMemory`] view of this region without
+    /// consuming `self`. Used in tests where the full `build()` /
+    /// `evolve()` pipeline is not available.
+    #[cfg(all(test, feature = "nanvix-unstable"))]
+    pub(crate) fn as_host_shared_memory(&self) -> HostSharedMemory {
+        let lock = Arc::new(RwLock::new(()));
+        HostSharedMemory {
+            region: self.region.clone(),
+            lock,
+        }
+    }
 }
 
 impl GuestSharedMemory {
@@ -888,9 +860,6 @@ impl SharedMemory for GuestSharedMemory {
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
         let mut excl = ExclusiveSharedMemory {
             region: self.region.clone(),
-            deferred_hshm: Arc::new(Mutex::new(None)),
-            #[cfg(feature = "nanvix-unstable")]
-            counter_taken: AtomicBool::new(false),
         };
         let ret = f(&mut excl);
         drop(excl);
@@ -1256,9 +1225,6 @@ impl SharedMemory for HostSharedMemory {
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
         let mut excl = ExclusiveSharedMemory {
             region: self.region.clone(),
-            deferred_hshm: Arc::new(Mutex::new(None)),
-            #[cfg(feature = "nanvix-unstable")]
-            counter_taken: AtomicBool::new(false),
         };
         let ret = f(&mut excl);
         drop(excl);

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -33,6 +33,8 @@ use crate::mem::memory_region::{DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegionFlags}
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::ExclusiveSharedMemory;
 #[cfg(feature = "nanvix-unstable")]
+use crate::mem::shared_mem::HostSharedMemory;
+#[cfg(feature = "nanvix-unstable")]
 use crate::mem::shared_mem::SharedMemory;
 use crate::sandbox::SandboxConfiguration;
 use crate::{MultiUseSandbox, Result, new_error};
@@ -70,11 +72,9 @@ pub(crate) struct SandboxRuntimeConfig {
 /// `decrement()` take `&self` rather than `&mut self`.
 ///
 /// The counter holds an `Arc<Mutex<Option<HostSharedMemory>>>` that is
-/// shared with [`ExclusiveSharedMemory`]. The `Option` is `None` until
-/// [`build()`](ExclusiveSharedMemory::build) populates it, at which point
-/// the counter can issue volatile writes via the proper protocol. This
-/// avoids adding lock overhead to `ExclusiveSharedMemory` during the
-/// exclusive setup phase.
+/// shared with [`UninitializedSandbox`]. The `Option` is `None` until
+/// [`evolve()`](UninitializedSandbox::evolve) populates it, at which point
+/// the counter can issue volatile writes via the proper protocol.
 ///
 /// Only one `GuestCounter` may be created per sandbox; a second call to
 /// [`UninitializedSandbox::guest_counter()`] returns an error.
@@ -85,7 +85,7 @@ pub struct GuestCounter {
 
 #[cfg(feature = "nanvix-unstable")]
 struct GuestCounterInner {
-    deferred_hshm: Arc<Mutex<Option<crate::mem::shared_mem::HostSharedMemory>>>,
+    deferred_hshm: Arc<Mutex<Option<HostSharedMemory>>>,
     offset: usize,
     value: u64,
 }
@@ -111,11 +111,12 @@ impl GuestCounter {
                 })?
                 .clone()
         };
-        inner.value = inner
+        let new_value = inner
             .value
             .checked_add(1)
             .ok_or_else(|| new_error!("GuestCounter overflow"))?;
-        shm.write::<u64>(inner.offset, inner.value)?;
+        shm.write::<u64>(inner.offset, new_value)?;
+        inner.value = new_value;
         Ok(())
     }
 
@@ -131,11 +132,12 @@ impl GuestCounter {
                 })?
                 .clone()
         };
-        inner.value = inner
+        let new_value = inner
             .value
             .checked_sub(1)
             .ok_or_else(|| new_error!("GuestCounter underflow"))?;
-        shm.write::<u64>(inner.offset, inner.value)?;
+        shm.write::<u64>(inner.offset, new_value)?;
+        inner.value = new_value;
         Ok(())
     }
 
@@ -170,6 +172,17 @@ pub struct UninitializedSandbox {
     // This is needed to convey the stack pointer between the snapshot
     // and the HyperlightVm creation
     pub(crate) stack_top_gva: u64,
+    /// Populated by [`evolve()`](Self::evolve) with a [`HostSharedMemory`]
+    /// view of scratch memory. Code that needs host-style volatile access
+    /// before `evolve()` (e.g. `GuestCounter`) can clone this `Arc` and
+    /// will see `Some` once `evolve()` completes.
+    #[cfg(feature = "nanvix-unstable")]
+    pub(crate) deferred_hshm: Arc<Mutex<Option<HostSharedMemory>>>,
+    /// Set to `true` once a [`GuestCounter`] has been handed out via
+    /// [`guest_counter()`](Self::guest_counter). Prevents creating
+    /// multiple counters that would have divergent cached values.
+    #[cfg(feature = "nanvix-unstable")]
+    counter_taken: std::sync::atomic::AtomicBool,
 }
 
 impl Debug for UninitializedSandbox {
@@ -267,10 +280,9 @@ impl UninitializedSandbox {
     /// the top of scratch memory, so both host and guest can locate it
     /// without an explicit GPA parameter.
     ///
-    /// The returned counter holds an `Arc` clone of the scratch memory's
+    /// The returned counter holds an `Arc` clone of the sandbox's
     /// `deferred_hshm`, so it will automatically gain access to the
-    /// [`HostSharedMemory`] once [`build()`](ExclusiveSharedMemory::build)
-    /// is called during [`evolve()`](Self::evolve).
+    /// [`HostSharedMemory`] once [`evolve()`](Self::evolve) completes.
     ///
     /// This method can only be called once; a second call returns an error
     /// because multiple counters would have divergent cached values.
@@ -280,12 +292,7 @@ impl UninitializedSandbox {
 
         use hyperlight_common::layout::SCRATCH_TOP_GUEST_COUNTER_OFFSET;
 
-        if self
-            .mgr
-            .scratch_mem
-            .counter_taken
-            .swap(true, Ordering::Relaxed)
-        {
+        if self.counter_taken.swap(true, Ordering::Relaxed) {
             return Err(new_error!(
                 "GuestCounter has already been created for this sandbox"
             ));
@@ -301,7 +308,7 @@ impl UninitializedSandbox {
         }
 
         let offset = scratch_size - SCRATCH_TOP_GUEST_COUNTER_OFFSET as usize;
-        let deferred_hshm = self.mgr.scratch_mem.deferred_hshm.clone();
+        let deferred_hshm = self.deferred_hshm.clone();
 
         Ok(GuestCounter {
             inner: Mutex::new(GuestCounterInner {
@@ -370,6 +377,10 @@ impl UninitializedSandbox {
             rt_cfg,
             load_info: snapshot.load_info(),
             stack_top_gva: snapshot.stack_top_gva(),
+            #[cfg(feature = "nanvix-unstable")]
+            deferred_hshm: Arc::new(Mutex::new(None)),
+            #[cfg(feature = "nanvix-unstable")]
+            counter_taken: std::sync::atomic::AtomicBool::new(false),
         };
 
         // If we were passed a writer for host print register it otherwise use the default.
@@ -449,6 +460,18 @@ impl UninitializedSandbox {
     ) -> Result<()> {
         self.register("HostPrint", print_func)
     }
+
+    /// Populate the deferred `HostSharedMemory` slot without running
+    /// the full `evolve()` pipeline. Used in tests where guest boot
+    /// is not available.
+    #[cfg(all(test, feature = "nanvix-unstable"))]
+    fn simulate_build(&self) {
+        let hshm = self.mgr.scratch_mem.as_host_shared_memory();
+        #[allow(clippy::unwrap_used)]
+        {
+            *self.deferred_hshm.lock().unwrap() = Some(hshm);
+        }
+    }
 }
 // Check to see if the current version of Windows is supported
 // Hyperlight is only supported on Windows 11 and Windows Server 2022 and later
@@ -1500,7 +1523,7 @@ mod tests {
         fn increment_decrement() {
             let mut sandbox = make_sandbox();
             let counter = sandbox.guest_counter().unwrap();
-            sandbox.mgr.scratch_mem.simulate_build();
+            sandbox.simulate_build();
 
             counter.increment().unwrap();
             assert_eq!(counter.value().unwrap(), 1);
@@ -1514,7 +1537,7 @@ mod tests {
         fn underflow_returns_error() {
             let mut sandbox = make_sandbox();
             let counter = sandbox.guest_counter().unwrap();
-            sandbox.mgr.scratch_mem.simulate_build();
+            sandbox.simulate_build();
 
             assert_eq!(counter.value().unwrap(), 0);
             let result = counter.decrement();

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -38,6 +38,21 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox};
 #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
 pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<MultiUseSandbox> {
     let (mut hshm, gshm) = u_sbox.mgr.build()?;
+
+    // Publish the HostSharedMemory for scratch so any pre-existing
+    // GuestCounter can begin issuing volatile writes.
+    #[cfg(feature = "nanvix-unstable")]
+    {
+        #[allow(clippy::unwrap_used)]
+        // The mutex can only be poisoned if a previous lock holder
+        // panicked.  Since we are the only writer at this point (the
+        // GuestCounter only reads inside `increment`/`decrement`),
+        // poisoning cannot happen.
+        {
+            *u_sbox.deferred_hshm.lock().unwrap() = Some(hshm.scratch_mem.clone());
+        }
+    }
+
     let mut vm = set_up_hypervisor_partition(
         gshm,
         &u_sbox.config,