70 lines
2.1 KiB
Rust
70 lines
2.1 KiB
Rust
//! @brief Example Rust-based BPF that tests out using a custom heap
|
|
|
|
use solana_program::{
|
|
account_info::AccountInfo,
|
|
entrypoint,
|
|
entrypoint::{ProgramResult, HEAP_LENGTH, HEAP_START_ADDRESS},
|
|
msg,
|
|
pubkey::Pubkey,
|
|
};
|
|
use std::{
|
|
alloc::{alloc, Layout},
|
|
mem::{align_of, size_of},
|
|
ptr::null_mut,
|
|
usize,
|
|
};
|
|
|
|
/// Developers can implement their own heap by defining their own
|
|
/// `#[global_allocator]`. The following implements a dummy for test purposes
|
|
/// but can be flushed out with whatever the developer sees fit.
|
|
struct BumpAllocator;
|
|
unsafe impl std::alloc::GlobalAlloc for BumpAllocator {
|
|
#[inline]
|
|
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
|
if layout.size() == usize::MAX - 0x42 {
|
|
// Return test value
|
|
0x42 as *mut u8
|
|
} else {
|
|
const POS_PTR: *mut usize = HEAP_START_ADDRESS as *mut usize;
|
|
const TOP_ADDRESS: usize = HEAP_START_ADDRESS + HEAP_LENGTH;
|
|
const BOTTOM_ADDRESS: usize = HEAP_START_ADDRESS + size_of::<*mut u8>();
|
|
|
|
let mut pos = *POS_PTR;
|
|
if pos == 0 {
|
|
// First time, set starting position
|
|
pos = TOP_ADDRESS;
|
|
}
|
|
pos = pos.saturating_sub(layout.size());
|
|
pos &= !(layout.align().saturating_sub(1));
|
|
if pos < BOTTOM_ADDRESS {
|
|
return null_mut();
|
|
}
|
|
*POS_PTR = pos;
|
|
pos as *mut u8
|
|
}
|
|
}
|
|
#[inline]
|
|
unsafe fn dealloc(&self, _: *mut u8, _: Layout) {
|
|
// I'm a bump allocator, I don't free
|
|
}
|
|
}
|
|
#[cfg(not(test))]
|
|
#[global_allocator]
|
|
static A: BumpAllocator = BumpAllocator;
|
|
|
|
entrypoint!(process_instruction);
|
|
#[allow(clippy::unnecessary_wraps)]
|
|
pub fn process_instruction(
|
|
_program_id: &Pubkey,
|
|
_accounts: &[AccountInfo],
|
|
_instruction_data: &[u8],
|
|
) -> ProgramResult {
|
|
msg!("Custom heap");
|
|
unsafe {
|
|
let layout = Layout::from_size_align(usize::MAX - 0x42, align_of::<u8>()).unwrap();
|
|
let ptr = alloc(layout);
|
|
assert_eq!(ptr as u64, 0x42);
|
|
}
|
|
Ok(())
|
|
}
|