tmp (x5)

Signed-off-by: danbugs <[email protected]>

Author: danbugs

src/hyperlight_common/src/flatbuffer_wrappers/hyperlight_peb.rs

@@ -25,5 +25,5 @@ pub mod guest_log_level;
 pub mod host_function_definition;
 /// cbindgen:ignore
 pub mod host_function_details;
-pub mod util;
 pub mod hyperlight_peb;
+pub mod util;

src/hyperlight_common/src/flatbuffer_wrappers/mod.rs

@@ -54,7 +54,5 @@ pub mod hyperlight {
         pub use self::log_level_generated::*;
         mod guest_log_data_generated;
         pub use self::guest_log_data_generated::*;
-        mod hyperlight_peb_generated;
-        pub use self::hyperlight_peb_generated::*;
     } // generated
 } // hyperlight

src/hyperlight_common/src/flatbuffers/hyperlight/generated/hyperlight_peb_generated.rs

@@ -13,18 +13,18 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-use alloc::vec::Vec;
 use core::arch::asm;
 
-use spin::Once;
 use hyperlight_common::flatbuffer_wrappers::hyperlight_peb::{HyperlightPEB, RunMode};
-use crate::{__security_cookie, HEAP_ALLOCATOR, PEB, RUNNING_MODE};
+use spin::Once;
+
 use crate::guest_function_call::dispatch_function;
