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

@@ -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"),
     }
 

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;

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:?}"
+        );
+    }
+}