Remove need for buffer validation by slightly modifying function signatures

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

Author: ludfjig

src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs

@@ -227,30 +227,6 @@ impl FunctionCall {
     }
 }
 
-#[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))]
-pub fn validate_guest_function_call_buffer(function_call_buffer: &[u8]) -> Result<()> {
-    let guest_function_call_fb = size_prefixed_root::<FbFunctionCall>(function_call_buffer)
-        .map_err(|e| anyhow::anyhow!("Error reading function call buffer: {:?}", e))?;
-    match guest_function_call_fb.function_call_type() {
-        FbFunctionCallType::guest => Ok(()),
-        other => {
-            bail!("Invalid function call type: {:?}", other);
-        }
-    }
-}
-
-#[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))]
-pub fn validate_host_function_call_buffer(function_call_buffer: &[u8]) -> Result<()> {
-    let host_function_call_fb = size_prefixed_root::<FbFunctionCall>(function_call_buffer)
-        .map_err(|e| anyhow::anyhow!("Error reading function call buffer: {:?}", e))?;
-    match host_function_call_fb.function_call_type() {
-        FbFunctionCallType::host => Ok(()),
-        other => {
-            bail!("Invalid function call type: {:?}", other);
-        }
-    }
-}
-
 impl TryFrom<&[u8]> for FunctionCall {
     type Error = Error;
     #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))]

src/hyperlight_host/src/func/guest_dispatch.rs

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use flatbuffers::FlatBufferBuilder;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
     ParameterValue, ReturnType, ReturnValue,
@@ -50,11 +49,8 @@ pub(crate) fn call_function_on_guest<WrapperGetterT: WrapperGetter>(
     );
 
     {
-        let mut builder = FlatBufferBuilder::new();
         let mem_mgr = wrapper_getter.get_mgr_wrapper_mut();
-        mem_mgr
-            .as_mut()
-            .write_guest_function_call(fc.encode(&mut builder))?;
+        mem_mgr.as_mut().write_guest_function_call(&fc)?;
     }
 
     let mut hv_handler = wrapper_getter.get_hv_handler().clone();

src/hyperlight_host/src/mem/mgr.rs

@@ -19,9 +19,8 @@ use std::cmp::Ordering;
 use std::str::from_utf8;
 use std::sync::{Arc, Mutex};
 
-use hyperlight_common::flatbuffer_wrappers::function_call::{
-    validate_guest_function_call_buffer, FunctionCall,
-};
+use flatbuffers::FlatBufferBuilder;
+use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
@@ -617,18 +616,20 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
     /// Writes a guest function call to memory
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(crate) fn write_guest_function_call(&mut self, buffer: &[u8]) -> Result<()> {
-        validate_guest_function_call_buffer(buffer).map_err(|e| {
-            new_error!(
-                "Guest function call buffer validation failed: {}",
-                e.to_string()
-            )
-        })?;
-
+    pub(crate) fn write_guest_function_call(&mut self, function_call: &FunctionCall) -> Result<()> {
+        match function_call.function_call_type() {
+            FunctionCallType::Host => {
+                log_then_return!(
+                    "Tried to serialize a host function call as a guest function call"
+                );
+            }
+            FunctionCallType::Guest => {}
+        };
+        let mut builder = FlatBufferBuilder::new();
         self.shared_mem.push_buffer(
             self.layout.input_data_buffer_offset,
             self.layout.sandbox_memory_config.get_input_data_size(),
-            buffer,
+            function_call.encode(&mut builder),
         )
     }