[common/input_output] making input_output stacks portable

If there's a mismatch in the guest's usize and the host's usize, our input and output
stacks would break. This commit fixes that by using u64 instead of usize for ptr accesses.

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

Author: danbugs

Diff for: src/hyperlight_common/src/host_calling.rs

@@ -1,14 +1,15 @@
-use crate::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType, ReturnValue};
-use crate::input_output::{InputDataSection, OutputDataSection};
+use alloc::string::ToString;
+use alloc::vec::Vec;
+use core::ffi::c_char;
+
 use anyhow::Result;
+
 use crate::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
+use crate::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType, ReturnValue};
+use crate::input_output::{InputDataSection, OutputDataSection};
 use crate::outb::{outb, OutBAction};
 use crate::PEB;
 
-use alloc::vec::Vec;
-use alloc::string::ToString;
-use core::ffi::c_char;
-
 /// 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> {
@@ -19,7 +20,10 @@ pub fn get_host_return_value<T: TryFrom<ReturnValue>>() -> Result<T> {
         .expect("Unable to deserialize a return value from host");
 
     T::try_from(return_value).map_err(|_| {
-        anyhow::anyhow!("Host return value was not a {} as expected", core::any::type_name::<T>())
+        anyhow::anyhow!(
+            "Host return value was not a {} as expected",
+            core::any::type_name::<T>()
+        )
     })
 }
 
