From ee3fc39f1c33a044a8f83cd538b133bf4015d421 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Mei=C3=9Fner?= <AlexanderMeissner@gmx.net>
Date: Mon, 28 Feb 2022 16:24:40 +0100
Subject: [PATCH] Refactor: Prepare stake_instruction.rs to remove
 `KeyedAccount`s (#23375)

* Adds instruction_account_indices to stake instruction.

* Uses instruction_account_indices in stake instruction.

* Replaces get_sysvar_with_account_check by get_sysvar_with_account_check2 in stake instruction.

* Adds check_number_of_instruction_accounts().
---
 programs/stake/src/stake_instruction.rs | 333 ++++++++++++++++++------
 1 file changed, 251 insertions(+), 82 deletions(-)

diff --git a/programs/stake/src/stake_instruction.rs b/programs/stake/src/stake_instruction.rs
index 315c3c5a7..7add8540d 100644
--- a/programs/stake/src/stake_instruction.rs
+++ b/programs/stake/src/stake_instruction.rs
@@ -7,12 +7,12 @@ use {
     crate::{config, stake_state::StakeAccount},
     log::*,
     solana_program_runtime::{
-        invoke_context::InvokeContext, sysvar_cache::get_sysvar_with_account_check,
+        invoke_context::InvokeContext, sysvar_cache::get_sysvar_with_account_check2,
     },
     solana_sdk::{
         feature_set,
         instruction::InstructionError,
-        keyed_account::{get_signers, keyed_account_at_index},
+        keyed_account::keyed_account_at_index,
         program_utils::limited_deserialize,
         stake::{
             instruction::StakeInstruction,
@@ -23,27 +23,118 @@ use {
     },
 };
 
+pub mod instruction_account_indices {
+    pub enum Initialize {
+        StakeAccount = 0,
+        Rent = 1,
+    }
+
+    pub enum Authorize {
+        StakeAccount = 0,
+        Clock = 1,
+        // CurrentAuthority = 2,
+        Custodian = 3,
+    }
+
+    pub enum AuthorizeWithSeed {
+        StakeAccount = 0,
+        AuthorityBase = 1,
+        Clock = 2,
+        Custodian = 3,
+    }
+
+    pub enum DelegateStake {
+        StakeAccount = 0,
+        VoteAccount = 1,
+        Clock = 2,
+        StakeHistory = 3,
+        ConfigAccount = 4,
+    }
+
+    pub enum Split {
+        StakeAccount = 0,
+        SplitTo = 1,
+    }
+
+    pub enum Merge {
+        StakeAccount = 0,
+        MergeFrom = 1,
+        Clock = 2,
+        StakeHistory = 3,
+    }
+
+    pub enum Withdraw {
+        StakeAccount = 0,
+        Recipient = 1,
+        Clock = 2,
+        StakeHistory = 3,
+        WithdrawAuthority = 4,
+        Custodian = 5,
+    }
+
+    pub enum Deactivate {
+        StakeAccount = 0,
+        Clock = 1,
+    }
+
+    pub enum SetLockup {
+        StakeAccount = 0,
+        // Clock = 1,
+    }
+
+    pub enum InitializeChecked {
+        StakeAccount = 0,
+        Rent = 1,
+        AuthorizedStaker = 2,
+        AuthorizedWithdrawer = 3,
+    }
+
+    pub enum AuthorizeChecked {
+        StakeAccount = 0,
+        Clock = 1,
+        // CurrentAuthority = 2,
+        Authorized = 3,
+        Custodian = 4,
+    }
+
+    pub enum AuthorizeCheckedWithSeed {
+        StakeAccount = 0,
+        AuthorityBase = 1,
+        Clock = 2,
+        Authorized = 3,
+        Custodian = 4,
+    }
+
+    pub enum SetLockupChecked {
+        StakeAccount = 0,
+        // Clock = 1,
+        Custodian = 2,
+    }
+}
+
 pub fn process_instruction(
     first_instruction_account: usize,
     data: &[u8],
     invoke_context: &mut InvokeContext,
 ) -> Result<(), InstructionError> {
+    let transaction_context = &invoke_context.transaction_context;
+    let instruction_context = transaction_context.get_current_instruction_context()?;
     let keyed_accounts = invoke_context.get_keyed_accounts()?;
 
     trace!("process_instruction: {:?}", data);
-    trace!("keyed_accounts: {:?}", keyed_accounts);
 
     let me = &keyed_account_at_index(keyed_accounts, first_instruction_account)?;
     if me.owner()? != id() {
         return Err(InstructionError::InvalidAccountOwner);
     }
 
-    let signers = get_signers(&keyed_accounts[first_instruction_account..]);
+    let signers = instruction_context.get_signers(transaction_context);
     match limited_deserialize(data)? {
         StakeInstruction::Initialize(authorized, lockup) => {
-            let rent = get_sysvar_with_account_check::rent(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?,
+            let rent = get_sysvar_with_account_check2::rent(
                 invoke_context,
+                instruction_context,
+                instruction_account_indices::Initialize::Rent as usize,
             )?;
             me.initialize(&authorized, &lockup, &rent)
         }
@@ -53,16 +144,18 @@ pub fn process_instruction(
                 .is_active(&feature_set::require_custodian_for_locked_stake_authorize::id());
 
             if require_custodian_for_locked_stake_authorize {
-                let clock = get_sysvar_with_account_check::clock(
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?,
+                let clock = get_sysvar_with_account_check2::clock(
                     invoke_context,
+                    instruction_context,
+                    instruction_account_indices::Authorize::Clock as usize,
                 )?;
-                let _current_authority =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?;
-                let custodian =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 3)
-                        .ok()
-                        .map(|ka| ka.unsigned_key());
+                let custodian = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::Authorize::Custodian as usize,
+                )
+                .ok()
+                .map(|ka| ka.unsigned_key());
 
                 me.authorize(
                     &signers,
@@ -84,21 +177,29 @@ pub fn process_instruction(
             }
         }
         StakeInstruction::AuthorizeWithSeed(args) => {
-            let authority_base =
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
+            instruction_context.check_number_of_instruction_accounts(2)?;
+            let authority_base = keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account
+                    + instruction_account_indices::AuthorizeWithSeed::AuthorityBase as usize,
+            )?;
             let require_custodian_for_locked_stake_authorize = invoke_context
                 .feature_set
                 .is_active(&feature_set::require_custodian_for_locked_stake_authorize::id());
 
             if require_custodian_for_locked_stake_authorize {
-                let clock = get_sysvar_with_account_check::clock(
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?,
+                let clock = get_sysvar_with_account_check2::clock(
                     invoke_context,
+                    instruction_context,
+                    instruction_account_indices::AuthorizeWithSeed::Clock as usize,
                 )?;
-                let custodian =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 3)
-                        .ok()
-                        .map(|ka| ka.unsigned_key());
+                let custodian = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeWithSeed::Custodian as usize,
+                )
+                .ok()
+                .map(|ka| ka.unsigned_key());
 
                 me.authorize_with_seed(
                     authority_base,
@@ -124,17 +225,28 @@ pub fn process_instruction(
             }
         }
         StakeInstruction::DelegateStake => {
-            let vote = keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
-            let clock = get_sysvar_with_account_check::clock(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?,
-                invoke_context,
+            instruction_context.check_number_of_instruction_accounts(2)?;
+            let vote = keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account
+                    + instruction_account_indices::DelegateStake::VoteAccount as usize,
             )?;
-            let stake_history = get_sysvar_with_account_check::stake_history(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 3)?,
+            let clock = get_sysvar_with_account_check2::clock(
                 invoke_context,
+                instruction_context,
+                instruction_account_indices::DelegateStake::Clock as usize,
+            )?;
+            let stake_history = get_sysvar_with_account_check2::stake_history(
+                invoke_context,
+                instruction_context,
+                instruction_account_indices::DelegateStake::StakeHistory as usize,
+            )?;
+            instruction_context.check_number_of_instruction_accounts(5)?;
+            let config_account = keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account
+                    + instruction_account_indices::DelegateStake::ConfigAccount as usize,
             )?;
-            let config_account =
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 4)?;
             if !config::check_id(config_account.unsigned_key()) {
                 return Err(InstructionError::InvalidArgument);
             }
@@ -143,20 +255,28 @@ pub fn process_instruction(
             me.delegate(vote, &clock, &stake_history, &config, &signers)
         }
         StakeInstruction::Split(lamports) => {
-            let split_stake =
-                &keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
+            instruction_context.check_number_of_instruction_accounts(2)?;
+            let split_stake = &keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account + instruction_account_indices::Split::SplitTo as usize,
+            )?;
             me.split(lamports, split_stake, &signers)
         }
         StakeInstruction::Merge => {
-            let source_stake =
-                &keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
-            let clock = get_sysvar_with_account_check::clock(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?,
-                invoke_context,
+            instruction_context.check_number_of_instruction_accounts(2)?;
+            let source_stake = &keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account + instruction_account_indices::Merge::MergeFrom as usize,
             )?;
-            let stake_history = get_sysvar_with_account_check::stake_history(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 3)?,
+            let clock = get_sysvar_with_account_check2::clock(
                 invoke_context,
+                instruction_context,
+                instruction_account_indices::Merge::Clock as usize,
+            )?;
+            let stake_history = get_sysvar_with_account_check2::stake_history(
+                invoke_context,
+                instruction_context,
+                instruction_account_indices::Merge::StakeHistory as usize,
             )?;
             me.merge(
                 invoke_context,
@@ -167,28 +287,46 @@ pub fn process_instruction(
             )
         }
         StakeInstruction::Withdraw(lamports) => {
-            let to = &keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
-            let clock = get_sysvar_with_account_check::clock(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?,
-                invoke_context,
+            instruction_context.check_number_of_instruction_accounts(2)?;
+            let to = &keyed_account_at_index(
+                keyed_accounts,
+                first_instruction_account
+                    + instruction_account_indices::Withdraw::Recipient as usize,
             )?;
-            let stake_history = get_sysvar_with_account_check::stake_history(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 3)?,
+            let clock = get_sysvar_with_account_check2::clock(
                 invoke_context,
+                instruction_context,
+                instruction_account_indices::Withdraw::Clock as usize,
             )?;
+            let stake_history = get_sysvar_with_account_check2::stake_history(
+                invoke_context,
+                instruction_context,
+                instruction_account_indices::Withdraw::StakeHistory as usize,
+            )?;
+            instruction_context.check_number_of_instruction_accounts(5)?;
             me.withdraw(
                 lamports,
                 to,
                 &clock,
                 &stake_history,
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 4)?,
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 5).ok(),
+                keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::Withdraw::WithdrawAuthority as usize,
+                )?,
+                keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::Withdraw::Custodian as usize,
+                )
+                .ok(),
             )
         }
         StakeInstruction::Deactivate => {
-            let clock = get_sysvar_with_account_check::clock(
-                keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?,
+            let clock = get_sysvar_with_account_check2::clock(
                 invoke_context,
+                instruction_context,
+                instruction_account_indices::Deactivate::Clock as usize,
             )?;
             me.deactivate(&clock, &signers)
         }
@@ -201,20 +339,29 @@ pub fn process_instruction(
                 .feature_set
                 .is_active(&feature_set::vote_stake_checked_instructions::id())
             {
+                instruction_context.check_number_of_instruction_accounts(4)?;
                 let authorized = Authorized {
-                    staker: *keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?
-                        .unsigned_key(),
+                    staker: *keyed_account_at_index(
+                        keyed_accounts,
+                        first_instruction_account
+                            + instruction_account_indices::InitializeChecked::AuthorizedStaker
+                                as usize,
+                    )?
+                    .unsigned_key(),
                     withdrawer: *keyed_account_at_index(
                         keyed_accounts,
-                        first_instruction_account + 3,
+                        first_instruction_account
+                            + instruction_account_indices::InitializeChecked::AuthorizedWithdrawer
+                                as usize,
                     )?
                     .signer_key()
                     .ok_or(InstructionError::MissingRequiredSignature)?,
                 };
 
-                let rent = get_sysvar_with_account_check::rent(
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?,
+                let rent = get_sysvar_with_account_check2::rent(
                     invoke_context,
+                    instruction_context,
+                    instruction_account_indices::InitializeChecked::Rent as usize,
                 )?;
                 me.initialize(&authorized, &Lockup::default(), &rent)
             } else {
@@ -226,20 +373,26 @@ pub fn process_instruction(
                 .feature_set
                 .is_active(&feature_set::vote_stake_checked_instructions::id())
             {
-                let clock = get_sysvar_with_account_check::clock(
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?,
+                let clock = get_sysvar_with_account_check2::clock(
                     invoke_context,
+                    instruction_context,
+                    instruction_account_indices::AuthorizeChecked::Clock as usize,
                 )?;
-                let _current_authority =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?;
-                let authorized_pubkey =
-                    &keyed_account_at_index(keyed_accounts, first_instruction_account + 3)?
-                        .signer_key()
-                        .ok_or(InstructionError::MissingRequiredSignature)?;
-                let custodian =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 4)
-                        .ok()
-                        .map(|ka| ka.unsigned_key());
+                instruction_context.check_number_of_instruction_accounts(4)?;
+                let authorized_pubkey = &keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeChecked::Authorized as usize,
+                )?
+                .signer_key()
+                .ok_or(InstructionError::MissingRequiredSignature)?;
+                let custodian = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeChecked::Custodian as usize,
+                )
+                .ok()
+                .map(|ka| ka.unsigned_key());
 
                 me.authorize(
                     &signers,
@@ -258,20 +411,34 @@ pub fn process_instruction(
                 .feature_set
                 .is_active(&feature_set::vote_stake_checked_instructions::id())
             {
-                let authority_base =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
-                let clock = get_sysvar_with_account_check::clock(
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 2)?,
-                    invoke_context,
+                instruction_context.check_number_of_instruction_accounts(2)?;
+                let authority_base = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeCheckedWithSeed::AuthorityBase
+                            as usize,
                 )?;
-                let authorized_pubkey =
-                    &keyed_account_at_index(keyed_accounts, first_instruction_account + 3)?
-                        .signer_key()
-                        .ok_or(InstructionError::MissingRequiredSignature)?;
-                let custodian =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 4)
-                        .ok()
-                        .map(|ka| ka.unsigned_key());
+                let clock = get_sysvar_with_account_check2::clock(
+                    invoke_context,
+                    instruction_context,
+                    instruction_account_indices::AuthorizeCheckedWithSeed::Clock as usize,
+                )?;
+                instruction_context.check_number_of_instruction_accounts(4)?;
+                let authorized_pubkey = &keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeCheckedWithSeed::Authorized
+                            as usize,
+                )?
+                .signer_key()
+                .ok_or(InstructionError::MissingRequiredSignature)?;
+                let custodian = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::AuthorizeCheckedWithSeed::Custodian as usize,
+                )
+                .ok()
+                .map(|ka| ka.unsigned_key());
 
                 me.authorize_with_seed(
                     authority_base,
@@ -292,9 +459,11 @@ pub fn process_instruction(
                 .feature_set
                 .is_active(&feature_set::vote_stake_checked_instructions::id())
             {
-                let custodian = if let Ok(custodian) =
-                    keyed_account_at_index(keyed_accounts, first_instruction_account + 2)
-                {
+                let custodian = if let Ok(custodian) = keyed_account_at_index(
+                    keyed_accounts,
+                    first_instruction_account
+                        + instruction_account_indices::SetLockupChecked::Custodian as usize,
+                ) {
                     Some(
                         *custodian
                             .signer_key()