Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(heaphook_rust): add aligned alloc #226

Merged
merged 6 commits into from
Oct 7, 2024
Merged

feat(heaphook_rust): add aligned alloc #226

Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
bb39c9b
feat(heaphook_rust): add aligned alloc
veqcc Oct 6, 2024
143c288
use unwrap()
veqcc Oct 7, 2024
e7055ed
refactor
veqcc Oct 7, 2024
debc6fb
rename
veqcc Oct 7, 2024
67382ea
delete
veqcc Oct 7, 2024
90f908e
avoid clippy suppression
veqcc Oct 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 97 additions & 15 deletions src/agnocast_heaphook_rust/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use rlsf::Tlsf;
use std::{
alloc::Layout,
cell::Cell,
collections::HashMap,
ffi::CStr,
mem::MaybeUninit,
os::raw::c_void,
Expand Down Expand Up @@ -46,6 +47,33 @@ static ORIGINAL_REALLOC: LazyLock<ReallocType> = LazyLock::new(|| {
}
});

type PosixMemalignType = unsafe extern "C" fn(*mut *mut c_void, usize, usize) -> i32;
static ORIGINAL_POSIX_MEMALIGN: LazyLock<PosixMemalignType> = LazyLock::new(|| {
let symbol: &CStr = CStr::from_bytes_with_nul(b"posix_memalign\0").unwrap();
unsafe {
let posix_memalign_ptr: *mut c_void = libc::dlsym(libc::RTLD_NEXT, symbol.as_ptr());
std::mem::transmute(posix_memalign_ptr)
}
});

type AlignedAllocType = unsafe extern "C" fn(usize, usize) -> *mut c_void;
static ORIGINAL_ALIGNED_ALLOC: LazyLock<AlignedAllocType> = LazyLock::new(|| {
let symbol: &CStr = CStr::from_bytes_with_nul(b"aligned_alloc\0").unwrap();
unsafe {
let aligned_alloc_ptr: *mut c_void = libc::dlsym(libc::RTLD_NEXT, symbol.as_ptr());
std::mem::transmute(aligned_alloc_ptr)
}
});

type MemalignType = unsafe extern "C" fn(usize, usize) -> *mut c_void;
static ORIGINAL_MEMALIGN: LazyLock<MemalignType> = LazyLock::new(|| {
let symbol: &CStr = CStr::from_bytes_with_nul(b"memalign\0").unwrap();
unsafe {
let memalign_ptr: *mut c_void = libc::dlsym(libc::RTLD_NEXT, symbol.as_ptr());
std::mem::transmute(memalign_ptr)
}
});

const FLLEN: usize = 28; // The maximum block size is (32 << 28) - 1 = 8_589_934_591 (nearly 8GiB)
const SLLEN: usize = 64; // The worst-case internal fragmentation is ((32 << 28) / 64 - 2) = 134_217_726 (nearly 128MiB)
type FLBitmap = u32; // FLBitmap should contain at least FLLEN bits
Expand Down Expand Up @@ -91,6 +119,9 @@ static TLSF: LazyLock<Mutex<TlsfType>> = LazyLock::new(|| {
Mutex::new(tlsf)
});

static ALIGNED_TO_ORIGINAL: LazyLock<Mutex<HashMap<usize, usize>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));

