fix(guest-bin): align user memory allocations

The guest allocation wrapper aligned only the total memory block but did
not guarantee that the pointer returned to the user was itself properly
aligned. This patch attempts to fix that.

It also adds alloc_aligned function which would be tricky to implement
outside the guest but will make implementing other posix like allocation
API much easier.

Signed-off-by: Tomasz Andrzejak <[email protected]>

Author: andreiltd

src/hyperlight_guest_bin/src/memory.rs

@@ -25,51 +25,71 @@ use hyperlight_guest::exit::abort_with_code;
 /*
     C-wrappers for Rust's registered global allocator.
 
-    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.
+    Each memory allocation via `malloc/calloc/realloc` is stored together with a `Header` containing
+    an `alloc::Layout` and offset information. The header 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:
+    Example: A call to malloc(64) with 16-byte alignment will allocate space for the header and 64 bytes:
 
     ----------------------------------------------------------------------------------------
-    | Layout { size: 64 + size_of::<Layout>(), ... }    |      64 bytes of memory         | ...
+    | allocated memory...     | Header { layout, offset }    |      64 bytes of memory    | ...
     ----------------------------------------------------------------------------------------
-                                                        ^
-                                                        |
-                                                        |
-                                                    ptr returned to caller
+    ^                         ^                              ^
+    |                         |                              |
+    raw_ptr                   header_ptr                     ptr returned to caller
+    (from allocator)          (at data_offset - header_size) (at data_offset)
+
+    The header contains:
+    - layout: The Layout used for the entire allocation (total_size, actual_align)
+    - offset: The offset from raw_ptr to the user data (data_offset)
 */
 
 // We assume the maximum alignment for any value is the alignment of u128.
-const MAX_ALIGN: usize = align_of::<u128>();
+const DEFAULT_ALIGN: usize = align_of::<u128>();
+
+#[repr(C)]
+struct Header {
+    layout: Layout,
+    offset: usize,
+}
 
 /// 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 {
+unsafe fn alloc_helper(size: usize, alignment: usize, zero: bool) -> *mut c_void {
     if size == 0 {
         return ptr::null_mut();
     }
 
-    // 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");
+    let actual_align = alignment.max(align_of::<Header>());
+    let header_size = size_of::<Header>();
+    let data_offset = header_size.next_multiple_of(actual_align);
+    let total_size = data_offset + size;
+
+    // Create layout for entire allocation
+    let layout =
+        Layout::from_size_align(total_size, actual_align).expect("Invalid layout parameters");
 
     unsafe {
         let raw_ptr = match zero {
             true => alloc::alloc::alloc_zeroed(layout),
             false => alloc::alloc::alloc(layout),
         };
+
         if raw_ptr.is_null() {
             abort_with_code(&[ErrorCode::MallocFailed as u8]);
-        } else {
-            let layout_ptr = raw_ptr as *mut Layout;
-            layout_ptr.write(layout);
-            layout_ptr.add(1) as *mut c_void
         }
+
+        // Place Header immediately before the user data region
+        let header_ptr = raw_ptr.add(data_offset - header_size).cast::<Header>();
+
+        header_ptr.write(Header {
+            layout,
+            offset: data_offset,
+        });
+        raw_ptr.add(data_offset) as *mut c_void
     }
 }
 
@@ -80,7 +100,7 @@ unsafe fn alloc_helper(size: usize, zero: bool) -> *mut c_void {
 /// The returned pointer must be freed with `memory::free` when it is no longer needed, otherwise memory will leak.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
-    unsafe { alloc_helper(size, false) }
+    unsafe { alloc_helper(size, DEFAULT_ALIGN, false) }
 }
 
 /// Allocates a block of memory for an array of `nmemb` elements, each of `size` bytes.
@@ -95,8 +115,22 @@ pub unsafe extern "C" fn calloc(nmemb: usize, size: usize) -> *mut c_void {
             .checked_mul(size)
             .expect("nmemb * size should not overflow in calloc");
 
-        alloc_helper(total_size, true)
+        alloc_helper(total_size, DEFAULT_ALIGN, true)
+    }
+}
+
+/// Allocates aligned memory.
+///
+/// # Safety
+/// The returned pointer must be freed with `free` when it is no longer needed.
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn aligned_alloc(alignment: usize, size: usize) -> *mut c_void {
+    // Validate alignment
+    if alignment == 0 || (alignment & (alignment - 1)) != 0 {
+        return ptr::null_mut();
     }
+
+    unsafe { alloc_helper(size, alignment, false) }
 }
 
 /// Frees the memory block pointed to by `ptr`.
@@ -105,12 +139,20 @@ pub unsafe extern "C" fn calloc(nmemb: usize, size: usize) -> *mut c_void {
 /// `ptr` must be a pointer to a memory block previously allocated by `memory::malloc`, `memory::calloc`, or `memory::realloc`.
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn free(ptr: *mut c_void) {
-    if !ptr.is_null() {
-        unsafe {
-            let block_start = (ptr as *const Layout).sub(1);
-            let layout = block_start.read();
-            alloc::alloc::dealloc(block_start as *mut u8, layout)
-        }
+    if ptr.is_null() {
+        return;
+    }
+
+    let user_ptr = ptr as *const u8;
+
+    unsafe {
+        // Read the Header just before the user data
+        let header_ptr = user_ptr.sub(size_of::<Header>()).cast::<Header>();
+        let Header { layout, offset } = header_ptr.read();
+
+        // Deallocate from the original base pointer
+        let raw_ptr = user_ptr.sub(offset) as *mut u8;
+        alloc::alloc::dealloc(raw_ptr, layout);
     }
 }
 
@@ -134,26 +176,25 @@ pub unsafe extern "C" fn realloc(ptr: *mut c_void, size: usize) -> *mut c_void {
         return ptr::null_mut();
     }
 
-    let total_new_size = size
-        .checked_add(size_of::<Layout>())
-        .expect("data and layout size should not overflow in realloc");
-
-    let block_start = unsafe { (ptr as *const Layout).sub(1) };
-    let old_layout = unsafe { block_start.read() };
-    let new_layout = Layout::from_size_align(total_new_size, MAX_ALIGN).unwrap();
+    unsafe {
+        let user_ptr = ptr as *const u8;
+        let header_ptr = user_ptr.sub(size_of::<Header>()).cast::<Header>();
+        let Header {
+            layout: old_layout,
+            offset,
+        } = header_ptr.read();
+        let old_user_size = old_layout.size() - offset;
+
+        let new_ptr = alloc_helper(size, old_layout.align(), false);
+        if new_ptr.is_null() {
+            return ptr::null_mut();
+        }
 
-    let new_block_start =
-        unsafe { alloc::alloc::realloc(block_start as *mut u8, old_layout, total_new_size) }
-            as *mut Layout;
+        // Copy the lesser of old and new sizes
+        let copy_size = old_user_size.min(size);
+        ptr::copy_nonoverlapping(user_ptr, new_ptr as *mut u8, copy_size);
 
-    if new_block_start.is_null() {
-        // Realloc failed
-        abort_with_code(&[ErrorCode::MallocFailed as u8]);
-    } else {
-        // Update the stored Layout, then return ptr to memory right after the Layout.
-        unsafe {
-            new_block_start.write(new_layout);
-            new_block_start.add(1) as *mut c_void
-        }
+        free(ptr);
+        new_ptr
     }
 }