Skip to content

Commit

Permalink
added check_account_infos util
Browse files Browse the repository at this point in the history
  • Loading branch information
@buffalojoec
buffalojoec committed Aug 14, 2023
1 parent 128b422 commit b63946f
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 122 deletions.
67 changes: 23 additions & 44 deletions libraries/tlv-account-resolution/src/account.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,43 +7,20 @@ use {
crate::{error::AccountResolutionError, seeds::Seed},
bytemuck::{Pod, Zeroable},
solana_program::{
account_info::AccountInfo,
instruction::{AccountMeta, Instruction},
program_error::ProgramError,
account_info::AccountInfo, instruction::AccountMeta, program_error::ProgramError,
pubkey::Pubkey,
},
spl_type_length_value::pod::PodBool,
};

/// De-escalate an account meta if necessary
fn de_escalate_account_meta(account_meta: &mut AccountMeta, account_metas: &[AccountMeta]) {
// This is a little tricky to read, but the idea is to see if
// this account is marked as writable or signer anywhere in
// the instruction at the start. If so, DON'T escalate it to
// be a writer or signer in the CPI
let maybe_highest_privileges = account_metas
.iter()
.filter(|&x| x.pubkey == account_meta.pubkey)
.map(|x| (x.is_signer, x.is_writable))
.reduce(|acc, x| (acc.0 || x.0, acc.1 || x.1));
// If `Some`, then the account was found somewhere in the instruction
if let Some((is_signer, is_writable)) = maybe_highest_privileges {
if !is_signer && is_signer != account_meta.is_signer {
// Existing account is *NOT* a signer already, but the CPI
// wants it to be, so de-escalate to not be a signer
account_meta.is_signer = false;
}
if !is_writable && is_writable != account_meta.is_writable {
// Existing account is *NOT* writable already, but the CPI
// wants it to be, so de-escalate to not be writable
account_meta.is_writable = false;
}
}
}

