hyperlight_guest_bin/

memory.rs

/*
Copyright 2025  The Hyperlight Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

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