Fix realloc by updating the stored Layout. Previously the old Layout was copied over but not updated to account for the new size, resulting in all kinds of UB and memory corruption

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_guest/build.rs

@@ -133,11 +133,6 @@ fn cargo_main() {
         cfg.define("__x86_64__", None);
         cfg.define("__LITTLE_ENDIAN__", None);
 
-        cfg.define("malloc", "hlmalloc");
-        cfg.define("calloc", "hlcalloc");
-        cfg.define("free", "hlfree");
-        cfg.define("realloc", "hlrealloc");
-
         // silence compiler warnings
         cfg.flag("-Wno-sign-compare");
         cfg.flag("-Wno-bitwise-op-parentheses");

src/hyperlight_guest/src/memory.rs

@@ -25,19 +25,42 @@ use crate::entrypoint::abort_with_code;
 
 extern crate alloc;
 
-#[no_mangle]
-pub extern "C" fn hlmalloc(size: usize) -> *mut c_void {
-    alloc_helper(size, false)
-}
+/*
+    C-wrappers for Rust's registered global allocator.
 
-pub fn alloc_helper(size: usize, zero: bool) -> *mut c_void {
-    // Allocate a block that includes space for both layout information and data
+    Each memory allocation via `malloc/calloc/realloc` is stored together with a `alloc::Layout` describing
+    the size and alignment of the allocation. This layout is stored just before the actual raw memory returned to the caller.
+
+    Example: A call to malloc(64) will allocate space for both an `alloc::Layout` and 64 bytes of memory:
+
+    ----------------------------------------------------------------------------------------
+    | Layout { size: 64 + size_of::<Layout>(), ... }    |      64 bytes of memory         | ...
+    ----------------------------------------------------------------------------------------
+                                                        ^
+                                                        |
+                                                        |
+                                                    ptr returned to caller
+*/
+
+// We assume the maximum alignment for any value is the alignment of u128.
+const MAX_ALIGN: usize = align_of::<u128>();
+
+/// Allocates a block of memory with the given size. The memory is only guaranteed to be initialized to 0s if `zero` is true, otherwise
+/// it may or may not be initialized.
+///
+/// # Safety
+/// The returned pointer must be freed with `memory::free` when it is no longer needed, otherwise memory will leak.
+unsafe fn alloc_helper(size: usize, zero: bool) -> *mut c_void {
     if size == 0 {
         return ptr::null_mut();
     }
 
-    let total_size = size + size_of::<Layout>();
-    let layout = Layout::from_size_align(total_size, align_of::<usize>()).unwrap();
+    // Allocate a block that includes space for both layout information and data
+    let total_size = size
+        .checked_add(size_of::<Layout>())
+        .expect("data and layout size should not overflow in alloc");
+    let layout = Layout::from_size_align(total_size, MAX_ALIGN).expect("Invalid layout");
+
     unsafe {
         let raw_ptr = match zero {
             true => alloc::alloc::alloc_zeroed(layout),
@@ -53,14 +76,36 @@ pub fn alloc_helper(size: usize, zero: bool) -> *mut c_void {
     }
 }
 
+/// Allocates a block of memory with the given size.
+/// The memory is not guaranteed to be initialized to 0s.
+///
+/// # Safety
+/// The returned pointer must be freed with `memory::free` when it is no longer needed, otherwise memory will leak.
+#[no_mangle]
+pub unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
+    alloc_helper(size, false)
+}
+
+/// Allocates a block of memory for an array of `nmemb` elements, each of `size` bytes.
+/// The memory is initialized to 0s.
+///
+/// # Safety
+/// The returned pointer must be freed with `memory::free` when it is no longer needed, otherwise memory will leak.
 #[no_mangle]
-pub extern "C" fn hlcalloc(n: usize, size: usize) -> *mut c_void {
-    let total_size = n * size;
+pub unsafe extern "C" fn calloc(nmemb: usize, size: usize) -> *mut c_void {
+    let total_size = nmemb
+        .checked_mul(size)
+        .expect("nmemb * size should not overflow in calloc");
+
     alloc_helper(total_size, true)
 }
 
+/// Frees the memory block pointed to by `ptr`.
+///
+/// # Safety
+/// `ptr` must be a pointer to a memory block previously allocated by `memory::malloc`, `memory::calloc`, or `memory::realloc`.
 #[no_mangle]
-pub extern "C" fn hlfree(ptr: *mut c_void) {
+pub unsafe extern "C" fn free(ptr: *mut c_void) {
     if !ptr.is_null() {
         unsafe {
             let block_start = (ptr as *const Layout).sub(1);
@@ -70,26 +115,43 @@ pub extern "C" fn hlfree(ptr: *mut c_void) {
     }
 }
 
+/// Changes the size of the memory block pointed to by `ptr` to `size` bytes. If the returned ptr is non-null,
+/// any usage of the old memory block is immediately undefined behavior.
+///
+/// # Safety
+/// `ptr` must be a pointer to a memory block previously allocated by `memory::malloc`, `memory::calloc`, or `memory::realloc`.
 #[no_mangle]
-pub extern "C" fn hlrealloc(ptr: *mut c_void, size: usize) -> *mut c_void {
+pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
     if ptr.is_null() {
         // If the pointer is null, treat as a malloc
-        return hlmalloc(size);
+        return malloc(size);
+    }
+
+    if size == 0 {
+        // If the size is 0, treat as a free and return null
+        free(ptr);
+        return ptr::null_mut();
     }
 
     unsafe {
+        let total_new_size = size
+            .checked_add(size_of::<Layout>())
+            .expect("data and layout size should not overflow in realloc");
+
         let block_start = (ptr as *const Layout).sub(1);
-        let layout = block_start.read();
-        let total_new_size = size + size_of::<Layout>();
+        let old_layout = block_start.read();
+        let new_layout = Layout::from_size_align(total_new_size, align_of::<usize>()).unwrap();
+        
         let new_block_start =
-            alloc::alloc::realloc(block_start as *mut u8, layout, total_new_size) as *mut Layout;
+            alloc::alloc::realloc(block_start as *mut u8, old_layout, total_new_size)
+                as *mut Layout;
 
         if new_block_start.is_null() {
             // Realloc failed
             abort_with_code(ErrorCode::MallocFailed as i32);
         } else {
-            // Return the pointer just after the layout information
-            // since old layout should still as it would have been copied
+            // Update the stored Layout, then return ptr to memory right after the Layout.
+            new_block_start.write(new_layout);
             new_block_start.add(1) as *mut c_void
         }
     }

src/hyperlight_host/tests/integration_test.rs

@@ -186,14 +186,14 @@ fn guest_panic() {
 }
 
 #[test]
-fn guest_hlmalloc() {
+fn guest_malloc() {
     // this test is rust-only
     let sbox1: SingleUseSandbox = new_uninit_rust().unwrap().evolve(Noop::default()).unwrap();
 
     let size_to_allocate = 2000;
     let res = sbox1
         .call_guest_function_by_name(
-            "TestHlMalloc", // uses hlmalloc
+            "TestMalloc",
             ReturnType::Int,
             Some(vec![ParameterValue::Int(size_to_allocate)]),
         )
@@ -202,7 +202,7 @@ fn guest_hlmalloc() {
 }
 
 #[test]
-fn guest_malloc() {
+fn guest_allocate_vec() {
     let sbox1: SingleUseSandbox = new_uninit().unwrap().evolve(Noop::default()).unwrap();
 
     let size_to_allocate = 2000;
@@ -220,14 +220,14 @@ fn guest_malloc() {
 
 // checks that malloc failures are captured correctly
 #[test]
-fn guest_hlmalloc_abort() {
+fn guest_malloc_abort() {
     let sbox1: SingleUseSandbox = new_uninit_rust().unwrap().evolve(Noop::default()).unwrap();
 
     let size = 20000000; // some big number that should fail when allocated
 
     let res = sbox1
         .call_guest_function_by_name(
-            "TestHlMalloc",
+            "TestMalloc",
             ReturnType::Int,
             Some(vec![ParameterValue::Int(size)]),
         )

src/tests/c_guests/c_simpleguest/main.c

@@ -19,7 +19,7 @@ double echo_double(double d) { return d; }
 
 hl_Vec *set_byte_array_to_zero(const hl_FunctionCall* params) {
   hl_Vec input = params->parameters[0].value.VecBytes;
-  uint8_t *x = hlmalloc(input.len);
+  uint8_t *x = malloc(input.len);
   for (uintptr_t i = 0; i < input.len; i++) {
     x[i] = 0;
   }
@@ -90,7 +90,7 @@ int small_var(void) {
 }
 
 int call_malloc(int32_t size) {
-  void *heap_memory = hlmalloc(size);
+  void *heap_memory = malloc(size);
   if (NULL == heap_memory) {
     hl_set_error(hl_ErrorCode_GuestError, "Malloc Failed");
   }
@@ -99,12 +99,12 @@ int call_malloc(int32_t size) {
 }
 
 int malloc_and_free(int32_t size) {
-  void *heap_memory = hlmalloc(size);
+  void *heap_memory = malloc(size);
   if (NULL == heap_memory) {
     hl_set_error(hl_ErrorCode_GuestError, "Malloc Failed");
   }
 
-  hlfree(heap_memory);
+  free(heap_memory);
 
   return size;
 }

src/tests/rust_guests/callbackguest/Cargo.lock

@@ -54,7 +54,7 @@ use hyperlight_guest::guest_function_register::register_function;
 use hyperlight_guest::host_function_call::{
     call_host_function, get_host_value_return_as_int, get_host_value_return_as_ulong,
 };
-use hyperlight_guest::memory::hlmalloc;
+use hyperlight_guest::memory::malloc;
 use hyperlight_guest::{logging, MIN_STACK_ADDRESS};
 use log::{error, LevelFilter};
 
@@ -623,11 +623,9 @@ fn execute_on_heap(_function_call: &FunctionCall) -> Result<Vec<u8>> {
     Ok(get_flatbuffer_result_from_string("fail"))
 }
 
-#[no_mangle]
-#[allow(improper_ctypes_definitions)]
-pub extern "C" fn test_rust_malloc(function_call: &FunctionCall) -> Result<Vec<u8>> {
+fn test_rust_malloc(function_call: &FunctionCall) -> Result<Vec<u8>> {
     if let ParameterValue::Int(code) = function_call.parameters.clone().unwrap()[0].clone() {
-        let ptr = hlmalloc(code as usize);
+        let ptr = unsafe { malloc(code as usize) };
         Ok(get_flatbuffer_result_from_int(ptr as i32))
     } else {
         Err(HyperlightGuestError::new(
@@ -1012,7 +1010,7 @@ pub extern "C" fn hyperlight_main() {
     register_function(guest_panic_def);
 
     let rust_malloc_def = GuestFunctionDefinition::new(
-        "TestHlMalloc".to_string(),
+        "TestMalloc".to_string(),
         Vec::from(&[ParameterType::Int]),
         ReturnType::Int,
         test_rust_malloc as i64,