gaspersic/solana
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Safer invoke context (#22898)

Browse Source 
* Safer invoke context

* feedback and rebase with master
This commit is contained in:
Jack May
2022-02-03 02:34:51 -08:00
committed by GitHub
parent bd1850df25
commit cc94a93b56
8 changed files with 150 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
program-runtime/src/invoke_context.rs
View File

@ -241,10 +241,15 @@ impl<'a> InvokeContext<'a> {
let mut sysvar_cache = SysvarCache::default(); let mut sysvar_cache = SysvarCache::default();
sysvar_cache.fill_missing_entries(|pubkey| { sysvar_cache.fill_missing_entries(|pubkey| {
(0..transaction_context.get_number_of_accounts()).find_map(|index| { (0..transaction_context.get_number_of_accounts()).find_map(|index| {
if transaction_context.get_key_of_account_at_index(index) == pubkey { if transaction_context
.get_key_of_account_at_index(index)
.unwrap()
== pubkey
{
Some( Some(
transaction_context transaction_context
.get_account_at_index(index) .get_account_at_index(index)
.unwrap()
.borrow() .borrow()
.clone(), .clone(),
) )
@ -283,10 +288,13 @@ impl<'a> InvokeContext<'a> {
return Err(InstructionError::CallDepth); return Err(InstructionError::CallDepth);
} }
let program_id = program_indices.last().map(|account_index| { let program_id = program_indices
.last()
.map(|account_index| {
self.transaction_context self.transaction_context
.get_key_of_account_at_index(*account_index) .get_key_of_account_at_index(*account_index)
}); })
.transpose()?;
if program_id.is_none() if program_id.is_none()
&& self && self
.feature_set .feature_set
@ -328,12 +336,13 @@ impl<'a> InvokeContext<'a> {
{ {
let account = self let account = self
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction) .get_account_at_index(instruction_account.index_in_transaction)?
.borrow() .borrow()
.clone(); .clone();
self.pre_accounts.push(PreAccount::new( self.pre_accounts.push(PreAccount::new(
self.transaction_context self.transaction_context.get_key_of_account_at_index(
.get_key_of_account_at_index(instruction_account.index_in_transaction), instruction_account.index_in_transaction,
)?,
account, account,
)); ));
return Ok(()); return Ok(());
@ -372,26 +381,26 @@ impl<'a> InvokeContext<'a> {
let keyed_accounts = program_indices let keyed_accounts = program_indices
.iter() .iter()
.map(|account_index| { .map(|account_index| {
( Ok((
false, false,
false, false,
self.transaction_context self.transaction_context
.get_key_of_account_at_index(*account_index), .get_key_of_account_at_index(*account_index)?,
self.transaction_context self.transaction_context
.get_account_at_index(*account_index), .get_account_at_index(*account_index)?,
) ))
}) })
.chain(instruction_accounts.iter().map(|instruction_account| { .chain(instruction_accounts.iter().map(|instruction_account| {
( Ok((
instruction_account.is_signer, instruction_account.is_signer,
instruction_account.is_writable, instruction_account.is_writable,
self.transaction_context self.transaction_context
.get_key_of_account_at_index(instruction_account.index_in_transaction), .get_key_of_account_at_index(instruction_account.index_in_transaction)?,
self.transaction_context self.transaction_context
.get_account_at_index(instruction_account.index_in_transaction), .get_account_at_index(instruction_account.index_in_transaction)?,
) ))
})) }))
.collect::<Vec<_>>(); .collect::<Result<Vec<_>, InstructionError>>()?;
// Unsafe will be removed together with the keyed_accounts // Unsafe will be removed together with the keyed_accounts
self.invoke_stack.push(StackFrame::new( self.invoke_stack.push(StackFrame::new(
@ -433,7 +442,7 @@ impl<'a> InvokeContext<'a> {
// Verify all executable accounts have zero outstanding refs // Verify all executable accounts have zero outstanding refs
for account_index in program_indices.iter() { for account_index in program_indices.iter() {
self.transaction_context self.transaction_context
.get_account_at_index(*account_index) .get_account_at_index(*account_index)?
.try_borrow_mut() .try_borrow_mut()
.map_err(|_| InstructionError::AccountBorrowOutstanding)?; .map_err(|_| InstructionError::AccountBorrowOutstanding)?;
} }
@ -446,7 +455,7 @@ impl<'a> InvokeContext<'a> {
// Verify account has no outstanding references // Verify account has no outstanding references
let _ = self let _ = self
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction) .get_account_at_index(instruction_account.index_in_transaction)?
.try_borrow_mut() .try_borrow_mut()
.map_err(|_| InstructionError::AccountBorrowOutstanding)?; .map_err(|_| InstructionError::AccountBorrowOutstanding)?;
} }
@ -454,7 +463,7 @@ impl<'a> InvokeContext<'a> {
pre_account_index = pre_account_index.saturating_add(1); pre_account_index = pre_account_index.saturating_add(1);
let account = self let account = self
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction) .get_account_at_index(instruction_account.index_in_transaction)?
.borrow(); .borrow();
pre_account pre_account
.verify( .verify(
@ -521,9 +530,9 @@ impl<'a> InvokeContext<'a> {
< transaction_context.get_number_of_accounts() < transaction_context.get_number_of_accounts()
{ {
let key = transaction_context let key = transaction_context
.get_key_of_account_at_index(instruction_account.index_in_transaction); .get_key_of_account_at_index(instruction_account.index_in_transaction)?;
let account = transaction_context let account = transaction_context
.get_account_at_index(instruction_account.index_in_transaction); .get_account_at_index(instruction_account.index_in_transaction)?;
let is_writable = if before_instruction_context_push { let is_writable = if before_instruction_context_push {
instruction_context instruction_context
.try_borrow_account( .try_borrow_account(
@ -607,7 +616,7 @@ impl<'a> InvokeContext<'a> {
for instruction_account in instruction_accounts.iter() { for instruction_account in instruction_accounts.iter() {
let account_length = self let account_length = self
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction) .get_account_at_index(instruction_account.index_in_transaction)?
.borrow() .borrow()
.data() .data()
.len(); .len();
@ -630,7 +639,7 @@ impl<'a> InvokeContext<'a> {
&& prev_size && prev_size
!= self != self
.transaction_context .transaction_context
.get_account_at_index(account_index) .get_account_at_index(account_index)?
.borrow() .borrow()
.data() .data()
.len() .len()
@ -798,8 +807,12 @@ impl<'a> InvokeContext<'a> {
*compute_units_consumed = 0; *compute_units_consumed = 0;
let program_id = program_indices let program_id = program_indices
.last() .last()
.map(|index| *self.transaction_context.get_key_of_account_at_index(*index)) .map(|index| {
.unwrap_or_else(native_loader::id); self.transaction_context
.get_key_of_account_at_index(*index)
.map(|pubkey| *pubkey)
})
.unwrap_or_else(|| Ok(native_loader::id()))?;
let stack_height = self let stack_height = self
.transaction_context .transaction_context
@ -997,7 +1010,10 @@ impl<'a> InvokeContext<'a> {
} }
// Get pubkey of account at index // Get pubkey of account at index
pub fn get_key_of_account_at_index(&self, index_in_transaction: usize) -> &Pubkey { pub fn get_key_of_account_at_index(
&self,
index_in_transaction: usize,
) -> Result<&Pubkey, InstructionError> {
self.transaction_context self.transaction_context
.get_key_of_account_at_index(index_in_transaction) .get_key_of_account_at_index(index_in_transaction)
} }
@ -1364,6 +1380,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(owned_index) .get_account_at_index(owned_index)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = (MAX_DEPTH + owned_index) as u8; .data_as_mut_slice()[0] = (MAX_DEPTH + owned_index) as u8;
invoke_context invoke_context
@ -1378,11 +1395,13 @@ mod tests {
let data = invoke_context let data = invoke_context
.transaction_context .transaction_context
.get_account_at_index(not_owned_index) .get_account_at_index(not_owned_index)
.unwrap()
.borrow_mut() .borrow_mut()
.data()[0]; .data()[0];
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(not_owned_index) .get_account_at_index(not_owned_index)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = (MAX_DEPTH + not_owned_index) as u8; .data_as_mut_slice()[0] = (MAX_DEPTH + not_owned_index) as u8;
assert_eq!( assert_eq!(
@ -1393,6 +1412,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(not_owned_index) .get_account_at_index(not_owned_index)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = data; .data_as_mut_slice()[0] = data;
@ -1468,6 +1488,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(1) .get_account_at_index(1)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = 1; .data_as_mut_slice()[0] = 1;
assert_eq!( assert_eq!(
@ -1477,6 +1498,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(1) .get_account_at_index(1)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = 0; .data_as_mut_slice()[0] = 0;
@ -1484,6 +1506,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(2) .get_account_at_index(2)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = 1; .data_as_mut_slice()[0] = 1;
assert_eq!( assert_eq!(
@ -1493,6 +1516,7 @@ mod tests {
invoke_context invoke_context
.transaction_context .transaction_context
.get_account_at_index(2) .get_account_at_index(2)
.unwrap()
.borrow_mut() .borrow_mut()
.data_as_mut_slice()[0] = 0; .data_as_mut_slice()[0] = 0;

5
program-test/src/lib.rs
View File

@ -271,7 +271,8 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
for instruction_account in instruction_accounts.iter() { for instruction_account in instruction_accounts.iter() {
let account_key = invoke_context let account_key = invoke_context
.transaction_context .transaction_context
.get_key_of_account_at_index(instruction_account.index_in_transaction); .get_key_of_account_at_index(instruction_account.index_in_transaction)
.unwrap();
let account_info_index = account_infos let account_info_index = account_infos
.iter() .iter()
.position(|account_info| account_info.unsigned_key() == account_key) .position(|account_info| account_info.unsigned_key() == account_key)
@ -281,6 +282,7 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
let mut account = invoke_context let mut account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction) .get_account_at_index(instruction_account.index_in_transaction)
.unwrap()
.borrow_mut(); .borrow_mut();
account.copy_into_owner_from_slice(account_info.owner.as_ref()); account.copy_into_owner_from_slice(account_info.owner.as_ref());
account.set_data_from_slice(&account_info.try_borrow_data().unwrap()); account.set_data_from_slice(&account_info.try_borrow_data().unwrap());
@ -309,6 +311,7 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow_mut(); .borrow_mut();
let account_info = &account_infos[account_info_index]; let account_info = &account_infos[account_info_index];
**account_info.try_borrow_mut_lamports().unwrap() = account.lamports(); **account_info.try_borrow_mut_lamports().unwrap() = account.lamports();

6
programs/bpf_loader/src/serialization.rs
View File

@ -493,6 +493,7 @@ mod tests {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow(); .borrow();
assert_eq!(account.lamports(), account_info.lamports()); assert_eq!(account.lamports(), account_info.lamports());
assert_eq!(account.data(), &account_info.data.borrow()[..]); assert_eq!(account.data(), &account_info.data.borrow()[..]);
@ -518,6 +519,7 @@ mod tests {
let mut account = invoke_context let mut account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow_mut(); .borrow_mut();
account.set_lamports(0); account.set_lamports(0);
account.set_data(vec![0; 0]); account.set_data(vec![0; 0]);
@ -536,6 +538,7 @@ mod tests {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow(); .borrow();
assert_eq!(&*account, original_account); assert_eq!(&*account, original_account);
} }
@ -567,6 +570,7 @@ mod tests {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow(); .borrow();
assert_eq!(account.lamports(), account_info.lamports()); assert_eq!(account.lamports(), account_info.lamports());
assert_eq!(account.data(), &account_info.data.borrow()[..]); assert_eq!(account.data(), &account_info.data.borrow()[..]);
@ -579,6 +583,7 @@ mod tests {
let mut account = invoke_context let mut account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow_mut(); .borrow_mut();
account.set_lamports(0); account.set_lamports(0);
account.set_data(vec![0; 0]); account.set_data(vec![0; 0]);
@ -596,6 +601,7 @@ mod tests {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(index_in_transaction) .get_account_at_index(index_in_transaction)
.unwrap()
.borrow(); .borrow();
assert_eq!(&*account, original_account); assert_eq!(&*account, original_account);
} }

21
programs/bpf_loader/src/syscalls.rs
View File

@ -2546,10 +2546,12 @@ where
for instruction_account in instruction_accounts.iter() { for instruction_account in instruction_accounts.iter() {
let account = invoke_context let account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(instruction_account.index_in_transaction); .get_account_at_index(instruction_account.index_in_transaction)
.map_err(SyscallError::InstructionError)?;
let account_key = invoke_context let account_key = invoke_context
.transaction_context .transaction_context
.get_key_of_account_at_index(instruction_account.index_in_transaction); .get_key_of_account_at_index(instruction_account.index_in_transaction)
.map_err(SyscallError::InstructionError)?;
if account.borrow().executable() { if account.borrow().executable() {
// Use the known account // Use the known account
accounts.push((instruction_account.index_in_transaction, None)); accounts.push((instruction_account.index_in_transaction, None));
@ -2725,6 +2727,7 @@ fn call<'a, 'b: 'a>(
let callee_account = invoke_context let callee_account = invoke_context
.transaction_context .transaction_context
.get_account_at_index(*callee_account_index) .get_account_at_index(*callee_account_index)
.map_err(SyscallError::InstructionError)?
.borrow(); .borrow();
*caller_account.lamports = callee_account.lamports(); *caller_account.lamports = callee_account.lamports();
*caller_account.owner = *callee_account.owner(); *caller_account.owner = *callee_account.owner();
@ -3103,16 +3106,20 @@ impl<'a, 'b> SyscallObject<BpfError> for SyscallGetProcessedSiblingInstruction<'
*program_id = *program_id =
instruction_context.get_program_id(invoke_context.transaction_context); instruction_context.get_program_id(invoke_context.transaction_context);
data.clone_from_slice(instruction_context.get_instruction_data()); data.clone_from_slice(instruction_context.get_instruction_data());
let account_metas = instruction_context let account_metas = question_mark!(
instruction_context
.get_instruction_accounts_metas() .get_instruction_accounts_metas()
.iter() .iter()
.map(|meta| AccountMeta { .map(|meta| Ok(AccountMeta {
pubkey: *invoke_context pubkey: *invoke_context
.get_key_of_account_at_index(meta.index_in_transaction), .get_key_of_account_at_index(meta.index_in_transaction)
.map_err(SyscallError::InstructionError)?,
is_signer: meta.is_signer, is_signer: meta.is_signer,
is_writable: meta.is_writable, is_writable: meta.is_writable,
}) }))
.collect::<Vec<_>>(); .collect::<Result<Vec<_>, EbpfError<BpfError>>>(),
result
);
accounts.clone_from_slice(account_metas.as_slice()); accounts.clone_from_slice(account_metas.as_slice());
} }
*data_len = instruction_context.get_instruction_data().len(); *data_len = instruction_context.get_instruction_data().len();

4
runtime/src/account_rent_state.rs
View File

@ -60,7 +60,9 @@ pub(crate) fn check_rent_state(
debug!( debug!(
"Account {:?} not rent exempt, state {:?}", "Account {:?} not rent exempt, state {:?}",
transaction_context.get_key_of_account_at_index(index), transaction_context.get_key_of_account_at_index(index),
transaction_context.get_account_at_index(index).borrow(), transaction_context
.get_account_at_index(index)
.map(|account| account.borrow()),
); );
return Err(TransactionError::InvalidRentPayingAccount); return Err(TransactionError::InvalidRentPayingAccount);
} }

11
runtime/src/bank/transaction_account_state_info.rs
View File

@ -22,7 +22,8 @@ impl Bank {
(0..message.account_keys_len()) (0..message.account_keys_len())
.map(|i| { .map(|i| {
let rent_state = if message.is_writable(i) { let rent_state = if message.is_writable(i) {
let account = transaction_context.get_account_at_index(i).borrow(); let state = if let Ok(account) = transaction_context.get_account_at_index(i) {
let account = account.borrow();
// Native programs appear to be RentPaying because they carry low lamport // Native programs appear to be RentPaying because they carry low lamport
// balances; however they will never be loaded as writable // balances; however they will never be loaded as writable
@ -35,6 +36,14 @@ impl Bank {
} else { } else {
None None
}; };
debug_assert!(
state.is_some(),
"message and transaction context out of sync, fatal"
);
state
} else {
None
};
TransactionAccountStateInfo { rent_state } TransactionAccountStateInfo { rent_state }
}) })
.collect() .collect()

48
runtime/src/message_processor.rs
View File

@ -103,6 +103,7 @@ impl MessageProcessor {
let mut mut_account_ref = invoke_context let mut mut_account_ref = invoke_context
.transaction_context .transaction_context
.get_account_at_index(account_index) .get_account_at_index(account_index)
.map_err(|_| TransactionError::InvalidAccountIndex)?
.borrow_mut(); .borrow_mut();
instructions::store_current_index( instructions::store_current_index(
mut_account_ref.data_as_mut_slice(), mut_account_ref.data_as_mut_slice(),
@ -243,8 +244,14 @@ mod tests {
let program_indices = vec![vec![2]]; let program_indices = vec![vec![2]];
let executors = Rc::new(RefCell::new(Executors::default())); let executors = Rc::new(RefCell::new(Executors::default()));
let account_metas = vec![ let account_metas = vec![
AccountMeta::new(*transaction_context.get_key_of_account_at_index(0), true), AccountMeta::new(
AccountMeta::new_readonly(*transaction_context.get_key_of_account_at_index(1), false), *transaction_context.get_key_of_account_at_index(0).unwrap(),
true,
),
AccountMeta::new_readonly(
*transaction_context.get_key_of_account_at_index(1).unwrap(),
false,
),
]; ];
let message = SanitizedMessage::Legacy(Message::new( let message = SanitizedMessage::Legacy(Message::new(
@ -253,7 +260,7 @@ mod tests {
&MockSystemInstruction::Correct, &MockSystemInstruction::Correct,
account_metas.clone(), account_metas.clone(),
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let sysvar_cache = SysvarCache::default(); let sysvar_cache = SysvarCache::default();
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
@ -276,6 +283,7 @@ mod tests {
assert_eq!( assert_eq!(
transaction_context transaction_context
.get_account_at_index(0) .get_account_at_index(0)
.unwrap()
.borrow() .borrow()
.lamports(), .lamports(),
100 100
@ -283,6 +291,7 @@ mod tests {
assert_eq!( assert_eq!(
transaction_context transaction_context
.get_account_at_index(1) .get_account_at_index(1)
.unwrap()
.borrow() .borrow()
.lamports(), .lamports(),
0 0
@ -294,7 +303,7 @@ mod tests {
&MockSystemInstruction::TransferLamports { lamports: 50 }, &MockSystemInstruction::TransferLamports { lamports: 50 },
account_metas.clone(), account_metas.clone(),
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
builtin_programs, builtin_programs,
@ -326,7 +335,7 @@ mod tests {
&MockSystemInstruction::ChangeData { data: 50 }, &MockSystemInstruction::ChangeData { data: 50 },
account_metas, account_metas,
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
builtin_programs, builtin_programs,
@ -439,9 +448,18 @@ mod tests {
let program_indices = vec![vec![2]]; let program_indices = vec![vec![2]];
let executors = Rc::new(RefCell::new(Executors::default())); let executors = Rc::new(RefCell::new(Executors::default()));
let account_metas = vec![ let account_metas = vec![
AccountMeta::new(*transaction_context.get_key_of_account_at_index(0), true), AccountMeta::new(
AccountMeta::new(*transaction_context.get_key_of_account_at_index(1), false), *transaction_context.get_key_of_account_at_index(0).unwrap(),
AccountMeta::new(*transaction_context.get_key_of_account_at_index(0), false), true,
),
AccountMeta::new(
*transaction_context.get_key_of_account_at_index(1).unwrap(),
false,
),
AccountMeta::new(
*transaction_context.get_key_of_account_at_index(0).unwrap(),
false,
),
]; ];
// Try to borrow mut the same account // Try to borrow mut the same account
@ -451,7 +469,7 @@ mod tests {
&MockSystemInstruction::BorrowFail, &MockSystemInstruction::BorrowFail,
account_metas.clone(), account_metas.clone(),
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let sysvar_cache = SysvarCache::default(); let sysvar_cache = SysvarCache::default();
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
@ -485,7 +503,7 @@ mod tests {
&MockSystemInstruction::MultiBorrowMut, &MockSystemInstruction::MultiBorrowMut,
account_metas.clone(), account_metas.clone(),
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
builtin_programs, builtin_programs,
@ -515,7 +533,7 @@ mod tests {
}, },
account_metas, account_metas,
)], )],
Some(transaction_context.get_key_of_account_at_index(0)), Some(transaction_context.get_key_of_account_at_index(0).unwrap()),
)); ));
let result = MessageProcessor::process_message( let result = MessageProcessor::process_message(
builtin_programs, builtin_programs,
@ -537,6 +555,7 @@ mod tests {
assert_eq!( assert_eq!(
transaction_context transaction_context
.get_account_at_index(0) .get_account_at_index(0)
.unwrap()
.borrow() .borrow()
.lamports(), .lamports(),
80 80
@ -544,12 +563,17 @@ mod tests {
assert_eq!( assert_eq!(
transaction_context transaction_context
.get_account_at_index(1) .get_account_at_index(1)
.unwrap()
.borrow() .borrow()
.lamports(), .lamports(),
20 20
); );
assert_eq!( assert_eq!(
transaction_context.get_account_at_index(0).borrow().data(), transaction_context
.get_account_at_index(0)
.unwrap()
.borrow()
.data(),
&vec![42] &vec![42]
); );
} }

18
sdk/src/transaction_context.rs
View File

@ -96,13 +96,23 @@ impl TransactionContext {
} }
/// Searches for an account by its key /// Searches for an account by its key
pub fn get_key_of_account_at_index(&self, index_in_transaction: usize) -> &Pubkey { pub fn get_key_of_account_at_index(
&self.account_keys[index_in_transaction] &self,
index_in_transaction: usize,
) -> Result<&Pubkey, InstructionError> {
self.account_keys
.get(index_in_transaction)
.ok_or(InstructionError::NotEnoughAccountKeys)
} }
/// Searches for an account by its key /// Searches for an account by its key
pub fn get_account_at_index(&self, index_in_transaction: usize) -> &RefCell<AccountSharedData> { pub fn get_account_at_index(
&self.accounts[index_in_transaction] &self,
index_in_transaction: usize,
) -> Result<&RefCell<AccountSharedData>, InstructionError> {
self.accounts
.get(index_in_transaction)
.ok_or(InstructionError::NotEnoughAccountKeys)
} }
/// Searches for an account by its key /// Searches for an account by its key