Unify HostFunctionXX traits into a single HostFunction (#464)

jprendes · web-flow · commit 2cd054cc4923 · 2025-05-07T18:02:25.000+01:00
* Unify all the HostFunctionXX traits into a single HostFunction trait

Signed-off-by: Jorge Prendes &lt;jorge.prendes@gmail.com&gt;

* Migrate code from HostFunctionXX to HostFunction

Signed-off-by: Jorge Prendes &lt;jorge.prendes@gmail.com&gt;

---------

Signed-off-by: Jorge Prendes &lt;jorge.prendes@gmail.com&gt;
diff --git a/src/hyperlight_host/benches/benchmarks.rs b/src/hyperlight_host/benches/benchmarks.rs
@@ -19,7 +19,7 @@ use std::time::Duration;
 
 use criterion::{criterion_group, criterion_main, Criterion};
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
-use hyperlight_host::func::HostFunction2;
+use hyperlight_host::func::HostFunction;
 use hyperlight_host::sandbox::{MultiUseSandbox, SandboxConfiguration, UninitializedSandbox};
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
diff --git a/src/hyperlight_host/examples/guest-debugging/main.rs b/src/hyperlight_host/examples/guest-debugging/main.rs
@@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex};
 use std::thread;
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
-use hyperlight_host::func::HostFunction0;
+use hyperlight_host::func::HostFunction;
 #[cfg(gdb)]
 use hyperlight_host::sandbox::config::DebugInfo;
 use hyperlight_host::sandbox::SandboxConfiguration;
diff --git a/src/hyperlight_host/examples/hello-world/main.rs b/src/hyperlight_host/examples/hello-world/main.rs
@@ -18,7 +18,7 @@ use std::sync::{Arc, Mutex};
 use std::thread;
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
-use hyperlight_host::func::HostFunction0;
+use hyperlight_host::func::HostFunction;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{MultiUseSandbox, UninitializedSandbox};
diff --git a/src/hyperlight_host/src/func/guest_dispatch.rs b/src/hyperlight_host/src/func/guest_dispatch.rs
@@ -112,7 +112,7 @@ mod tests {
 
     use super::*;
     use crate::func::call_ctx::MultiUseGuestCallContext;
-    use crate::func::host_functions::HostFunction0;
+    use crate::func::host_functions::HostFunction;
     use crate::sandbox::is_hypervisor_present;
     use crate::sandbox::uninitialized::GuestBinary;
     use crate::sandbox_state::sandbox::EvolvableSandbox;
diff --git a/src/hyperlight_host/src/func/host_functions.rs b/src/hyperlight_host/src/func/host_functions.rs
@@ -20,7 +20,6 @@ use std::sync::{Arc, Mutex};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterType, ParameterValue, ReturnType,
 };
-use paste::paste;
 use tracing::{instrument, Span};
 
 use super::{HyperlightFunction, SupportedParameterType, SupportedReturnType};
@@ -54,145 +53,35 @@ impl HostFunctionDefinition {
     }
 }
 
