//
// hardened-malloc: Global allocator using GrapheneOS allocator
// src/lib.rs: Global allocator definition
//
// Copyright (c) 2025, 2026 Ali Polatel <alip@chesswob.org>
// Based in part upon hardened_malloc-rs/src/lib.rs which is
//   Copyright (c) strawberry <strawberry@puppygock.gay>
//   SPDX-License-Identifier: Apache-2.0 OR MIT
//
// SPDX-License-Identifier: MIT

#![no_std]

use core::{
    alloc::{GlobalAlloc, Layout},
    ffi::c_void,
};

// POSIX
pub use hardened_malloc_sys::posix_memalign;
// C standard
pub use hardened_malloc_sys::{aligned_alloc, calloc, free, malloc, realloc};
// hardened_malloc extensions
pub use hardened_malloc_sys::{free_sized, malloc_object_size, malloc_object_size_fast};

// From: hardened-malloc/h_malloc.c:
// static const size_t min_align = 16;
// Keep in sync!
const MIN_ALIGN: usize = 16;

pub struct HardenedMalloc;

unsafe impl GlobalAlloc for HardenedMalloc {
    #[inline(never)]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            return malloc(layout.size()) as *mut u8;
        }

        aligned_alloc(layout.align(), layout.size()) as *mut u8
    }

    #[inline(never)]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        if layout.align() <= MIN_ALIGN {
            return calloc(layout.size(), 1) as *mut u8;
        }

        let ptr = self.alloc(layout);
        if ptr.is_null() {
            return core::ptr::null_mut();
        }

        ptr.write_bytes(0, layout.size());
        ptr
    }

    #[inline(never)]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // Deallocate memory using free_sized as necessary.
        // Pass the size of the allocation to ensure proper sized deallocation.
        if layout.align() <= MIN_ALIGN {
            free_sized(ptr as *mut c_void, layout.size());
        } else {
            free(ptr as *mut c_void);
        }
    }

    #[inline(never)]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, size: usize) -> *mut u8 {
        // Quoting realloc(3p):
        // Applications should only free the [old] space if errno(3) was changed.
        if ptr.is_null() {
            let layout = Layout::from_size_align(size, layout.align());
            return match layout {
                Ok(layout) => self.alloc(layout),
                Err(_) => core::ptr::null_mut(),
            };
        }

        if size == 0 {
            self.dealloc(ptr, layout);
            return core::ptr::null_mut();
        }

        if layout.align() <= MIN_ALIGN {
            return realloc(ptr as *mut c_void, size) as *mut u8;
        }

        let new_ptr = aligned_alloc(layout.align(), size);
        if new_ptr.is_null() {
            return core::ptr::null_mut();
        }

        let size = size.min(layout.size());
        core::ptr::copy_nonoverlapping(ptr, new_ptr as *mut u8, size);

        free(ptr as *mut c_void);
        new_ptr as *mut u8
    }
}
