refactor: make guest dispatch API version-resilient

- Rename guest_dispatch_function to guest_dispatch_function_v2 to force
  linker errors for out-of-date downstream code
- Add guest_dispatch! macro so users don't define the symbol by hand;
  future signature changes will produce compile errors automatically
- Parameterize GuestFunctionDefinition<F: Copy> over the function
  pointer type, eliminating transmutes in the C API dispatch path
- Remove redundant clone in print_output_with_host_print now that
  FunctionCall is passed by value

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

Author: ludfjig

README.md

@@ -97,11 +97,10 @@ pub extern "C" fn hyperlight_main() {
     // any initialization code goes here
 }
 
-#[no_mangle]
-pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
+hyperlight_guest_bin::guest_dispatch!(|function_call| {
     let function_name = function_call.function_name;
     bail!(ErrorCode::GuestFunctionNotFound => "{function_name}");
-}
+});
 ```
 
 Build the guest using [cargo-hyperlight](https://github.com/hyperlight-dev/cargo-hyperlight):

docs/how-to-build-a-hyperlight-guest-binary.md

@@ -9,8 +9,7 @@ binary can be used with Hyperlight:
   `pub fn hyperlight_main()`
 - Hyperlight expects 
   `hl_Vec* c_guest_dispatch_function(const hl_FunctionCall *functioncall)` or
-  `pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>>`
-  to be defined in the binary so that in case the host calls a function that is
+  the `guest_dispatch!` macro (Rust) to be used in the binary so that in case the host calls a function that is
   not registered by the guest, this function is called instead.
 - to be callable by the host, a function needs to be registered by the guest in
   the `hyperlight_main` function.

src/hyperlight_component_util/src/guest.rs

@@ -210,7 +210,7 @@ fn emit_export_extern_decl<'a, 'b, 'c>(
             let marshal_result = emit_hl_marshal_result(s, ret.clone(), &ft.result);
             let trait_path = s.cur_trait_path();
             quote! {
-                fn #n<T: Guest>(fc: &::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall) -> ::hyperlight_guest::error::Result<::alloc::vec::Vec<u8>> {
+                fn #n<T: Guest>(fc: ::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall) -> ::hyperlight_guest::error::Result<::alloc::vec::Vec<u8>> {
                     <T as Guest>::with_guest_state(|state| {
                         #(#pds)*
                         #(#get_instance)*
@@ -223,7 +223,7 @@ fn emit_export_extern_decl<'a, 'b, 'c>(
                         ::alloc::string::ToString::to_string(#fname),
                         ::alloc::vec![#(#pts),*],
                         ::hyperlight_common::flatbuffer_wrappers::function_types::ReturnType::VecBytes,
-                        #n::<T> as usize
+                        #n::<T>
                     )
                 );
             }

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -27,8 +27,6 @@ use tracing::{Span, instrument};
 
 use crate::{GUEST_HANDLE, REGISTERED_GUEST_FUNCTIONS};
 
-type GuestFunc = fn(FunctionCall) -> Result<Vec<u8>>;
-
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>> {
     // Validate this is a Guest Function Call
@@ -59,23 +57,19 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
         // Verify that the function call has the correct parameter types and length.
         registered_function_definition.verify_parameters(&function_call_parameter_types)?;
 
-        let p_function = unsafe {
-            let function_pointer = registered_function_definition.function_pointer;
-            core::mem::transmute::<usize, GuestFunc>(function_pointer)
-        };
-
-        p_function(function_call)
+        (registered_function_definition.function_pointer)(function_call)
     } else {
-        // The given function is not registered. The guest should implement a function called guest_dispatch_function to handle this.
+        // The given function is not registered. The guest should implement a function called
+        // guest_dispatch_function_v2 to handle this. Use the `guest_dispatch!` macro to define it.
 
         // TODO: ideally we would define a default implementation of this with weak linkage so the guest is not required
         // to implement the function but its seems that weak linkage is an unstable feature so for now its probably better
         // to not do that.
         unsafe extern "Rust" {
-            fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>>;
+            fn guest_dispatch_function_v2(function_call: FunctionCall) -> Result<Vec<u8>>;
         }
 
-        unsafe { guest_dispatch_function(function_call) }
+        unsafe { guest_dispatch_function_v2(function_call) }
     }
 }
 

src/hyperlight_guest_bin/src/guest_function/definition.rs

@@ -28,17 +28,23 @@ use hyperlight_common::func::{
 };
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 
-/// The definition of a function exposed from the guest to the host
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub struct GuestFunctionDefinition {
+/// The function pointer type for Rust guest functions.
+pub type GuestFunc = fn(FunctionCall) -> Result<Vec<u8>>;
+
+/// The definition of a function exposed from the guest to the host.
+///
+/// The type parameter `F` is the function pointer type. For Rust guests this
+/// is [`GuestFunc`]; the C API uses its own `CGuestFunc` type.
+#[derive(Debug, Clone)]
+pub struct GuestFunctionDefinition<F: Copy> {
     /// The function name
     pub function_name: String,
     /// The type of the parameter values for the host function call.
     pub parameter_types: Vec<ParameterType>,
     /// The type of the return value from the host function call
     pub return_type: ReturnType,
-    /// The function pointer to the guest function
-    pub function_pointer: usize,
+    /// The function pointer to the guest function.
+    pub function_pointer: F,
 }
 
 /// Trait for functions that can be converted to a `fn(FunctionCall) -> Result<Vec<u8>>`
@@ -57,7 +63,7 @@ where
     fn into_guest_function(self) -> fn(FunctionCall) -> Result<Vec<u8>>;
 }
 
-/// Trait for functions that can be converted to a `GuestFunctionDefinition`
+/// Trait for functions that can be converted to a `GuestFunctionDefinition<GuestFunc>`
 pub trait AsGuestFunctionDefinition<Output, Args>
 where
     Self: Function<Output, Args, HyperlightGuestError>,
@@ -66,7 +72,10 @@ where
     Args: ParameterTuple,
 {
     /// Get the `GuestFunctionDefinition` for this function
-    fn as_guest_function_definition(&self, name: impl Into<String>) -> GuestFunctionDefinition;
+    fn as_guest_function_definition(
+        &self,
+        name: impl Into<String>,
+    ) -> GuestFunctionDefinition<GuestFunc>;
 }
 
 fn into_flatbuffer_result(value: ReturnValue) -> Vec<u8> {
@@ -141,17 +150,19 @@ macro_rules! impl_host_function {
     };
 }
 
-impl<F, Args, Output> AsGuestFunctionDefinition<Output, Args> for F
+impl<Func, Args, Output> AsGuestFunctionDefinition<Output, Args> for Func
 where
-    F: IntoGuestFunction<Output, Args>,
+    Func: IntoGuestFunction<Output, Args>,
     Args: ParameterTuple,
     Output: SupportedReturnType,
 {
-    fn as_guest_function_definition(&self, name: impl Into<String>) -> GuestFunctionDefinition {
+    fn as_guest_function_definition(
+        &self,
+        name: impl Into<String>,
+    ) -> GuestFunctionDefinition<GuestFunc> {
         let parameter_types = Args::TYPE.to_vec();
         let return_type = Output::TYPE;
         let function_pointer = self.into_guest_function();
-        let function_pointer = function_pointer as usize;
 
         GuestFunctionDefinition {
             function_name: name.into(),
@@ -164,13 +175,13 @@ where
 
 for_each_tuple!(impl_host_function);
 
-impl GuestFunctionDefinition {
+impl<F: Copy> GuestFunctionDefinition<F> {
     /// Create a new `GuestFunctionDefinition`.
     pub fn new(
         function_name: String,
         parameter_types: Vec<ParameterType>,
         return_type: ReturnType,
-        function_pointer: usize,
+        function_pointer: F,
     ) -> Self {
         Self {
             function_name,
@@ -180,12 +191,12 @@ impl GuestFunctionDefinition {
         }
     }
 
-    /// Create a new `GuestFunctionDefinition` from a function that implements
-    /// `AsGuestFunctionDefinition`.
+    /// Create a new `GuestFunctionDefinition<GuestFunc>` from a function that
+    /// implements `AsGuestFunctionDefinition`.
     pub fn from_fn<Output, Args>(
         function_name: String,
         function: impl AsGuestFunctionDefinition<Output, Args>,
-    ) -> Self
+    ) -> GuestFunctionDefinition<GuestFunc>
     where
         Args: ParameterTuple,
         Output: SupportedReturnType,

src/hyperlight_guest_bin/src/guest_function/register.rs

@@ -19,19 +19,19 @@ use alloc::string::String;
 
 use hyperlight_common::func::{ParameterTuple, SupportedReturnType};
 
-use super::definition::GuestFunctionDefinition;
+use super::definition::{GuestFunc, GuestFunctionDefinition};
 use crate::REGISTERED_GUEST_FUNCTIONS;
 use crate::guest_function::definition::AsGuestFunctionDefinition;
 
 /// Represents the functions that the guest exposes to the host.
-#[derive(Debug, Default, Clone)]
-pub struct GuestFunctionRegister {
+#[derive(Debug, Clone)]
+pub struct GuestFunctionRegister<F: Copy = GuestFunc> {
     /// Currently registered guest functions
-    guest_functions: BTreeMap<String, GuestFunctionDefinition>,
+    guest_functions: BTreeMap<String, GuestFunctionDefinition<F>>,
 }
 
-impl GuestFunctionRegister {
-    /// Create a new `GuestFunctionDetails`.
+impl<F: Copy> GuestFunctionRegister<F> {
+    /// Create a new `GuestFunctionRegister`.
     pub const fn new() -> Self {
         Self {
             guest_functions: BTreeMap::new(),
@@ -44,12 +44,19 @@ impl GuestFunctionRegister {
     /// otherwise the previous `GuestFunctionDefinition` is returned.
     pub fn register(
         &mut self,
-        guest_function: GuestFunctionDefinition,
-    ) -> Option<GuestFunctionDefinition> {
+        guest_function: GuestFunctionDefinition<F>,
+    ) -> Option<GuestFunctionDefinition<F>> {
         self.guest_functions
             .insert(guest_function.function_name.clone(), guest_function)
     }
 
+    /// Gets a `GuestFunctionDefinition` by its `name` field.
+    pub fn get(&self, function_name: &str) -> Option<&GuestFunctionDefinition<F>> {
+        self.guest_functions.get(function_name)
+    }
+}
+
+impl GuestFunctionRegister<GuestFunc> {
     pub fn register_fn<Output, Args>(
         &mut self,
         name: impl Into<String>,
@@ -61,14 +68,9 @@ impl GuestFunctionRegister {
         let gfd = f.as_guest_function_definition(name);
         self.register(gfd);
     }
-
-    /// Gets a `GuestFunctionDefinition` by its `name` field.
-    pub fn get(&self, function_name: &str) -> Option<&GuestFunctionDefinition> {
-        self.guest_functions.get(function_name)
-    }
 }
 
-pub fn register_function(function_definition: GuestFunctionDefinition) {
+pub fn register_function(function_definition: GuestFunctionDefinition<GuestFunc>) {
     unsafe {
         // This is currently safe, because we are single threaded, but we
         // should find a better way to do this, see issue #808

src/hyperlight_guest_bin/src/host_comm.rs

@@ -80,9 +80,9 @@ pub fn read_n_bytes_from_user_memory(num: u64) -> Result<Vec<u8>> {
 ///
 /// This function requires memory to be setup to be used. In particular, the
 /// existence of the input and output memory regions.
-pub fn print_output_with_host_print(function_call: &FunctionCall) -> Result<Vec<u8>> {
+pub fn print_output_with_host_print(function_call: FunctionCall) -> Result<Vec<u8>> {
     let handle = unsafe { GUEST_HANDLE };
-    if let ParameterValue::String(message) = function_call.parameters.clone().unwrap()[0].clone() {
+    if let ParameterValue::String(message) = function_call.parameters.unwrap().remove(0) {
         let res = handle.call_host_function::<i32>(
             "HostPrint",
             Some(Vec::from(&[ParameterValue::String(message.to_string())])),

src/hyperlight_guest_bin/src/lib.rs

@@ -51,6 +51,41 @@ pub mod host_comm;
 pub mod memory;
 pub mod paging;
 
+/// Defines the fallback dispatch function for guest function calls that were
+/// not registered via [`#[guest_function]`](crate::guest_function) or
+/// [`register_function`](crate::guest_function::register::register_function).
+///
+/// Registered guest functions are dispatched directly and do not go through
+/// this function. Only calls whose name has no registration end up here.
+///
+/// Use this macro instead of defining the symbol by hand so that signature
+/// changes in future versions of `hyperlight_guest_bin` produce a compile
+/// error.
+///
+/// The body receives a [`FunctionCall`](hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall)
+/// by value and must return `Result<Vec<u8>>`.
+///
+/// The macro expansion references `hyperlight_common`, `hyperlight_guest`,
+/// and `alloc`. The calling crate must have these as direct dependencies
+/// (or `extern crate alloc` for `no_std` crates).
+///
+/// ```ignore
+/// guest_dispatch!(|function_call| {
+///     bail!(ErrorCode::GuestFunctionNotFound => "{}", function_call.function_name);
+/// });
+/// ```
+#[macro_export]
+macro_rules! guest_dispatch {
+    (|$fc:ident| $body:expr) => {
+        #[unsafe(no_mangle)]
+        pub fn guest_dispatch_function_v2(
+            $fc: ::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall,
+        ) -> ::hyperlight_guest::error::Result<::alloc::vec::Vec<u8>> {
+            $body
+        }
+    };
+}
+
 // Globals
 #[cfg(feature = "mem_profile")]
 struct ProfiledLockedHeap<const ORDER: usize>(LockedHeap<ORDER>);

