diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs index 03ed0ef397..990d4d203c 100644 --- a/program-test/src/lib.rs +++ b/program-test/src/lib.rs @@ -395,6 +395,11 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs { .set_return_data(caller, data.to_vec()) .unwrap(); } + + fn sol_get_stack_height(&self) -> u64 { + let invoke_context = get_invoke_context(); + invoke_context.get_stack_height().try_into().unwrap() + } } pub fn find_file(filename: &str) -> Option { diff --git a/program-test/tests/cpi.rs b/program-test/tests/cpi.rs index b0833c95c9..cacfc5e969 100644 --- a/program-test/tests/cpi.rs +++ b/program-test/tests/cpi.rs @@ -3,7 +3,7 @@ use { solana_sdk::{ account_info::{next_account_info, AccountInfo}, entrypoint::{ProgramResult, MAX_PERMITTED_DATA_INCREASE}, - instruction::{AccountMeta, Instruction}, + instruction::{get_stack_height, AccountMeta, Instruction}, msg, program::invoke, pubkey::Pubkey, @@ -226,3 +226,73 @@ async fn cpi_create_account() { .await .unwrap(); } + +// Process instruction to invoke into another program +fn invoker_stack_height( + _program_id: &Pubkey, + accounts: &[AccountInfo], + _input: &[u8], +) -> ProgramResult { + // if we can call `msg!` successfully, then InvokeContext exists as required + msg!("Processing invoker instruction before CPI"); + let stack_height = get_stack_height(); + assert_eq!(stack_height, 1); + let account_info_iter = &mut accounts.iter(); + let invoked_program_info = next_account_info(account_info_iter)?; + invoke( + &Instruction::new_with_bytes(*invoked_program_info.key, &[], vec![]), + &[invoked_program_info.clone()], + )?; + msg!("Processing invoker instruction after CPI"); + Ok(()) +} + +// Process instruction to be invoked by another program +#[allow(clippy::unnecessary_wraps)] +fn invoked_stack_height( + _program_id: &Pubkey, + _accounts: &[AccountInfo], + _input: &[u8], +) -> ProgramResult { + let stack_height = get_stack_height(); + assert_eq!(stack_height, 2); + Ok(()) +} + +#[tokio::test] +async fn stack_height() { + let invoker_stack_height_program_id = Pubkey::new_unique(); + let invoked_stack_height_program_id = Pubkey::new_unique(); + let mut program_test = ProgramTest::new( + "invoker_stack_height", + invoker_stack_height_program_id, + processor!(invoker_stack_height), + ); + program_test.add_program( + "invoked_stack_height", + invoked_stack_height_program_id, + processor!(invoked_stack_height), + ); + + let mut context = program_test.start_with_context().await; + let instructions = vec![Instruction::new_with_bytes( + invoker_stack_height_program_id, + &[], + vec![AccountMeta::new_readonly( + invoked_stack_height_program_id, + false, + )], + )]; + let transaction = Transaction::new_signed_with_payer( + &instructions, + Some(&context.payer.pubkey()), + &[&context.payer], + context.last_blockhash, + ); + + context + .banks_client + .process_transaction(transaction) + .await + .unwrap(); +}