Persist host function metadata in Snapshot

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

Author: ludfjig

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,
         )
     }
 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -207,13 +207,19 @@ impl MultiUseSandbox {
             .get_snapshot_sregs()
             .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
         let entrypoint = self.vm.get_entrypoint();
+        let host_functions = (&*self.host_funcs.try_lock().map_err(|e| {
+            crate::new_error!("Error locking host_funcs at {}:{}: {}", file!(), line!(), e)
+        })?)
+            .into();
+
         let memory_snapshot = self.mem_mgr.snapshot(
             self.id,
             mapped_regions_vec,
             &root_pt_gpas,
             stack_top_gpa,
             sregs,
             entrypoint,
+            host_functions,
         )?;
         let snapshot = Arc::new(memory_snapshot);
         self.snapshot = Some(snapshot.clone());

src/hyperlight_host/src/sandbox/snapshot/mod.rs

@@ -17,6 +17,7 @@ limitations under the License.
 use std::collections::{BTreeMap, HashMap};
 use std::sync::atomic::{AtomicU64, Ordering};
 
+use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
 use hyperlight_common::layout::{scratch_base_gpa, scratch_base_gva};
 use hyperlight_common::vmem;
 use hyperlight_common::vmem::{
@@ -115,6 +116,16 @@ pub struct Snapshot {
     /// restored sandbox's guest-visible counter so the guest can tell
     /// which snapshot it is currently a clone of.
     snapshot_generation: u64,
+
+    /// Names and signatures of host functions registered on the
+    /// sandbox at the time this snapshot was taken. Persisted to disk
+    /// so that [`crate::MultiUseSandbox::from_snapshot`] can reject
+    /// a `HostFunctions` set that is missing required functions or
+    /// has mismatched signatures.
+    ///
+    /// Empty for snapshots created via test-only constructors that
+    /// bypass the normal sandbox path.
+    host_functions: HostFunctionDetails,
 }
 impl core::convert::AsRef<Snapshot> for Snapshot {
     fn as_ref(&self) -> &Self {
@@ -423,6 +434,9 @@ impl Snapshot {
             sregs: None,
             entrypoint: NextAction::Initialise(load_addr + entrypoint_va - base_va),
             snapshot_generation: 0,
+            host_functions: HostFunctionDetails {
+                host_functions: None,
+            },
         })
     }
 
@@ -447,6 +461,7 @@ impl Snapshot {
         sregs: CommonSpecialRegisters,
         entrypoint: NextAction,
         snapshot_generation: u64,
+        host_functions: HostFunctionDetails,
     ) -> Result<Self> {
         let mut phys_seen = HashMap::<u64, usize>::new();
         let scratch_gva = scratch_base_gva(layout.get_scratch_size());
@@ -610,6 +625,7 @@ impl Snapshot {
             sregs: Some(sregs),
             entrypoint,
             snapshot_generation,
+            host_functions,
         })
     }
 
@@ -663,6 +679,65 @@ impl Snapshot {
     pub(crate) fn entrypoint(&self) -> NextAction {
         self.entrypoint
     }
+
+    /// Validate that `provided` is a superset of the host functions
+    /// recorded in this snapshot: every function that was registered
+    /// at snapshot time must also be present in `provided` with a
+    /// matching signature. Extras in `provided` are allowed.
+    ///
+    /// Snapshots created in-memory before this field was wired in
+    /// (or by test-only constructors) carry an empty function set,
+    /// in which case any `provided` set is accepted.
+    pub(crate) fn validate_host_functions(&self, provided: &crate::HostFunctions) -> Result<()> {
+        let required = match &self.host_functions.host_functions {
+            Some(v) => v,
+            None => return Ok(()),
+        };
+        if required.is_empty() {
+            return Ok(());
+        }
+
+        // Build a HostFunctionDetails view of the provided registry
+        // using the existing `From<&FunctionRegistry>` impl.
+        let provided_details: HostFunctionDetails = provided.inner().into();
+        let provided_funcs = provided_details.host_functions.unwrap_or_default();
+
+        let mut missing: Vec<String> = Vec::new();
+        let mut signature_mismatches: Vec<String> = Vec::new();
+
+        for req in required {
+            match provided_funcs
+                .iter()
+                .find(|f| f.function_name == req.function_name)
+            {
+                None => missing.push(req.function_name.clone()),
+                Some(found)
+                    if found.parameter_types != req.parameter_types
+                        || found.return_type != req.return_type =>
+                {
+                    signature_mismatches.push(format!(
+                        "{}: snapshot has {:?} -> {:?}, registered {:?} -> {:?}",
+                        req.function_name,
+                        req.parameter_types,
+                        req.return_type,
+                        found.parameter_types,
+                        found.return_type,
+                    ));
+                }
+                Some(_) => {}
+            }
+        }
+
+        if missing.is_empty() && signature_mismatches.is_empty() {
+            return Ok(());
+        }
+
+        Err(crate::new_error!(
+            "snapshot host function mismatch: missing={:?}, signature_mismatches={:?}",
+            missing,
+            signature_mismatches
+        ))
+    }
 }
 
 impl PartialEq for Snapshot {
@@ -674,6 +749,7 @@ impl PartialEq for Snapshot {
 #[cfg(test)]
 #[cfg(not(feature = "i686-guest"))]
 mod tests {
+    use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
     use hyperlight_common::vmem::{self, BasicMapping, Mapping, MappingKind, PAGE_SIZE};
 
     use crate::hypervisor::regs::CommonSpecialRegisters;
@@ -747,6 +823,7 @@ mod tests {
             default_sregs(),
             super::NextAction::None,
             1,
+            HostFunctionDetails::default(),
         )
         .unwrap();
 
@@ -764,6 +841,7 @@ mod tests {
             default_sregs(),
             super::NextAction::None,
             2,
+            HostFunctionDetails::default(),
         )
         .unwrap();