-macro_rules! host_function {
-    // Special case for zero parameters
-    (0) => {
-        paste! {
-            /// Trait for registering a host function with zero parameters.
-            pub trait HostFunction0<'a, R: SupportedReturnType<R>> {
-                /// Register the host function with the given name in the sandbox.
-                fn register(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                ) -> Result<()>;
-
-                /// Register the host function with the given name in the sandbox, allowing extra syscalls.
-                #[cfg(all(feature = "seccomp", target_os = "linux"))]
-                fn register_with_extra_allowed_syscalls(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                    extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
-                ) -> Result<()>;
-            }
-
-            impl<'a, T, R> HostFunction0<'a, R> for Arc<Mutex<T>>
-            where
-                T: FnMut() -> Result<R> + Send + 'static,
-                R: SupportedReturnType<R>,
-            {
-                #[instrument(
-                    err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
-                )]
-                fn register(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                ) -> Result<()> {
-                    register_host_function_0(self.clone(), sandbox, name, None)
-                }
-
-                #[cfg(all(feature = "seccomp", target_os = "linux"))]
-                #[instrument(
-                    err(Debug), skip(self, sandbox, extra_allowed_syscalls),
-                    parent = Span::current(), level = "Trace"
-                )]
-                fn register_with_extra_allowed_syscalls(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                    extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
-                ) -> Result<()> {
-                    register_host_function_0(self.clone(), sandbox, name, Some(extra_allowed_syscalls))
-                }
-            }
-
-            fn register_host_function_0<T, R>(
-                self_: Arc<Mutex<T>>,
-                sandbox: &mut UninitializedSandbox,
-                name: &str,
-                extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
-            ) -> Result<()>
-            where
-                T: FnMut() -> Result<R> + Send + 'static,
-                R: SupportedReturnType<R>,
-            {
-                let cloned = self_.clone();
-                let func = Box::new(move |_: Vec<ParameterValue>| {
-                    let result = cloned
-                        .try_lock()
-                        .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?()?;
-                    Ok(result.get_hyperlight_value())
-                });
-
-                if let Some(_eas) = extra_allowed_syscalls {
-                    if cfg!(all(feature = "seccomp", target_os = "linux")) {
-                        // Register with extra allowed syscalls
-                        #[cfg(all(feature = "seccomp", target_os = "linux"))]
-                        {
-                            sandbox
-                                .host_funcs
-                                .try_lock()
-                                .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-                                .register_host_function_with_syscalls(
-                                    &HostFunctionDefinition::new(name.to_string(), None, R::get_hyperlight_type()),
-                                    HyperlightFunction::new(func),
-                                    _eas,
-                                )?;
-                        }
-                    } else {
-                        // Log and return an error
-                        log_then_return!("Extra allowed syscalls are only supported on Linux with seccomp enabled");
-                    }
-                } else {
-                    // Register without extra allowed syscalls
-                    sandbox
-                        .host_funcs
-                        .try_lock()
-                        .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-                        .register_host_function(
-                            &HostFunctionDefinition::new(name.to_string(), None, R::get_hyperlight_type()),
-                            HyperlightFunction::new(func),
-                        )?;
-                }
+/// Trait for registering a host function
+pub trait HostFunction<R, Args> {
+    /// Register the host function with the given name in the sandbox.
+    fn register(&self, sandbox: &mut UninitializedSandbox, name: &str) -> Result<()>;
+
+    /// Register the host function with the given name in the sandbox, allowing extra syscalls.
+    #[cfg(all(feature = "seccomp", target_os = "linux"))]
+    fn register_with_extra_allowed_syscalls(
+        &self,
+        sandbox: &mut UninitializedSandbox,
+        name: &str,
+        extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
+    ) -> Result<()>;
+}
 
-                Ok(())
-            }
-        }
+macro_rules! impl_host_function {
+    (@count) => { 0 };
+    (@count $P:ident $(, $R:ident)*) => {
+        impl_host_function!(@count $($R),*) + 1
     };
-    // General case for one or more parameters
-    ($N:expr, $($P:ident),+) => {
-        paste! {
-            /// Trait for registering a host function with $N parameters.
-            pub trait [<HostFunction $N>]<'a, $($P,)* R>
+    (@impl $($P:ident),*) => {
+        const _: () = {
+            impl<R $(, $P)*, F> HostFunction<R, ($($P,)*)> for Arc<Mutex<F>>
             where
-                $($P: SupportedParameterType<$P> + Clone + 'a,)*
+                F: FnMut($($P),*) -> Result<R> + Send + 'static,
+                $($P: SupportedParameterType<$P> + Clone,)*
                 R: SupportedReturnType<R>,
             {
                 /// Register the host function with the given name in the sandbox.
-                fn register(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                ) -> Result<()>;
-
-                /// Register the host function with the given name in the sandbox, allowing extra syscalls.
-                #[cfg(all(feature = "seccomp", target_os = "linux"))]
-                fn register_with_extra_allowed_syscalls(
-                    &self,
-                    sandbox: &mut UninitializedSandbox,
-                    name: &str,
-                    extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
-                ) -> Result<()>;
-            }
-
-            impl<'a, T, $($P,)* R> [<HostFunction $N>]<'a, $($P,)* R> for Arc<Mutex<T>>
-            where
-                T: FnMut($($P),*) -> Result<R> + Send + 'static,
-                $($P: SupportedParameterType<$P> + Clone + 'a,)*
-                R: SupportedReturnType<R>,
-            {
                 #[instrument(
                     err(Debug), skip(self, sandbox), parent = Span::current(), level = "Trace"
                 )]
@@ -201,9 +90,10 @@ macro_rules! host_function {
                     sandbox: &mut UninitializedSandbox,
                     name: &str,
                 ) -> Result<()> {
-                    [<register_host_function_ $N>](self.clone(), sandbox, name, None)
+                    register_host_function(self.clone(), sandbox, name, None)
                 }
 
+                /// Register the host function with the given name in the sandbox, allowing extra syscalls.
                 #[cfg(all(feature = "seccomp", target_os = "linux"))]
                 #[instrument(
                     err(Debug), skip(self, sandbox, extra_allowed_syscalls),
@@ -215,31 +105,28 @@ macro_rules! host_function {
                     name: &str,
                     extra_allowed_syscalls: Vec<ExtraAllowedSyscall>,
                 ) -> Result<()> {
-                    [<register_host_function_ $N>](self.clone(), sandbox, name, Some(extra_allowed_syscalls))
+                    register_host_function(self.clone(), sandbox, name, Some(extra_allowed_syscalls))
                 }
             }
 
-            fn [<register_host_function_ $N>]<'a, T, $($P,)* R>(
+            fn register_host_function<T, $($P,)* R>(
                 self_: Arc<Mutex<T>>,
                 sandbox: &mut UninitializedSandbox,
                 name: &str,
                 extra_allowed_syscalls: Option<Vec<ExtraAllowedSyscall>>,
             ) -> Result<()>
             where
                 T: FnMut($($P),*) -> Result<R> + Send + 'static,
-                $($P: SupportedParameterType<$P> + Clone + 'a,)*
+                $($P: SupportedParameterType<$P> + Clone,)*
                 R: SupportedReturnType<R>,
             {
+                const N: usize = impl_host_function!(@count $($P),*);
                 let cloned = self_.clone();
                 let func = Box::new(move |args: Vec<ParameterValue>| {
-                    if args.len() != $N {
-                        log_then_return!(UnexpectedNoOfArguments(args.len(), $N));
-                    }
-
-                    let mut args_iter = args.into_iter();
-                    $(
-                        let $P = $P::get_inner(args_iter.next().unwrap())?;
-                    )*
+                    let ($($P,)*) = match <[ParameterValue; N]>::try_from(args) {
+                        Ok([$($P,)*]) => ($($P::get_inner($P)?,)*),
+                        Err(args) => { log_then_return!(UnexpectedNoOfArguments(args.len(), N)); }
+                    };
 
                     let result = cloned
                         .try_lock()
@@ -292,18 +179,15 @@ macro_rules! host_function {
 
                 Ok(())
             }
-        }
+        };
+    };
+    () => {
+        impl_host_function!(@impl);
+    };
+    ($P:ident $(, $R:ident)*) => {
+        impl_host_function!($($R),*);
+        impl_host_function!(@impl $P $(, $R)*);
     };
 }
 
-host_function!(0);
-host_function!(1, P1);
-host_function!(2, P1, P2);
-host_function!(3, P1, P2, P3);
-host_function!(4, P1, P2, P3, P4);
-host_function!(5, P1, P2, P3, P4, P5);
-host_function!(6, P1, P2, P3, P4, P5, P6);
-host_function!(7, P1, P2, P3, P4, P5, P6, P7);
-host_function!(8, P1, P2, P3, P4, P5, P6, P7, P8);
-host_function!(9, P1, P2, P3, P4, P5, P6, P7, P8, P9);
-host_function!(10, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10);
+impl_host_function!(P1, P2, P3, P4, P5, P6, P7, P8, P9, P10);
diff --git a/src/hyperlight_host/src/func/mod.rs b/src/hyperlight_host/src/func/mod.rs
@@ -77,25 +77,5 @@ impl HyperlightFunction {
     }
 }
 
-/// Re-export for `HostFunction0` trait
-pub use host_functions::HostFunction0;
-/// Re-export for `HostFunction1` trait
-pub use host_functions::HostFunction1;
-/// Re-export for `HostFunction10` trait
-pub use host_functions::HostFunction10;
-/// Re-export for `HostFunction2` trait
-pub use host_functions::HostFunction2;
-/// Re-export for `HostFunction3` trait
-pub use host_functions::HostFunction3;
-/// Re-export for `HostFunction4` trait
-pub use host_functions::HostFunction4;
-/// Re-export for `HostFunction5` trait
-pub use host_functions::HostFunction5;
-/// Re-export for `HostFunction6` trait
-pub use host_functions::HostFunction6;
-/// Re-export for `HostFunction7` trait
-pub use host_functions::HostFunction7;
-/// Re-export for `HostFunction8` trait
-pub use host_functions::HostFunction8;
-/// Re-export for `HostFunction9` trait
-pub use host_functions::HostFunction9;
+/// Re-export for `HostFunction` trait
+pub use host_functions::HostFunction;
diff --git a/src/hyperlight_host/src/sandbox/uninitialized.rs b/src/hyperlight_host/src/sandbox/uninitialized.rs
@@ -30,7 +30,7 @@ use super::mem_mgr::MemMgrWrapper;
 use super::run_options::SandboxRunOptions;
 use super::uninitialized_evolve::evolve_impl_multi_use;
 use crate::error::HyperlightError::GuestBinaryShouldBeAFile;
-use crate::func::host_functions::HostFunction1;
+use crate::func::host_functions::HostFunction;
 use crate::mem::exe::ExeInfo;
 use crate::mem::mgr::{SandboxMemoryManager, STACK_COOKIE_LEN};
 use crate::mem::shared_mem::ExclusiveSharedMemory;
@@ -128,7 +128,7 @@ impl UninitializedSandbox {
         guest_binary: GuestBinary,
         cfg: Option<SandboxConfiguration>,
         sandbox_run_options: Option<SandboxRunOptions>,
-        host_print_writer: Option<&dyn HostFunction1<String, i32>>,
+        host_print_writer: Option<&dyn HostFunction<i32, (String,)>>,
     ) -> Result<Self> {
         log_build_details();
 
@@ -355,7 +355,7 @@ mod tests {
     use tracing_core::Subscriber;
     use uuid::Uuid;
 
-    use crate::func::{HostFunction1, HostFunction2};
+    use crate::func::HostFunction;
     use crate::sandbox::uninitialized::GuestBinary;
     use crate::sandbox::SandboxConfiguration;
     use crate::sandbox_state::sandbox::EvolvableSandbox;
diff --git a/src/hyperlight_host/tests/common/mod.rs b/src/hyperlight_host/tests/common/mod.rs
@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use hyperlight_host::func::HostFunction1;
+use hyperlight_host::func::HostFunction;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
 use hyperlight_host::{GuestBinary, MultiUseSandbox, Result, UninitializedSandbox};
@@ -46,7 +46,7 @@ pub fn new_uninit_rust() -> Result<UninitializedSandbox> {
 }
 
 pub fn get_simpleguest_sandboxes(
-    writer: Option<&dyn HostFunction1<String, i32>>, // An optional writer to make sure correct info is passed to the host printer
+    writer: Option<&dyn HostFunction<i32, (String,)>>, // An optional writer to make sure correct info is passed to the host printer
 ) -> Vec<MultiUseSandbox> {
     let elf_path = get_c_or_rust_simpleguest_path();
     let exe_path = format!("{elf_path}.exe");
@@ -99,7 +99,7 @@ pub fn get_simpleguest_sandboxes(
 }
 
 pub fn get_callbackguest_uninit_sandboxes(
-    writer: Option<&dyn HostFunction1<String, i32>>, // An optional writer to make sure correct info is passed to the host printer
+    writer: Option<&dyn HostFunction<i32, (String,)>>, // An optional writer to make sure correct info is passed to the host printer
 ) -> Vec<UninitializedSandbox> {
     let elf_path = get_c_or_rust_callbackguest_path();
     let exe_path = format!("{elf_path}.exe");
diff --git a/src/hyperlight_host/tests/sandbox_host_tests.rs b/src/hyperlight_host/tests/sandbox_host_tests.rs
@@ -18,7 +18,7 @@ use core::f64;
 use std::sync::{Arc, Mutex};
 
 use common::new_uninit;
-use hyperlight_host::func::{HostFunction1, ParameterValue, ReturnType, ReturnValue};
+use hyperlight_host::func::{HostFunction, ParameterValue, ReturnType, ReturnValue};
 use hyperlight_host::sandbox::SandboxConfiguration;
 use hyperlight_host::sandbox_state::sandbox::EvolvableSandbox;
 use hyperlight_host::sandbox_state::transition::Noop;