src/hyperlight_guest_capi/src/dispatch.rs

@@ -18,7 +18,6 @@ use alloc::boxed::Box;
 use alloc::slice;
 use alloc::vec::Vec;
 use core::ffi::{CStr, c_char};
-use core::mem;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ReturnType};
@@ -29,7 +28,8 @@ use hyperlight_guest_bin::guest_function::register::GuestFunctionRegister;
 use hyperlight_guest_bin::host_comm::call_host_function_without_returning_result;
 
 use crate::types::{FfiFunctionCall, FfiVec};
-static mut REGISTERED_C_GUEST_FUNCTIONS: GuestFunctionRegister = GuestFunctionRegister::new();
+static mut REGISTERED_C_GUEST_FUNCTIONS: GuestFunctionRegister<CGuestFunc> =
+    GuestFunctionRegister::new();
 
 type CGuestFunc = extern "C" fn(&FfiFunctionCall) -> Box<FfiVec>;
 
@@ -40,7 +40,7 @@ unsafe extern "C" {
 }
 
 #[unsafe(no_mangle)]
-pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
+pub fn guest_dispatch_function_v2(function_call: FunctionCall) -> Result<Vec<u8>> {
     // Use &raw const to get an immutable reference to the static HashMap
     // this is to avoid the clippy warning "shared reference to mutable static"
     if let Some(registered_func) =
@@ -55,10 +55,7 @@ pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
         registered_func.verify_parameters(&function_call_parameter_types)?;
 
         let ffi_func_call = FfiFunctionCall::from_function_call(function_call)?;
-
-        let guest_func =
-            unsafe { mem::transmute::<usize, CGuestFunc>(registered_func.function_pointer) };
-        let function_result = guest_func(&ffi_func_call);
+        let function_result = (registered_func.function_pointer)(&ffi_func_call);
 
         unsafe { Ok(FfiVec::into_vec(*function_result)) }
     } else {
@@ -94,8 +91,7 @@ pub extern "C" fn hl_register_function_definition(
 
     let func_params = unsafe { slice::from_raw_parts(params_type, param_no).to_vec() };
 
-    let func_def =
-        GuestFunctionDefinition::new(func_name, func_params, return_type, func_ptr as usize);
+    let func_def = GuestFunctionDefinition::new(func_name, func_params, return_type, func_ptr);
 
     // Use &raw mut to get a mutable raw pointer, then dereference it
     // this is to avoid the clippy warning "shared reference to mutable static"

src/tests/rust_guests/simpleguest/src/main.rs

@@ -44,7 +44,7 @@ use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest::exit::{abort_with_code, abort_with_code_and_message};
 use hyperlight_guest_bin::exception::arch::{Context, ExceptionInfo};
-use hyperlight_guest_bin::guest_function::definition::GuestFunctionDefinition;
+use hyperlight_guest_bin::guest_function::definition::{GuestFunc, GuestFunctionDefinition};
 use hyperlight_guest_bin::guest_function::register::register_function;
 use hyperlight_guest_bin::host_comm::{
     call_host_function, call_host_function_without_returning_result, get_host_return_value_raw,
@@ -652,11 +652,11 @@ fn call_host_expect_error(hostfuncname: String) -> Result<()> {
 #[no_mangle]
 #[instrument(skip_all, parent = Span::current(), level= "Trace")]
 pub extern "C" fn hyperlight_main() {
-    let print_output_def = GuestFunctionDefinition::new(
+    let print_output_def = GuestFunctionDefinition::<GuestFunc>::new(
         "PrintOutputWithHostPrint".to_string(),
         Vec::from(&[ParameterType::String]),
         ReturnType::Int,
-        print_output_with_host_print as usize,
+        print_output_with_host_print,
     );
     register_function(print_output_def);
 }
@@ -814,9 +814,7 @@ fn fuzz_host_function(func: FunctionCall) -> Result<Vec<u8>> {
     }
 }
 
-#[no_mangle]
-#[instrument(skip_all, parent = Span::current(), level= "Trace")]
-pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
+hyperlight_guest_bin::guest_dispatch!(|function_call| {
     // This test checks the stack behavior of the input/output buffer
     // by calling the host before serializing the function call.
     // If the stack is not working correctly, the input or output buffer will be
@@ -858,4 +856,4 @@ pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
     }
 
     Ok(get_flatbuffer_result(99))
-}
+});