Refactor HostFuncsWrapper so that the helper functions are private methods

Signed-off-by: Jorge Prendes <[email protected]>

Author: jprendes

src/hyperlight_host/src/sandbox/host_funcs.rs

@@ -40,7 +40,7 @@ impl HostFuncsWrapper {
         name: String,
         func: HyperlightFunction,
     ) -> Result<()> {
-        register_host_function_helper(self, name, func, None)
+        self.register_host_function_helper(name, func, None)
     }
 
     /// Register a host function with the sandbox, with a list of extra syscalls
@@ -53,7 +53,7 @@ impl HostFuncsWrapper {
         func: HyperlightFunction,
         extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
     ) -> Result<()> {
-        register_host_function_helper(self, name, func, Some(extra_allowed_syscalls))
+        self.register_host_function_helper(name, func, Some(extra_allowed_syscalls))
     }
 
     /// Assuming a host function called `"HostPrint"` exists, and takes a
@@ -63,11 +63,7 @@ impl HostFuncsWrapper {
     /// and `Err` otherwise.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(super) fn host_print(&mut self, msg: String) -> Result<i32> {
-        let res = call_host_func_impl(
-            &self.functions_map,
-            "HostPrint",
-            vec![ParameterValue::String(msg)],
-        )?;
+        let res = self.call_host_func_impl("HostPrint", vec![ParameterValue::String(msg)])?;
         res.try_into()
             .map_err(|_| HostFunctionNotFound("HostPrint".to_string()))
     }
@@ -84,97 +80,40 @@ impl HostFuncsWrapper {
         name: &str,
         args: Vec<ParameterValue>,
     ) -> Result<ReturnValue> {
-        call_host_func_impl(&self.functions_map, name, args)
+        self.call_host_func_impl(name, args)
     }
-}
-
-fn register_host_function_helper(
-    self_: &mut HostFuncsWrapper,
-    name: String,
-    func: HyperlightFunction,
-    extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
-) -> Result<()> {
-    if let Some(_syscalls) = extra_allowed_syscalls {
-        #[cfg(all(feature = "seccomp", target_os = "linux"))]
-        self_.functions_map.insert(name, (func, Some(_syscalls)));
 
+    fn register_host_function_helper(
+        &mut self,
+        name: String,
+        func: HyperlightFunction,
+        extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
+    ) -> Result<()> {
         #[cfg(not(all(feature = "seccomp", target_os = "linux")))]
-        return Err(new_error!(
-            "Extra syscalls are only supported on Linux with seccomp"
-        ));
-    } else {
-        self_.functions_map.insert(name, (func, None));
+        if extra_allowed_syscalls.is_some() {
+            return Err(new_error!(
+                "Extra syscalls are only supported on Linux with seccomp"
+            ));
+        }
+        self.functions_map
+            .insert(name, (func, extra_allowed_syscalls));
+        Ok(())
     }
 
-    Ok(())
-}
-
-#[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-fn call_host_func_impl(
-    host_funcs: &HashMap<String, (HyperlightFunction, Option<Vec<ExtraAllowedSyscall>>)>,
-    name: &str,
-    args: Vec<ParameterValue>,
-) -> Result<ReturnValue> {
-    // Inner function containing the common logic
-    fn call_func(
-        host_funcs: &HashMap<String, (HyperlightFunction, Option<Vec<ExtraAllowedSyscall>>)>,
-        name: &str,
-        args: Vec<ParameterValue>,
-    ) -> Result<ReturnValue> {
-        let func_with_syscalls = host_funcs
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
+    fn call_host_func_impl(&self, name: &str, args: Vec<ParameterValue>) -> Result<ReturnValue> {
+        let FunctionEntry {
+            function,
+            extra_allowed_syscalls,
+        } = self
+            .functions_map
             .get(name)
             .ok_or_else(|| HostFunctionNotFound(name.to_string()))?;
 
-        let func = func_with_syscalls.0.clone();
-
-        #[cfg(all(feature = "seccomp", target_os = "linux"))]
-        {
-            let syscalls = func_with_syscalls.1.clone();
-            let seccomp_filter =
-                crate::seccomp::guest::get_seccomp_filter_for_host_function_worker_thread(
-                    syscalls,
-                )?;
-            seccompiler::apply_filter(&seccomp_filter)?;
-        }
-
-        crate::metrics::maybe_time_and_emit_host_call(name, || func.call(args))
-    }
-
-    cfg_if::cfg_if! {
-        if #[cfg(all(feature = "seccomp", target_os = "linux"))] {
-            // Clone variables for the thread
-            let host_funcs_cloned = host_funcs.clone();
-            let name_cloned = name.to_string();
-            let args_cloned = args.clone();
-
-            // Create a new thread when seccomp is enabled on Linux
-            let join_handle = std::thread::Builder::new()
-                .name(format!("Host Function Worker Thread for: {:?}", name_cloned))
-                .spawn(move || {
-                    // We have a `catch_unwind` here because, if a disallowed syscall is issued,
-                    // we handle it by panicking. This is to avoid returning execution to the
-                    // offending host function—for two reasons: (1) if a host function is issuing
-                    // disallowed syscalls, it could be unsafe to return to, and (2) returning
-                    // execution after trapping the disallowed syscall can lead to UB (e.g., try
-                    // running a host function that attempts to sleep without `SYS_clock_nanosleep`,
-                    // you'll block the syscall but panic in the aftermath).
-                    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| call_func(&host_funcs_cloned, &name_cloned, args_cloned))) {
-                        Ok(val) => val,
-                        Err(err) => {
-                            if let Some(crate::HyperlightError::DisallowedSyscall) = err.downcast_ref::<crate::HyperlightError>() {
-                                return Err(crate::HyperlightError::DisallowedSyscall)
-                            }
-
-                            crate::log_then_return!("Host function {} panicked", name_cloned);
-                        }
-                    }
-                })?;
-
-            join_handle.join().map_err(|_| new_error!("Error joining thread executing host function"))?
-        } else {
-            // Directly call the function without creating a new thread
-            call_func(host_funcs, name, args)
-        }
+        // Create a new thread when seccomp is enabled on Linux
+        maybe_with_seccomp(name, extra_allowed_syscalls.as_deref(), || {
+            crate::metrics::maybe_time_and_emit_host_call(name, || function.call(args))
+        })
     }
 }
 
