diff --git a/programs/vest_api/src/vest_processor.rs b/programs/vest_api/src/vest_processor.rs index 321002bac..2ed760d47 100644 --- a/programs/vest_api/src/vest_processor.rs +++ b/programs/vest_api/src/vest_processor.rs @@ -10,6 +10,7 @@ use solana_config_api::get_config_data; use solana_sdk::{ account::{Account, KeyedAccount}, instruction::InstructionError, + instruction_processor_utils::next_keyed_account, pubkey::Pubkey, }; @@ -58,6 +59,7 @@ pub fn process_instruction( keyed_accounts: &mut [KeyedAccount], data: &[u8], ) -> Result<(), InstructionError> { + let keyed_accounts_iter = &mut keyed_accounts.iter_mut(); let instruction = deserialize(data).map_err(|_| InstructionError::InvalidInstructionData)?; match instruction { @@ -68,10 +70,7 @@ pub fn process_instruction( date_pubkey, total_lamports, } => { - let contract_keyed_account = match keyed_accounts { - [ka0] => ka0, - _ => return Err(InstructionError::InvalidArgument), - }; + let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?; let contract_account = &mut contract_keyed_account.account; let vest_state = VestState { terminator_pubkey, @@ -84,10 +83,8 @@ pub fn process_instruction( vest_state.serialize(&mut contract_account.data) } VestInstruction::SetPayee(payee_pubkey) => { - let (contract_keyed_account, old_payee_keyed_account) = match keyed_accounts { - [ka0, ka1] => (ka0, ka1), - _ => return Err(InstructionError::InvalidArgument), - }; + let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?; + let old_payee_keyed_account = next_keyed_account(keyed_accounts_iter)?; let contract_account = &mut contract_keyed_account.account; let mut vest_state = VestState::deserialize(&contract_account.data)?; parse_signed_account(old_payee_keyed_account, &vest_state.payee_pubkey)?; @@ -95,11 +92,9 @@ pub fn process_instruction( vest_state.serialize(&mut contract_account.data) } VestInstruction::RedeemTokens => { - let (contract_keyed_account, date_keyed_account, payee_keyed_account) = - match keyed_accounts { - [ka0, ka1, ka2] => (ka0, ka1, ka2), - _ => return Err(InstructionError::InvalidArgument), - }; + let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?; + let date_keyed_account = next_keyed_account(keyed_accounts_iter)?; + let payee_keyed_account = next_keyed_account(keyed_accounts_iter)?; let contract_account = &mut contract_keyed_account.account; let mut vest_state = VestState::deserialize(&contract_account.data)?; let current_date = parse_date_account(date_keyed_account, &vest_state.date_pubkey)?; @@ -109,12 +104,9 @@ pub fn process_instruction( vest_state.serialize(&mut contract_account.data) } VestInstruction::Terminate => { - let (contract_keyed_account, terminator_keyed_account, payee_keyed_account) = - match keyed_accounts { - [ka0, ka1] => (ka0, ka1, None), - [ka0, ka1, ka2] => (ka0, ka1, Some(ka2)), - _ => return Err(InstructionError::InvalidArgument), - }; + let contract_keyed_account = next_keyed_account(keyed_accounts_iter)?; + let terminator_keyed_account = next_keyed_account(keyed_accounts_iter)?; + let payee_keyed_account = keyed_accounts_iter.next(); let contract_account = &mut contract_keyed_account.account; let mut vest_state = VestState::deserialize(&contract_account.data)?; let terminator_account = @@ -335,7 +327,7 @@ mod tests { .send_message(&[&alice_keypair], message) .unwrap_err() .unwrap(), - TransactionError::InstructionError(1, InstructionError::InvalidArgument) + TransactionError::InstructionError(1, InstructionError::NotEnoughAccountKeys) ); } diff --git a/sdk/src/instruction.rs b/sdk/src/instruction.rs index 026c89d60..cf8b55de6 100644 --- a/sdk/src/instruction.rs +++ b/sdk/src/instruction.rs @@ -67,6 +67,9 @@ pub enum InstructionError { /// Rent_epoch account changed, but shouldn't have RentEpochModified, + /// The instruction expected additional account keys + NotEnoughAccountKeys, + /// CustomError allows on-chain programs to implement program-specific error types and see /// them returned by the Solana runtime. A CustomError may be any type that is represented /// as or serialized to a u32 integer. diff --git a/sdk/src/instruction_processor_utils.rs b/sdk/src/instruction_processor_utils.rs index 853d7b8e0..8da5eafbc 100644 --- a/sdk/src/instruction_processor_utils.rs +++ b/sdk/src/instruction_processor_utils.rs @@ -38,6 +38,11 @@ where } } +/// Return the next KeyedAccount or a NotEnoughAccountKeys instruction error +pub fn next_keyed_account(iter: &mut I) -> Result { + iter.next().ok_or(InstructionError::NotEnoughAccountKeys) +} + pub trait DecodeError { fn decode_custom_error_to_enum(custom: u32) -> Option where