Make get_host_value_return_as generic

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_guest/src/host_function_call.rs

@@ -14,16 +14,17 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use alloc::format;
 use alloc::string::ToString;
 use alloc::vec::Vec;
 use core::arch::global_asm;
-use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterValue, ReturnType, ReturnValue,
 };
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 use hyperlight_common::mem::RunMode;
 
 use crate::error::{HyperlightGuestError, Result};
@@ -39,94 +40,20 @@ pub enum OutBAction {
     Abort = 102,
 }
 
-pub fn get_host_value_return_as_void() -> Result<()> {
+/// Get a return value from a host function call.
+/// This usually requires a host function to be called first using `call_host_function`.
+pub fn get_host_return_value<T: TryFrom<ReturnValue>>() -> Result<T> {
     let return_value = try_pop_shared_input_data_into::<ReturnValue>()
         .expect("Unable to deserialize a return value from host");
-    if let ReturnValue::Void = return_value {
-        Ok(())
-    } else {
-        Err(HyperlightGuestError::new(
-            ErrorCode::GuestError,
-            "Host return value was not void as expected".to_string(),
-        ))
-    }
-}
-
-pub fn get_host_value_return_as_int() -> Result<i32> {
-    let return_value = try_pop_shared_input_data_into::<ReturnValue>()
-        .expect("Unable to deserialize return value from host");
-
-    // check that return value is an int and return
-    if let ReturnValue::Int(i) = return_value {
-        Ok(i)
-    } else {
-        Err(HyperlightGuestError::new(
-            ErrorCode::GuestError,
-            "Host return value was not an int as expected".to_string(),
-        ))
-    }
-}
-
-pub fn get_host_value_return_as_uint() -> Result<u32> {
-    let return_value = try_pop_shared_input_data_into::<ReturnValue>()
-        .expect("Unable to deserialize return value from host");
-
-    // check that return value is an int and return
-    if let ReturnValue::UInt(ui) = return_value {
-        Ok(ui)
-    } else {
-        Err(HyperlightGuestError::new(
-            ErrorCode::GuestError,
-            "Host return value was not a uint as expected".to_string(),
-        ))
-    }
-}
-
-pub fn get_host_value_return_as_long() -> Result<i64> {
-    let return_value = try_pop_shared_input_data_into::<ReturnValue>()
-        .expect("Unable to deserialize return value from host");
-
-    // check that return value is an int and return
-    if let ReturnValue::Long(l) = return_value {
-        Ok(l)
-    } else {
-        Err(HyperlightGuestError::new(
-            ErrorCode::GuestError,
-            "Host return value was not a long as expected".to_string(),
-        ))
-    }
-}
-
-pub fn get_host_value_return_as_ulong() -> Result<u64> {
-    let return_value = try_pop_shared_input_data_into::<ReturnValue>()
-        .expect("Unable to deserialize return value from host");
-
-    // check that return value is an int and return
-    if let ReturnValue::ULong(ul) = return_value {
-        Ok(ul)
-    } else {
-        Err(HyperlightGuestError::new(
+    T::try_from(return_value).map_err(|_| {
+        HyperlightGuestError::new(
             ErrorCode::GuestError,
-            "Host return value was not a ulong as expected".to_string(),
-        ))
-    }
-}
-
-// TODO: Make this generic, return a Result<T, ErrorCode>
-
-pub fn get_host_value_return_as_vecbytes() -> Result<Vec<u8>> {
-    let return_value = try_pop_shared_input_data_into::<ReturnValue>()
-        .expect("Unable to deserialize return value from host");
-
-    // check that return value is an Vec<u8> and return
-    if let ReturnValue::VecBytes(v) = return_value {
-        Ok(v)
-    } else {
-        Err(HyperlightGuestError::new(
-            ErrorCode::GuestError,
-            "Host return value was not an VecBytes as expected".to_string(),
-        ))
-    }
+            format!(
+                "Host return value was not a {} as expected",
+                core::any::type_name::<T>()
+            ),
+        )
+    })
 }
 
 // TODO: Make this generic, return a Result<T, ErrorCode> this should allow callers to call this function and get the result type they expect
@@ -194,7 +121,7 @@ pub fn print_output_as_guest_function(function_call: &FunctionCall) -> Result<Ve
             Some(Vec::from(&[ParameterValue::String(message.to_string())])),
             ReturnType::Int,
         )?;
