Gate aligned program heap (#11808)

programs/bpf/benches/bpf_loader.rs

@@ -6,6 +6,7 @@ use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
 use solana_rbpf::EbpfVm;
 use solana_sdk::{
     account::Account,
+    bpf_loader,
     entrypoint_native::{ComputeMeter, InvokeContext, Logger, ProcessInstruction},
     instruction::{CompiledInstruction, InstructionError},
     message::Message,
@@ -80,7 +81,9 @@ fn bench_program_alu(bencher: &mut Bencher) {
     let mut invoke_context = MockInvokeContext::default();
 
     let elf = load_elf().unwrap();
-    let (mut vm, _) = solana_bpf_loader_program::create_vm(&elf, &[], &mut invoke_context).unwrap();
+    let (mut vm, _) =
+        solana_bpf_loader_program::create_vm(&bpf_loader::id(), &elf, &[], &mut invoke_context)
+            .unwrap();
 
     println!("Interpreted:");
     assert_eq!(

programs/bpf_loader/src/lib.rs

@@ -60,6 +60,7 @@ pub enum BPFError {
 impl UserDefinedError for BPFError {}
 
 pub fn create_vm<'a>(
+    loader_id: &Pubkey,
     prog: &'a [u8],
     parameter_accounts: &'a [KeyedAccount<'a>],
     invoke_context: &'a mut dyn InvokeContext,
@@ -68,7 +69,8 @@ pub fn create_vm<'a>(
     vm.set_verifier(bpf_verifier::check)?;
     vm.set_elf(&prog)?;
 
-    let heap_region = syscalls::register_syscalls(&mut vm, parameter_accounts, invoke_context)?;
+    let heap_region =
+        syscalls::register_syscalls(loader_id, &mut vm, parameter_accounts, invoke_context)?;
 
     Ok((vm, heap_region))
 }
@@ -148,14 +150,18 @@ pub fn process_instruction(
         {
             let compute_meter = invoke_context.get_compute_meter();
             let program_account = program.try_account_ref_mut()?;
-            let (mut vm, heap_region) =
-                match create_vm(&program_account.data, &parameter_accounts, invoke_context) {
-                    Ok(info) => info,
-                    Err(e) => {
-                        log!(logger, "Failed to create BPF VM: {}", e);
-                        return Err(BPFLoaderError::VirtualMachineCreationFailed.into());
-                    }
-                };
+            let (mut vm, heap_region) = match create_vm(
+                program_id,
+                &program_account.data,
+                &parameter_accounts,
+                invoke_context,
+            ) {
+                Ok(info) => info,
+                Err(e) => {
+                    log!(logger, "Failed to create BPF VM: {}", e);
+                    return Err(BPFLoaderError::VirtualMachineCreationFailed.into());
+                }
+            };
 
             log!(logger, "Call BPF program {}", program.unsigned_key());
             let instruction_meter = ThisInstructionMeter { compute_meter };

programs/bpf_loader/src/syscalls.rs

@@ -92,6 +92,7 @@ use crate::allocator_bump::BPFAllocator;
 const DEFAULT_HEAP_SIZE: usize = 32 * 1024;
 
 pub fn register_syscalls<'a>(
+    loader_id: &Pubkey,
     vm: &mut EbpfVm<'a, BPFError>,
     callers_keyed_accounts: &'a [KeyedAccount<'a>],
     invoke_context: &'a mut dyn InvokeContext,
@@ -147,6 +148,7 @@ pub fn register_syscalls<'a>(
     vm.register_syscall_with_context_ex(
         "sol_alloc_free_",
         Box::new(SyscallSolAllocFree {
+            aligned: *loader_id != bpf_loader_deprecated::id(),
             allocator: BPFAllocator::new(heap, MM_HEAP_START),
         }),
     )?;
@@ -347,6 +349,7 @@ impl SyscallObject<BPFError> for SyscallLogU64 {
 /// information about that memory (start address and size) is passed
 /// to the VM to use for enforcement.
 pub struct SyscallSolAllocFree {
+    aligned: bool,
     allocator: BPFAllocator,
 }
 impl SyscallObject<BPFError> for SyscallSolAllocFree {
@@ -360,7 +363,12 @@ impl SyscallObject<BPFError> for SyscallSolAllocFree {
         _ro_regions: &[MemoryRegion],
         _rw_regions: &[MemoryRegion],
     ) -> Result<u64, EbpfError<BPFError>> {
-        let layout = match Layout::from_size_align(size as usize, align_of::<u128>()) {
+        let align = if self.aligned {
+            align_of::<u128>()
+        } else {
+            align_of::<u8>()
+        };
+        let layout = match Layout::from_size_align(size as usize, align) {
             Ok(layout) => layout,
             Err(_) => return Ok(0),
         };
@@ -1195,6 +1203,7 @@ mod tests {
             let ro_regions = &[MemoryRegion::default()];
             let rw_regions = &[MemoryRegion::new_from_slice(&heap, MM_HEAP_START)];
             let mut syscall = SyscallSolAllocFree {
+                aligned: true,
                 allocator: BPFAllocator::new(heap, MM_HEAP_START),
             };
             assert_ne!(
@@ -1216,12 +1225,35 @@ mod tests {
                 0
             );
         }
-        // many small allocs
+        // many small unaligned allocs
         {
             let heap = vec![0_u8; 100];
             let ro_regions = &[MemoryRegion::default()];
             let rw_regions = &[MemoryRegion::new_from_slice(&heap, MM_HEAP_START)];
             let mut syscall = SyscallSolAllocFree {
+                aligned: false,
+                allocator: BPFAllocator::new(heap, MM_HEAP_START),
+            };
+            for _ in 0..100 {
+                assert_ne!(
+                    syscall.call(1, 0, 0, 0, 0, ro_regions, rw_regions).unwrap(),
+                    0
+                );
+            }
+            assert_eq!(
+                syscall
+                    .call(100, 0, 0, 0, 0, ro_regions, rw_regions)
+                    .unwrap(),
+                0
+            );
+        }
+        // many small aligned allocs
+        {
+            let heap = vec![0_u8; 100];
+            let ro_regions = &[MemoryRegion::default()];
+            let rw_regions = &[MemoryRegion::new_from_slice(&heap, MM_HEAP_START)];
+            let mut syscall = SyscallSolAllocFree {
+                aligned: true,
                 allocator: BPFAllocator::new(heap, MM_HEAP_START),
             };
             for _ in 0..12 {
@@ -1244,6 +1276,7 @@ mod tests {
             let ro_regions = &[MemoryRegion::default()];
             let rw_regions = &[MemoryRegion::new_from_slice(&heap, MM_HEAP_START)];
             let mut syscall = SyscallSolAllocFree {
+                aligned: true,
                 allocator: BPFAllocator::new(heap, MM_HEAP_START),
             };
             let address = syscall