Fix broken realloc, rename and mark hlmalloc/hlfree/hlcalloc/hlrealloc unsafe

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 extern "C" fn hlcalloc(n: usize, size: usize) -> *mut c_void {
-    let total_size = n * size;
+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 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, MAX_ALIGN).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/simpleguest/src/main.rs

@@ -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,