fix(guest-bin): align user memory allocations (#753)

* 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]>

* Add comment on alloc_helper invariants

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

---------

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

Author: andreiltd

src/hyperlight_guest_bin/src/memory.rs

@@ -40,36 +40,51 @@ use hyperlight_guest::exit::abort_with_code;
 */
 
 // 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>();
+const HEADER_LEN: usize = size_of::<Header>();
+
+#[repr(transparent)]
+// A header that stores the layout information for the allocated memory block.
+struct Header(Layout);
 
 /// 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.
 ///
+/// # Invariants
+/// `alignment` must be non-zero and a power of two
+///
 /// # 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 data_offset = HEADER_LEN.next_multiple_of(actual_align);
+
+    let Some(total_size) = data_offset.checked_add(size) else {
+        abort_with_code(&[ErrorCode::MallocFailed as u8]);
+    };
+
+    // 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_LEN).cast::<Header>();
+        header_ptr.write(Header(layout));
+        raw_ptr.add(data_offset) as *mut c_void
     }
 }
 
@@ -80,7 +95,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,22 +110,45 @@ 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`.
 ///
 /// # Safety
 /// `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(HEADER_LEN).cast::<Header>();
+        let layout = header_ptr.read().0;
+
+        // Deallocate from the original base pointer
+        let offset = HEADER_LEN.next_multiple_of(layout.align());
+        let raw_ptr = user_ptr.sub(offset) as *mut u8;
+        alloc::alloc::dealloc(raw_ptr, layout);
     }
 }
 
@@ -134,26 +172,24 @@ 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 user_ptr = ptr as *const u8;
 
-    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 header_ptr = user_ptr.sub(HEADER_LEN).cast::<Header>();
 
-    let new_block_start =
-        unsafe { alloc::alloc::realloc(block_start as *mut u8, old_layout, total_new_size) }
-            as *mut Layout;
+        let old_layout = header_ptr.read().0;
+        let old_offset = HEADER_LEN.next_multiple_of(old_layout.align());
+        let old_user_size = old_layout.size() - old_offset;
 
-    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
+        let new_ptr = alloc_helper(size, old_layout.align(), false);
+        if new_ptr.is_null() {
+            return ptr::null_mut();
         }
+
+        let copy_size = old_user_size.min(size);
+        ptr::copy_nonoverlapping(user_ptr, new_ptr as *mut u8, copy_size);
+
+        free(ptr);
+        new_ptr
     }
 }