Create sandboxes directly from snapshots

src/hyperlight_host/src/func/host_functions.rs

@@ -52,7 +52,8 @@ impl Registerable for UninitializedSandbox {
             return_type: Output::TYPE,
         };
 
-        (*hfs).register_host_function(name.to_string(), entry)
+        (*hfs).register_host_function(name.to_string(), entry);
+        Ok(())
     }
 }
 
@@ -92,7 +93,31 @@ impl Registerable for crate::MultiUseSandbox {
             return_type: Output::TYPE,
         };
 
-        (*hfs).register_host_function(name.to_string(), entry)
+        (*hfs).register_host_function(name.to_string(), entry);
+
+        // Registration mutates the host-function set captured in
+        // snapshots. Invalidate the cached snapshot so the next
+        // `snapshot()` call reflects the updated registry.
+        self.snapshot = None;
+        Ok(())
+    }
+}
+
+impl Registerable for crate::HostFunctions {
+    fn register_host_function<Args: ParameterTuple, Output: SupportedReturnType>(
+        &mut self,
+        name: &str,
+        hf: impl Into<HostFunction<Output, Args>>,
+    ) -> Result<()> {
+        let entry = FunctionEntry {
+            function: hf.into().into(),
+            parameter_types: Args::TYPE,
+            return_type: Output::TYPE,
+        };
+
+        self.inner_mut()
+            .register_host_function(name.to_string(), entry);
+        Ok(())
     }
 }
 
@@ -236,7 +261,7 @@ pub(crate) fn register_host_function<Args: ParameterTuple, Output: SupportedRetu
         .host_funcs
         .try_lock()
         .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-        .register_host_function(name.to_string(), entry)?;
+        .register_host_function(name.to_string(), entry);
 
     Ok(())
 }

src/hyperlight_host/src/hypervisor/gdb/arch.rs

