solana/sdk/src/transaction_context.rs

//! Successors of instruction_context_context::StackFrame, KeyedAccount and AccountInfo

use crate::{
    account::{AccountSharedData, ReadableAccount, WritableAccount},
    instruction::InstructionError,
    lamports::LamportsError,
    pubkey::Pubkey,
};
use std::cell::{RefCell, RefMut};

pub type TransactionAccount = (Pubkey, AccountSharedData);

#[derive(Clone, Debug)]
pub struct InstructionAccount {
    pub index_in_transaction: usize,
    pub index_in_caller: usize,
    pub is_signer: bool,
    pub is_writable: bool,
}

/// 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<Pubkey>,
    accounts: Vec<RefCell<AccountSharedData>>,
    instruction_context_capacity: usize,
    instruction_context_stack: Vec<InstructionContext>,
    return_data: (Pubkey, Vec<u8>),
}

impl TransactionContext {
    /// Constructs a new TransactionContext
    pub fn new(
        transaction_accounts: Vec<TransactionAccount>,
        instruction_context_capacity: usize,
    ) -> Self {
        let (account_keys, accounts) = transaction_accounts
            .into_iter()
            .map(|(key, account)| (key, RefCell::new(account)))
            .unzip();
        Self {
            account_keys,
            accounts,
            instruction_context_capacity,
            instruction_context_stack: Vec::with_capacity(instruction_context_capacity),
            return_data: (Pubkey::default(), Vec::new()),
        }
    }

    /// Used by the bank in the runtime to write back the processed accounts
    pub fn deconstruct(self) -> Vec<TransactionAccount> {
        self.account_keys
            .into_iter()
            .zip(
                self.accounts
                    .into_iter()
                    .map(|account| account.into_inner()),
            )
            .collect()
    }

    /// Used in mock_process_instruction
    pub fn deconstruct_without_keys(self) -> Result<Vec<AccountSharedData>, InstructionError> {
        if !self.instruction_context_stack.is_empty() {
            return Err(InstructionError::CallDepth);
        }
        Ok(self
            .accounts
            .into_iter()
            .map(|account| account.into_inner())
            .collect())
    }

    /// Returns the total number of accounts loaded in this Transaction
    pub fn get_number_of_accounts(&self) -> usize {
        self.accounts.len()
    }

    /// Searches for an account by its key
    pub fn get_key_of_account_at_index(&self, index_in_transaction: usize) -> &Pubkey {
        &self.account_keys[index_in_transaction]
    }

    /// Searches for an account by its key
    pub fn get_account_at_index(&self, index_in_transaction: usize) -> &RefCell<AccountSharedData> {
        &self.accounts[index_in_transaction]
    }

    /// Searches for an account by its key
    pub fn find_index_of_account(&self, pubkey: &Pubkey) -> Option<usize> {
        self.account_keys.iter().position(|key| key == pubkey)
    }

    /// Searches for a program account by its key
    pub fn find_index_of_program_account(&self, pubkey: &Pubkey) -> Option<usize> {
        self.account_keys.iter().rposition(|key| key == pubkey)
    }

    /// Gets an InstructionContext by its height in the stack
    pub fn get_instruction_context_at(
        &self,
        level: usize,
    ) -> Result<&InstructionContext, InstructionError> {
        if level >= self.instruction_context_stack.len() {
            return Err(InstructionError::CallDepth);
        }
        Ok(&self.instruction_context_stack[level])
    }

    /// Gets the max height of the InstructionContext stack
    pub fn get_instruction_context_capacity(&self) -> usize {
        self.instruction_context_capacity
    }

