feat: restrict guest capabilities during snapshot checkpoint

CHANGELOG.md

@@ -60,6 +60,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   rejects snapshots without transport state.
 * Running snapshots checkpoint dirty virtqueues before capture. Ordinary calls
   keep their deferred result path.
+* Snapshot checkpoints reject host function calls. Violations poison the
+  sandbox.
 * Use the reclaimed stack pages to raise the default G2H and H2G pools to 12
   and 8 pages.
 * `hyperlight_guest_bin::exception::arch`, previously available on

src/hyperlight_host/src/mem/mgr.rs

@@ -133,6 +133,15 @@ impl ReadonlySharedMemory {
 }
 pub(crate) use unused_hack::SnapshotSharedMemory;
 
+/// Guest operations the host accepts during a VM entry.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum ExecutionMode {
+    /// Guest initialization and guest function calls.
+    Unrestricted,
+    /// Snapshot checkpoint preparation. Host function calls are rejected.
+    Restricted,
+}
+
 /// A struct that is responsible for laying out and managing the memory
 /// for a given `Sandbox`.
 pub(crate) struct SandboxMemoryManager<S: SharedMemory> {
@@ -151,6 +160,8 @@ pub(crate) struct SandboxMemoryManager<S: SharedMemory> {
     pub(crate) original_entrypoint: u64,
     /// Buffer for accumulating guest abort messages
     pub(crate) abort_buffer: Vec<u8>,
+    /// Guest operations the host accepts during the current VM entry.
+    pub(crate) execution_mode: ExecutionMode,
     /// Generation counter: how many snapshots have been taken from
     /// this sandbox's execution path from init to here. Incremented
     /// on each `snapshot` call; on `restore_snapshot` we inherit the
@@ -174,6 +185,7 @@ impl<S: Clone + SharedMemory> Clone for SandboxMemoryManager<S> {
             next_action: self.next_action,
             original_entrypoint: self.original_entrypoint,
             abort_buffer: self.abort_buffer.clone(),
+            execution_mode: self.execution_mode,
             snapshot_count: self.snapshot_count,
             g2h_consumer: None,
             h2g_consumer: None,
@@ -314,6 +326,7 @@ where
             next_action,
             original_entrypoint: 0,
             abort_buffer: Vec::new(),
+            execution_mode: ExecutionMode::Unrestricted,
             snapshot_count: 0,
             g2h_consumer: None,
             h2g_consumer: None,
@@ -368,6 +381,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             next_action: self.next_action,
             original_entrypoint: self.original_entrypoint,
             abort_buffer: self.abort_buffer,
+            execution_mode: self.execution_mode,
             snapshot_count: self.snapshot_count,
             g2h_consumer: None,
             h2g_consumer: None,
@@ -380,6 +394,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             next_action: self.next_action,
             original_entrypoint: self.original_entrypoint,
             abort_buffer: Vec::new(), // Guest doesn't need abort buffer
+            execution_mode: self.execution_mode,
             snapshot_count: self.snapshot_count,
             g2h_consumer: None,
             h2g_consumer: None,
@@ -691,6 +706,10 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
         self.g2h_consumer = None;
         self.h2g_consumer = None;
+        self.execution_mode = ExecutionMode::Unrestricted;
+
+        // A panic during restricted dispatch can skip the mode reset.
+        self.execution_mode = ExecutionMode::Unrestricted;
 
         let gsnapshot = if *snapshot.memory() == self.shared_mem {
             // If the snapshot memory is already the correct memory,

src/hyperlight_host/src/sandbox/initialized.rs

@@ -21,7 +21,7 @@ use crate::func::{ParameterTuple, SupportedReturnType};
 use crate::hypervisor::InterruptHandle;
 use crate::hypervisor::hyperlight_vm::{HyperlightVm, HyperlightVmError};
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
-use crate::mem::mgr::SandboxMemoryManager;
+use crate::mem::mgr::{ExecutionMode, SandboxMemoryManager};
 use crate::mem::shared_mem::{HostSharedMemory, SharedMemory as _};
 use crate::metrics::{
     METRIC_GUEST_ERROR, METRIC_GUEST_ERROR_LABEL_CODE, maybe_time_and_emit_guest_call,
@@ -477,7 +477,7 @@ impl Sandbox {
     fn checkpoint_transport_for_snapshot(&mut self) -> Result<()> {
         self.with_guest_execution(|sbox| {
             sbox.mem_mgr.begin_snapshot_checkpoint()?;
-            sbox.dispatch_guest_call()?;
+            sbox.dispatch_guest_call_restricted()?;
 
             sbox.mem_mgr.finish_snapshot_checkpoint()?;
 
@@ -1028,6 +1028,15 @@ impl Sandbox {
             })
     }
 
+    /// Dispatch in restricted mode, then restore the previous mode.
+    fn dispatch_guest_call_restricted(&mut self) -> Result<()> {
+        let mode = std::mem::replace(&mut self.mem_mgr.execution_mode, ExecutionMode::Restricted);
+        let res = self.dispatch_guest_call();
+        self.mem_mgr.execution_mode = mode;
+
+        res
+    }
+
     /// Returns a handle for interrupting guest execution.
     ///
     /// # Examples
@@ -1258,6 +1267,7 @@ mod tests {
     #[cfg(not(gdb))]
     use crate::hypervisor::hyperlight_vm::test_support::VmOperation;
     use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
+    use crate::mem::mgr::ExecutionMode;
     use crate::mem::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory as _};
     use crate::sandbox::SandboxConfiguration;
     use crate::sandbox::uninitialized::{GuestBlob, GuestEnvironment};
@@ -1606,6 +1616,45 @@ mod tests {
         assert_eq!(sandbox.status(), SandboxStatus::Ready);
     }
 
+    #[test]
+    fn restricted_dispatch_rejects_host_calls() {
+        let mut sandbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
+            .build()
+            .unwrap();
+
+        let call = FunctionCall::new(
+            "PrintOutput".into(),
+            Some(vec![ParameterValue::String("checkpoint".into())]),
+            FunctionCallType::Guest,
+            i32::TYPE,
+        );
+        sandbox.mem_mgr.write_guest_function_call(&call).unwrap();
+
+        let error = sandbox.dispatch_guest_call_restricted().unwrap_err();
+
+        assert!(
+            error.to_string().contains("during a snapshot checkpoint"),
+            "{error}"
+        );
+        assert_eq!(sandbox.status(), SandboxStatus::Poisoned);
+        assert_eq!(sandbox.mem_mgr.execution_mode, ExecutionMode::Unrestricted);
+    }
+
+    #[test]
+    fn restore_resets_execution_mode() {
+        let mut sandbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())
+            .build()
+            .unwrap();
+
+        let snapshot = sandbox.snapshot().unwrap();
+        sandbox.mem_mgr.execution_mode = ExecutionMode::Restricted;
+        sandbox.restore(snapshot).unwrap();
+
+        sandbox
+            .call::<i32>("PrintOutput", "restored".to_string())
+            .unwrap();
+    }
+
     #[test]
     fn guest_execution_preserves_pre_entry_cancellation() {
         let mut sandbox = SandboxBuilder::from_file(simple_guest_as_pathbuf())

src/hyperlight_host/src/sandbox/outb.rs

@@ -16,7 +16,7 @@ use tracing::{Span, instrument};
 use super::host_funcs::FunctionRegistry;
 #[cfg(feature = "mem_profile")]
 use crate::hypervisor::regs::CommonRegisters;
-use crate::mem::mgr::SandboxMemoryManager;
+use crate::mem::mgr::{ExecutionMode, SandboxMemoryManager};
 use crate::mem::shared_mem::HostSharedMemory;
 use crate::mem::virtq;
 #[cfg(feature = "mem_profile")]
@@ -36,6 +36,8 @@ pub enum HandleOutbError {
     InvalidPort(String),
     #[error("Failed to read host function call: {0}")]
     ReadHostFunctionCall(String),
+    #[error("Guest called a host function during a snapshot checkpoint")]
+    HostCallDuringCheckpoint,
     #[error("Failed to acquire lock at {0}:{1} - {2}")]
     LockFailed(&'static str, u32, String),
     #[error("Failed to write host function response: {0}")]
@@ -212,6 +214,7 @@ fn outb_virtq_call(
     host_funcs: &Arc<Mutex<FunctionRegistry>>,
 ) -> Result<(), HandleOutbError> {
     let max_recv_len = mem_mgr.layout.get_g2h_queue_dims().pool_len();
+    let mode = mem_mgr.execution_mode;
 
     let Some(consumer) = mem_mgr.g2h_consumer.as_mut() else {
         return Err(HandleOutbError::ReadHostFunctionCall(
@@ -233,6 +236,10 @@ fn outb_virtq_call(
         let header = virtq::read_message_header(&mut request)
             .map_err(|error| HandleOutbError::ReadHostFunctionCall(error.to_string()))?;
 
+        if mode == ExecutionMode::Restricted && header.kind == MsgKind::Request {
+            return Err(HandleOutbError::HostCallDuringCheckpoint);
+        }
+
         match header.kind {
             MsgKind::Request => break (request, reply, header),
             MsgKind::Log => {