From ea06ab28409a82be00d1aa50440387725b9dd5c0 Mon Sep 17 00:00:00 2001 From: Alessandro Decina Date: Thu, 27 Jul 2023 20:27:04 +0700 Subject: [PATCH] cpi: add tests that check that CPI updates all accounts at once --- programs/sbf/rust/invoke/src/processor.rs | 12 +++- programs/sbf/tests/programs.rs | 77 +++++++++++++++-------- 2 files changed, 60 insertions(+), 29 deletions(-) diff --git a/programs/sbf/rust/invoke/src/processor.rs b/programs/sbf/rust/invoke/src/processor.rs index 00be6e4acc358b..81f81ec0678465 100644 --- a/programs/sbf/rust/invoke/src/processor.rs +++ b/programs/sbf/rust/invoke/src/processor.rs @@ -828,7 +828,7 @@ fn process_instruction( assert_eq!(account.data_len(), rc_box_size); // update the serialized length so we don't error out early with AccountDataSizeChanged - unsafe { *serialized_len_ptr = account.data_len() as u64 }; + unsafe { *serialized_len_ptr = rc_box_size as u64 }; // hack to avoid dropping the RcBox, which is supposed to be on the // heap but we put it into account data. If we don't do this, @@ -847,13 +847,17 @@ fn process_instruction( const INVOKE_PROGRAM_INDEX: usize = 3; const REALLOC_PROGRAM_INDEX: usize = 4; let account = &accounts[ARGUMENT_INDEX]; + let target_account_index = instruction_data[1] as usize; + let target_account = &accounts[target_account_index]; let realloc_program_id = accounts[REALLOC_PROGRAM_INDEX].key; let invoke_program_id = accounts[INVOKE_PROGRAM_INDEX].key; account.realloc(0, false).unwrap(); account.assign(realloc_program_id); + target_account.realloc(0, false).unwrap(); + target_account.assign(realloc_program_id); let rc_box_addr = - account.data.borrow_mut().as_mut_ptr() as *mut RcBox>; + target_account.data.borrow_mut().as_mut_ptr() as *mut RcBox>; let rc_box_size = mem::size_of::>>(); unsafe { std::ptr::write( @@ -876,6 +880,8 @@ fn process_instruction( ); } + let serialized_len_ptr = + unsafe { account.data.borrow_mut().as_mut_ptr().offset(-8) as *mut u64 }; // Place a RcBox> in the account data. This // allows us to test having CallerAccount::ref_to_len_in_vm in an // account region. @@ -897,6 +903,7 @@ fn process_instruction( *realloc_program_id, &[ (accounts[ARGUMENT_INDEX].key, true, false), + (target_account.key, true, false), (realloc_program_id, false, false), (invoke_program_id, false, false), ], @@ -906,6 +913,7 @@ fn process_instruction( ) .unwrap(); + unsafe { *serialized_len_ptr = rc_box_size as u64 }; // hack to avoid dropping the RcBox, which is supposed to be on the // heap but we put it into account data. If we don't do this, // dropping the Rc will cause diff --git a/programs/sbf/tests/programs.rs b/programs/sbf/tests/programs.rs index 9436573d824d11..de839fcb99fb0b 100644 --- a/programs/sbf/tests/programs.rs +++ b/programs/sbf/tests/programs.rs @@ -4066,34 +4066,57 @@ fn test_cpi_account_ownership_writability() { } ); - // Similar to the test above where we try to make CPI write into account - // data. This variant is for when direct mapping is enabled. - let account = AccountSharedData::new(42, 0, &invoke_program_id); - bank.store_account(&account_keypair.pubkey(), &account); - let instruction_data = vec![TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE, 42, 42, 42]; - let instruction = Instruction::new_with_bytes( - invoke_program_id, - &instruction_data, - account_metas.clone(), - ); - let result = bank_client.send_and_confirm_instruction(&mint_keypair, instruction); - if direct_mapping { - assert_eq!( - result.unwrap_err().unwrap(), - TransactionError::InstructionError( - 0, - InstructionError::ExternalAccountDataModified - ) + // We're going to try and make CPI write ref_to_len_in_vm into a 2nd + // account, so we add an extra one here. + let account2_keypair = Keypair::new(); + let mut account_metas = account_metas.clone(); + account_metas.push(AccountMeta::new(account2_keypair.pubkey(), false)); + + for target_account in [1, account_metas.len() as u8 - 1] { + // Similar to the test above where we try to make CPI write into account + // data. This variant is for when direct mapping is enabled. + let account = AccountSharedData::new(42, 0, &invoke_program_id); + bank.store_account(&account_keypair.pubkey(), &account); + let account = AccountSharedData::new(42, 0, &invoke_program_id); + bank.store_account(&account2_keypair.pubkey(), &account); + let instruction_data = vec![ + TEST_FORBID_LEN_UPDATE_AFTER_OWNERSHIP_CHANGE, + target_account, + 42, + 42, + ]; + let instruction = Instruction::new_with_bytes( + invoke_program_id, + &instruction_data, + account_metas.clone(), ); - } else { - // we expect this to succeed as after updating `ref_to_len_in_vm`, - // CPI will sync the actual account data between the callee and the - // caller, _always_ writing over the location pointed by - // `ref_to_len_in_vm`. To verify this, we check that the account - // data is in fact all zeroes like it is in the callee. - result.unwrap(); - let account = bank.get_account(&account_keypair.pubkey()).unwrap(); - assert_eq!(account.data(), vec![0; 40]); + let message = Message::new(&[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); + if direct_mapping { + assert_eq!( + result.unwrap_err(), + TransactionError::InstructionError( + 0, + InstructionError::ProgramFailedToComplete + ) + ); + // We haven't moved the data pointer, but ref_to_len_vm _is_ in + // the account data vm range and that's not allowed either. + assert!( + logs.iter().any(|log| log.contains("Invalid pointer")), + "{logs:?}" + ); + } else { + // we expect this to succeed as after updating `ref_to_len_in_vm`, + // CPI will sync the actual account data between the callee and the + // caller, _always_ writing over the location pointed by + // `ref_to_len_in_vm`. To verify this, we check that the account + // data is in fact all zeroes like it is in the callee. + result.unwrap(); + let account = bank.get_account(&account_keypair.pubkey()).unwrap(); + assert_eq!(account.data(), vec![0; 40]); + } } } }