From 6ec65987f7aaa9454442b77976554d96cd193a6b Mon Sep 17 00:00:00 2001 From: Alessandro Decina Date: Wed, 19 Jun 2024 15:21:24 +0000 Subject: [PATCH] programs/sbf: add test for stack/heap zeroing Add TEST_STACK_HEAP_ZEROED which tests that stack and heap regions are zeroed across reuse from the memory pool. --- programs/sbf/rust/invoke/src/lib.rs | 52 +++++++++++++++++++ programs/sbf/rust/invoke_dep/src/lib.rs | 1 + programs/sbf/tests/programs.rs | 68 ++++++++++++++++++++++++- 3 files changed, 119 insertions(+), 2 deletions(-) diff --git a/programs/sbf/rust/invoke/src/lib.rs b/programs/sbf/rust/invoke/src/lib.rs index 82461c2bc5595c..7aaa65d3c31982 100644 --- a/programs/sbf/rust/invoke/src/lib.rs +++ b/programs/sbf/rust/invoke/src/lib.rs @@ -1393,6 +1393,58 @@ fn process_instruction<'a>( account.data.borrow_mut()[write_offset] ^= 0xe5; } } + TEST_STACK_HEAP_ZEROED => { + msg!("TEST_STACK_HEAP_ZEROED"); + const MM_STACK_START: u64 = 0x200000000; + const MM_HEAP_START: u64 = 0x300000000; + const ZEROS: [u8; 256 * 1024] = [0; 256 * 1024]; + const FORTYTWOS: [u8; 256 * 1024] = [42; 256 * 1024]; + const STACK_FRAME_SIZE: usize = 4096; + const MAX_CALL_DEPTH: usize = 64; + + // Check that the heap is always zeroed. + // + // At this point the code up to here will have allocated some values on the heap. The + // bump allocator writes the current heap pointer to the start of the memory region. We + // read it to find the slice of unallocated memory and check that it's zeroed. We then + // fill this memory with a sentinel value, and in the next nested invocation check that + // it's been zeroed as expected. + let heap_len = usize::from_le_bytes(instruction_data[1..9].try_into().unwrap()); + let heap = unsafe { slice::from_raw_parts_mut(MM_HEAP_START as *mut u8, heap_len) }; + let pos = usize::from_le_bytes(heap[0..8].try_into().unwrap()) + .saturating_sub(MM_HEAP_START as usize); + assert!(heap[8..pos] == ZEROS[8..pos], "heap not zeroed"); + heap[8..pos].copy_from_slice(&FORTYTWOS[8..pos]); + + // Check that the stack is zeroed too. + // + // We don't know in which frame we are now, so we skip a few (10) frames at the start + // which might have been used by the current call stack. We check that the memory for + // the 10..MAX_CALL_DEPTH frames is zeroed. Then we write a sentinel value, and in the + // next nested invocation check that it's been zeroed. + let stack = + unsafe { slice::from_raw_parts_mut(MM_STACK_START as *mut u8, 0x100000000) }; + for i in 10..MAX_CALL_DEPTH { + let stack = &mut stack[i * STACK_FRAME_SIZE..][..STACK_FRAME_SIZE]; + assert!(stack == &ZEROS[..STACK_FRAME_SIZE], "stack not zeroed"); + stack.fill(42); + } + + // Recurse to check that the stack and heap are zeroed. + // + // We recurse until we go over max CPI depth and error out. Stack and heap allocations + // are reused across CPI, by going over max depth we ensure that it's impossible to get + // non-zeroed regions through execution. + invoke( + &create_instruction( + *program_id, + &[(program_id, false, false)], + instruction_data.to_vec(), + ), + accounts, + ) + .unwrap(); + } _ => panic!("unexpected program data"), } diff --git a/programs/sbf/rust/invoke_dep/src/lib.rs b/programs/sbf/rust/invoke_dep/src/lib.rs index f33515f602bd8f..066e900b7f9d2e 100644 --- a/programs/sbf/rust/invoke_dep/src/lib.rs +++ b/programs/sbf/rust/invoke_dep/src/lib.rs @@ -40,6 +40,7 @@ pub const TEST_CPI_INVALID_DATA_POINTER: u8 = 37; pub const TEST_CPI_CHANGE_ACCOUNT_DATA_MEMORY_ALLOCATION: u8 = 38; pub const TEST_WRITE_ACCOUNT: u8 = 39; pub const TEST_CALLEE_ACCOUNT_UPDATES: u8 = 40; +pub const TEST_STACK_HEAP_ZEROED: u8 = 41; pub const MINT_INDEX: usize = 0; pub const ARGUMENT_INDEX: usize = 1; diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs index 9fdcee23aa4749..0956a27943f180 100644 --- a/programs/sbf/tests/programs.rs +++ b/programs/sbf/tests/programs.rs @@ -1,4 +1,3 @@ -#![cfg(any(feature = "sbf_c", feature = "sbf_rust"))] #![allow(clippy::clone_on_copy)] #![allow(clippy::needless_range_loop)] #![allow(clippy::needless_borrow)] @@ -7,7 +6,6 @@ #![allow(clippy::unnecessary_cast)] #![allow(clippy::uninlined_format_args)] -#[cfg(feature = "sbf_rust")] use { itertools::izip, solana_account_decoder::parse_bpf_loader::{ @@ -4678,3 +4676,69 @@ fn test_update_callee_account() { }); } } + +#[test] +fn test_stack_heap_zeroed() { + solana_logger::setup(); + + let GenesisConfigInfo { + genesis_config, + mint_keypair, + .. + } = create_genesis_config(100_123_456_789); + + let bank = Bank::new_for_tests(&genesis_config); + + let (bank, bank_forks) = bank.wrap_with_bank_forks_for_tests(); + let mut bank_client = BankClient::new_shared(bank); + let authority_keypair = Keypair::new(); + + let (bank, invoke_program_id) = load_upgradeable_program_and_advance_slot( + &mut bank_client, + bank_forks.as_ref(), + &mint_keypair, + &authority_keypair, + "solana_sbf_rust_invoke", + ); + + let account_keypair = Keypair::new(); + let mint_pubkey = mint_keypair.pubkey(); + let account_metas = vec![ + AccountMeta::new(mint_pubkey, true), + AccountMeta::new(account_keypair.pubkey(), false), + AccountMeta::new_readonly(invoke_program_id, false), + ]; + + // Check multiple heap sizes. It's generally a good idea, and also it's needed to ensure that + // pooled heap and stack values are reused - and therefore zeroed - across executions. + for heap_len in [32usize * 1024, 64 * 1024, 128 * 1024, 256 * 1024] { + // TEST_STACK_HEAP_ZEROED will recursively check that stack and heap are zeroed until it + // reaches max CPI invoke depth. We make it fail at max depth so we're sure that there's no + // legit way to access non-zeroed stack and heap regions. + let mut instruction_data = vec![TEST_STACK_HEAP_ZEROED]; + instruction_data.extend_from_slice(&heap_len.to_le_bytes()); + + let instruction = Instruction::new_with_bytes( + invoke_program_id, + &instruction_data, + account_metas.clone(), + ); + + let message = Message::new( + &[ + ComputeBudgetInstruction::set_compute_unit_limit(1_400_000), + ComputeBudgetInstruction::request_heap_frame(heap_len as u32), + instruction, + ], + Some(&mint_pubkey), + ); + let tx = Transaction::new(&[&mint_keypair], message.clone(), bank.last_blockhash()); + let (result, _, logs) = process_transaction_and_record_inner(&bank, tx); + assert!(result.is_err(), "{result:?}"); + assert!( + logs.iter() + .any(|log| log.contains("Cross-program invocation call depth too deep")), + "{logs:?}" + ); + } +}