fn tlsf_allocate(size: usize) -> *mut c_void {
let layout: Layout = Layout::from_size_align(size, ALIGNMENT).unwrap_or_else(|error| {
panic!("{}: size={}, alignment={}", error, size, ALIGNMENT);
Expand Down Expand Up @@ -126,6 +157,16 @@ fn tlsf_deallocate(ptr: std::ptr::NonNull<u8>) {
unsafe { tlsf.deallocate(ptr, ALIGNMENT) }
}

fn tlsf_aligned_alloc(alignment: usize, size: usize) -> *mut c_void {
let addr: usize = tlsf_allocate(size + alignment) as usize;
let aligned_addr: usize = addr + alignment - (addr % alignment);

let mut aligned_to_original = ALIGNED_TO_ORIGINAL.lock().unwrap();
aligned_to_original.insert(aligned_addr, addr);

aligned_addr as *mut c_void
}

thread_local! {
static HOOKED : Cell<bool> = const { Cell::new(false) }
}
Expand Down Expand Up @@ -158,7 +199,18 @@ pub extern "C" fn free(ptr: *mut c_void) {
unsafe { ORIGINAL_FREE(ptr) }
} else {
hooked.set(true);
tlsf_deallocate(non_null_ptr);

let mut aligned_to_original = ALIGNED_TO_ORIGINAL.lock().unwrap();

if let Some(original_addr) = aligned_to_original.get(&ptr_addr) {
let original_ptr: std::ptr::NonNull<u8> =
std::ptr::NonNull::new(*original_addr as *mut c_void as *mut u8).unwrap();
aligned_to_original.remove(&ptr_addr);
tlsf_deallocate(original_ptr);
} else {
tlsf_deallocate(non_null_ptr);
}

hooked.set(false);
}
});
Expand Down Expand Up @@ -196,7 +248,16 @@ pub extern "C" fn realloc(ptr: *mut c_void, new_size: usize) -> *mut c_void {
if !(0x40000000000..=0x50000000000).contains(&ptr_addr) {
unsafe { ORIGINAL_REALLOC(ptr, new_size) }
} else {
tlsf_reallocate(non_null_ptr, new_size)
let mut aligned_to_original = ALIGNED_TO_ORIGINAL.lock().unwrap();
if let Some(original_addr) = aligned_to_original.get(&ptr_addr) {
let original_ptr: std::ptr::NonNull<u8> =
std::ptr::NonNull::new(*original_addr as *mut c_void as *mut u8)
.unwrap();
aligned_to_original.remove(&ptr_addr);
tlsf_reallocate(original_ptr, new_size)
} else {
tlsf_reallocate(non_null_ptr, new_size)
}
}
} else {
tlsf_allocate(new_size)
Expand All @@ -209,25 +270,46 @@ pub extern "C" fn realloc(ptr: *mut c_void, new_size: usize) -> *mut c_void {
}

#[no_mangle]
pub extern "C" fn posix_memalign(
_memptr: *mut *mut c_void,
_alignment: usize,
_size: usize,
) -> i32 {
eprintln!("TODO: posix_memalign is not supported");
0
#[allow(clippy::not_unsafe_ptr_arg_deref)]
pub extern "C" fn posix_memalign(memptr: *mut *mut c_void, alignment: usize, size: usize) -> i32 {
HOOKED.with(|hooked: &Cell<bool>| {
if hooked.get() {
unsafe { ORIGINAL_POSIX_MEMALIGN(memptr, alignment, size) }
} else {
hooked.set(true);
unsafe { *memptr = tlsf_aligned_alloc(alignment, size) };
hooked.set(false);
0
}
})
}

#[no_mangle]
pub extern "C" fn aligned_alloc(_alignment: usize, _size: usize) -> *mut c_void {
eprintln!("TODO: aligned_alloc is not supported");
std::ptr::null_mut()
pub extern "C" fn aligned_alloc(alignment: usize, size: usize) -> *mut c_void {
HOOKED.with(|hooked: &Cell<bool>| {
if hooked.get() {
unsafe { ORIGINAL_ALIGNED_ALLOC(alignment, size) }
} else {
hooked.set(true);
let ret = tlsf_aligned_alloc(alignment, size);
hooked.set(false);
ret
}
})
}

#[no_mangle]
pub extern "C" fn memalign(_alignment: usize, _size: usize) -> *mut c_void {
eprintln!("TODO: memalign is not supported");
std::ptr::null_mut()
pub extern "C" fn memalign(alignment: usize, size: usize) -> *mut c_void {
HOOKED.with(|hooked: &Cell<bool>| {
if hooked.get() {
unsafe { ORIGINAL_MEMALIGN(alignment, size) }
} else {
hooked.set(true);
let ret = tlsf_aligned_alloc(alignment, size);
hooked.set(false);
ret
}
})
}

#[no_mangle]
Expand Down
Loading