+use crate::{__security_cookie, HEAP_ALLOCATOR, PEB};
 
 #[inline(never)]
 pub fn halt() {
     unsafe {
-        if RUNNING_MODE == RunMode::Hypervisor {
+        if let RunMode::Hypervisor = PEB.clone().unwrap().run_mode {
             asm!("hlt", options(nostack))
         }
     }
@@ -72,50 +72,26 @@ static INIT: Once = Once::new();
 #[no_mangle]
 pub extern "win64" fn entrypoint(
     hyperlight_peb_ptr: u64,
-    hyperlight_peb_size: u64,
+    _hyperlight_peb_size: u64,
     seed: u64,
     _max_log_level: u64,
 ) {
-    INIT.call_once(|| {
-        unsafe {
-            let peb_slice: &mut [u8] = core::slice::from_raw_parts_mut(
-                hyperlight_peb_ptr as *mut u8,
-                hyperlight_peb_size as usize,
-            );
-            let mut peb: HyperlightPEB = peb_slice.try_into().unwrap();
-            let run_mode = peb.run_mode.clone().unwrap();
-
-            __security_cookie = seed;
-
-            let heap_start = peb.guest_heap_data_ptr.unwrap() as usize;
-            let heap_size = peb.guest_heap_data_size.unwrap() as usize;
-            HEAP_ALLOCATOR
-                .try_lock()
-                .expect("Failed to access HEAP_ALLOCATOR")
-                .init(heap_start, heap_size);
-
-            match run_mode {
-                RunMode::Hypervisor => {
-                    RUNNING_MODE = RunMode::Hypervisor;
-                }
-                RunMode::InProcessLinux | RunMode::InProcessWindows => {
-                    RUNNING_MODE = run_mode;
-                }
-                _ => {
-                    panic!("Invalid runmode in PEB");
-                }
-            }
-
-            dispatch_function();
-
-            peb.guest_function_dispatch_ptr = Some(dispatch_function as usize as u64);
-            PEB = Some(peb.clone());
-            let peb_serialized: Vec<u8> = peb.try_into().unwrap();
-            core::slice::from_raw_parts_mut(hyperlight_peb_ptr as *mut u8, peb_serialized.len())
-                .copy_from_slice(&peb_serialized);
-
-            hyperlight_main();
-        }
+    INIT.call_once(|| unsafe {
+        let peb = hyperlight_peb_ptr as *mut HyperlightPEB;
+
+        HEAP_ALLOCATOR
+            .try_lock()
+            .expect("Failed to access HEAP_ALLOCATOR")
+            .init((*peb).guest_heap_data_ptr as usize, (*peb).guest_heap_data_size as usize);
+
+        (*peb).set_default_memory_layout();
+        (*peb).guest_function_dispatch_ptr = dispatch_function as usize as u64;
+
+        __security_cookie = seed;
+
+        PEB = Some((*peb).clone());
+
+        hyperlight_main();
     });
 
     halt();
@@ -165,4 +141,4 @@ pub extern "win64" fn entrypoint(
 //             .init(heap_start, heap_size);
 //
 //
-//         reset_error();
+//         reset_error();

src/hyperlight_common/src/flatbuffers/mod.rs

@@ -35,20 +35,19 @@ pub(crate) fn write_error(error_code: ErrorCode, message: Option<&str>) {
         .expect("Invalid guest_error_buffer, could not be converted to a Vec<u8>");
 
     unsafe {
-        assert!(!peb.guest_error_data_ptr.is_some());
+        assert!(!peb.guest_error_data_ptr != 0);
         let len = guest_error_buffer.len();
-        if guest_error_buffer.len() > peb.guest_error_data_size.unwrap() as usize {
+        if guest_error_buffer.len() > peb.guest_error_data_size as usize {
             error!(
                 "Guest error buffer is too small to hold the error message: size {} buffer size {} message may be truncated",
                 guest_error_buffer.len(),
-                peb.guest_error_data_size.unwrap() as usize
+                peb.guest_error_data_size as usize
             );
             // get the length of the message
             let message_len = message.map_or("".to_string(), |m| m.to_string()).len();
             // message is too long, truncate it
             let truncate_len = message_len
-                - (guest_error_buffer.len()
-                - peb.guest_error_data_size.unwrap() as usize);
+                - (guest_error_buffer.len() - peb.guest_error_data_size as usize);
             let truncated_message = message
                 .map_or("".to_string(), |m| m.to_string())
                 .chars()
@@ -67,7 +66,7 @@ pub(crate) fn write_error(error_code: ErrorCode, message: Option<&str>) {
         // but, because copy_nonoverlapping doesn't return anything, we can't do that.
         // Instead, we do the prior asserts/checks to check the destination pointer isn't null
         // and that there is enough space in the destination buffer for the copy.
-        let dest_ptr = peb.guest_error_data_ptr.unwrap() as *mut u8;
+        let dest_ptr = peb.guest_error_data_ptr as *mut u8;
         core::ptr::copy_nonoverlapping(guest_error_buffer.as_ptr(), dest_ptr, len);
     }
 }
@@ -76,9 +75,9 @@ pub(crate) fn reset_error() {
     unsafe {
         let peb = PEB.clone().unwrap();
         core::ptr::write_bytes(
-            peb.guest_error_data_ptr.unwrap() as *mut u8,
+            peb.guest_error_data_ptr as *mut u8,
             0,
-            peb.guest_error_data_size.unwrap() as usize,
+            peb.guest_error_data_size as usize,
         );
     }
 }

src/hyperlight_guest/src/entrypoint.rs

@@ -21,6 +21,7 @@ use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, Functi
 use hyperlight_common::flatbuffer_wrappers::function_types::ParameterType;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::input_output::{InputDataSection, OutputDataSection};
+
 use crate::entrypoint::halt;
 use crate::error::{HyperlightGuestError, Result};
 use crate::guest_error::{reset_error, set_error};
@@ -86,22 +87,27 @@ fn internal_dispatch_function() -> Result<()> {
 
     let peb = unsafe { PEB.clone().unwrap() };
 
-    let input_data_ptr = peb.input_data_ptr.unwrap();
-    let input_data_size = peb.input_data_size.unwrap();
-    let input_data_section = InputDataSection::new(input_data_ptr as *mut u8, input_data_size as usize);
+    let input_data_ptr = peb.input_data_ptr;
+    let input_data_size = peb.input_data_size;
+    let input_data_section =
+        InputDataSection::new(input_data_ptr as *mut u8, input_data_size as usize);
 
-    let output_data_ptr = peb.output_data_ptr.unwrap();
-    let output_data_size = peb.output_data_size.unwrap();
-    let output_data_section = OutputDataSection::new(output_data_ptr as *mut u8, output_data_size as usize);
+    let output_data_ptr = peb.output_data_ptr;
+    let output_data_size = peb.output_data_size;
+    let output_data_section =
+        OutputDataSection::new(output_data_ptr as *mut u8, output_data_size as usize);
 
-    let function_call = input_data_section.try_pop_shared_input_data_into::<FunctionCall>()
+    let function_call = input_data_section
+        .try_pop_shared_input_data_into::<FunctionCall>()
         .expect("Function call deserialization failed");
 
     let result_vec = call_guest_function(function_call).inspect_err(|e| {
         set_error(e.kind.clone(), e.message.as_str());
     })?;
 
-    Ok(output_data_section.push_shared_output_data(result_vec).unwrap())
+    Ok(output_data_section
+        .push_shared_output_data(result_vec)
+        .unwrap())
 }
 
 // This is implemented as a separate function to make sure that epilogue in the internal_dispatch_function is called before the halt()

src/hyperlight_guest/src/guest_error.rs

@@ -19,7 +19,7 @@ limitations under the License.
 // Deps
 use buddy_system_allocator::LockedHeap;
 use guest_function_register::GuestFunctionRegister;
-use hyperlight_common::flatbuffer_wrappers::hyperlight_peb::{HyperlightPEB, RunMode};
+use hyperlight_common::flatbuffer_wrappers::hyperlight_peb::HyperlightPEB;
 
 extern crate alloc;
 
@@ -105,7 +105,6 @@ pub static mut MIN_STACK_ADDRESS: u64 = 0;
 // pub(crate) static mut OUTB_PTR_WITH_CONTEXT: Option<
 //     extern "win64" fn(*mut core::ffi::c_void, u16, u8),
 // > = None;
-pub static mut RUNNING_MODE: RunMode = RunMode::None;
 
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
     GuestFunctionRegister::new();

src/hyperlight_guest/src/guest_function_call.rs

@@ -72,7 +72,13 @@ impl MultiUseGuestCallContext {
         // !Send (and !Sync), we also don't need to worry about
         // synchronization
 
-        call_function_on_guest(&mut self.sbox.hv_handler, &mut self.sbox.mem_mgr, func_name, func_ret_type, args)
+        call_function_on_guest(
+            &mut self.sbox.hv_handler,
+            &mut self.sbox.mem_mgr,
+            func_name,
+            func_ret_type,
+            args,
+        )
     }
 
     /// Close out the context and get back the internally-stored

src/hyperlight_guest/src/lib.rs

@@ -22,10 +22,10 @@ use tracing::{instrument, Span};
 
 use super::guest_err::check_for_guest_error;
 use crate::hypervisor::hypervisor_handler::{HypervisorHandler, HypervisorHandlerAction};
-use crate::HyperlightError::GuestExecutionHungOnHostFunctionCall;
-use crate::{HyperlightError, Result};
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::shared_mem::HostSharedMemory;
+use crate::HyperlightError::GuestExecutionHungOnHostFunctionCall;
+use crate::{HyperlightError, Result};
 
 /// Call a guest function by name, using the given `wrapper_getter`.
 #[instrument(
@@ -54,7 +54,9 @@ pub(crate) fn call_function_on_guest(
         .try_into()
         .map_err(|_| HyperlightError::Error("Failed to serialize FunctionCall".to_string()))?;
 
-    mem_mgr.write_guest_function_call(&buffer)?;
+    let input_data_region = mem_mgr.read_hyperlight_peb()?.get_input_data_region();
+
+    mem_mgr.write_guest_function_call(input_data_region, &buffer)?;
 
     match hv_handler.execute_hypervisor_handler_action(
         HypervisorHandlerAction::DispatchCallFromHost(function_name.to_string()),
@@ -63,9 +65,7 @@ pub(crate) fn call_function_on_guest(
         Err(e) => match e {
             HyperlightError::HypervisorHandlerMessageReceiveTimedout() => {
                 timedout = true;
-                match hv_handler.terminate_hypervisor_handler_execution_and_reinitialise(
-                    mem_mgr,
-                )? {
+                match hv_handler.terminate_hypervisor_handler_execution_and_reinitialise(mem_mgr)? {
                     HyperlightError::HypervisorHandlerExecutionCancelAttemptOnFinishedExecution() =>
                         {}
                     // ^^^ do nothing, we just want to actually get the Flatbuffer return value
@@ -81,8 +81,10 @@ pub(crate) fn call_function_on_guest(
     // mem_mgr.check_stack_guard()?; // <- wrapper around mem_mgr `check_for_stack_guard`
     check_for_guest_error(mem_mgr)?;
 
+    let output_data_region = hv_handler.get_output_data_region();
+
     mem_mgr
-        .get_guest_function_call_result()
+        .get_guest_function_call_result(output_data_region)
         .map_err(|e| {
             if timedout {
                 // if we timed-out, but still got here