refactor: enforce single GuestCounter, harden encapsulation

- Add counter_taken AtomicBool to ExclusiveSharedMemory; guest_counter()
  returns an error on a second call instead of silently creating a
  duplicate with divergent cached state.
- Initialize value to 0 instead of reading from the offset — scratch
  memory is always zero at that point.
- Revert HostSharedMemory fields to private and add a #[cfg(test)]
  simulate_build() method so tests don't leak shm internals.
- Add only_one_counter_allowed test.

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

Author: danbugs

src/hyperlight_host/src/mem/shared_mem.rs

@@ -20,6 +20,8 @@ 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 hyperlight_common::mem::PAGE_SIZE_USIZE;
@@ -142,6 +144,11 @@ pub struct ExclusiveSharedMemory {
     /// 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,
 }
 unsafe impl Send for ExclusiveSharedMemory {}
 
@@ -326,8 +333,8 @@ unsafe impl Send for GuestSharedMemory {}
 /// becoming completely divorced from the view of the VM.
 #[derive(Clone, Debug)]
 pub struct HostSharedMemory {
-    pub(crate) region: Arc<HostMapping>,
-    pub(crate) lock: Arc<RwLock<()>>,
+    region: Arc<HostMapping>,
+    lock: Arc<RwLock<()>>,
 }
 unsafe impl Send for HostSharedMemory {}
 
@@ -430,6 +437,8 @@ impl ExclusiveSharedMemory {
                 size: total_size,
             }),
             deferred_hshm: Arc::new(Mutex::new(None)),
+            #[cfg(feature = "nanvix-unstable")]
+            counter_taken: AtomicBool::new(false),
         })
     }
 
@@ -549,6 +558,8 @@ impl ExclusiveSharedMemory {
                 handle,
             }),
             deferred_hshm: Arc::new(Mutex::new(None)),
+            #[cfg(feature = "nanvix-unstable")]
+            counter_taken: AtomicBool::new(false),
         })
     }
 
@@ -661,7 +672,13 @@ impl ExclusiveSharedMemory {
         };
         // Publish the HostSharedMemory so any pre-existing GuestCounter
         // can begin issuing volatile writes via the proper protocol.
