Local program allocator (#12679) (#12767)

(cherry picked from commit 630eb3b907)

Co-authored-by: Jack May <jack@solana.com>

programs/bpf/Cargo.lock

@@ -1791,6 +1791,13 @@ dependencies = [
  "solana-sdk",
 ]
 
+[[package]]
+name = "solana-bpf-rust-custom-heap"
+version = "1.3.16"
+dependencies = [
+ "solana-sdk",
+]
+
 [[package]]
 name = "solana-bpf-rust-dep-crate"
 version = "1.3.16"

programs/bpf/Cargo.toml

@@ -38,6 +38,7 @@ members = [
     "rust/128bit_dep",
     "rust/alloc",
     "rust/call_depth",
+    "rust/custom_heap",
     "rust/dep_crate",
     "rust/deprecated_loader",
     "rust/dup_accounts",

programs/bpf/build.rs

@@ -68,6 +68,7 @@ fn main() {
             "128bit",
             "alloc",
             "call_depth",
+            "custom_heap",
             "dep_crate",
             "deprecated_loader",
             "dup_accounts",

programs/bpf/rust/custom_heap/Cargo.toml

@@ -0,0 +1,27 @@
+
+# Note: This crate must be built using do.sh
+
+[package]
+name = "solana-bpf-rust-custom-heap"
+version = "1.3.16"
+description = "Solana BPF test program written in Rust"
+authors = ["Solana Maintainers <maintainers@solana.foundation>"]
+repository = "https://github.com/solana-labs/solana"
+license = "Apache-2.0"
+homepage = "https://solana.com/"
+edition = "2018"
+
+[dependencies]
+solana-sdk = { path = "../../../../sdk/", version = "1.3.16", default-features = false }
+
+[features]
+custom-heap = []
+program = ["custom-heap", "solana-sdk/program"]
+default = ["program", "solana-sdk/default"]
+
+[lib]
+name = "solana_bpf_rust_custom_heap"
+crate-type = ["cdylib"]
+
+[package.metadata.docs.rs]
+targets = ["x86_64-unknown-linux-gnu"]

programs/bpf/rust/custom_heap/Xargo.toml

@@ -0,0 +1,2 @@
+[target.bpfel-unknown-unknown.dependencies.std]
+features = []

programs/bpf/rust/custom_heap/src/lib.rs

@@ -0,0 +1,68 @@
+//! @brief Example Rust-based BPF that tests out using a custom heap
+
+use solana_sdk::{
+    account_info::AccountInfo,
+    entrypoint,
+    entrypoint::{ProgramResult, HEAP_LENGTH, HEAP_START_ADDRESS},
+    info,
+    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);
+fn process_instruction(
+    _program_id: &Pubkey,
+    _accounts: &[AccountInfo],
+    _instruction_data: &[u8],
+) -> ProgramResult {
+    info!("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(())
+}

programs/bpf/tests/programs.rs

@@ -145,6 +145,7 @@ fn test_program_bpf_sanity() {
         programs.extend_from_slice(&[
             ("solana_bpf_rust_128bit", true),
             ("solana_bpf_rust_alloc", true),
+            ("solana_bpf_rust_custom_heap", true),
             ("solana_bpf_rust_dep_crate", true),
             ("solana_bpf_rust_external_spend", false),
             ("solana_bpf_rust_iter", true),
@@ -699,12 +700,12 @@ fn assert_instruction_count() {
             ("solana_bpf_rust_128bit", 543),
             ("solana_bpf_rust_alloc", 19082),
             ("solana_bpf_rust_dep_crate", 2),
-            ("solana_bpf_rust_external_spend", 485),
+            ("solana_bpf_rust_external_spend", 538),
             ("solana_bpf_rust_iter", 723),
             ("solana_bpf_rust_many_args", 231),
-            ("solana_bpf_rust_noop", 459),
-            ("solana_bpf_rust_param_passing", 54),
-            ("solana_bpf_rust_sanity", 2223),
+            ("solana_bpf_rust_noop", 512),
+            ("solana_bpf_rust_param_passing", 46),
+            ("solana_bpf_rust_sanity", 1989),
         ]);
     }
 

sdk/src/entrypoint.rs

@@ -5,8 +5,10 @@ extern crate alloc;
 use crate::{account_info::AccountInfo, program_error::ProgramError, pubkey::Pubkey};
 use alloc::vec::Vec;
 use std::{
+    alloc::Layout,
     cell::RefCell,
     mem::{align_of, size_of},
+    ptr::null_mut,
     rc::Rc,
     // Hide Result from bindgen gets confused about generics in non-generic type declarations
     result::Result as ResultGeneric,
@@ -17,24 +19,39 @@ pub type ProgramResult = ResultGeneric<(), ProgramError>;
 
 /// User implemented function to process an instruction
 ///
-/// program_id: Program ID of the currently executing program
-/// accounts: Accounts passed as part of the instruction
-/// instruction_data: Instruction data
+/// program_id: Program ID of the currently executing program accounts: Accounts
+/// passed as part of the instruction instruction_data: Instruction data
 pub type ProcessInstruction =
     fn(program_id: &Pubkey, accounts: &[AccountInfo], instruction_data: &[u8]) -> ProgramResult;
 
 /// Programs indicate success with a return value of 0
 pub const SUCCESS: u64 = 0;
 
-/// Declare the entry point of the program.
+/// Start address of the memory region used for program heap.
+pub const HEAP_START_ADDRESS: usize = 0x300000000;
+/// Length of the heap memory region used for program heap.
+pub const HEAP_LENGTH: usize = 32 * 1024;
+
+/// Declare the entry point of the program and use the default local heap
+/// implementation
 ///
-/// Deserialize the program input arguments and call
-/// the user defined `process_instruction` function.
-/// Users must call this macro otherwise an entry point for
-/// their program will not be created.
+/// Deserialize the program input arguments and call the user defined
+/// `process_instruction` function. Users must call this macro otherwise an
+/// entry point for their program will not be created.
+///
+/// If the program defines the feature `custom-heap` then the default heap
+/// implementation will not be included and the program is free to implement
+/// their own `#[global_allocator]`
 #[macro_export]
 macro_rules! entrypoint {
     ($process_instruction:ident) => {
+        #[cfg(all(not(feature = "custom-heap"), not(test)))]
+        #[global_allocator]
+        static A: $crate::entrypoint::BumpAllocator = $crate::entrypoint::BumpAllocator {
+            start: $crate::entrypoint::HEAP_START_ADDRESS,
+            len: $crate::entrypoint::HEAP_LENGTH,
+        };
+
         /// # Safety
         #[no_mangle]
         pub unsafe extern "C" fn entrypoint(input: *mut u8) -> u64 {
@@ -48,6 +65,35 @@ macro_rules! entrypoint {
     };
 }
 
+/// The bump allocator used as the default rust heap when running programs.
+pub struct BumpAllocator {
+    pub start: usize,
+    pub len: usize,
+}
+unsafe impl std::alloc::GlobalAlloc for BumpAllocator {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let pos_ptr = self.start as *mut usize;
+
+        let mut pos = *pos_ptr;
+        if pos == 0 {
+            // First time, set starting position
+            pos = self.start + self.len;
+        }
+        pos = pos.saturating_sub(layout.size());
+        pos &= !(layout.align().wrapping_sub(1));
+        if pos < self.start + size_of::<*mut u8>() {
+            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
+    }
+}
+
 /// Maximum number of bytes a program may add to an account during a single realloc
 pub const MAX_PERMITTED_DATA_INCREASE: usize = 1_024 * 10;
 
@@ -142,3 +188,70 @@ pub unsafe fn deserialize<'a>(input: *mut u8) -> (&'a Pubkey, Vec<AccountInfo<'a
 
     (program_id, accounts, instruction_data)
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use std::alloc::GlobalAlloc;
+
+    #[test]
+    fn test_bump_allocator() {
+        // alloc the entire
+        {
+            let heap = vec![0u8; 128];
+            let allocator = BumpAllocator {
+                start: heap.as_ptr() as *const _ as usize,
+                len: heap.len(),
+            };
+            for i in 0..128 - size_of::<*mut u8>() {
+                let ptr = unsafe {
+                    allocator.alloc(Layout::from_size_align(1, size_of::<u8>()).unwrap())
+                };
+                assert_eq!(
+                    ptr as *const _ as usize,
+                    heap.as_ptr() as *const _ as usize + heap.len() - 1 - i
+                );
+            }
+            assert_eq!(null_mut(), unsafe {
+                allocator.alloc(Layout::from_size_align(1, 1).unwrap())
+            });
+        }
+        // check alignment
+        {
+            let heap = vec![0u8; 128];
+            let allocator = BumpAllocator {
+                start: heap.as_ptr() as *const _ as usize,
+                len: heap.len(),
+            };
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(1, size_of::<u8>()).unwrap()) };
+            assert_eq!(0, ptr.align_offset(size_of::<u8>()));
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(1, size_of::<u16>()).unwrap()) };
+            assert_eq!(0, ptr.align_offset(size_of::<u16>()));
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(1, size_of::<u32>()).unwrap()) };
+            assert_eq!(0, ptr.align_offset(size_of::<u32>()));
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(1, size_of::<u64>()).unwrap()) };
+            assert_eq!(0, ptr.align_offset(size_of::<u64>()));
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(1, size_of::<u128>()).unwrap()) };
+            assert_eq!(0, ptr.align_offset(size_of::<u128>()));
+            let ptr = unsafe { allocator.alloc(Layout::from_size_align(1, 64).unwrap()) };
+            assert_eq!(0, ptr.align_offset(64));
+        }
+        // alloc entire block (minus the pos ptr)
+        {
+            let heap = vec![0u8; 128];
+            let allocator = BumpAllocator {
+                start: heap.as_ptr() as *const _ as usize,
+                len: heap.len(),
+            };
+            let ptr =
+                unsafe { allocator.alloc(Layout::from_size_align(120, size_of::<u8>()).unwrap()) };
+            assert_ne!(ptr, null_mut());
+            assert_eq!(0, ptr.align_offset(size_of::<u64>()));
+        }
+    }
+}