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,96 +80,77 @@ 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));
-    }
-
-    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
-            .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)?;
+        if extra_allowed_syscalls.is_some() {
+            return Err(new_error!(
+                "Extra syscalls are only supported on Linux with seccomp"
+            ));
         }
-
-        crate::metrics::maybe_time_and_emit_host_call(name, || func.call(args))
+        self.functions_map
+            .insert(name, (func, extra_allowed_syscalls));
+        Ok(())
     }
 
-    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();
-
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
+    fn call_host_func_impl(&self, name: &str, args: Vec<ParameterValue>) -> Result<ReturnValue> {
+        // Inner function containing the common logic
+        let do_call = || {
+            let (func, syscalls) = self
+                .functions_map
+                .get(name)
+                .ok_or_else(|| HostFunctionNotFound(name.to_string()))?;
+
+            #[cfg(all(feature = "seccomp", target_os = "linux"))]
+            {
+                let seccomp_filter =
+                    crate::seccomp::guest::get_seccomp_filter_for_host_function_worker_thread(
+                        syscalls.clone(),
+                    )?;
+                seccompiler::apply_filter(&seccomp_filter)?;
+            }
+
+            crate::metrics::maybe_time_and_emit_host_call(name, || func.call(args))
+        };
+
+        if cfg!(all(feature = "seccomp", target_os = "linux")) {
             // 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 || {
+            std::thread::scope(|s| {
+                s.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))) {
+                    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(do_call)) {
                         Ok(val) => val,
                         Err(err) => {
-                            if let Some(crate::HyperlightError::DisallowedSyscall) = err.downcast_ref::<crate::HyperlightError>() {
-                                return Err(crate::HyperlightError::DisallowedSyscall)
+                            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);
+                            crate::log_then_return!("Host function {} panicked", name);
                         }
                     }
-                })?;
-
-            join_handle.join().map_err(|_| new_error!("Error joining thread executing host function"))?
+                })
+                .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)
+            do_call()
         }
     }
 }