From edb20d6909236c44f0bb520bdb777f6a7496ac5a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= Date: Thu, 30 Dec 2021 15:46:36 +0100 Subject: [PATCH] Splits index of InstructionAccount into index_in_transaction and index_in_caller. (#22165) --- program-runtime/src/invoke_context.rs | 167 +++++++++------ program-test/src/lib.rs | 51 ++--- programs/bpf_loader/src/serialization.rs | 7 +- programs/bpf_loader/src/syscalls.rs | 34 ++- rbpf-cli/src/main.rs | 3 +- runtime/src/message_processor.rs | 11 +- sdk/src/transaction_context.rs | 256 +++++++++++++++++------ 7 files changed, 341 insertions(+), 188 deletions(-) diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs index 3bd73dec5..026d3496e 100644 --- a/program-runtime/src/invoke_context.rs +++ b/program-runtime/src/invoke_context.rs @@ -257,16 +257,19 @@ impl<'a> InvokeContext<'a> { } self.pre_accounts = Vec::with_capacity(instruction_accounts.len()); - let mut work = |_index_in_instruction: usize, entry: &InstructionAccount| { - if entry.index < self.transaction_context.get_number_of_accounts() { + let mut work = |_index_in_instruction: usize, + instruction_account: &InstructionAccount| { + if instruction_account.index_in_transaction + < self.transaction_context.get_number_of_accounts() + { let account = self .transaction_context - .get_account_at_index(entry.index) + .get_account_at_index(instruction_account.index_in_transaction) .borrow() .clone(); self.pre_accounts.push(PreAccount::new( self.transaction_context - .get_key_of_account_at_index(entry.index), + .get_key_of_account_at_index(instruction_account.index_in_transaction), account, )); return Ok(()); @@ -308,9 +311,9 @@ impl<'a> InvokeContext<'a> { instruction_account.is_signer, instruction_account.is_writable, self.transaction_context - .get_key_of_account_at_index(instruction_account.index), + .get_key_of_account_at_index(instruction_account.index_in_transaction), self.transaction_context - .get_account_at_index(instruction_account.index), + .get_account_at_index(instruction_account.index_in_transaction), ) })) .collect::>(); @@ -361,7 +364,7 @@ impl<'a> InvokeContext<'a> { // Verify account has no outstanding references let _ = self .transaction_context - .get_account_at_index(instruction_account.index) + .get_account_at_index(instruction_account.index_in_transaction) .try_borrow_mut() .map_err(|_| InstructionError::AccountBorrowOutstanding)?; } @@ -369,7 +372,7 @@ impl<'a> InvokeContext<'a> { pre_account_index = pre_account_index.saturating_add(1); let account = self .transaction_context - .get_account_at_index(instruction_account.index) + .get_account_at_index(instruction_account.index_in_transaction) .borrow(); pre_account .verify( @@ -428,10 +431,13 @@ impl<'a> InvokeContext<'a> { // Verify the per-account instruction results let (mut pre_sum, mut post_sum) = (0_u128, 0_u128); let mut work = |index_in_instruction: usize, instruction_account: &InstructionAccount| { - if instruction_account.index < transaction_context.get_number_of_accounts() { - let key = - transaction_context.get_key_of_account_at_index(instruction_account.index); - let account = transaction_context.get_account_at_index(instruction_account.index); + if instruction_account.index_in_transaction + < transaction_context.get_number_of_accounts() + { + let key = transaction_context + .get_key_of_account_at_index(instruction_account.index_in_transaction); + let account = transaction_context + .get_account_at_index(instruction_account.index_in_transaction); let is_writable = if let Some(caller_write_privileges) = caller_write_privileges { caller_write_privileges[index_in_instruction] } else { @@ -502,11 +508,11 @@ impl<'a> InvokeContext<'a> { for instruction_account in instruction_accounts.iter() { let account_length = self .transaction_context - .get_account_at_index(instruction_account.index) + .get_account_at_index(instruction_account.index_in_transaction) .borrow() .data() .len(); - prev_account_sizes.push((instruction_account.index, account_length)); + prev_account_sizes.push((instruction_account.index_in_transaction, account_length)); } self.process_instruction( @@ -552,10 +558,11 @@ impl<'a> InvokeContext<'a> { // Finds the index of each account in the instruction by its pubkey. // Then normalizes / unifies the privileges of duplicate accounts. // Note: This works like visit_each_account_once() and is an O(n^2) algorithm too. + let caller_keyed_accounts = self.get_keyed_accounts()?; let mut deduplicated_instruction_accounts: Vec = Vec::new(); let mut duplicate_indicies = Vec::with_capacity(instruction.accounts.len()); for account_meta in instruction.accounts.iter() { - let account_index = self + let index_in_transaction = self .transaction_context .find_index_of_account(&account_meta.pubkey) .ok_or_else(|| { @@ -566,37 +573,21 @@ impl<'a> InvokeContext<'a> { ); InstructionError::MissingAccount })?; - if let Some(duplicate_index) = deduplicated_instruction_accounts - .iter() - .position(|instruction_account| instruction_account.index == account_index) + if let Some(duplicate_index) = + deduplicated_instruction_accounts + .iter() + .position(|instruction_account| { + instruction_account.index_in_transaction == index_in_transaction + }) { duplicate_indicies.push(duplicate_index); let instruction_account = &mut deduplicated_instruction_accounts[duplicate_index]; instruction_account.is_signer |= account_meta.is_signer; instruction_account.is_writable |= account_meta.is_writable; } else { - duplicate_indicies.push(deduplicated_instruction_accounts.len()); - deduplicated_instruction_accounts.push(InstructionAccount { - index: account_index, - is_signer: account_meta.is_signer, - is_writable: account_meta.is_writable, - }); - } - } - let instruction_accounts = duplicate_indicies - .into_iter() - .map(|duplicate_index| deduplicated_instruction_accounts[duplicate_index].clone()) - .collect(); - - // Check for privilege escalation - let caller_keyed_accounts = self.get_instruction_keyed_accounts()?; - let caller_write_privileges = instruction - .accounts - .iter() - .map(|account_meta| { - let keyed_account = caller_keyed_accounts + let index_in_caller = caller_keyed_accounts .iter() - .find(|keyed_account| *keyed_account.unsigned_key() == account_meta.pubkey) + .position(|keyed_account| *keyed_account.unsigned_key() == account_meta.pubkey) .ok_or_else(|| { ic_msg!( self, @@ -605,24 +596,47 @@ impl<'a> InvokeContext<'a> { ); InstructionError::MissingAccount })?; + duplicate_indicies.push(deduplicated_instruction_accounts.len()); + deduplicated_instruction_accounts.push(InstructionAccount { + index_in_transaction, + index_in_caller, + is_signer: account_meta.is_signer, + is_writable: account_meta.is_writable, + }); + } + } + let instruction_accounts: Vec = duplicate_indicies + .into_iter() + .map(|duplicate_index| deduplicated_instruction_accounts[duplicate_index].clone()) + .collect(); + + // Check for privilege escalation + let caller_write_privileges = instruction_accounts + .iter() + .map(|instruction_account| { + let keyed_account = &caller_keyed_accounts[instruction_account.index_in_caller]; // Readonly in caller cannot become writable in callee - if account_meta.is_writable && !keyed_account.is_writable() { + if instruction_account.is_writable && !keyed_account.is_writable() { ic_msg!( self, "{}'s writable privilege escalated", - account_meta.pubkey, + keyed_account.unsigned_key(), ); return Err(InstructionError::PrivilegeEscalation); } // To be signed in the callee, // it must be either signed in the caller or by the program - if account_meta.is_signer + if instruction_account.is_signer && !(keyed_account.signer_key().is_some() - || signers.contains(&account_meta.pubkey)) + || signers.contains(keyed_account.unsigned_key())) { - ic_msg!(self, "{}'s signer privilege escalated", account_meta.pubkey); + ic_msg!( + self, + "{}'s signer privilege escalated", + keyed_account.unsigned_key() + ); return Err(InstructionError::PrivilegeEscalation); } @@ -720,7 +734,7 @@ impl<'a> InvokeContext<'a> { data: instruction_data.to_vec(), accounts: instruction_accounts .iter() - .map(|instruction_account| instruction_account.index as u8) + .map(|instruction_account| instruction_account.index_in_transaction as u8) .collect(), }; instruction_recorder @@ -907,16 +921,21 @@ pub struct MockInvokeContextPreparation { pub fn prepare_mock_invoke_context( transaction_accounts: Vec, instruction_accounts: Vec, + program_indices: &[usize], ) -> MockInvokeContextPreparation { let instruction_accounts = instruction_accounts .iter() - .map(|account_meta| InstructionAccount { - index: transaction_accounts + .map(|account_meta| { + let index_in_transaction = transaction_accounts .iter() .position(|(key, _account)| *key == account_meta.pubkey) - .unwrap_or(transaction_accounts.len()), - is_signer: account_meta.is_signer, - is_writable: account_meta.is_writable, + .unwrap_or(transaction_accounts.len()); + InstructionAccount { + index_in_transaction, + index_in_caller: program_indices.len().saturating_add(index_in_transaction), + is_signer: account_meta.is_signer, + is_writable: account_meta.is_writable, + } }) .collect(); MockInvokeContextPreparation { @@ -950,7 +969,8 @@ pub fn with_mock_invoke_context R>( is_signer: false, is_writable: false, }]; - let preparation = prepare_mock_invoke_context(transaction_accounts, instruction_accounts); + let preparation = + prepare_mock_invoke_context(transaction_accounts, instruction_accounts, &program_indices); let transaction_context = TransactionContext::new( preparation.transaction_accounts, ComputeBudget::default().max_invoke_depth, @@ -972,9 +992,10 @@ pub fn mock_process_instruction_with_sysvars( sysvars: &[(Pubkey, Vec)], process_instruction: ProcessInstructionWithContext, ) -> Vec { - let mut preparation = prepare_mock_invoke_context(transaction_accounts, instruction_accounts); + program_indices.insert(0, transaction_accounts.len()); + let mut preparation = + prepare_mock_invoke_context(transaction_accounts, instruction_accounts, &program_indices); let processor_account = AccountSharedData::new(0, 0, &solana_sdk::native_loader::id()); - program_indices.insert(0, preparation.transaction_accounts.len()); preparation .transaction_accounts .push((*loader_id, processor_account)); @@ -1023,7 +1044,7 @@ fn visit_each_account_once( // Note: This is an O(n^2) algorithm, // but performed on a very small slice and requires no heap allocations for before in instruction_accounts[..index].iter() { - if before.index == instruction_account.index { + if before.index_in_transaction == instruction_account.index_in_transaction { continue 'root; // skip dups } } @@ -1058,7 +1079,7 @@ mod tests { let mut work = |index_in_instruction: usize, entry: &InstructionAccount| { unique_entries += 1; index_sum_a += index_in_instruction; - index_sum_b += entry.index; + index_sum_b += entry.index_in_transaction; Ok(()) }; visit_each_account_once(accounts, &mut work).unwrap(); @@ -1070,22 +1091,26 @@ mod tests { (3, 3, 19), do_work(&[ InstructionAccount { - index: 7, + index_in_transaction: 7, + index_in_caller: 0, is_signer: false, is_writable: false, }, InstructionAccount { - index: 3, + index_in_transaction: 3, + index_in_caller: 1, is_signer: false, is_writable: false, }, InstructionAccount { - index: 9, + index_in_transaction: 9, + index_in_caller: 2, is_signer: false, is_writable: false, }, InstructionAccount { - index: 3, + index_in_transaction: 3, + index_in_caller: 1, is_signer: false, is_writable: false, }, @@ -1180,7 +1205,8 @@ mod tests { AccountSharedData::new(index as u64, 1, &invoke_stack[index]), )); instruction_accounts.push(InstructionAccount { - index, + index_in_transaction: index, + index_in_caller: 1 + index, is_signer: false, is_writable: true, }); @@ -1191,12 +1217,13 @@ mod tests { AccountSharedData::new(1, 1, &solana_sdk::pubkey::Pubkey::default()), )); instruction_accounts.push(InstructionAccount { - index, + index_in_transaction: index, + index_in_caller: 1 + index, is_signer: false, is_writable: false, }); } - let transaction_context = TransactionContext::new(accounts, 1); + let transaction_context = TransactionContext::new(accounts, MAX_DEPTH); let mut invoke_context = InvokeContext::new_mock(&transaction_context, &[]); // Check call depth increases and has a limit @@ -1217,12 +1244,14 @@ mod tests { let not_owned_index = owned_index - 1; let instruction_accounts = vec![ InstructionAccount { - index: not_owned_index, + index_in_transaction: not_owned_index, + index_in_caller: 1 + not_owned_index, is_signer: false, is_writable: true, }, InstructionAccount { - index: owned_index, + index_in_transaction: owned_index, + index_in_caller: 1 + owned_index, is_signer: false, is_writable: true, }, @@ -1318,8 +1347,9 @@ mod tests { let instruction_accounts = metas .iter() .enumerate() - .map(|(account_index, account_meta)| InstructionAccount { - index: account_index, + .map(|(index_in_transaction, account_meta)| InstructionAccount { + index_in_transaction, + index_in_caller: program_indices.len() + index_in_transaction, is_signer: account_meta.is_signer, is_writable: account_meta.is_writable, }) @@ -1436,8 +1466,9 @@ mod tests { let instruction_accounts = metas .iter() .enumerate() - .map(|(account_index, account_meta)| InstructionAccount { - index: account_index, + .map(|(index_in_transaction, account_meta)| InstructionAccount { + index_in_transaction, + index_in_caller: program_indices.len() + index_in_transaction, is_signer: account_meta.is_signer, is_writable: account_meta.is_writable, }) diff --git a/program-test/src/lib.rs b/program-test/src/lib.rs index 02c2be435..f179d1a67 100644 --- a/program-test/src/lib.rs +++ b/program-test/src/lib.rs @@ -249,30 +249,30 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs { .prepare_instruction(instruction, &signers) .unwrap(); - // Convert AccountInfos into Accounts - let mut accounts = Vec::with_capacity(instruction_accounts.len()); + // Copy caller's account_info modifications into invoke_context accounts + let mut account_indices = Vec::with_capacity(instruction_accounts.len()); for instruction_account in instruction_accounts.iter() { let account_key = invoke_context .transaction_context - .get_key_of_account_at_index(instruction_account.index); - let account_info = account_infos + .get_key_of_account_at_index(instruction_account.index_in_transaction); + let account_info_index = account_infos .iter() - .find(|account_info| account_info.unsigned_key() == account_key) + .position(|account_info| account_info.unsigned_key() == account_key) .ok_or(InstructionError::MissingAccount) .unwrap(); - { - let mut account = invoke_context - .transaction_context - .get_account_at_index(instruction_account.index) - .borrow_mut(); - account.copy_into_owner_from_slice(account_info.owner.as_ref()); - account.set_data_from_slice(&account_info.try_borrow_data().unwrap()); - account.set_lamports(account_info.lamports()); - account.set_executable(account_info.executable); - account.set_rent_epoch(account_info.rent_epoch); - } + let account_info = &account_infos[account_info_index]; + let mut account = invoke_context + .transaction_context + .get_account_at_index(instruction_account.index_in_transaction) + .borrow_mut(); + account.copy_into_owner_from_slice(account_info.owner.as_ref()); + account.set_data_from_slice(&account_info.try_borrow_data().unwrap()); + account.set_lamports(account_info.lamports()); + account.set_executable(account_info.executable); + account.set_rent_epoch(account_info.rent_epoch); if instruction_account.is_writable { - accounts.push((instruction_account.index, account_info)); + account_indices + .push((instruction_account.index_in_transaction, account_info_index)); } } @@ -285,22 +285,23 @@ impl solana_sdk::program_stubs::SyscallStubs for SyscallStubs { ) .map_err(|err| ProgramError::try_from(err).unwrap_or_else(|err| panic!("{}", err)))?; - // Copy writeable account modifications back into the caller's AccountInfos - for (account_index, account_info) in accounts.into_iter() { + // Copy invoke_context accounts modifications into caller's account_info + for (index_in_transaction, account_info_index) in account_indices.into_iter() { let account = invoke_context .transaction_context - .get_account_at_index(account_index); - let account_borrow = account.borrow(); - **account_info.try_borrow_mut_lamports().unwrap() = account_borrow.lamports(); + .get_account_at_index(index_in_transaction) + .borrow_mut(); + let account_info = &account_infos[account_info_index]; + **account_info.try_borrow_mut_lamports().unwrap() = account.lamports(); let mut data = account_info.try_borrow_mut_data()?; - let new_data = account_borrow.data(); - if account_info.owner != account_borrow.owner() { + let new_data = account.data(); + if account_info.owner != account.owner() { // TODO Figure out a better way to allow the System Program to set the account owner #[allow(clippy::transmute_ptr_to_ptr)] #[allow(mutable_transmutes)] let account_info_mut = unsafe { transmute::<&Pubkey, &mut Pubkey>(account_info.owner) }; - *account_info_mut = *account_borrow.owner(); + *account_info_mut = *account.owner(); } // TODO: Figure out how to allow the System Program to resize the account data assert!( diff --git a/programs/bpf_loader/src/serialization.rs b/programs/bpf_loader/src/serialization.rs index 994e060e4..db6c4297d 100644 --- a/programs/bpf_loader/src/serialization.rs +++ b/programs/bpf_loader/src/serialization.rs @@ -451,8 +451,11 @@ mod tests { ]; let instruction_data = vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]; let program_indices = [0]; - let preparation = - prepare_mock_invoke_context(transaction_accounts.clone(), instruction_accounts); + let preparation = prepare_mock_invoke_context( + transaction_accounts.clone(), + instruction_accounts, + &program_indices, + ); let transaction_context = TransactionContext::new(preparation.transaction_accounts, 1); let mut invoke_context = InvokeContext::new_mock(&transaction_context, &[]); invoke_context diff --git a/programs/bpf_loader/src/syscalls.rs b/programs/bpf_loader/src/syscalls.rs index 15de33ac6..3e382666a 100644 --- a/programs/bpf_loader/src/syscalls.rs +++ b/programs/bpf_loader/src/syscalls.rs @@ -2203,8 +2203,13 @@ where F: Fn(&T, &InvokeContext) -> Result, EbpfError>, { let keyed_accounts = invoke_context - .get_instruction_keyed_accounts() + .get_keyed_accounts() .map_err(SyscallError::InstructionError)?; + let number_of_program_accounts = keyed_accounts.len() + - invoke_context + .get_instruction_keyed_accounts() + .map_err(SyscallError::InstructionError)? + .len(); let mut accounts = Vec::with_capacity(instruction_accounts.len().saturating_add(1)); let program_account_index = program_indices @@ -2217,13 +2222,13 @@ where for instruction_account in instruction_accounts.iter() { let account = invoke_context .transaction_context - .get_account_at_index(instruction_account.index); + .get_account_at_index(instruction_account.index_in_transaction); let account_key = invoke_context .transaction_context - .get_key_of_account_at_index(instruction_account.index); + .get_key_of_account_at_index(instruction_account.index_in_transaction); if account.borrow().executable() { // Use the known account - accounts.push((instruction_account.index, None)); + accounts.push((instruction_account.index_in_transaction, None)); } else if let Some(caller_account_index) = account_info_keys.iter().position(|key| *key == account_key) { @@ -2238,20 +2243,11 @@ where account.set_rent_epoch(caller_account.rent_epoch); } let caller_account = if instruction_account.is_writable { - if let Some(orig_data_len_index) = keyed_accounts - .iter() - .position(|keyed_account| keyed_account.unsigned_key() == account_key) - .map(|index| { - // index starts at first instruction account - index - keyed_accounts.len().saturating_sub(orig_data_lens.len()) - }) - .and_then(|index| { - if index >= orig_data_lens.len() { - None - } else { - Some(index) - } - }) + let orig_data_len_index = instruction_account + .index_in_caller + .saturating_sub(number_of_program_accounts); + if keyed_accounts[instruction_account.index_in_caller].unsigned_key() == account_key + && orig_data_len_index < orig_data_lens.len() { caller_account.original_data_len = orig_data_lens[orig_data_len_index]; } else { @@ -2269,7 +2265,7 @@ where } else { None }; - accounts.push((instruction_account.index, caller_account)); + accounts.push((instruction_account.index_in_transaction, caller_account)); } else { ic_msg!( invoke_context, diff --git a/rbpf-cli/src/main.rs b/rbpf-cli/src/main.rs index e43939d31..79b6fdd60 100644 --- a/rbpf-cli/src/main.rs +++ b/rbpf-cli/src/main.rs @@ -209,8 +209,9 @@ native machine code before execting it in the virtual machine.", input.instruction_data } }; - let preparation = prepare_mock_invoke_context(transaction_accounts, instruction_accounts); let program_indices = [0, 1]; + let preparation = + prepare_mock_invoke_context(transaction_accounts, instruction_accounts, &program_indices); let transaction_context = TransactionContext::new(preparation.transaction_accounts, 1); let mut invoke_context = InvokeContext::new_mock(&transaction_context, &[]); invoke_context diff --git a/runtime/src/message_processor.rs b/runtime/src/message_processor.rs index 4c9d4beb7..2796e1fc2 100644 --- a/runtime/src/message_processor.rs +++ b/runtime/src/message_processor.rs @@ -117,12 +117,13 @@ impl MessageProcessor { let instruction_accounts = instruction .accounts .iter() - .map(|account_index| { - let account_index = *account_index as usize; + .map(|index_in_transaction| { + let index_in_transaction = *index_in_transaction as usize; InstructionAccount { - index: account_index, - is_signer: message.is_signer(account_index), - is_writable: message.is_writable(account_index), + index_in_transaction, + index_in_caller: program_indices.len().saturating_add(index_in_transaction), + is_signer: message.is_signer(index_in_transaction), + is_writable: message.is_writable(index_in_transaction), } }) .collect::>(); diff --git a/sdk/src/transaction_context.rs b/sdk/src/transaction_context.rs index 555a9226d..1e0badc2f 100644 --- a/sdk/src/transaction_context.rs +++ b/sdk/src/transaction_context.rs @@ -3,6 +3,7 @@ use crate::{ account::{AccountSharedData, ReadableAccount, WritableAccount}, instruction::InstructionError, + lamports::LamportsError, pubkey::Pubkey, }; use std::cell::{RefCell, RefMut}; @@ -11,7 +12,8 @@ pub type TransactionAccount = (Pubkey, AccountSharedData); #[derive(Clone, Debug)] pub struct InstructionAccount { - pub index: usize, + pub index_in_transaction: usize, + pub index_in_caller: usize, pub is_signer: bool, pub is_writable: bool, } @@ -19,6 +21,7 @@ pub struct InstructionAccount { /// Loaded transaction shared between runtime and programs. /// /// This context is valid for the entire duration of a transaction being processed. +#[derive(Debug)] pub struct TransactionContext { account_keys: Vec, accounts: Vec>, @@ -98,12 +101,12 @@ impl TransactionContext { /// Gets an InstructionContext by its height in the stack pub fn get_instruction_context_at( &self, - instruction_context_height: usize, + level: usize, ) -> Result<&InstructionContext, InstructionError> { - if instruction_context_height >= self.instruction_context_stack.len() { + if level >= self.instruction_context_stack.len() { return Err(InstructionError::CallDepth); } - Ok(&self.instruction_context_stack[instruction_context_height]) + Ok(&self.instruction_context_stack[level]) } /// Gets the max height of the InstructionContext stack @@ -111,66 +114,36 @@ impl TransactionContext { self.instruction_context_capacity } - /// Gets the height of the current InstructionContext - pub fn get_instruction_context_height(&self) -> usize { - self.instruction_context_stack.len().saturating_sub(1) + /// Gets the level of the next InstructionContext + pub fn get_instruction_context_stack_height(&self) -> usize { + self.instruction_context_stack.len() } /// Returns the current InstructionContext pub fn get_current_instruction_context(&self) -> Result<&InstructionContext, InstructionError> { - self.get_instruction_context_at(self.get_instruction_context_height()) - } - - /// Gets the last program account of the current InstructionContext - pub fn try_borrow_program_account(&self) -> Result { - let instruction_context = self.get_current_instruction_context()?; - instruction_context.try_borrow_account( - self, - instruction_context - .number_of_program_accounts - .saturating_sub(1), - ) - } - - /// Gets an instruction account of the current InstructionContext - pub fn try_borrow_instruction_account( - &self, - index_in_instruction: usize, - ) -> Result { - let instruction_context = self.get_current_instruction_context()?; - instruction_context.try_borrow_account( - self, - instruction_context - .number_of_program_accounts - .saturating_add(index_in_instruction), - ) + let level = self + .instruction_context_stack + .len() + .checked_sub(1) + .ok_or(InstructionError::CallDepth)?; + self.get_instruction_context_at(level) } /// Pushes a new InstructionContext pub fn push( &mut self, - number_of_program_accounts: usize, + program_accounts: &[usize], instruction_accounts: &[InstructionAccount], - instruction_data: Vec, + instruction_data: &[u8], ) -> Result<(), InstructionError> { if self.instruction_context_stack.len() >= self.instruction_context_capacity { return Err(InstructionError::CallDepth); } - let mut result = InstructionContext { - instruction_data, - number_of_program_accounts, - account_indices: Vec::with_capacity(instruction_accounts.len()), - account_is_signer: Vec::with_capacity(instruction_accounts.len()), - account_is_writable: Vec::with_capacity(instruction_accounts.len()), - }; - for instruction_account in instruction_accounts.iter() { - result.account_indices.push(instruction_account.index); - result.account_is_signer.push(instruction_account.is_signer); - result - .account_is_writable - .push(instruction_account.is_writable); - } - self.instruction_context_stack.push(result); + self.instruction_context_stack.push(InstructionContext { + program_accounts: program_accounts.to_vec(), + instruction_accounts: instruction_accounts.to_vec(), + instruction_data: instruction_data.to_vec(), + }); Ok(()) } @@ -183,6 +156,20 @@ impl TransactionContext { Ok(()) } + /// Returns the key of the current InstructionContexts program account + pub fn get_program_key(&self) -> Result<&Pubkey, InstructionError> { + let instruction_context = self.get_current_instruction_context()?; + let program_account = instruction_context.try_borrow_program_account(self)?; + Ok(&self.account_keys[program_account.index_in_transaction]) + } + + /// Returns the owner of the current InstructionContexts program account + pub fn get_loader_key(&self) -> Result { + let instruction_context = self.get_current_instruction_context()?; + let program_account = instruction_context.try_borrow_program_account(self)?; + Ok(*program_account.get_owner()) + } + /// Gets the return data of the current InstructionContext or any above pub fn get_return_data(&self) -> (&Pubkey, &[u8]) { (&self.return_data.0, &self.return_data.1) @@ -190,8 +177,7 @@ impl TransactionContext { /// Set the return data of the current InstructionContext pub fn set_return_data(&mut self, data: Vec) -> Result<(), InstructionError> { - let pubkey = *self.try_borrow_program_account()?.get_key(); - self.return_data = (pubkey, data); + self.return_data = (*self.get_program_key()?, data); Ok(()) } } @@ -199,30 +185,29 @@ impl TransactionContext { /// Loaded instruction shared between runtime and programs. /// /// This context is valid for the entire duration of a (possibly cross program) instruction being processed. +#[derive(Debug)] pub struct InstructionContext { - number_of_program_accounts: usize, - account_indices: Vec, - account_is_signer: Vec, - account_is_writable: Vec, + program_accounts: Vec, + instruction_accounts: Vec, instruction_data: Vec, } impl InstructionContext { /// Number of program accounts pub fn get_number_of_program_accounts(&self) -> usize { - self.number_of_program_accounts + self.program_accounts.len() } /// Number of accounts in this Instruction (without program accounts) pub fn get_number_of_instruction_accounts(&self) -> usize { - self.account_indices - .len() - .saturating_sub(self.number_of_program_accounts) + self.instruction_accounts.len() } - /// Total number of accounts in this Instruction (with program accounts) - pub fn get_total_number_of_accounts(&self) -> usize { - self.account_indices.len() + /// Number of accounts in this Instruction + pub fn get_number_of_accounts(&self) -> usize { + self.program_accounts + .len() + .saturating_add(self.instruction_accounts.len()) } /// Data parameter for the programs `process_instruction` handler @@ -230,16 +215,49 @@ impl InstructionContext { &self.instruction_data } + /// Searches for a program account by its key + pub fn find_index_of_program_account( + &self, + transaction_context: &TransactionContext, + pubkey: &Pubkey, + ) -> Option { + self.program_accounts + .iter() + .position(|index_in_transaction| { + &transaction_context.account_keys[*index_in_transaction] == pubkey + }) + } + + /// Searches for an account by its key + pub fn find_index_of_account( + &self, + transaction_context: &TransactionContext, + pubkey: &Pubkey, + ) -> Option { + self.instruction_accounts + .iter() + .position(|instruction_account| { + &transaction_context.account_keys[instruction_account.index_in_transaction] + == pubkey + }) + .map(|index| index.saturating_add(self.program_accounts.len())) + } + /// Tries to borrow an account from this Instruction pub fn try_borrow_account<'a, 'b: 'a>( &'a self, transaction_context: &'b TransactionContext, index_in_instruction: usize, ) -> Result, InstructionError> { - if index_in_instruction >= self.account_indices.len() { + let index_in_transaction = if index_in_instruction < self.program_accounts.len() { + self.program_accounts[index_in_instruction] + } else if index_in_instruction < self.get_number_of_accounts() { + self.instruction_accounts + [index_in_instruction.saturating_sub(self.program_accounts.len())] + .index_in_transaction + } else { return Err(InstructionError::NotEnoughAccountKeys); - } - let index_in_transaction = self.account_indices[index_in_instruction]; + }; if index_in_transaction >= transaction_context.accounts.len() { return Err(InstructionError::MissingAccount); } @@ -254,9 +272,35 @@ impl InstructionContext { account, }) } + + /// Gets the last program account of the current InstructionContext + pub fn try_borrow_program_account<'a, 'b: 'a>( + &'a self, + transaction_context: &'b TransactionContext, + ) -> Result, InstructionError> { + self.try_borrow_account( + transaction_context, + self.program_accounts.len().saturating_sub(1), + ) + } + + /// Gets an instruction account of the current InstructionContext + pub fn try_borrow_instruction_account<'a, 'b: 'a>( + &'a self, + transaction_context: &'b TransactionContext, + index_in_instruction: usize, + ) -> Result, InstructionError> { + self.try_borrow_account( + transaction_context, + self.program_accounts + .len() + .saturating_add(index_in_instruction), + ) + } } /// Shared account borrowed from the TransactionContext and an InstructionContext. +#[derive(Debug)] pub struct BorrowedAccount<'a> { transaction_context: &'a TransactionContext, instruction_context: &'a InstructionContext, @@ -266,6 +310,16 @@ pub struct BorrowedAccount<'a> { } impl<'a> BorrowedAccount<'a> { + /// Returns the index of this account (transaction wide) + pub fn get_index_in_transaction(&self) -> usize { + self.index_in_transaction + } + + /// Returns the index of this account (instruction wide) + pub fn get_index_in_instruction(&self) -> usize { + self.index_in_instruction + } + /// Returns the public key of this account (transaction wide) pub fn get_key(&self) -> &Pubkey { &self.transaction_context.account_keys[self.index_in_transaction] @@ -299,6 +353,24 @@ impl<'a> BorrowedAccount<'a> { Ok(()) } + /// Adds lamports to this account (transaction wide) + pub fn checked_add_lamports(&mut self, lamports: u64) -> Result<(), InstructionError> { + self.set_lamports( + self.get_lamports() + .checked_add(lamports) + .ok_or(LamportsError::ArithmeticOverflow)?, + ) + } + + /// Subtracts lamports from this account (transaction wide) + pub fn checked_sub_lamports(&mut self, lamports: u64) -> Result<(), InstructionError> { + self.set_lamports( + self.get_lamports() + .checked_sub(lamports) + .ok_or(LamportsError::ArithmeticUnderflow)?, + ) + } + /// Returns a read-only slice of the account data (transaction wide) pub fn get_data(&self) -> &[u8] { self.account.data() @@ -312,6 +384,40 @@ impl<'a> BorrowedAccount<'a> { Ok(self.account.data_as_mut_slice()) } + /// Guarded alternative to `get_data_mut()?.copy_from_slice()` which checks if the account size matches + pub fn copy_from_slice(&mut self, data: &[u8]) -> Result<(), InstructionError> { + if !self.is_writable() { + return Err(InstructionError::Immutable); + } + let account_data = self.account.data_as_mut_slice(); + if data.len() != account_data.len() { + return Err(InstructionError::AccountDataSizeChanged); + } + account_data.copy_from_slice(data); + Ok(()) + } + + /// Deserializes the account data into a state + pub fn get_state(&self) -> Result { + self.account + .deserialize_data() + .map_err(|_| InstructionError::InvalidAccountData) + } + + /// Serializes a state into the account data + pub fn set_state(&mut self, state: &T) -> Result<(), InstructionError> { + if !self.is_writable() { + return Err(InstructionError::Immutable); + } + let data = self.account.data_as_mut_slice(); + let serialized_size = + bincode::serialized_size(state).map_err(|_| InstructionError::GenericError)?; + if serialized_size > data.len() as u64 { + return Err(InstructionError::AccountDataTooSmall); + } + bincode::serialize_into(&mut *data, state).map_err(|_| InstructionError::GenericError) + } + /*pub fn realloc(&self, new_len: usize, zero_init: bool) { // TODO }*/ @@ -332,11 +438,25 @@ impl<'a> BorrowedAccount<'a> { /// Returns whether this account is a signer (instruction wide) pub fn is_signer(&self) -> bool { - self.instruction_context.account_is_signer[self.index_in_instruction] + if self.index_in_instruction < self.instruction_context.program_accounts.len() { + false + } else { + self.instruction_context.instruction_accounts[self + .index_in_instruction + .saturating_sub(self.instruction_context.program_accounts.len())] + .is_signer + } } /// Returns whether this account is writable (instruction wide) pub fn is_writable(&self) -> bool { - self.instruction_context.account_is_writable[self.index_in_instruction] + if self.index_in_instruction < self.instruction_context.program_accounts.len() { + false + } else { + self.instruction_context.instruction_accounts[self + .index_in_instruction + .saturating_sub(self.instruction_context.program_accounts.len())] + .is_writable + } } }