-        let res_i = get_host_value_return_as_int()?;
+        let res_i = get_host_return_value::<i32>()?;
         Ok(get_flatbuffer_result(res_i))
     } else {
         Err(HyperlightGuestError::new(

src/hyperlight_guest_capi/src/flatbuffer.rs

@@ -2,10 +2,7 @@ use alloc::boxed::Box;
 use core::ffi::{c_char, CStr};
 
 use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
-use hyperlight_guest::host_function_call::{
-    get_host_value_return_as_int, get_host_value_return_as_long, get_host_value_return_as_uint,
-    get_host_value_return_as_ulong,
-};
+use hyperlight_guest::host_function_call::get_host_return_value;
 
 use crate::types::FfiVec;
 
@@ -83,23 +80,23 @@ pub extern "C" fn hl_flatbuffer_result_from_Bytes(data: *const u8, len: usize) -
 
 #[no_mangle]
 pub extern "C" fn hl_get_host_return_value_as_Int() -> i32 {
-    get_host_value_return_as_int().expect("Unable to get host return value as int")
+    get_host_return_value().expect("Unable to get host return value as int")
 }
 
 #[no_mangle]
 pub extern "C" fn hl_get_host_return_value_as_UInt() -> u32 {
-    get_host_value_return_as_uint().expect("Unable to get host return value as uint")
+    get_host_return_value().expect("Unable to get host return value as uint")
 }
 
 // the same for long, ulong
 #[no_mangle]
 pub extern "C" fn hl_get_host_return_value_as_Long() -> i64 {
-    get_host_value_return_as_long().expect("Unable to get host return value as long")
+    get_host_return_value().expect("Unable to get host return value as long")
 }
 
 #[no_mangle]
 pub extern "C" fn hl_get_host_return_value_as_ULong() -> u64 {
-    get_host_value_return_as_ulong().expect("Unable to get host return value as ulong")
+    get_host_return_value().expect("Unable to get host return value as ulong")
 }
 
 // TODO add bool, float, double, string, vecbytes

src/tests/rust_guests/callbackguest/src/main.rs

@@ -35,7 +35,7 @@ use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest::guest_function_definition::GuestFunctionDefinition;
 use hyperlight_guest::guest_function_register::register_function;
 use hyperlight_guest::host_function_call::{
-    call_host_function, get_host_value_return_as_int, print_output_as_guest_function,
+    call_host_function, get_host_return_value, print_output_as_guest_function,
 };
 use hyperlight_guest::logging::log_message;
 
@@ -51,7 +51,7 @@ fn send_message_to_host_method(
         ReturnType::Int,
     )?;
 
-    let result = get_host_value_return_as_int()?;
+    let result = get_host_return_value::<i32>()?;
 
     Ok(get_flatbuffer_result(result))
 }

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

@@ -45,9 +45,7 @@ use hyperlight_guest::entrypoint::{abort_with_code, abort_with_code_and_message}
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest::guest_function_definition::GuestFunctionDefinition;
 use hyperlight_guest::guest_function_register::register_function;
-use hyperlight_guest::host_function_call::{
-    call_host_function, get_host_value_return_as_int, get_host_value_return_as_ulong,
-};
+use hyperlight_guest::host_function_call::{call_host_function, get_host_return_value};
 use hyperlight_guest::memory::malloc;
 use hyperlight_guest::{logging, MIN_STACK_ADDRESS};
 use log::{error, LevelFilter};
@@ -94,7 +92,7 @@ fn print_output(message: &str) -> Result<Vec<u8>> {
         Some(Vec::from(&[ParameterValue::String(message.to_string())])),
         ReturnType::Int,
     )?;
-    let result = get_host_value_return_as_int()?;
+    let result = get_host_return_value::<i32>()?;
     Ok(get_flatbuffer_result(result))
 }
 
@@ -681,7 +679,7 @@ fn violate_seccomp_filters(function_call: &FunctionCall) -> Result<Vec<u8>> {
     if function_call.parameters.is_none() {
         call_host_function("MakeGetpidSyscall", None, ReturnType::ULong)?;
 
-        let res = get_host_value_return_as_ulong()?;
+        let res = get_host_return_value::<u64>()?;
 
         Ok(get_flatbuffer_result(res))
     } else {
@@ -703,7 +701,7 @@ fn add(function_call: &FunctionCall) -> Result<Vec<u8>> {
             ReturnType::Int,
         )?;
 
-        let res = get_host_value_return_as_int()?;
+        let res = get_host_return_value::<i32>()?;
 
         Ok(get_flatbuffer_result(res))
     } else {
@@ -1130,7 +1128,7 @@ pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
         Some(Vec::from(&[ParameterValue::String(message.to_string())])),
         ReturnType::Int,
     )?;
-    let result = get_host_value_return_as_int()?;
+    let result = get_host_return_value::<i32>()?;
     let function_name = function_call.function_name.clone();
     let param_len = function_call.parameters.clone().unwrap_or_default().len();
     let call_type = function_call.function_call_type().clone();