WIP

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ## [Prerelease] - Unreleased
 
 ### Added
+* Add `RngRestorePolicy` to `SandboxConfiguration` to preserve or reseed the guest libc PRNG on snapshot restore.
 
 ### Changed
 * **Breaking:** Guest MSR state is now saved and restored across snapshots.

src/hyperlight_common/src/layout.rs

@@ -27,6 +27,7 @@ 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_SNAPSHOT_GENERATION_OFFSET: u64 = 0x20;
+pub const SCRATCH_TOP_LIBC_RNG_SEED_OFFSET: u64 = 0x28;
 pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x30;
 
 pub fn scratch_base_gpa(size: usize) -> u64 {

src/hyperlight_guest/src/layout.rs

@@ -35,4 +35,8 @@ pub fn snapshot_generation_gva() -> *mut u64 {
     use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET};
     (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET + 1) as *mut u64
 }
+pub fn libc_rng_seed_gva() -> *mut u64 {
+    use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_LIBC_RNG_SEED_OFFSET};
+    (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_LIBC_RNG_SEED_OFFSET + 1) as *mut u64
+}
 pub use arch::{scratch_base_gpa, scratch_base_gva};

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -86,6 +86,9 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
 }
 
 pub(crate) fn internal_dispatch_function() {
+    #[cfg(feature = "libc")]
+    crate::refresh_libc_rng();
+
     // Read the current TSC to report it to the host with the spans/events
     // This helps calculating the timestamps relative to the guest call
     #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))]

src/hyperlight_guest_bin/src/lib.rs

@@ -210,6 +210,23 @@ unsafe extern "C" {
     fn srand(seed: u32);
 }
 
+#[cfg(feature = "libc")]
+fn fold_libc_seed(seed: u64) -> u32 {
+    (seed ^ (seed >> 32)) as u32
+}
+
+#[cfg(feature = "libc")]
+pub(crate) fn refresh_libc_rng() {
+    let seed_ptr = hyperlight_guest::layout::libc_rng_seed_gva();
+    let seed = unsafe { seed_ptr.read_volatile() };
+    if seed != 0 {
+        unsafe {
+            seed_ptr.write_volatile(0);
+            srand(fold_libc_seed(seed));
+        }
+    }
+}
+
 #[tracing::instrument(skip_all, parent = tracing::Span::current(), level= "Trace")]
 extern "C" fn hyperlight_main_default() {
     // no-op
@@ -254,8 +271,7 @@ pub(crate) extern "C" fn generic_init(
 
     #[cfg(feature = "libc")]
     unsafe {
-        let srand_seed = (((peb_address << 8) ^ (_seed >> 4)) >> 32) as u32;
-        srand(srand_seed);
+        srand(fold_libc_seed(_seed));
     }
 
     unsafe {

src/hyperlight_host/src/mem/mgr.rs

@@ -531,6 +531,13 @@ impl SandboxMemoryManager<HostSharedMemory> {
         self.scratch_mem.write::<u64>(base_offset, value)
     }
 
+    pub(crate) fn request_libc_rng_reseed(&mut self, seed: u64) -> Result<()> {
+        self.update_scratch_bookkeeping_item(
+            hyperlight_common::layout::SCRATCH_TOP_LIBC_RNG_SEED_OFFSET,
+            seed,
+        )
+    }
+
     fn update_scratch_bookkeeping(&mut self) -> Result<()> {
         use hyperlight_common::layout::*;
         let scratch_size = self.scratch_mem.mem_size();
@@ -555,6 +562,7 @@ impl SandboxMemoryManager<HostSharedMemory> {
             SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET,
             self.snapshot_count,
         )?;
+        self.update_scratch_bookkeeping_item(SCRATCH_TOP_LIBC_RNG_SEED_OFFSET, 0)?;
 
         // Initialise the guest input and output data buffers in
         // scratch memory. TODO: remove the need for this.

src/hyperlight_host/src/sandbox/config.rs

@@ -41,6 +41,19 @@ pub enum GuestMsrError {
     },
 }
 
+/// Controls Hyperlight-managed guest randomness across snapshot restores.
+#[derive(Copy, Clone, Debug, Default, Eq, PartialEq)]
+#[repr(u8)]
+pub enum RngRestorePolicy {
+    /// Reseed the guest libc PRNG before execution resumes.
+    ///
+    /// Snapshots taken before Hyperlight v0.17.0 ignore this policy and preserve their PRNG state.
+    Refresh,
+    /// Preserve the guest libc PRNG state captured by the snapshot.
+    #[default]
+    Preserve,
+}
+
 /// The complete set of configuration needed to create a Sandbox
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 #[repr(C)]
@@ -92,6 +105,8 @@ pub struct SandboxConfiguration {
     /// Number of valid entries in `guest_msrs`.
     #[cfg(target_arch = "x86_64")]
     guest_msrs_count: usize,
+    /// How snapshot restores handle the guest libc PRNG
+    rng_restore_policy: RngRestorePolicy,
 }
 
 impl SandboxConfiguration {
@@ -137,6 +152,7 @@ impl SandboxConfiguration {
             scratch_size,
             interrupt_retry_delay,
             interrupt_vcpu_sigrtmin_offset,
+            rng_restore_policy: RngRestorePolicy::default(),
             #[cfg(gdb)]
             guest_debug_info,
             #[cfg(crashdump)]
@@ -299,6 +315,16 @@ impl SandboxConfiguration {
         self.scratch_size = scratch_size;
     }
 
+    /// Set how snapshot restores handle the guest libc PRNG.
+    pub fn set_rng_restore_policy(&mut self, policy: RngRestorePolicy) {
+        self.rng_restore_policy = policy;
+    }
+
+    /// Get how snapshot restores handle the guest libc PRNG.
+    pub fn get_rng_restore_policy(&self) -> RngRestorePolicy {
+        self.rng_restore_policy
+    }
+
     #[cfg(crashdump)]
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn get_guest_core_dump(&self) -> bool {
@@ -347,7 +373,16 @@ impl Default for SandboxConfiguration {
 mod tests {
     #[cfg(target_arch = "x86_64")]
     use super::GuestMsrError;
-    use super::SandboxConfiguration;
+    use super::{RngRestorePolicy, SandboxConfiguration};
+
+    #[test]
+    fn rng_restore_policy_defaults_to_preserve_and_can_be_refreshed() {
+        let mut config = SandboxConfiguration::default();
+        assert_eq!(config.get_rng_restore_policy(), RngRestorePolicy::Preserve);
+
+        config.set_rng_restore_policy(RngRestorePolicy::Refresh);
+        assert_eq!(config.get_rng_restore_policy(), RngRestorePolicy::Refresh);
+    }
 
     #[test]
     #[cfg(target_arch = "x86_64")]

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -42,6 +42,16 @@ use crate::metrics::{
 };
 use crate::{HyperlightError, Result, log_then_return};
 
+fn libc_rng_seed() -> u64 {
+    loop {
+        let seed = rand::random::<u64>();
+        // Zero means no reseed request in scratch memory.
+        if seed != 0 {
+            return seed;
+        }
+    }
+}
+
 /// A fully initialized sandbox that can execute guest functions multiple times.
 ///
 /// Guest functions can be called repeatedly while maintaining state between calls.
@@ -95,6 +105,7 @@ pub struct MultiUseSandbox {
     /// Given (snapshot_mem, scratch_mem, cr3), returns a list of root GPAs.
     /// If not set, only CR3 is used as the single root.
     pt_root_finder: Option<PtRootFinder>,
+    rng_restore_policy: crate::sandbox::RngRestorePolicy,
 }
 
 /// Callback for discovering page table roots from guest memory.
@@ -120,6 +131,7 @@ impl MultiUseSandbox {
         host_funcs: Arc<Mutex<FunctionRegistry>>,
         mgr: SandboxMemoryManager<HostSharedMemory>,
         vm: HyperlightVm,
+        rng_restore_policy: crate::sandbox::RngRestorePolicy,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> MultiUseSandbox {
         Self {
@@ -131,6 +143,7 @@ impl MultiUseSandbox {
             dbg_mem_access_fn,
             snapshot: None,
             pt_root_finder: None,
+            rng_restore_policy,
         }
     }
 
@@ -210,8 +223,6 @@ impl MultiUseSandbox {
         host_funcs: crate::HostFunctions,
         config: Option<crate::sandbox::SandboxConfiguration>,
     ) -> Result<Self> {
-        use rand::RngExt;
-
         use crate::mem::ptr::RawPtr;
         use crate::sandbox::uninitialized_evolve::set_up_hypervisor_partition;
 
@@ -275,10 +286,7 @@ impl MultiUseSandbox {
             load_info,
         )?;
 
-        let seed = {
-            let mut rng = rand::rng();
-            rng.random::<u64>()
-        };
+        let seed = libc_rng_seed();
         let peb_addr = RawPtr::from(u64::try_from(hshm.layout.peb_address())?);
 
         #[cfg(gdb)]
@@ -296,6 +304,12 @@ impl MultiUseSandbox {
         )
         .map_err(crate::hypervisor::hyperlight_vm::HyperlightVmError::Initialize)?;
 
+        if config.get_rng_restore_policy() == crate::sandbox::RngRestorePolicy::Refresh
+            && matches!(snapshot.next_action(), super::snapshot::NextAction::Call(_))
+        {
+            hshm.request_libc_rng_reseed(seed)?;
+        }
+
         // If the snapshot was taken from an already-initialized guest
         // (NextAction::Call), apply the captured special registers so
         // the guest resumes in the correct CPU state.
@@ -326,6 +340,7 @@ impl MultiUseSandbox {
             host_funcs,
             hshm,
             vm,
+            config.get_rng_restore_policy(),
             #[cfg(gdb)]
             dbg_mem_wrapper,
         );
@@ -553,6 +568,9 @@ impl MultiUseSandbox {
         }
 
         let (gsnapshot, gscratch) = self.mem_mgr.restore_snapshot(&snapshot)?;
+        if self.rng_restore_policy == crate::sandbox::RngRestorePolicy::Refresh {
+            self.mem_mgr.request_libc_rng_reseed(libc_rng_seed())?;
+        }
         if let Some(gsnapshot) = gsnapshot {
             self.vm
                 .update_snapshot_mapping(gsnapshot)

src/hyperlight_host/src/sandbox/mod.rs

@@ -43,8 +43,8 @@ pub(crate) mod trace;
 
 /// Trait used by the macros to paper over the differences between hyperlight and hyperlight-wasm
 pub use callable::Callable;
-/// Re-export for `SandboxConfiguration` type
-pub use config::SandboxConfiguration;
+/// Re-export for sandbox configuration types
+pub use config::{RngRestorePolicy, SandboxConfiguration};
 /// Re-export for the `MultiUseSandbox` type
 pub use initialized_multi_use::{MultiUseSandbox, PtRootFinder};
 /// Re-export for `GuestBinary` type

src/hyperlight_host/src/sandbox/snapshot/file_tests.rs

@@ -20,7 +20,7 @@ limitations under the License.
 
 use std::sync::Arc;
 
-use hyperlight_testing::simple_guest_as_pathbuf;
+use hyperlight_testing::{c_simple_guest_as_pathbuf, simple_guest_as_pathbuf};
 use serde_json::Value;
 use sha2::{Digest as _, Sha256};
 
@@ -37,6 +37,24 @@ fn create_test_sandbox() -> MultiUseSandbox {
         .unwrap()
 }
 
+fn create_c_test_sandbox() -> MultiUseSandbox {
+    create_c_test_sandbox_with_config(None)
+}
+
+fn create_c_test_sandbox_with_config(
+    config: Option<crate::sandbox::SandboxConfiguration>,
+) -> MultiUseSandbox {
+    let path = c_simple_guest_as_pathbuf();
+    UninitializedSandbox::new(GuestBinary::FilePath(path), config)
+        .unwrap()
+        .evolve()
+        .unwrap()
+}
+
+fn random_sequence(sandbox: &mut MultiUseSandbox) -> [i32; 4] {
+    std::array::from_fn(|_| sandbox.call("NextRandom", ()).unwrap())
+}
+
 fn create_snapshot() -> Arc<Snapshot> {
     let mut sbox = create_test_sandbox();
     sbox.snapshot().unwrap()
@@ -3175,6 +3193,62 @@ fn from_snapshot_silently_ignores_layout_overrides() {
     assert_eq!(new_snap.layout().get_scratch_size(), original_scratch);
 }
 
+#[test]
+fn from_snapshot_preserves_guest_libc_rng_by_default() {
+    let mut sandbox = create_c_test_sandbox();
+    let snapshot = sandbox.snapshot().unwrap();
+
+    let mut first =
+        MultiUseSandbox::from_snapshot(snapshot.clone(), HostFunctions::default(), None).unwrap();
+    let mut second =
+        MultiUseSandbox::from_snapshot(snapshot, HostFunctions::default(), None).unwrap();
+
+    assert_eq!(random_sequence(&mut first), random_sequence(&mut second));
+}
+
+#[test]
+fn from_snapshot_reseeds_guest_libc_rng_when_requested() {
+    use crate::sandbox::{RngRestorePolicy, SandboxConfiguration};
+
+    let mut sandbox = create_c_test_sandbox();
+    let snapshot = sandbox.snapshot().unwrap();
+    let mut config = SandboxConfiguration::default();
+    config.set_rng_restore_policy(RngRestorePolicy::Refresh);
+    let mut first =
+        MultiUseSandbox::from_snapshot(snapshot.clone(), HostFunctions::default(), Some(config))
+            .unwrap();
+    let mut second =
+        MultiUseSandbox::from_snapshot(snapshot, HostFunctions::default(), Some(config)).unwrap();
+
+    assert_ne!(random_sequence(&mut first), random_sequence(&mut second));
+}
+
+#[test]
+fn restore_preserves_guest_libc_rng_by_default() {
+    let mut sandbox = create_c_test_sandbox();
+    let snapshot = sandbox.snapshot().unwrap();
+    let expected = random_sequence(&mut sandbox);
+
+    sandbox.restore(snapshot).unwrap();
+
+    assert_eq!(random_sequence(&mut sandbox), expected);
+}
+
+#[test]
+fn restore_reseeds_guest_libc_rng_when_requested() {
+    use crate::sandbox::{RngRestorePolicy, SandboxConfiguration};
+
+    let mut config = SandboxConfiguration::default();
+    config.set_rng_restore_policy(RngRestorePolicy::Refresh);
+    let mut sandbox = create_c_test_sandbox_with_config(Some(config));
+    let snapshot = sandbox.snapshot().unwrap();
+    let captured = random_sequence(&mut sandbox);
+
+    sandbox.restore(snapshot).unwrap();
+
+    assert_ne!(random_sequence(&mut sandbox), captured);
+}
+
 /// `from_snapshot` honors `guest_core_dump=true` so that
 /// `generate_crashdump_to_dir` writes a file.
 #[test]