Use new API in host and guest

Signed-off-by: Jorge Prendes <jorge.prendes@gmail.com>

Author: jprendes

src/hyperlight_guest/src/guest_handle/host_comm.rs

@@ -18,15 +18,17 @@ use alloc::format;
 use alloc::string::ToString;
 use alloc::vec::Vec;
 
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     FunctionCallResult, ParameterValue, ReturnType, ReturnValue,
 };
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
 use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel;
-use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
+use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
+use hyperlight_common::flatbuffer_wrappers::util::{
+    decode, encode, encode_extend, estimate_flatbuffer_capacity,
+};
 use hyperlight_common::outb::OutBAction;
 use tracing::instrument;
 
@@ -124,10 +126,15 @@ impl GuestHandle {
             return_type,
         );
 
-        let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity);
-
-        let host_function_call_buffer = host_function_call.encode(&mut builder);
-        self.push_shared_output_data(host_function_call_buffer)?;
+        let host_function_call_buffer = Vec::with_capacity(estimated_capacity);
+        let host_function_call_buffer =
+            encode_extend(&host_function_call, host_function_call_buffer).map_err(|e| {
+                HyperlightGuestError::new(
+                    ErrorCode::GuestError,
+                    format!("Error serializing host function call to flatbuffer: {}", e),
+                )
+            })?;
+        self.push_shared_output_data(&host_function_call_buffer)?;
 
         unsafe {
             out32(OutBAction::CallFunction as u16, 0);
@@ -173,9 +180,8 @@ impl GuestHandle {
                 line,
             );
 
-            let bytes: Vec<u8> = guest_log_data
-                .try_into()
-                .expect("Failed to convert GuestLogData to bytes");
+            let bytes: Vec<u8> =
+                encode(&guest_log_data).expect("Failed to convert GuestLogData to bytes");
 
             self.push_shared_output_data(&bytes)
                 .expect("Unable to push log data to shared output data");

src/hyperlight_guest/src/guest_handle/io.rs

@@ -20,6 +20,7 @@ use core::any::type_name;
 use core::slice::from_raw_parts_mut;
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
+use hyperlight_common::flatbuffer_wrappers::util::{Deserialize, decode};
 use tracing::instrument;
 
 use super::handle::GuestHandle;
@@ -30,7 +31,7 @@ impl GuestHandle {
     #[instrument(skip_all, level = "Trace")]
     pub fn try_pop_shared_input_data_into<T>(&self) -> Result<T>
     where
-        T: for<'a> TryFrom<&'a [u8]>,
+        T: for<'a> Deserialize<'a>,
     {
         let peb_ptr = self.peb().unwrap();
         let input_stack_size = unsafe { (*peb_ptr).input_stack.size as usize };
@@ -69,7 +70,7 @@ impl GuestHandle {
         let buffer = &idb[last_element_offset_rel as usize..];
 
         // convert the buffer to T
-        let type_t = match T::try_from(buffer) {
+        let type_t = match decode::<T>(buffer) {
             Ok(t) => Ok(t),
             Err(_e) => {
                 return Err(HyperlightGuestError::new(

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -17,10 +17,10 @@ limitations under the License.
 use alloc::format;
 use alloc::vec::Vec;
 
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterType};
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
+use hyperlight_common::flatbuffer_wrappers::util::encode;
 use hyperlight_guest::bail;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use tracing::instrument;
@@ -115,10 +115,9 @@ pub(crate) fn internal_dispatch_function() {
         Err(err) => {
             let guest_error = Err(GuestError::new(err.kind, err.message));
             let fcr = FunctionCallResult::new(guest_error);
-            let mut builder = FlatBufferBuilder::new();
-            let data = fcr.encode(&mut builder);
+            let data = encode(&fcr).unwrap();
             handle
-                .push_shared_output_data(data)
+                .push_shared_output_data(&data)
                 .expect("Failed to serialize function call result");
         }
     }

src/hyperlight_guest_capi/src/error.rs

@@ -16,9 +16,9 @@ limitations under the License.
 
 use core::ffi::{CStr, c_char};
 
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult;
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
+use hyperlight_common::flatbuffer_wrappers::util::encode;
 use hyperlight_guest_bin::GUEST_HANDLE;
 
 use crate::alloc::borrow::ToOwned;
@@ -33,12 +33,11 @@ pub extern "C" fn hl_set_error(err: ErrorCode, message: *const c_char) {
             .to_owned(),
     ));
     let fcr = FunctionCallResult::new(guest_error);
-    let mut builder = FlatBufferBuilder::new();
-    let data = fcr.encode(&mut builder);
+    let data = encode(&fcr).unwrap();
     unsafe {
         #[allow(static_mut_refs)] // we are single threaded
         GUEST_HANDLE
-            .push_shared_output_data(data)
+            .push_shared_output_data(&data)
             .expect("Failed to set error")
     }
 }

src/hyperlight_host/benches/benchmarks.rs

@@ -19,10 +19,11 @@ use std::thread;
 use std::time::{Duration, Instant};
 
 use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
-use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
+use hyperlight_common::flatbuffer_wrappers::util::{
+    decode, encode, encode_extend, estimate_flatbuffer_capacity,
+};
 use hyperlight_host::GuestBinary;
 use hyperlight_host::mem::shared_mem::ExclusiveSharedMemory;
 use hyperlight_host::sandbox::{MultiUseSandbox, SandboxConfiguration, UninitializedSandbox};
@@ -437,18 +438,17 @@ fn function_call_serialization_benchmark(c: &mut Criterion) {
                 function_call.function_name.as_str(),
                 function_call.parameters.as_deref().unwrap_or(&[]),
             );
-            let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity);
-            let serialized: &[u8] = function_call.encode(&mut builder);
+            let serialized = Vec::with_capacity(estimated_capacity);
+            let serialized = encode_extend(&function_call, serialized).unwrap();
             std::hint::black_box(serialized);
         });
     });
 
     group.bench_function("deserialize_function_call", |b| {
-        let mut builder = FlatBufferBuilder::new();
-        let bytes = function_call.clone().encode(&mut builder);
+        let bytes = encode(&function_call).unwrap();
 
         b.iter(|| {
-            let deserialized: FunctionCall = bytes.try_into().unwrap();
+            let deserialized: FunctionCall = decode(&bytes).unwrap();
             std::hint::black_box(deserialized);
         });
     });

src/hyperlight_host/src/mem/mgr.rs

@@ -16,12 +16,12 @@ limitations under the License.
 #[cfg(feature = "nanvix-unstable")]
 use std::mem::offset_of;
 
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{
     FunctionCall, validate_guest_function_call_buffer,
 };
 use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
+use hyperlight_common::flatbuffer_wrappers::util::encode;
 use hyperlight_common::vmem::{self, PAGE_TABLE_SIZE};
 #[cfg(all(feature = "crashdump", not(feature = "i686-guest")))]
 use hyperlight_common::vmem::{BasicMapping, MappingKind};
@@ -451,13 +451,16 @@ impl SandboxMemoryManager<HostSharedMemory> {
         &mut self,
         res: &FunctionCallResult,
     ) -> Result<()> {
-        let mut builder = FlatBufferBuilder::new();
-        let data = res.encode(&mut builder);
+        let data = encode(res).map_err(|_| {
+            new_error!(
+                "write_response_from_host_function_call: failed to convert FunctionCallResult to Vec<u8>"
+            )
+        })?;
 
         self.scratch_mem.push_buffer(
             self.layout.get_input_data_buffer_scratch_host_offset(),
             self.layout.input_data_size,
-            data,
+            &data,
         )
     }
 

src/hyperlight_host/src/mem/shared_mem.rs

@@ -22,6 +22,7 @@ use std::mem::{align_of, size_of};
 use std::ptr::null_mut;
 use std::sync::{Arc, RwLock};
 
+use hyperlight_common::flatbuffer_wrappers::util::{Deserialize, decode};
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use tracing::{Span, instrument};
 #[cfg(target_os = "windows")]
@@ -1130,17 +1131,11 @@ impl HostSharedMemory {
         Ok(())
     }
 
-    /// Pops the given given buffer into a `T` and returns it.
-    /// NOTE! the data must be a size-prefixed flatbuffer, and
-    /// buffer_start_offset must point to the beginning of the buffer
-    pub fn try_pop_buffer_into<T>(
+    fn try_pop_buffer_raw(
         &mut self,
         buffer_start_offset: usize,
         buffer_size: usize,
-    ) -> Result<T>
-    where
-        T: for<'b> TryFrom<&'b [u8]>,
-    {
+    ) -> Result<Vec<u8>> {
         // get the stackpointer
         let stack_pointer_rel = self.read::<u64>(buffer_start_offset)? as usize;
 
@@ -1180,17 +1175,15 @@ impl HostSharedMemory {
 
         // Get the size of the flatbuffer buffer from memory
         let fb_buffer_size = {
-            let raw_prefix = self.read::<u32>(last_element_offset_abs)?;
-            // flatbuffer byte arrays are prefixed by 4 bytes indicating
-            // the remaining size; add 4 for the prefix itself.
-            let total = raw_prefix.checked_add(4).ok_or_else(|| {
-                new_error!(
-                    "Corrupt buffer size prefix: value {} overflows when adding 4-byte header.",
-                    raw_prefix
-                )
-            })?;
-            usize::try_from(total)
-        }?;
+            stack_pointer_rel - last_element_offset_rel - 8
+            /*
+            let size_i32 = self.read::<u32>(last_element_offset_abs)? + 4;
+            // ^^^ flatbuffer byte arrays are prefixed by 4 bytes
+            // indicating its size, so, to get the actual size, we need
+            // to add 4.
+            usize::try_from(size_i32)
+            */
+        };
 
         if fb_buffer_size > max_element_size {
             return Err(new_error!(
@@ -1203,12 +1196,6 @@ impl HostSharedMemory {
         let mut result_buffer = vec![0; fb_buffer_size];
 
         self.copy_to_slice(&mut result_buffer, last_element_offset_abs)?;
-        let to_return = T::try_from(result_buffer.as_slice()).map_err(|_e| {
-            new_error!(
-                "pop_buffer_into: failed to convert buffer to {}",
-                type_name::<T>()
-            )
-        })?;
 
         // update the stack pointer to point to the element we just popped off since that is now free
         self.write::<u64>(buffer_start_offset, last_element_offset_rel as u64)?;
@@ -1217,6 +1204,27 @@ impl HostSharedMemory {
         let num_bytes_to_zero = stack_pointer_rel - last_element_offset_rel;
         self.fill(0, last_element_offset_abs, num_bytes_to_zero)?;
 
+        Ok(result_buffer)
+    }
+
+    /// Pops the given given buffer into a `T` and returns it.
+    /// NOTE! the data must be a size-prefixed flatbuffer, and
+    /// buffer_start_offset must point to the beginning of the buffer
+    pub fn try_pop_buffer_into<T>(
+        &mut self,
+        buffer_start_offset: usize,
+        buffer_size: usize,
+    ) -> Result<T>
+    where
+        T: for<'b> Deserialize<'b>,
+    {
+        let result_buffer = self.try_pop_buffer_raw(buffer_start_offset, buffer_size)?;
+        let to_return: T = decode(result_buffer.as_slice()).map_err(|_e| {
+            new_error!(
+                "pop_buffer_into: failed to convert buffer to {}",
+                type_name::<T>()
+            )
+        })?;
         Ok(to_return)
     }
 }
@@ -1730,16 +1738,6 @@ mod tests {
     mod try_pop_buffer_bounds {
         use super::*;
 
-        #[derive(Debug, PartialEq)]
-        struct RawBytes(Vec<u8>);
-
-        impl TryFrom<&[u8]> for RawBytes {
-            type Error = String;
-            fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> {
-                Ok(RawBytes(value.to_vec()))
-            }
-        }
-
         /// Create a buffer with stack pointer initialized to 8 (empty).
         fn make_buffer(mem_size: usize) -> super::super::HostSharedMemory {
             let eshm = ExclusiveSharedMemory::new(mem_size).unwrap();
@@ -1760,8 +1758,8 @@ mod tests {
             data.extend_from_slice(payload);
 
             hshm.push_buffer(0, mem_size, &data).unwrap();
-            let result: RawBytes = hshm.try_pop_buffer_into(0, mem_size).unwrap();
-            assert_eq!(result.0, data);
+            let result = hshm.try_pop_buffer_raw(0, mem_size).unwrap();
+            assert_eq!(result, data);
         }
 
         #[test]
@@ -1778,7 +1776,7 @@ mod tests {
             // Corrupt size prefix at element start (offset 8) to near u32::MAX.
             hshm.write::<u32>(8, 0xFFFF_FFFBu32).unwrap(); // +4 = 0xFFFF_FFFF
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains("Corrupt buffer size prefix: flatbuffer claims 4294967295 bytes but the element slot is only 9 bytes"),
@@ -1801,7 +1799,7 @@ mod tests {
             // Corrupt back-pointer (offset 16) to 0 (before valid range).
             hshm.write::<u64>(16, 0u64).unwrap();
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains(
@@ -1826,7 +1824,7 @@ mod tests {
             // Corrupt back-pointer (offset 16) to 9999 (past stack pointer 24).
             hshm.write::<u64>(16, 9999u64).unwrap();
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains(
@@ -1851,7 +1849,7 @@ mod tests {
             // Corrupt size prefix: claim 5 bytes (total 9), exceeding the 8-byte slot.
             hshm.write::<u32>(8, 5u32).unwrap(); // fb_buffer_size = 5 + 4 = 9
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains("Corrupt buffer size prefix: flatbuffer claims 9 bytes but the element slot is only 8 bytes"),
@@ -1876,7 +1874,7 @@ mod tests {
             // stack_pointer_rel = 24. Set back-pointer to 23 (> 24 - 16 = 8, so rejected).
             hshm.write::<u64>(16, 23u64).unwrap();
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains(
@@ -1902,7 +1900,7 @@ mod tests {
             // Write 0xFFFF_FFFD as size prefix: checked_add(4) returns None.
             hshm.write::<u32>(8, 0xFFFF_FFFDu32).unwrap();
 
-            let result: Result<RawBytes> = hshm.try_pop_buffer_into(0, mem_size);
+            let result = hshm.try_pop_buffer_raw(0, mem_size);
             let err_msg = format!("{}", result.unwrap_err());
             assert!(
                 err_msg.contains("Corrupt buffer size prefix: value 4294967293 overflows when adding 4-byte header"),