hyperlight-guest-bin 0.17.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2025 The Hyperlight Authors.

use core::alloc::Layout;
use core::ffi::c_void;
use core::mem::{align_of, size_of};
use core::ptr;

use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
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.

    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 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, alignment: usize, zero: bool) -> *mut c_void {
    if size == 0 {
        return ptr::null_mut();
    }

    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]);
        }

        // 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
    }
}

/// 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.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn malloc(size: usize) -> *mut c_void {
    unsafe { alloc_helper(size, DEFAULT_ALIGN, 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.
#[unsafe(no_mangle)]
pub unsafe extern "C" fn calloc(nmemb: usize, size: usize) -> *mut c_void {
    unsafe {
        let total_size = nmemb
            .checked_mul(size)
            .expect("nmemb * size should not overflow in calloc");

        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() {
        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);
    }
}

/// 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`.
#[unsafe(no_mangle)]
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 unsafe { malloc(size) };
    }

    if size == 0 {
        // If the size is 0, treat as a free and return null
        unsafe {
            free(ptr);
        }
        return ptr::null_mut();
    }

    let user_ptr = ptr as *const u8;

    unsafe {
        let header_ptr = user_ptr.sub(HEADER_LEN).cast::<Header>();

        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;

        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
    }
}