@@ -44,8 +48,7 @@ pub fn call_host_function(
 
     let output_data_section: OutputDataSection =
         unsafe { (*PEB).clone() }.get_output_data_region().into();
-    output_data_section
-        .push_shared_output_data(host_function_call_buffer)?;
+    output_data_section.push_shared_output_data(host_function_call_buffer)?;
 
     outb(OutBAction::CallFunction as u16, 0);
 
@@ -66,4 +69,4 @@ pub fn print(message: &str) {
 #[no_mangle]
 pub unsafe extern "C" fn _putchar(c: c_char) {
     outb(OutBAction::DebugPrint as u16, c as u8);
-}
+}

Diff for: src/hyperlight_common/src/input_output.rs

@@ -13,135 +13,145 @@ 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 anyhow::{bail, Result};
 use alloc::vec::Vec;
-use core::any::type_name;
 use core::slice::from_raw_parts_mut;
 
+use anyhow::{bail, Result};
+
 pub struct InputDataSection {
     ptr: *mut u8,
-    len: usize,
+    len: u64,
 }
 
 impl InputDataSection {
-    pub fn new(ptr: *mut u8, len: usize) -> Self {
+    pub fn new(ptr: *mut u8, len: u64) -> Self {
         InputDataSection { ptr, len }
     }
 
     pub fn try_pop_shared_input_data_into<T>(&self) -> Result<T>
     where
         T: for<'a> TryFrom<&'a [u8]>,
     {
-        let input_data_buffer = unsafe { from_raw_parts_mut(self.ptr, self.len) };
+        let shared_buffer_size = self.len as usize;
 
-        if input_data_buffer.is_empty() {
-            bail!("Got a 0-size buffer in pop_shared_input_data_into");
+        let idb = unsafe { from_raw_parts_mut(self.ptr, shared_buffer_size) };
+
+        if idb.is_empty() {
+            bail!("Got a 0-size buffer in try_pop_shared_input_data_into");
         }
 
         // get relative offset to next free address
-        let stack_ptr_rel: usize = usize::from_le_bytes(
-            input_data_buffer[..8]
-                .try_into()
-                .expect("Shared input buffer too small"),
-        );
-
-        if stack_ptr_rel > self.len || stack_ptr_rel < 16 {
-            bail!("Invalid stack pointer: {} in pop_shared_input_data_into", stack_ptr_rel);
+        let stack_ptr_rel: u64 = u64::from_le_bytes(match idb[..8].try_into() {
+            Ok(bytes) => bytes,
+            Err(_) => bail!("shared input buffer too small"),
+        });
+
+        if stack_ptr_rel as usize > shared_buffer_size || stack_ptr_rel < 16 {
+            bail!(
+                "Invalid stack pointer: {} in try_pop_shared_input_data_into",
+                stack_ptr_rel
+            );
         }
 
         // go back 8 bytes and read. This is the offset to the element on top of stack
-        let last_element_offset_rel = usize::from_le_bytes(
-            input_data_buffer[stack_ptr_rel - 8..stack_ptr_rel]
-                .try_into()
-                .expect("Invalid stack pointer in pop_shared_input_data_into"),
+        let last_element_offset_rel = u64::from_le_bytes(
+            match idb[stack_ptr_rel as usize - 8..stack_ptr_rel as usize].try_into() {
+                Ok(bytes) => bytes,
+                Err(_) => bail!("Invalid stack pointer in pop_shared_input_data_into"),
+            },
         );
 
-        let buffer = &input_data_buffer[last_element_offset_rel..];
+        let buffer = &idb[last_element_offset_rel as usize..];
 
         // convert the buffer to T
         let type_t = match T::try_from(buffer) {
             Ok(t) => Ok(t),
-            Err(_e) => {
-                bail!("Unable to convert buffer to {}", type_name::<T>());
-            }
+            Err(_e) => bail!("failed to convert buffer to type T in pop_shared_input_data_into"),
         };
 
         // update the stack pointer to point to the element we just popped of since that is now free
-        input_data_buffer[..8].copy_from_slice(&last_element_offset_rel.to_le_bytes());
+        idb[..8].copy_from_slice(&last_element_offset_rel.to_le_bytes());
 
         // zero out popped off buffer
-        input_data_buffer[last_element_offset_rel..stack_ptr_rel].fill(0);
+        idb[last_element_offset_rel as usize..stack_ptr_rel as usize].fill(0);
 
         type_t
     }
 }
 
 pub struct OutputDataSection {
     pub ptr: *mut u8,
-    pub len: usize,
+    pub len: u64,
 }
 
 impl OutputDataSection {
-    pub fn new(ptr: *mut u8, len: usize) -> Self {
+    const STACK_PTR_SIZE: usize = size_of::<u64>();
+
+    pub fn new(ptr: *mut u8, len: u64) -> Self {
         OutputDataSection { ptr, len }
     }
 
     pub fn push_shared_output_data(&self, data: Vec<u8>) -> Result<()> {
-        let output_data_buffer = unsafe { from_raw_parts_mut(self.ptr, self.len) };
+        let shared_buffer_size = self.len as usize;
+        let odb: &mut [u8] = unsafe { from_raw_parts_mut(self.ptr, shared_buffer_size) };
 
-        if output_data_buffer.is_empty() {
-            bail!("Got a 0-size buffer in push_shared_output_data");
+        if odb.len() < Self::STACK_PTR_SIZE {
+            bail!("shared output buffer is too small");
         }
 
         // get offset to next free address on the stack
-        let mut stack_ptr_rel: usize = usize::from_le_bytes(
-            output_data_buffer[..8]
-                .try_into()
-                .expect("Shared output buffer too small"),
-        );
-
+        let mut stack_ptr_rel: u64 =
+            u64::from_le_bytes(match odb[..Self::STACK_PTR_SIZE].try_into() {
+                Ok(bytes) => bytes,
+                Err(_) => bail!("failed to get stack pointer in shared output buffer"),
+            });
+
+        // if stack_ptr_rel is 0, it means this is the first time we're using the output buffer, so
+        // we want to offset it by 8 as to not overwrite the stack_ptr location.
         if stack_ptr_rel == 0 {
             stack_ptr_rel = 8;
         }
 
         // check if the stack pointer is within the bounds of the buffer.
         // It can be equal to the size, but never greater
         // It can never be less than 8. An empty buffer's stack pointer is 8
-        if stack_ptr_rel > self.len || stack_ptr_rel < 8 {
-            bail!("Invalid stack pointer: {} in push_shared_output_data", stack_ptr_rel);
+        if stack_ptr_rel as usize > shared_buffer_size {
+            bail!("invalid stack pointer in shared output buffer");
         }
 
         // check if there is enough space in the buffer
-        let size_required = data.len() + 8; // the data plus the pointer pointing to the data
-        let size_available = self.len - stack_ptr_rel;
+        let size_required: usize = data.len() + 8; // the data plus the pointer pointing to the data
+        let size_available: usize = shared_buffer_size - stack_ptr_rel as usize;
         if size_required > size_available {
-            bail!("Not enough space in shared output buffer. Required: {}, Available: {}", size_required, size_available);
+            bail!("not enough space in shared output buffer");
         }
 
         // write the actual data
-        output_data_buffer[stack_ptr_rel..stack_ptr_rel + data.len()].copy_from_slice(&data);
+        odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(&data);
 
         // write the offset to the newly written data, to the top of the stack
-        let bytes = stack_ptr_rel.to_le_bytes();
-        output_data_buffer[stack_ptr_rel + data.len()..stack_ptr_rel + data.len() + 8]
+        let bytes: [u8; Self::STACK_PTR_SIZE] = stack_ptr_rel.to_le_bytes();
+        odb[stack_ptr_rel as usize + data.len()
+            ..stack_ptr_rel as usize + data.len() + Self::STACK_PTR_SIZE]
             .copy_from_slice(&bytes);
 
         // update stack pointer to point to next free address
-        let new_stack_ptr_rel = stack_ptr_rel + data.len() + 8;
-        output_data_buffer[0..8].copy_from_slice(&new_stack_ptr_rel.to_le_bytes());
+        let new_stack_ptr_rel: u64 =
+            (stack_ptr_rel as usize + data.len() + Self::STACK_PTR_SIZE) as u64;
+        odb[0..Self::STACK_PTR_SIZE].copy_from_slice(&new_stack_ptr_rel.to_le_bytes());
 
         Ok(())
     }
 }
 
 impl From<(u64, u64)> for InputDataSection {
     fn from((ptr, len): (u64, u64)) -> Self {
-        InputDataSection::new(ptr as *mut u8, len as usize)
+        InputDataSection::new(ptr as *mut u8, len)
     }
 }
 
 impl From<(u64, u64)> for OutputDataSection {
     fn from((ptr, len): (u64, u64)) -> Self {
-        OutputDataSection::new(ptr as *mut u8, len as usize)
+        OutputDataSection::new(ptr as *mut u8, len)
     }
 }

Diff for: src/hyperlight_common/src/outb.rs

@@ -1,5 +1,7 @@
 use core::arch;
+
 use anyhow::{bail, Result};
+
 use crate::hyperlight_peb::RunMode;
 use crate::RUNNING_MODE;
 
@@ -30,7 +32,9 @@ impl TryFrom<u16> for OutBAction {
 
 /// Issues an OUTB instruction to the specified port with the given value.
 fn hloutb(port: u16, val: u8) {
-    unsafe { arch::asm!("out dx, al", in("dx") port, in("al") val, options(preserves_flags, nomem, nostack)); }
+    unsafe {
+        arch::asm!("out dx, al", in("dx") port, in("al") val, options(preserves_flags, nomem, nostack));
+    }
 }
 
 pub fn outb(port: u16, value: u8) {

Diff for: src/hyperlight_guest/src/entrypoint.rs

@@ -18,10 +18,11 @@ use core::ffi::{c_char, CStr};
 use core::ptr::copy_nonoverlapping;
 
 use hyperlight_common::hyperlight_peb::{HyperlightPEB, RunMode};
-use log::LevelFilter;
-use spin::Once;
 use hyperlight_common::outb::{outb, OutBAction};
 use hyperlight_common::{PEB, RUNNING_MODE};
+use log::LevelFilter;
+use spin::Once;
+
 use crate::gdt::load_gdt;
 use crate::guest_error::reset_error;
 use crate::guest_function_call::dispatch_function;

Diff for: src/hyperlight_guest/src/guest_error.rs

@@ -19,9 +19,10 @@ use alloc::vec::Vec;
 use core::ffi::{c_char, CStr};
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError};
-use log::error;
 use hyperlight_common::outb::{outb, OutBAction};
 use hyperlight_common::PEB;
+use log::error;
+
 use crate::entrypoint::halt;
 
 pub(crate) fn write_error(error_code: ErrorCode, message: Option<&str>) {

Diff for: src/hyperlight_guest/src/lib.rs

@@ -96,4 +96,3 @@ pub static mut MIN_STACK_ADDRESS: u64 = 0;
 
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
     GuestFunctionRegister::new();
-

Diff for: src/hyperlight_host/src/sandbox/outb.rs

@@ -19,10 +19,11 @@ use std::sync::{Arc, Mutex};
 use hyperlight_common::flatbuffer_wrappers::function_types::ParameterValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
+use hyperlight_common::outb::OutBAction;
 use log::{Level, Record};
 use tracing::{instrument, Span};
 use tracing_log::format_trace;
-use hyperlight_common::outb::OutBAction;
+
 use super::host_funcs::HostFuncsWrapper;
 use crate::hypervisor::handlers::{OutBHandler, OutBHandlerFunction, OutBHandlerWrapper};
 use crate::mem::mgr::SandboxMemoryManager;

Diff for: src/tests/rust_guests/simpleguest/src/main.rs

@@ -40,6 +40,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel;
 use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
+use hyperlight_common::host_calling::{call_host_function, get_host_return_value, print};
 use hyperlight_common::PAGE_SIZE;
 use hyperlight_guest::entrypoint::{abort_with_code, abort_with_code_and_message};
 use hyperlight_guest::error::{HyperlightGuestError, Result};
@@ -48,7 +49,6 @@ use hyperlight_guest::guest_function_register::register_function;
 use hyperlight_guest::memory::malloc;
 use hyperlight_guest::{logging, MIN_STACK_ADDRESS};
 use log::{error, LevelFilter};
-use hyperlight_common::host_calling::{call_host_function, get_host_return_value, print};
 
 extern crate hyperlight_guest;