-        *self.deferred_hshm.lock().unwrap() = Some(hshm.clone());
+        // 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 {
@@ -671,6 +688,19 @@ 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 {
@@ -859,6 +889,8 @@ impl SharedMemory for GuestSharedMemory {
         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);
@@ -1225,6 +1257,8 @@ impl SharedMemory for HostSharedMemory {
         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

@@ -76,8 +76,8 @@ pub(crate) struct SandboxRuntimeConfig {
 /// avoids adding lock overhead to `ExclusiveSharedMemory` during the
 /// exclusive setup phase.
 ///
-/// The constructor takes `&mut self` on `UninitializedSandbox`, which
-/// prevents creating multiple instances simultaneously.
+/// Only one `GuestCounter` may be created per sandbox; a second call to
+/// [`UninitializedSandbox::guest_counter()`] returns an error.
 #[cfg(feature = "nanvix-unstable")]
 pub struct GuestCounter {
     inner: Mutex<GuestCounterInner>,
@@ -272,13 +272,25 @@ impl UninitializedSandbox {
     /// [`HostSharedMemory`] once [`build()`](ExclusiveSharedMemory::build)
     /// is called during [`evolve()`](Self::evolve).
     ///
-    /// Takes `&mut self` to prevent creating multiple instances
-    /// simultaneously (multiple instances would have divergent cached
-    /// values).
+    /// This method can only be called once; a second call returns an error
+    /// because multiple counters would have divergent cached values.
     #[cfg(feature = "nanvix-unstable")]
     pub fn guest_counter(&mut self) -> Result<GuestCounter> {
+        use std::sync::atomic::Ordering;
+
         use hyperlight_common::layout::SCRATCH_TOP_GUEST_COUNTER_OFFSET;
 
+        if self
+            .mgr
+            .scratch_mem
+            .counter_taken
+            .swap(true, Ordering::Relaxed)
+        {
+            return Err(new_error!(
+                "GuestCounter has already been created for this sandbox"
+            ));
+        }
+
         let scratch_size = self.mgr.scratch_mem.mem_size();
         if (SCRATCH_TOP_GUEST_COUNTER_OFFSET as usize) > scratch_size {
             return Err(new_error!(
@@ -290,13 +302,12 @@ impl UninitializedSandbox {
 
         let offset = scratch_size - SCRATCH_TOP_GUEST_COUNTER_OFFSET as usize;
         let deferred_hshm = self.mgr.scratch_mem.deferred_hshm.clone();
-        let value = self.mgr.scratch_mem.read_u64(offset)?;
 
         Ok(GuestCounter {
             inner: Mutex::new(GuestCounterInner {
                 deferred_hshm,
                 offset,
-                value,
+                value: 0,
             }),
         })
     }
@@ -1449,12 +1460,9 @@ mod tests {
 
     #[cfg(feature = "nanvix-unstable")]
     mod guest_counter_tests {
-        use std::sync::{Arc, RwLock};
-
         use hyperlight_testing::simple_guest_as_string;
 
         use crate::UninitializedSandbox;
-        use crate::mem::shared_mem::HostSharedMemory;
         use crate::sandbox::uninitialized::GuestBinary;
 
         fn make_sandbox() -> UninitializedSandbox {
@@ -1465,24 +1473,21 @@ mod tests {
             .expect("Failed to create sandbox")
         }
 
-        /// Simulate `build()` for the scratch region by storing a
-        /// `HostSharedMemory` into the deferred slot.
-        fn simulate_build(sandbox: &UninitializedSandbox) {
-            let lock = Arc::new(RwLock::new(()));
-            let hshm = HostSharedMemory {
-                region: sandbox.mgr.scratch_mem.region.clone(),
-                lock,
-            };
-            *sandbox.mgr.scratch_mem.deferred_hshm.lock().unwrap() = Some(hshm);
-        }
-
         #[test]
         fn create_guest_counter() {
             let mut sandbox = make_sandbox();
             let counter = sandbox.guest_counter();
             assert!(counter.is_ok());
         }
 
+        #[test]
+        fn only_one_counter_allowed() {
+            let mut sandbox = make_sandbox();
+            let _c1 = sandbox.guest_counter().unwrap();
+            let c2 = sandbox.guest_counter();
+            assert!(c2.is_err());
+        }
+
         #[test]
         fn fails_before_build() {
             let mut sandbox = make_sandbox();
@@ -1495,22 +1500,21 @@ mod tests {
         fn increment_decrement() {
             let mut sandbox = make_sandbox();
             let counter = sandbox.guest_counter().unwrap();
-            simulate_build(&sandbox);
+            sandbox.mgr.scratch_mem.simulate_build();
 
-            let initial = counter.value().unwrap();
             counter.increment().unwrap();
-            assert_eq!(counter.value().unwrap(), initial + 1);
+            assert_eq!(counter.value().unwrap(), 1);
             counter.increment().unwrap();
-            assert_eq!(counter.value().unwrap(), initial + 2);
+            assert_eq!(counter.value().unwrap(), 2);
             counter.decrement().unwrap();
-            assert_eq!(counter.value().unwrap(), initial + 1);
+            assert_eq!(counter.value().unwrap(), 1);
         }
 
         #[test]
         fn underflow_returns_error() {
             let mut sandbox = make_sandbox();
             let counter = sandbox.guest_counter().unwrap();
-            simulate_build(&sandbox);
+            sandbox.mgr.scratch_mem.simulate_build();
 
             assert_eq!(counter.value().unwrap(), 0);
             let result = counter.decrement();