    /// 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> {
        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,
        program_accounts: &[usize],
        instruction_accounts: &[InstructionAccount],
        instruction_data: &[u8],
    ) -> Result<(), InstructionError> {
        if self.instruction_context_stack.len() >= self.instruction_context_capacity {
            return Err(InstructionError::CallDepth);
        }
        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(())
    }

    /// Pops the current InstructionContext
    pub fn pop(&mut self) -> Result<(), InstructionError> {
        if self.instruction_context_stack.is_empty() {
            return Err(InstructionError::CallDepth);
        }
        self.instruction_context_stack.pop();
        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<Pubkey, InstructionError> {
        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)
    }

    /// Set the return data of the current InstructionContext
    pub fn set_return_data(&mut self, data: Vec<u8>) -> Result<(), InstructionError> {
        self.return_data = (*self.get_program_key()?, data);
        Ok(())
    }
}

/// 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 {
    program_accounts: Vec<usize>,
    instruction_accounts: Vec<InstructionAccount>,
    instruction_data: Vec<u8>,
}

impl InstructionContext {
    /// Number of program accounts
    pub fn get_number_of_program_accounts(&self) -> usize {
        self.program_accounts.len()
    }

    /// Number of accounts in this Instruction (without program accounts)
    pub fn get_number_of_instruction_accounts(&self) -> usize {
        self.instruction_accounts.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
    pub fn get_instruction_data(&self) -> &[u8] {
        &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<usize> {
        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<usize> {
        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<BorrowedAccount<'a>, InstructionError> {
        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);
        };
        if index_in_transaction >= transaction_context.accounts.len() {
            return Err(InstructionError::MissingAccount);
        }
        let account = transaction_context.accounts[index_in_transaction]
            .try_borrow_mut()
            .map_err(|_| InstructionError::AccountBorrowFailed)?;
        Ok(BorrowedAccount {
            transaction_context,
            instruction_context: self,
            index_in_transaction,
            index_in_instruction,
            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<BorrowedAccount<'a>, 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<BorrowedAccount<'a>, 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,
    index_in_transaction: usize,
    index_in_instruction: usize,
    account: RefMut<'a, AccountSharedData>,
}

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]
    }

    /// Returns the owner of this account (transaction wide)
    pub fn get_owner(&self) -> &Pubkey {
        self.account.owner()
    }

    /// Assignes the owner of this account (transaction wide)
    pub fn set_owner(&mut self, pubkey: Pubkey) -> Result<(), InstructionError> {
        if !self.is_writable() {
            return Err(InstructionError::Immutable);
        }
        self.account.set_owner(pubkey);
        Ok(())
    }

    /// Returns the number of lamports of this account (transaction wide)
    pub fn get_lamports(&self) -> u64 {
        self.account.lamports()
    }

    /// Overwrites the number of lamports of this account (transaction wide)
    pub fn set_lamports(&mut self, lamports: u64) -> Result<(), InstructionError> {
        if !self.is_writable() {
            return Err(InstructionError::Immutable);
        }
        self.account.set_lamports(lamports);
        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()
    }

    /// Returns a writable slice of the account data (transaction wide)
    pub fn get_data_mut(&mut self) -> Result<&mut [u8], InstructionError> {
        if !self.is_writable() {
            return Err(InstructionError::Immutable);
        }
        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<T: serde::de::DeserializeOwned>(&self) -> Result<T, InstructionError> {
        self.account
            .deserialize_data()
            .map_err(|_| InstructionError::InvalidAccountData)
    }

    /// Serializes a state into the account data
    pub fn set_state<T: serde::Serialize>(&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
    }*/

    /// Returns whether this account is executable (transaction wide)
    pub fn is_executable(&self) -> bool {
        self.account.executable()
    }

    /// Configures whether this account is executable (transaction wide)
    pub fn set_executable(&mut self, is_executable: bool) -> Result<(), InstructionError> {
        if !self.is_writable() {
            return Err(InstructionError::Immutable);
        }
        self.account.set_executable(is_executable);
        Ok(())
    }

    /// Returns whether this account is a signer (instruction wide)
    pub fn is_signer(&self) -> bool {
        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 {
        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
        }
    }
}