Allow rewriting and reading of user-data

Signed-off-by: Quetzal Bradley <qbradley@microsoft.com>

Author: qbradley

src/hyperlight_guest/src/guest_handle/host_comm.rs

@@ -36,9 +36,7 @@ use crate::exit::out32;
 impl GuestHandle {
     /// Get user memory region as bytes.
     pub fn read_n_bytes_from_user_memory(&self, num: u64) -> Result<Vec<u8>> {
-        let peb_ptr = self.peb().unwrap();
-        let user_memory_region_ptr = unsafe { (*peb_ptr).init_data.ptr as *mut u8 };
-        let user_memory_region_size = unsafe { (*peb_ptr).init_data.size };
+        let (user_memory_region_ptr, user_memory_region_size) = self.get_user_memory();
 
         if num > user_memory_region_size {
             Err(HyperlightGuestError::new(
@@ -57,6 +55,14 @@ impl GuestHandle {
         }
     }
 
+    pub fn get_user_memory(&self) -> (*const u8, u64) {
+        let peb_ptr = self.peb().unwrap();
+        let user_memory_region_ptr = unsafe { (*peb_ptr).init_data.ptr as *const u8 };
+        let user_memory_region_size = unsafe { (*peb_ptr).init_data.size };
+
+        (user_memory_region_ptr, user_memory_region_size)
+    }
+
     /// Get a return value from a host function call.
     /// This usually requires a host function to be called first using
     /// `call_host_function_internal`.

src/hyperlight_guest_bin/src/host_comm.rs

@@ -70,6 +70,11 @@ pub fn read_n_bytes_from_user_memory(num: u64) -> Result<Vec<u8>> {
     handle.read_n_bytes_from_user_memory(num)
 }
 
+pub fn get_user_memory() -> (*const u8, u64) {
+    let handle = unsafe { GUEST_HANDLE };
+    handle.get_user_memory()
+}
+
 /// Print a message using the host's print function.
 ///
 /// This function requires memory to be setup to be used. In particular, the

src/hyperlight_guest_capi/src/dispatch.rs

@@ -26,7 +26,9 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest_bin::guest_function::definition::GuestFunctionDefinition;
 use hyperlight_guest_bin::guest_function::register::GuestFunctionRegister;
-use hyperlight_guest_bin::host_comm::call_host_function_without_returning_result;
+use hyperlight_guest_bin::host_comm::{
+    call_host_function_without_returning_result, get_user_memory,
+};
 
 use crate::types::{FfiFunctionCall, FfiVec};
 static mut REGISTERED_C_GUEST_FUNCTIONS: GuestFunctionRegister = GuestFunctionRegister::new();
@@ -114,3 +116,15 @@ pub extern "C" fn hl_call_host_function(function_call: &FfiFunctionCall) {
     let _ = call_host_function_without_returning_result(&func_name, Some(parameters), return_type)
         .expect("Failed to call host function");
 }
+
+#[unsafe(no_mangle)]
+pub extern "C" fn hl_get_user_memory_ptr() -> *const u8 {
+    let (user_memory_region_ptr, _) = get_user_memory();
+    user_memory_region_ptr
+}
+
+#[unsafe(no_mangle)]
+pub extern "C" fn hl_get_user_memory_size() -> u64 {
+    let (_, user_memory_region_size) = get_user_memory();
+    user_memory_region_size
+}

src/hyperlight_host/src/func/call_ctx.rs

@@ -95,6 +95,22 @@ impl MultiUseGuestCallContext {
     pub fn map_file_cow(&mut self, fp: &std::path::Path, guest_base: u64) -> Result<u64> {
         self.sbox.map_file_cow(fp, guest_base)
     }
+
+    /// Write new data to init data
+    pub fn rewrite_init_data(&mut self, user_memory: &[u8]) -> Result<()> {
+        use crate::sandbox::WrapperGetter;
+        let mem_mgr = self.sbox.get_mgr_wrapper_mut();
+        mem_mgr.rewrite_init_data(user_memory)?;
+        Ok(())
+    }
+
+    /// Read init data
+    pub fn read_init_data(&mut self, user_memory: &mut [u8]) -> Result<()> {
+        use crate::sandbox::WrapperGetter;
+        let mem_mgr = self.sbox.get_mgr_wrapper_mut();
+        mem_mgr.read_init_data(user_memory)?;
+        Ok(())
+    }
 }
 
 impl Callable for MultiUseGuestCallContext {

src/hyperlight_host/src/mem/layout.rs

@@ -91,7 +91,7 @@ pub(crate) struct SandboxMemoryLayout {
     pub(super) stack_size: usize,
     /// The heap size of this sandbox.
     pub(super) heap_size: usize,
-    init_data_size: usize,
+    pub(crate) init_data_size: usize,
 
     /// The following fields are offsets to the actual PEB struct fields.
     /// They are used when writing the PEB struct itself
@@ -114,7 +114,7 @@ pub(crate) struct SandboxMemoryLayout {
     guest_heap_buffer_offset: usize,
     guard_page_offset: usize,
     guest_user_stack_buffer_offset: usize, // the lowest address of the user stack
-    init_data_offset: usize,
+    pub(crate) init_data_offset: usize,
 
     // other
     pub(crate) peb_address: usize,

src/hyperlight_host/src/sandbox/mem_mgr.rs

@@ -120,4 +120,24 @@ impl MemMgrWrapper<HostSharedMemory> {
         self.unwrap_mgr()
             .check_stack_guard(*self.get_stack_cookie())
     }
+
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub(crate) fn rewrite_init_data(&mut self, user_memory: &[u8]) -> Result<()> {
+        let mgr = self.unwrap_mgr_mut();
+        let layout = mgr.layout;
+        let shared_mem = mgr.get_shared_mem_mut();
+        assert!(user_memory.len() <= layout.init_data_size);
+        shared_mem.copy_from_slice(user_memory, layout.init_data_offset)?;
+        Ok(())
+    }
+
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub(crate) fn read_init_data(&mut self, user_memory: &mut [u8]) -> Result<()> {
+        let mgr = self.unwrap_mgr_mut();
+        let layout = mgr.layout;
+        let shared_mem = mgr.get_shared_mem_mut();
+        assert!(user_memory.len() <= layout.init_data_size);
+        shared_mem.copy_to_slice(user_memory, layout.init_data_offset)?;
+        Ok(())
+    }
 }