/// Resolve a program-derived address (PDA) from the instruction data
/// and the accounts that have already been resolved
fn resolve_pda(seeds: &[Seed], instruction: &Instruction) -> Result<Pubkey, ProgramError> {
fn resolve_pda(
seeds: &[Seed],
accounts: &[AccountMeta],
instruction_data: &[u8],
program_id: &Pubkey,
) -> Result<Pubkey, ProgramError> {
let mut pda_seeds: Vec<&[u8]> = vec![];
for config in seeds {
match config {
Expand All @@ -52,19 +29,18 @@ fn resolve_pda(seeds: &[Seed], instruction: &Instruction) -> Result<Pubkey, Prog
Seed::InstructionData { index, length } => {
let arg_start = *index as usize;
let arg_end = arg_start + *length as usize;
pda_seeds.push(&instruction.data[arg_start..arg_end]);
pda_seeds.push(&instruction_data[arg_start..arg_end]);
}
Seed::AccountKey { index } => {
let account_index = *index as usize;
let account_meta = instruction
.accounts
let account_meta = accounts
.get(account_index)
.ok_or::<ProgramError>(AccountResolutionError::AccountNotFound.into())?;
pda_seeds.push(account_meta.pubkey.as_ref());
}
}
}
Ok(Pubkey::find_program_address(&pda_seeds, &instruction.program_id).0)
Ok(Pubkey::find_program_address(&pda_seeds, program_id).0)
}

/// `Pod` type for defining a required account in a validation account.
Expand Down Expand Up @@ -118,21 +94,24 @@ impl ExtraAccountMeta {

/// Resolve an `ExtraAccountMeta` into an `AccountMeta`, potentially
/// resolving a program-derived address (PDA) if necessary
pub fn resolve(&self, instruction: &Instruction) -> Result<AccountMeta, ProgramError> {
let mut account_meta = match self.discriminator {
0 => AccountMeta::try_from(self)?,
pub fn resolve(
&self,
accounts: &[AccountMeta],
instruction_data: &[u8],
program_id: &Pubkey,
) -> Result<AccountMeta, ProgramError> {
match self.discriminator {
0 => AccountMeta::try_from(self),
1 => {
let seeds = Seed::unpack_address_config(&self.address_config)?;
AccountMeta {
pubkey: resolve_pda(&seeds, instruction)?,
Ok(AccountMeta {
pubkey: resolve_pda(&seeds, accounts, instruction_data, program_id)?,
is_signer: self.is_signer.into(),
is_writable: self.is_writable.into(),
}
})
}
_ => return Err(ProgramError::InvalidAccountData),
};
de_escalate_account_meta(&mut account_meta, &instruction.accounts);
Ok(account_meta)
_ => Err(ProgramError::InvalidAccountData),
}
}
}

Expand Down
77 changes: 75 additions & 2 deletions libraries/tlv-account-resolution/src/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
use {
crate::{account::ExtraAccountMeta, error::AccountResolutionError},
solana_program::{
account_info::AccountInfo, instruction::Instruction, program_error::ProgramError,
account_info::AccountInfo,
instruction::{AccountMeta, Instruction},
program_error::ProgramError,
pubkey::Pubkey,
},
spl_discriminator::SplDiscriminate,
spl_type_length_value::{
Expand All @@ -12,6 +15,41 @@ use {
},
};

/// De-escalate an account meta if necessary
fn de_escalate_account_meta(account_meta: &mut AccountMeta, account_metas: &[AccountMeta]) {
// This is a little tricky to read, but the idea is to see if
// this account is marked as writable or signer anywhere in
// the instruction at the start. If so, DON'T escalate it to
// be a writer or signer in the CPI
let maybe_highest_privileges = account_metas
.iter()
.filter(|&x| x.pubkey == account_meta.pubkey)
.map(|x| (x.is_signer, x.is_writable))
.reduce(|acc, x| (acc.0 || x.0, acc.1 || x.1));
// If `Some`, then the account was found somewhere in the instruction
if let Some((is_signer, is_writable)) = maybe_highest_privileges {
if !is_signer && is_signer != account_meta.is_signer {
// Existing account is *NOT* a signer already, but the CPI
// wants it to be, so de-escalate to not be a signer
account_meta.is_signer = false;
}
if !is_writable && is_writable != account_meta.is_writable {
// Existing account is *NOT* writable already, but the CPI
// wants it to be, so de-escalate to not be writable
account_meta.is_writable = false;
}
}
}

/// Helper to convert an `AccountInfo` to an `AccountMeta`
fn account_meta_from_info(account_info: &AccountInfo) -> AccountMeta {
AccountMeta {
pubkey: *account_info.key,
is_signer: account_info.is_signer,
is_writable: account_info.is_writable,
}
}

/// Stateless helper for storing additional accounts required for an
/// instruction.
///
Expand Down Expand Up @@ -105,6 +143,35 @@ impl ExtraAccountMetaList {
.saturating_add(PodSlice::<ExtraAccountMeta>::size_of(num_items)?))
}

/// Checks provided account infos against validation data, using
/// instruction data and program ID to resolve any dynamic PDAs
/// if necessary.
pub fn check_account_infos<T: SplDiscriminate>(
account_infos: &[AccountInfo],
instruction_data: &[u8],
program_id: &Pubkey,
data: &[u8],
) -> Result<(), ProgramError> {
let state = TlvStateBorrowed::unpack(data).unwrap();
let extra_meta_list = ExtraAccountMetaList::unpack_with_tlv_state::<T>(&state)?;
let extra_account_metas = extra_meta_list.data();

let provided_metas = account_infos
.iter()
.map(account_meta_from_info)
.collect::<Vec<_>>();

for config in extra_account_metas.iter() {
let meta = config.resolve(&provided_metas, instruction_data, program_id)?;
provided_metas
.iter()
.find(|&x| *x == meta)
.ok_or(AccountResolutionError::IncorrectAccount)?;
}

Ok(())
}

/// Add the additional account metas to an existing instruction
pub fn add_to_instruction<T: SplDiscriminate>(
instruction: &mut Instruction,
Expand All @@ -115,7 +182,13 @@ impl ExtraAccountMetaList {
let extra_account_metas = PodSlice::<ExtraAccountMeta>::unpack(bytes)?;

for extra_meta in extra_account_metas.data().iter() {
instruction.accounts.push(extra_meta.resolve(instruction)?);
let mut meta = extra_meta.resolve(
&instruction.accounts,
&instruction.data,
&instruction.program_id,
)?;
de_escalate_account_meta(&mut meta, &instruction.accounts);
instruction.accounts.push(meta);
}
Ok(())
}
Expand Down
76 changes: 5 additions & 71 deletions token/transfer-hook-example/src/processor.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ use {
solana_program::{
account_info::{next_account_info, AccountInfo},
entrypoint::ProgramResult,
instruction::AccountMeta,
msg,
program::invoke_signed,
program_error::ProgramError,
Expand All @@ -24,7 +23,6 @@ use {
get_extra_account_metas_address, get_extra_account_metas_address_and_bump_seed,
instruction::{ExecuteInstruction, TransferHookInstruction},
},
spl_type_length_value::state::TlvStateBorrowed,
};

fn check_token_account_is_transferring(account_info: &AccountInfo) -> Result<(), ProgramError> {
Expand All @@ -38,27 +36,6 @@ fn check_token_account_is_transferring(account_info: &AccountInfo) -> Result<(),
}
}

fn check_extra_meta(
account_info: &AccountInfo,
extra_meta: &AccountMeta,
) -> Result<(), ProgramError> {
if !(&extra_meta.pubkey == account_info.key
&& extra_meta.is_signer == account_info.is_signer
&& extra_meta.is_writable == account_info.is_writable)
{
return Err(TransferHookError::IncorrectAccount.into());
}
Ok(())
}

fn next_extra_account_meta<'a>(
extra_account_metas_iter: &mut impl Iterator<Item = &'a ExtraAccountMeta>,
) -> Result<&'a ExtraAccountMeta, ProgramError> {
extra_account_metas_iter
.next()
.ok_or(ProgramError::NotEnoughAccountKeys)
}

/// Processes an [Execute](enum.TransferHookInstruction.html) instruction.
pub fn process_execute(
program_id: &Pubkey,
Expand All @@ -85,55 +62,12 @@ pub fn process_execute(
}

let data = extra_account_metas_info.try_borrow_data()?;
let state = TlvStateBorrowed::unpack(&data).unwrap();
let extra_account_metas =
ExtraAccountMetaList::unpack_with_tlv_state::<ExecuteInstruction>(&state)?;
let account_metas = extra_account_metas.data();

// if incorrect number of are provided, error
if account_metas.len() != account_info_iter.len() {
return Err(TransferHookError::IncorrectAccount.into());
}

// Let's assume that they're provided in the correct order
let extra_account_metas_iter = &mut account_metas.iter();
check_extra_meta(
next_account_info(account_info_iter)?,
&AccountMeta::try_from(next_extra_account_meta(extra_account_metas_iter)?)?,
)?;
check_extra_meta(
next_account_info(account_info_iter)?,
&AccountMeta::try_from(next_extra_account_meta(extra_account_metas_iter)?)?,
)?;
let next_config = next_extra_account_meta(extra_account_metas_iter)?;
check_extra_meta(
next_account_info(account_info_iter)?,
&AccountMeta {
pubkey: Pubkey::find_program_address(
&[b"seed-prefix", source_account_info.key.as_ref()],
program_id,
)
.0,
is_signer: bool::from(next_config.is_signer),
is_writable: bool::from(next_config.is_writable),
},
)?;
let next_config = next_extra_account_meta(extra_account_metas_iter)?;
check_extra_meta(
next_account_info(account_info_iter)?,
&AccountMeta {
pubkey: Pubkey::find_program_address(
&[&amount.to_le_bytes(), destination_account_info.key.as_ref()],
program_id,
)
.0,
is_signer: bool::from(next_config.is_signer),
is_writable: bool::from(next_config.is_writable),
},
)?;
check_extra_meta(
next_account_info(account_info_iter)?,
&AccountMeta::try_from(next_extra_account_meta(extra_account_metas_iter)?)?,
ExtraAccountMetaList::check_account_infos::<ExecuteInstruction>(
accounts,
&TransferHookInstruction::Execute { amount }.pack(),
program_id,
&data,
)?;

Ok(())
Expand Down
12 changes: 7 additions & 5 deletions token/transfer-hook-example/tests/functional.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,8 @@ async fn success_execute() {

let extra_account_metas_address = get_extra_account_metas_address(&mint_address, &program_id);

let writable_pubkey = Pubkey::new_unique();

let init_extra_account_metas = [
ExtraAccountMeta::new_with_pubkey(&sysvar::instructions::id(), false, false).unwrap(),
ExtraAccountMeta::new_with_pubkey(&mint_authority_pubkey, true, false).unwrap(),
Expand All @@ -183,7 +185,7 @@ async fn success_execute() {
true,
)
.unwrap(),
ExtraAccountMeta::new_with_pubkey(&extra_account_metas_address, false, true).unwrap(),
ExtraAccountMeta::new_with_pubkey(&writable_pubkey, false, true).unwrap(),
];

let extra_pda_1 = Pubkey::find_program_address(
Expand All @@ -208,7 +210,7 @@ async fn success_execute() {
AccountMeta::new_readonly(mint_authority_pubkey, true),
AccountMeta::new(extra_pda_1, false),
AccountMeta::new(extra_pda_2, false),
AccountMeta::new(extra_account_metas_address, false),
AccountMeta::new(writable_pubkey, false),
];

let mut context = program_test.start_with_context().await;
Expand Down Expand Up @@ -280,7 +282,7 @@ async fn success_execute() {
AccountMeta::new_readonly(mint_authority_pubkey, true),
AccountMeta::new(extra_pda_1, false),
AccountMeta::new(extra_pda_2, false),
AccountMeta::new(wallet.pubkey(), false),
AccountMeta::new(Pubkey::new_unique(), false),
];
let transaction = Transaction::new_signed_with_payer(
&[execute_with_extra_account_metas(
Expand Down Expand Up @@ -327,7 +329,7 @@ async fn success_execute() {
AccountMeta::new_readonly(mint_authority_pubkey, true),
AccountMeta::new(extra_pda_1, false),
AccountMeta::new(wrong_pda_2, false),
AccountMeta::new(extra_account_metas_address, false),
AccountMeta::new(writable_pubkey, false),
];
let transaction = Transaction::new_signed_with_payer(
&[execute_with_extra_account_metas(
Expand Down Expand Up @@ -366,7 +368,7 @@ async fn success_execute() {
AccountMeta::new_readonly(mint_authority_pubkey, false),
AccountMeta::new(extra_pda_1, false),
AccountMeta::new(extra_pda_2, false),
AccountMeta::new(extra_account_metas_address, false),
AccountMeta::new(writable_pubkey, false),
];
let transaction = Transaction::new_signed_with_payer(
&[execute_with_extra_account_metas(
Expand Down

0 comments on commit b63946f

Please sign in to comment.