@@ -197,3 +136,55 @@ pub(super) fn default_writer_func(s: String) -> Result<i32> {
         }
     }
 }
+
+#[cfg(all(feature = "seccomp", target_os = "linux"))]
+fn maybe_with_seccomp<T: Send>(
+    name: &str,
+    syscalls: Option<&[ExtraAllowedSyscall]>,
+    f: impl FnOnce() -> Result<T> + Send,
+) -> Result<T> {
+    use crate::seccomp::guest::get_seccomp_filter_for_host_function_worker_thread;
+
+    // Use a scoped thread so that we can pass around references without having to clone them.
+    crossbeam::thread::scope(|s| {
+        s.builder()
+            .name(format!("Host Function Worker Thread for: {name:?}",))
+            .spawn(move |_| {
+                let seccomp_filter = get_seccomp_filter_for_host_function_worker_thread(syscalls)?;
+                seccompiler::apply_filter(&seccomp_filter)?;
+
+                // We have a `catch_unwind` here because, if a disallowed syscall is issued,
+                // we handle it by panicking. This is to avoid returning execution to the
+                // offending host function—for two reasons: (1) if a host function is issuing
+                // disallowed syscalls, it could be unsafe to return to, and (2) returning
+                // execution after trapping the disallowed syscall can lead to UB (e.g., try
+                // running a host function that attempts to sleep without `SYS_clock_nanosleep`,
+                // you'll block the syscall but panic in the aftermath).
+                match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
+                    Ok(val) => val,
+                    Err(err) => {
+                        if let Some(crate::HyperlightError::DisallowedSyscall) =
+                            err.downcast_ref::<crate::HyperlightError>()
+                        {
+                            return Err(crate::HyperlightError::DisallowedSyscall);
+                        }
+
+                        crate::log_then_return!("Host function {} panicked", name);
+                    }
+                }
+            })?
+            .join()
+            .map_err(|_| new_error!("Error joining thread executing host function"))?
+    })
+    .unwrap() // we've already joined the spawned thread, so no error can happen here
+}
+
+#[cfg(not(all(feature = "seccomp", target_os = "linux")))]
+fn maybe_with_seccomp<T: Send>(
+    _name: &str,
+    _syscalls: Option<&[ExtraAllowedSyscall]>,
+    f: impl FnOnce() -> Result<T> + Send,
+) -> Result<T> {
+    // No seccomp, just call the function
+    f()
+}

src/hyperlight_host/src/seccomp/guest.rs

@@ -69,14 +69,15 @@ fn syscalls_allowlist() -> Result<Vec<(i64, Vec<SeccompRule>)>> {
 /// (e.g., `KVM_SET_USER_MEMORY_REGION`, `KVM_GET_API_VERSION`, `KVM_CREATE_VM`,
 /// or `KVM_CREATE_VCPU`).
 pub(crate) fn get_seccomp_filter_for_host_function_worker_thread(
-    extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
+    extra_allowed_syscalls: Option<&[ExtraAllowedSyscall]>,
 ) -> Result<BpfProgram> {
     let mut allowed_syscalls = syscalls_allowlist()?;
 
     if let Some(extra_allowed_syscalls) = extra_allowed_syscalls {
         allowed_syscalls.extend(
             extra_allowed_syscalls
-                .into_iter()
+                .iter()
+                .copied()
                 .map(|syscall| (syscall, vec![])),
         );