[guest,tests/rust_guests] change guest library to conform to new APIs in common library and use CGM

The host now defines a single undifferentiated memory region—the custom guest memory region (or, CGM)—that
the guest can address as it sees fit. This commit changes the guest library to leverage that and perform
said address using the APIs defined in the common library.

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

Author: danbugs

src/hyperlight_guest/src/chkstk.rs

@@ -17,7 +17,7 @@ limitations under the License.
 use core::arch::global_asm;
 use core::mem::size_of;
 
-use hyperlight_common::mem::RunMode;
+use hyperlight_common::hyperlight_peb::RunMode;
 
 use crate::guest_error::{set_invalid_runmode_error, set_stack_allocate_error};
 use crate::{MIN_STACK_ADDRESS, RUNNING_MODE};

src/hyperlight_guest/src/entrypoint.rs

@@ -13,12 +13,11 @@ 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 core::arch::asm;
-use core::ffi::{c_char, c_void, CStr};
+use core::ffi::{c_char, CStr};
 use core::ptr::copy_nonoverlapping;
 
-use hyperlight_common::mem::{HyperlightPEB, RunMode};
+use hyperlight_common::hyperlight_peb::{HyperlightPEB, RunMode};
 use log::LevelFilter;
 use spin::Once;
 
@@ -28,10 +27,7 @@ use crate::guest_function_call::dispatch_function;
 use crate::guest_logger::init_logger;
 use crate::host_function_call::{outb, OutBAction};
 use crate::idtr::load_idt;