@@ -61,9 +61,8 @@ pub(crate) const DR6_HW_BP_FLAGS_MASK: u64 = 0x0F << DR6_HW_BP_FLAGS_POS;
 /// Determine the reason the vCPU stopped
 /// This is done by checking the DR6 register and the exception id
 pub(crate) fn vcpu_stop_reason(
-    vm: &mut dyn DebuggableVm,
+    vm: &dyn DebuggableVm,
     dr6: u64,
-    entrypoint: u64,
     exception: u32,
 ) -> std::result::Result<VcpuStopReason, VcpuStopReasonError> {
     let CommonRegisters { rip, .. } = vm.regs()?;
@@ -81,10 +80,6 @@ pub(crate) fn vcpu_stop_reason(
         // Check page 19-4 Vol. 3B of Intel 64 and IA-32
         // Architectures Software Developer's Manual
         if DR6_HW_BP_FLAGS_MASK & dr6 != 0 {
-            if rip == entrypoint {
-                vm.remove_hw_breakpoint(entrypoint)?;
-                return Ok(VcpuStopReason::EntryPointBp);
-            }
             return Ok(VcpuStopReason::HwBp);
         }
     }
@@ -98,12 +93,10 @@ pub(crate) fn vcpu_stop_reason(
         r"The vCPU exited because of an unknown reason:
         rip: {:?}
         dr6: {:?}
-        entrypoint: {:?}
         exception: {:?}
         ",
         rip,
         dr6,
-        entrypoint,
         exception,
     );
 

src/hyperlight_host/src/hypervisor/gdb/event_loop.rs

@@ -59,7 +59,6 @@ impl run_blocking::BlockingEventLoop for GdbBlockingEventLoop {
                     // Resume execution if unknown reason for stop
                     let stop_response = match stop_reason {
                         VcpuStopReason::DoneStep => BaseStopReason::DoneStep,
-                        VcpuStopReason::EntryPointBp => BaseStopReason::HwBreak(()),
                         VcpuStopReason::SwBp => BaseStopReason::SwBreak(()),
                         VcpuStopReason::HwBp => BaseStopReason::HwBreak(()),
                         // This is a consequence of the GDB client sending an interrupt signal

src/hyperlight_host/src/hypervisor/gdb/mod.rs

@@ -171,10 +171,6 @@ impl DebugMemoryAccess {
 pub enum VcpuStopReason {
     Crash,
     DoneStep,
-    /// Hardware breakpoint inserted by the hypervisor so the guest can be stopped
-    /// at the entry point. This is used to avoid the guest from executing
-    /// the entry point code before the debugger is connected
-    EntryPointBp,
     HwBp,
     SwBp,
     Interrupt,

src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs

@@ -389,6 +389,13 @@ pub(crate) struct HyperlightVm {
     pub(super) gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
     #[cfg(gdb)]
     pub(super) sw_breakpoints: HashMap<u64, u8>, // addr -> original instruction
+    /// One-shot hw breakpoint installed at the entry address when gdb is
+    /// enabled, so the gdb stub gets a `VcpuStopped` to enter its event
+    /// loop on the first vCPU run after construction. Cleared by the
+    /// `VmExit::Debug` arm of `run` the first time a `HwBp` stop fires
+    /// at the entry address.
+    #[cfg(gdb)]
+    pub(super) one_shot_entry_bp: Option<u64>,
     #[cfg(feature = "mem_profile")]
     pub(super) trace_info: MemTraceInfo,
     #[cfg(crashdump)]
@@ -598,17 +605,28 @@ impl HyperlightVm {
             match exit_reason {
                 #[cfg(gdb)]
                 Ok(VmExit::Debug { dr6, exception }) => {
-                    let initialise = match self.entrypoint {
-                        NextAction::Initialise(initialise) => initialise,
-                        _ => 0,
-                    };
-                    // Handle debug event (breakpoints)
+                    // Classify the debug exit. `vcpu_stop_reason` is a
+                    // pure classifier and has no side effects on the VM.
                     let stop_reason = crate::hypervisor::gdb::arch::vcpu_stop_reason(
-                        self.vm.as_mut(),
+                        self.vm.as_ref(),
                         dr6,
-                        initialise,
                         exception,
                     )?;
+                    // Remove the one-shot entry breakpoint installed by
+                    // `HyperlightVm::new` the first time it fires so it
+                    // does not interfere with later user-installed
+                    // breakpoints at the same address.
+                    if matches!(stop_reason, VcpuStopReason::HwBp)
+                        && let Some(entry_addr) = self.one_shot_entry_bp
+                    {
+                        let rip = self.vm.regs().map_err(VcpuStopReasonError::GetRegs)?.rip;
+                        if rip == entry_addr {
+                            self.vm
+                                .remove_hw_breakpoint(entry_addr)
+                                .map_err(VcpuStopReasonError::RemoveHwBreakpoint)?;
+                            self.one_shot_entry_bp = None;
+                        }
+                    }
                     if let Err(e) = self.handle_debug(dbg_mem_access_fn.clone(), stop_reason) {
                         break Err(e.into());
                     }

src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs

@@ -168,6 +168,8 @@ impl HyperlightVm {
             gdb_conn,
             #[cfg(gdb)]
             sw_breakpoints: HashMap::new(),
+            #[cfg(gdb)]
+            one_shot_entry_bp: None,
             #[cfg(feature = "mem_profile")]
             trace_info,
             #[cfg(crashdump)]
@@ -182,12 +184,21 @@ impl HyperlightVm {
         #[cfg(gdb)]
         if ret.gdb_conn.is_some() {
             ret.send_dbg_msg(DebugResponse::InterruptHandle(ret.interrupt_handle.clone()))?;
-            // Add breakpoint to the entry point address, if we are going to initialise
+            // Add breakpoint at the entry point address. The breakpoint
+            // is removed on first hit by the run loop. Tracked via
+            // `one_shot_entry_bp` so it does not interfere with later
+            // user-installed breakpoints at the same address.
             ret.vm.set_debug(true).map_err(VmError::Debug)?;
-            if let NextAction::Initialise(initialise) = entrypoint {
+            let entry_addr = match entrypoint {
+                NextAction::Initialise(addr) | NextAction::Call(addr) => Some(addr),
+                #[cfg(test)]
+                NextAction::None => None,
+            };
+            if let Some(addr) = entry_addr {
                 ret.vm
-                    .add_hw_breakpoint(initialise)
+                    .add_hw_breakpoint(addr)
                     .map_err(CreateHyperlightVmError::AddHwBreakpoint)?;
+                ret.one_shot_entry_bp = Some(addr);
             }
         }
 

src/hyperlight_host/src/lib.rs

@@ -91,6 +91,9 @@ pub use hypervisor::virtual_machine::is_hypervisor_present;
 pub use sandbox::MultiUseSandbox;
 /// The re-export for the `UninitializedSandbox` type
 pub use sandbox::UninitializedSandbox;
+/// A collection of host functions that can be supplied to a sandbox
+/// constructor (e.g. [`MultiUseSandbox::from_snapshot`]).
+pub use sandbox::host_funcs::HostFunctions;
 /// The re-export for the `GuestBinary` type
 pub use sandbox::uninitialized::GuestBinary;
 /// The re-export for the `GuestCounter` type

src/hyperlight_host/src/mem/layout.rs

@@ -226,16 +226,16 @@ pub(crate) struct SandboxMemoryLayout {
     /// The size of the guest code section.
     pub(crate) code_size: usize,
     /// The size of the init data section (guest blob).
-    init_data_size: usize,
+    pub(crate) init_data_size: usize,
     /// Permission flags for the init data region.
     #[cfg_attr(feature = "i686-guest", allow(unused))]
-    init_data_permissions: Option<MemoryRegionFlags>,
+    pub(crate) init_data_permissions: Option<MemoryRegionFlags>,
     /// The size of the scratch region in physical memory.
-    scratch_size: usize,
+    pub(crate) scratch_size: usize,
     /// The size of the snapshot region in physical memory.
-    snapshot_size: usize,
+    pub(crate) snapshot_size: usize,
     /// The size of the page tables (None if not yet set).
-    pt_size: Option<usize>,
+    pub(crate) pt_size: Option<usize>,
 }
 
 impl Debug for SandboxMemoryLayout {
@@ -295,7 +295,7 @@ impl SandboxMemoryLayout {
     /// Both the scratch region and the snapshot region are bounded by
     /// this size. The value is arbitrary but chosen to be large enough
     /// for most workloads while preventing accidental resource exhaustion.
-    const MAX_MEMORY_SIZE: usize = (16 * 1024 * 1024 * 1024) - Self::BASE_ADDRESS; // 16 GiB - BASE_ADDRESS
+    pub(crate) const MAX_MEMORY_SIZE: usize = (16 * 1024 * 1024 * 1024) - Self::BASE_ADDRESS; // 16 GiB - BASE_ADDRESS
 
     /// The base address of the sandbox's memory.
     pub(crate) const BASE_ADDRESS: usize = 0x1000;

src/hyperlight_host/src/mem/mgr.rs

@@ -22,6 +22,7 @@ use hyperlight_common::flatbuffer_wrappers::function_call::{
 };
 use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
+use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
 use hyperlight_common::vmem::{self, PAGE_TABLE_SIZE};
 #[cfg(all(feature = "crashdump", not(feature = "i686-guest")))]
 use hyperlight_common::vmem::{BasicMapping, MappingKind};
@@ -298,6 +299,7 @@ where
     }
 
     /// Create a snapshot with the given mapped regions
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn snapshot(
         &mut self,
         sandbox_id: u64,
@@ -306,6 +308,7 @@ where
         rsp_gva: u64,
         sregs: CommonSpecialRegisters,
         entrypoint: NextAction,
+        host_functions: HostFunctionDetails,
     ) -> Result<Snapshot> {
         self.snapshot_count += 1;
         Snapshot::new(
@@ -320,6 +323,7 @@ where
             sregs,
             entrypoint,
             self.snapshot_count,
+            host_functions,
         )
     }
 }
@@ -330,7 +334,13 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         let shared_mem = s.memory().to_mgr_snapshot_mem()?;
         let scratch_mem = ExclusiveSharedMemory::new(s.layout().get_scratch_size())?;
         let entrypoint = s.entrypoint();
-        Ok(Self::new(layout, shared_mem, scratch_mem, entrypoint))
+        let mut mgr = Self::new(layout, shared_mem, scratch_mem, entrypoint);
+        // Inherit the snapshot's generation number for the same
+        // reason `restore_snapshot` does: the guest-visible counter
+        // reflects "which snapshot is the sandbox currently a clone
+        // of", not "how many snapshots this partition has taken".
+        mgr.snapshot_count = s.snapshot_generation();
+        Ok(mgr)
     }
 
     /// Wraps ExclusiveSharedMemory::build