[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

src/hyperlight_common/src/input_output.rs

@@ -15,133 +15,141 @@ limitations under the License.
 */
 use anyhow::{bail, Result};
 use alloc::vec::Vec;
-use core::any::type_name;
 use core::slice::from_raw_parts_mut;
 
 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"),
-        );
+        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 > self.len || stack_ptr_rel < 16 {
-            bail!("Invalid stack pointer: {} in pop_shared_input_data_into", stack_ptr_rel);
+        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)
     }
-}
+}