-use crate::{
-    __security_cookie, HEAP_ALLOCATOR, MIN_STACK_ADDRESS, OS_PAGE_SIZE, OUTB_PTR,
-    OUTB_PTR_WITH_CONTEXT, P_PEB, RUNNING_MODE,
-};
+use crate::{__security_cookie, HEAP_ALLOCATOR, MIN_STACK_ADDRESS, PEB, RUNNING_MODE};
 
 #[inline(never)]
 pub fn halt() {
@@ -57,10 +53,9 @@ pub fn abort_with_code(code: i32) -> ! {
 /// # Safety
 /// This function is unsafe because it dereferences a raw pointer.
 pub unsafe fn abort_with_code_and_message(code: i32, message_ptr: *const c_char) -> ! {
-    let peb_ptr = P_PEB.unwrap();
     copy_nonoverlapping(
         message_ptr,
-        (*peb_ptr).guestPanicContextData.guestPanicContextDataBuffer as *mut c_char,
+        (*PEB).guest_panic_context_ptr as *mut c_char,
         CStr::from_ptr(message_ptr).count_bytes() + 1, // +1 for null terminator
     );
     outb(OutBAction::Abort as u16, code as u8);
@@ -77,79 +72,64 @@ static INIT: Once = Once::new();
 // Note: entrypoint cannot currently have a stackframe >4KB, as that will invoke __chkstk on msvc
 //       target without first having setup global `RUNNING_MODE` variable, which __chkstk relies on.
 #[no_mangle]
-pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_level: u64) {
-    if peb_address == 0 {
-        panic!("PEB address is null");
-    }
-
-    INIT.call_once(|| {
-        unsafe {
-            P_PEB = Some(peb_address as *mut HyperlightPEB);
-            let peb_ptr = P_PEB.unwrap();
-            __security_cookie = peb_address ^ seed;
-
-            let srand_seed = ((peb_address << 8 ^ seed >> 4) >> 32) as u32;
-
-            // Set the seed for the random number generator for C code using rand;
-            srand(srand_seed);
-
-            // set up the logger
-            let max_log_level = LevelFilter::iter()
-                .nth(max_log_level as usize)
-                .expect("Invalid log level");
-            init_logger(max_log_level);
-
-            match (*peb_ptr).runMode {
-                RunMode::Hypervisor => {
-                    RUNNING_MODE = RunMode::Hypervisor;
-                    // This static is to make it easier to implement the __chkstk function in assembly.
-                    // It also means that should we change the layout of the struct in the future, we
-                    // don't have to change the assembly code.
-                    MIN_STACK_ADDRESS = (*peb_ptr).gueststackData.minUserStackAddress;
-
-                    // Setup GDT and IDT
-                    load_gdt();
-                    load_idt();
-                }
-                RunMode::InProcessLinux | RunMode::InProcessWindows => {
-                    RUNNING_MODE = (*peb_ptr).runMode;
-
-                    OUTB_PTR = {
-                        let outb_ptr: extern "win64" fn(u16, u8) =
-                            core::mem::transmute((*peb_ptr).pOutb);
-                        Some(outb_ptr)
-                    };
-
-                    if (*peb_ptr).pOutbContext.is_null() {
-                        panic!("OutbContext is null");
-                    }
-
-                    OUTB_PTR_WITH_CONTEXT = {
-                        let outb_ptr_with_context: extern "win64" fn(*mut c_void, u16, u8) =
-                            core::mem::transmute((*peb_ptr).pOutb);
-                        Some(outb_ptr_with_context)
-                    };
-                }
-                _ => {
-                    panic!("Invalid runmode in PEB");
-                }
-            }
-
-            let heap_start = (*peb_ptr).guestheapData.guestHeapBuffer as usize;
-            let heap_size = (*peb_ptr).guestheapData.guestHeapSize as usize;
-            HEAP_ALLOCATOR
-                .try_lock()
-                .expect("Failed to access HEAP_ALLOCATOR")
-                .init(heap_start, heap_size);
-
-            OS_PAGE_SIZE = ops as u32;
-
-            (*peb_ptr).guest_function_dispatch_ptr = dispatch_function as usize as u64;
-
-            reset_error();
-
-            hyperlight_main();
+pub extern "win64" fn entrypoint(
+    peb_address: u64,
+    // TODO(danbugs:297): remove the extra arg
+    _: u64,
+    seed: u64,
+    max_log_level: u64,
+) {
+    INIT.call_once(|| unsafe {
+        PEB = peb_address as *mut HyperlightPEB;
+        RUNNING_MODE = (*PEB).clone().run_mode;
+
+        // The guest receives an undifferentiated block of memory that it can address as it sees fit.
+        // This 'addressing' is done by writing to the PEB the guest's memory layout via this function,
+        // or by directly altering the PEB. `set_default_memory_layout` will configure the PEB to
+        // with a memory layout that is compatible with the expectations of guests that use the
+        // `hyperlight_guest` library (e.g., simpleguest, and callbackguest).
+        (*PEB).set_default_memory_layout();
+
+        // The guest sets the address to a "guest function dispatch" function, which is a function
+        // that is called by the host to dispatch calls to guest functions.
+        (*PEB).guest_function_dispatch_ptr = dispatch_function as usize as u64;
+
+        // Set up the guest heap
+        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,
+            );
+
+        __security_cookie = seed;
+
+        // Set the seed for the random number generator for C code using rand;
+        let srand_seed = ((peb_address << 8 ^ seed >> 4) >> 32) as u32;
+        srand(srand_seed);
+
+        // Set up the logger
+        let max_log_level = LevelFilter::iter()
+            .nth(max_log_level as usize)
+            .expect("Invalid log level");
+        init_logger(max_log_level);
+
+        if (*PEB).run_mode == RunMode::Hypervisor {
+            // This static is to make it easier to implement the __chkstk function in assembly.
+            // It also means that, should we change the layout of the struct in the future, we
+            // don't have to change the assembly code. Plus, while this could be accessible via
+            // the PEB, we don't want to expose it entirely to user code.
+            MIN_STACK_ADDRESS = (*PEB).guest_stack_data_ptr;
+
+            // Setup GDT and IDT
+            load_gdt();
+            load_idt();
         }
+
+        reset_error();
+
+        hyperlight_main();
     });
 
     halt();

src/hyperlight_guest/src/guest_error.rs

@@ -23,9 +23,10 @@ use log::error;
 
 use crate::entrypoint::halt;
 use crate::host_function_call::{outb, OutBAction};
-use crate::P_PEB;
+use crate::PEB;
 
 pub(crate) fn write_error(error_code: ErrorCode, message: Option<&str>) {
+    let peb = unsafe { (*PEB).clone() };
     let guest_error = GuestError::new(
         error_code.clone(),
         message.map_or("".to_string(), |m| m.to_string()),
@@ -35,20 +36,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!(!(*P_PEB.unwrap()).guestErrorData.guestErrorBuffer.is_null());
+        assert_ne!(!peb.guest_error_data_ptr, 0);
         let len = guest_error_buffer.len();
-        if guest_error_buffer.len() > (*P_PEB.unwrap()).guestErrorData.guestErrorSize 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(),
-                (*P_PEB.unwrap()).guestErrorData.guestErrorSize 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()
-                    - (*P_PEB.unwrap()).guestErrorData.guestErrorSize as usize);
+            let truncate_len =
+                message_len - (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,18 +67,17 @@ 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 = (*P_PEB.unwrap()).guestErrorData.guestErrorBuffer 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);
     }
 }
 
 pub(crate) fn reset_error() {
     unsafe {
-        let peb_ptr = P_PEB.unwrap();
         core::ptr::write_bytes(
-            (*peb_ptr).guestErrorData.guestErrorBuffer,
+            (*PEB).guest_error_data_ptr as *mut u8,
             0,
-            (*peb_ptr).guestErrorData.guestErrorSize as usize,
+            (*PEB).guest_error_data_size as usize,
         );
     }
 }

src/hyperlight_guest/src/guest_function_call.rs

@@ -20,13 +20,12 @@ use alloc::vec::Vec;
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 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};
-use crate::shared_input_data::try_pop_shared_input_data_into;
-use crate::shared_output_data::push_shared_output_data;
-use crate::REGISTERED_GUEST_FUNCTIONS;
+use crate::{PEB, REGISTERED_GUEST_FUNCTIONS};
 
 type GuestFunc = fn(&FunctionCall) -> Result<Vec<u8>>;
 
@@ -86,14 +85,22 @@ fn internal_dispatch_function() -> Result<()> {
     #[cfg(debug_assertions)]
     log::trace!("internal_dispatch_function");
 
-    let function_call = try_pop_shared_input_data_into::<FunctionCall>()
+    let peb = unsafe { (*PEB).clone() };
+
+    let input_data_section: InputDataSection = peb.get_input_data_region().into();
+    let output_data_section: OutputDataSection = peb.get_output_data_region().into();
+
+    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());
     })?;
 
-    push_shared_output_data(result_vec)
+    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/host_error.rs

@@ -14,25 +14,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use core::ffi::c_void;
 use core::slice::from_raw_parts;
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
 
-use crate::P_PEB;
+use crate::PEB;
 
 pub(crate) fn check_for_host_error() {
     unsafe {
-        let peb_ptr = P_PEB.unwrap();
-        let guest_error_buffer_ptr = (*peb_ptr).guestErrorData.guestErrorBuffer as *mut u8;
-        let guest_error_buffer_size = (*peb_ptr).guestErrorData.guestErrorSize as usize;
+        // TODO(danbugs:297): shouldn't this be `host_error_data_*`?
+        let guest_error_buffer_ptr = (*PEB).guest_error_data_ptr as *mut u8;
+        let guest_error_buffer_size = (*PEB).guest_error_data_size as usize;
 
         let guest_error_buffer = from_raw_parts(guest_error_buffer_ptr, guest_error_buffer_size);
 
         if !guest_error_buffer.is_empty() {
             let guest_error = GuestError::try_from(guest_error_buffer).expect("Invalid GuestError");
             if guest_error.code != ErrorCode::NoError {
-                (*peb_ptr).outputdata.outputDataBuffer = usize::MAX as *mut c_void;
+                (*PEB).output_data_ptr = u64::MAX;
                 panic!(
                     "Guest Error: {:?} - {}",
                     guest_error.code, guest_error.message