Refactor: Make program_id always last in program chain (#20598)

* Replaces program_id field in InvokeContextStackFrame by index.

* Swaps order of program account and programdata account.

* Removes program_id parameter from InvokeContext::push().

program-runtime/src/instruction_processor.rs

@@ -485,7 +485,7 @@ impl InstructionProcessor {
             );
             return Err(InstructionError::AccountNotExecutable);
         }
-        let mut program_indices = vec![program_account_index];
+        let mut program_indices = vec![];
         if program_account.borrow().owner() == &bpf_loader_upgradeable::id() {
             if let UpgradeableLoaderState::Program {
                 programdata_address,
@@ -512,6 +512,7 @@ impl InstructionProcessor {
                 return Err(InstructionError::MissingAccount);
             }
         }
+        program_indices.push(program_account_index);
 
         Ok((message, caller_write_privileges, program_indices))
     }
@@ -593,8 +594,6 @@ impl InstructionProcessor {
             .get(0)
             .ok_or(InstructionError::GenericError)?;
 
-        let program_id = instruction.program_id(&message.account_keys);
-
         // Verify the calling program hasn't misbehaved
         invoke_context.verify_and_update(instruction, account_indices, caller_write_privileges)?;
 
@@ -602,13 +601,7 @@ impl InstructionProcessor {
         invoke_context.set_return_data(Vec::new())?;
 
         // Invoke callee
-        invoke_context.push(
+        invoke_context.push(message, instruction, program_indices, Some(account_indices))?;
-            program_id,
-            message,
-            instruction,
-            program_indices,
-            Some(account_indices),
-        )?;
 
         let mut instruction_processor = InstructionProcessor::default();
         for (program_id, process_instruction) in invoke_context.get_programs().iter() {

programs/bpf_loader/src/lib.rs

@@ -199,72 +199,71 @@ fn process_instruction_common(
     use_jit: bool,
 ) -> Result<(), InstructionError> {
     let logger = invoke_context.get_logger();
-    let keyed_accounts = invoke_context.get_keyed_accounts()?;
+    let program_id = invoke_context.get_caller()?;
 
+    let keyed_accounts = invoke_context.get_keyed_accounts()?;
     let first_account = keyed_account_at_index(keyed_accounts, first_instruction_account)?;
-    if first_account.executable()? {
+    let second_account = keyed_account_at_index(keyed_accounts, first_instruction_account + 1);
+    let (program, next_first_instruction_account) = if first_account.unsigned_key() == program_id {
+        (first_account, first_instruction_account)
+    } else if second_account
+        .as_ref()
+        .map(|keyed_account| keyed_account.unsigned_key() == program_id)
+        .unwrap_or(false)
+    {
+        (second_account?, first_instruction_account + 1)
+    } else {
+        if first_account.executable()? {
+            ic_logger_msg!(logger, "BPF loader is executable");
+            return Err(InstructionError::IncorrectProgramId);
+        }
+        (first_account, first_instruction_account)
+    };
+
+    if program.executable()? {
         debug_assert_eq!(
             first_instruction_account,
             1 - (invoke_context.invoke_depth() > 1) as usize,
         );
 
-        let program_id = invoke_context.get_caller()?;
+        if !check_loader_id(&program.owner()?) {
-        if first_account.unsigned_key() != program_id {
-            ic_logger_msg!(logger, "Program id mismatch");
-            return Err(InstructionError::IncorrectProgramId);
-        }
-
-        let (first_instruction_account, program_data_offset) =
-            if bpf_loader_upgradeable::check_id(&first_account.owner()?) {
-                if let UpgradeableLoaderState::Program {
-                    programdata_address,
-                } = first_account.state()?
-                {
-                    let programdata =
-                        keyed_account_at_index(keyed_accounts, first_instruction_account + 1)?;
-                    if programdata_address != *programdata.unsigned_key() {
-                        ic_logger_msg!(
-                            logger,
-                            "Wrong ProgramData account for this Program account"
-                        );
-                        return Err(InstructionError::InvalidArgument);
-                    }
-                    if !matches!(
-                        programdata.state()?,
-                        UpgradeableLoaderState::ProgramData {
-                            slot: _,
-                            upgrade_authority_address: _,
-                        }
-                    ) {
-                        ic_logger_msg!(logger, "Program has been closed");
-                        return Err(InstructionError::InvalidAccountData);
-                    }
-                    (
-                        first_instruction_account + 1,
-                        UpgradeableLoaderState::programdata_data_offset()?,
-                    )
-                } else {
-                    ic_logger_msg!(logger, "Invalid Program account");
-                    return Err(InstructionError::InvalidAccountData);
-                }
-            } else {
-                (first_instruction_account, 0)
-            };
-
-        let keyed_accounts = invoke_context.get_keyed_accounts()?;
-        let program = keyed_account_at_index(keyed_accounts, first_instruction_account)?;
-        let loader_id = &program.owner()?;
-
-        if !check_loader_id(loader_id) {
             ic_logger_msg!(logger, "Executable account not owned by the BPF loader");
             return Err(InstructionError::IncorrectProgramId);
         }
 
+        let program_data_offset = if bpf_loader_upgradeable::check_id(&program.owner()?) {
+            if let UpgradeableLoaderState::Program {
+                programdata_address,
+            } = program.state()?
+            {
+                if programdata_address != *first_account.unsigned_key() {
+                    ic_logger_msg!(logger, "Wrong ProgramData account for this Program account");
+                    return Err(InstructionError::InvalidArgument);
+                }
+                if !matches!(
+                    first_account.state()?,
+                    UpgradeableLoaderState::ProgramData {
+                        slot: _,
+                        upgrade_authority_address: _,
+                    }
+                ) {
+                    ic_logger_msg!(logger, "Program has been closed");
+                    return Err(InstructionError::InvalidAccountData);
+                }
+                UpgradeableLoaderState::programdata_data_offset()?
+            } else {
+                ic_logger_msg!(logger, "Invalid Program account");
+                return Err(InstructionError::InvalidAccountData);
+            }
+        } else {
+            0
+        };
+
         let executor = match invoke_context.get_executor(program_id) {
             Some(executor) => executor,
             None => {
                 let executor = create_executor(
-                    first_instruction_account,
+                    next_first_instruction_account,
                     program_data_offset,
                     invoke_context,
                     use_jit,
@@ -275,37 +274,32 @@ fn process_instruction_common(
             }
         };
         executor.execute(
-            first_instruction_account,
+            next_first_instruction_account,
             instruction_data,
             invoke_context,
             use_jit,
-        )?
+        )
     } else {
         debug_assert_eq!(first_instruction_account, 1);
-
-        let program_id = invoke_context.get_caller()?;
-        if !check_loader_id(program_id) {
-            ic_logger_msg!(logger, "Invalid BPF loader id");
-            return Err(InstructionError::IncorrectProgramId);
-        }
-
         if bpf_loader_upgradeable::check_id(program_id) {
             process_loader_upgradeable_instruction(
                 first_instruction_account,
                 instruction_data,
                 invoke_context,
                 use_jit,
-            )?;
+            )
-        } else {
+        } else if bpf_loader::check_id(program_id) || bpf_loader_deprecated::check_id(program_id) {
             process_loader_instruction(
                 first_instruction_account,
                 instruction_data,
                 invoke_context,
                 use_jit,
-            )?;
+            )
+        } else {
+            ic_logger_msg!(logger, "Invalid BPF loader id");
+            Err(InstructionError::IncorrectProgramId)
         }
     }
-    Ok(())
 }
 
 fn process_loader_upgradeable_instruction(
@@ -3374,8 +3368,8 @@ mod tests {
 
         // Try to invoke closed account
         let keyed_accounts: Vec<KeyedAccountTuple> = vec![
-            (false, false, &program_address, &program_account),
             (false, false, &programdata_address, &programdata_account),
+            (false, false, &program_address, &program_account),
         ];
         assert_eq!(
             Err(InstructionError::InvalidAccountData),

runtime/src/accounts.rs

@@ -429,7 +429,7 @@ impl Accounts {
 
             // Add loader to chain
             let program_owner = *program.owner();
-
+            account_indices.insert(0, program_account_index);
             if bpf_loader_upgradeable::check_id(&program_owner) {
                 // The upgradeable loader requires the derived ProgramData account
                 if let Ok(UpgradeableLoaderState::Program {
@@ -457,7 +457,6 @@ impl Accounts {
                 }
             }
 
-            account_indices.insert(0, program_account_index);
             program_id = program_owner;
         }
         Ok(account_indices)
@@ -1923,8 +1922,8 @@ mod tests {
         let result = loaded_accounts[0].0.as_ref().unwrap();
         assert_eq!(result.accounts[..2], accounts[..2]);
         assert_eq!(result.accounts[result.program_indices[0][0]], accounts[5]);
-        assert_eq!(result.accounts[result.program_indices[0][1]], accounts[3]);
+        assert_eq!(result.accounts[result.program_indices[0][1]], accounts[4]);
-        assert_eq!(result.accounts[result.program_indices[0][2]], accounts[4]);
+        assert_eq!(result.accounts[result.program_indices[0][2]], accounts[3]);
     }
 
     #[test]

runtime/src/message_processor.rs

@@ -116,7 +116,6 @@ impl<'a> ThisInvokeContext<'a> {
 impl<'a> InvokeContext for ThisInvokeContext<'a> {
     fn push(
         &mut self,
-        key: &Pubkey,
         message: &Message,
         instruction: &CompiledInstruction,
         program_indices: &[usize],
@@ -126,6 +125,23 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
             return Err(InstructionError::CallDepth);
         }
 
+        if let Some(index_of_program_id) = program_indices.last() {
+            let program_id = &self.accounts[*index_of_program_id].0;
+            let contains = self
+                .invoke_stack
+                .iter()
+                .any(|frame| frame.program_id() == Some(program_id));
+            let is_last = if let Some(last_frame) = self.invoke_stack.last() {
+                last_frame.program_id() == Some(program_id)
+            } else {
+                false
+            };
+            if contains && !is_last {
+                // Reentrancy not allowed unless caller is calling itself
+                return Err(InstructionError::ReentrancyNotAllowed);
+            }
+        }
+
         if self.invoke_stack.is_empty() {
             self.pre_accounts = Vec::with_capacity(instruction.accounts.len());
             let mut work = |_unique_index: usize, account_index: usize| {
@@ -140,17 +156,6 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
             instruction.visit_each_account(&mut work)?;
         }
 
-        let contains = self.invoke_stack.iter().any(|frame| frame.key == *key);
-        let is_last = if let Some(last_frame) = self.invoke_stack.last() {
-            last_frame.key == *key
-        } else {
-            false
-        };
-        if contains && !is_last {
-            // Reentrancy not allowed unless caller is calling itself
-            return Err(InstructionError::ReentrancyNotAllowed);
-        }
-
         // Create the KeyedAccounts that will be passed to the program
         let demote_program_write_locks = self
             .feature_set
@@ -180,8 +185,9 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
                 )
             }))
             .collect::<Vec<_>>();
+
         self.invoke_stack.push(InvokeContextStackFrame::new(
-            *key,
+            program_indices.len(),
             create_keyed_accounts_unified(keyed_accounts.as_slice()),
         ));
         Ok(())
@@ -259,11 +265,11 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
         write_privileges: &[bool],
     ) -> Result<(), InstructionError> {
         let do_support_realloc = self.feature_set.is_active(&do_support_realloc::id());
-        let stack_frame = self
+        let program_id = self
             .invoke_stack
             .last()
+            .and_then(|frame| frame.program_id())
             .ok_or(InstructionError::CallDepth)?;
-        let program_id = &stack_frame.key;
         let rent = &self.rent;
         let logger = &self.logger;
         let accounts = &self.accounts;
@@ -325,7 +331,7 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
     fn get_caller(&self) -> Result<&Pubkey, InstructionError> {
         self.invoke_stack
             .last()
-            .map(|frame| &frame.key)
+            .and_then(|frame| frame.program_id())
             .ok_or(InstructionError::CallDepth)
     }
     fn remove_first_keyed_account(&mut self) -> Result<(), InstructionError> {
@@ -549,7 +555,7 @@ impl MessageProcessor {
 
             invoke_context.set_instruction_index(instruction_index);
             let result = invoke_context
-                .push(program_id, message, instruction, program_indices, None)
+                .push(message, instruction, program_indices, None)
                 .and_then(|_| {
                     instruction_processor
                         .process_instruction(&instruction.data, &mut invoke_context)?;
@@ -693,10 +699,9 @@ mod tests {
 
         // Check call depth increases and has a limit
         let mut depth_reached = 0;
-        for program_id in invoke_stack.iter() {
+        for _ in 0..invoke_stack.len() {
             if Err(InstructionError::CallDepth)
                 == invoke_context.push(
-                    program_id,
                     &message,
                     &message.instructions[0],
                     &[MAX_DEPTH + depth_reached],
@@ -796,13 +801,7 @@ mod tests {
             &fee_calculator,
         );
         invoke_context
-            .push(
+            .push(&message, &message.instructions[0], &[0], None)
-                &accounts[0].0,
-                &message,
-                &message.instructions[0],
-                &[0],
-                None,
-            )
             .unwrap();
         assert!(invoke_context
             .verify(&message, &message.instructions[0], &[0])
@@ -1253,13 +1252,7 @@ mod tests {
             &fee_calculator,
         );
         invoke_context
-            .push(
+            .push(&message, &caller_instruction, &program_indices[..1], None)
-                &caller_program_id,
-                &message,
-                &caller_instruction,
-                &program_indices[..1],
-                None,
-            )
             .unwrap();
 
         // not owned account modified by the caller (before the invoke)
@@ -1317,13 +1310,7 @@ mod tests {
                 Instruction::new_with_bincode(callee_program_id, &case.0, metas.clone());
             let message = Message::new(&[callee_instruction], None);
             invoke_context
-                .push(
+                .push(&message, &caller_instruction, &program_indices[..1], None)
-                    &caller_program_id,
-                    &message,
-                    &caller_instruction,
-                    &program_indices[..1],
-                    None,
-                )
                 .unwrap();
             let caller_write_privileges = message
                 .account_keys
@@ -1410,13 +1397,7 @@ mod tests {
             &fee_calculator,
         );
         invoke_context
-            .push(
+            .push(&message, &caller_instruction, &program_indices, None)
-                &caller_program_id,
-                &message,
-                &caller_instruction,
-                &program_indices,
-                None,
-            )
             .unwrap();
 
         // not owned account modified by the invoker
@@ -1469,13 +1450,7 @@ mod tests {
                 Instruction::new_with_bincode(callee_program_id, &case.0, metas.clone());
             let message = Message::new(&[callee_instruction.clone()], None);
             invoke_context
-                .push(
+                .push(&message, &caller_instruction, &program_indices, None)
-                    &caller_program_id,
-                    &message,
-                    &caller_instruction,
-                    &program_indices,
-                    None,
-                )
                 .unwrap();
             assert_eq!(
                 InstructionProcessor::native_invoke(

sdk/src/process_instruction.rs

@@ -31,23 +31,29 @@ pub type ProcessInstructionWithContext =
     fn(usize, &[u8], &mut dyn InvokeContext) -> Result<(), InstructionError>;
 
 pub struct InvokeContextStackFrame<'a> {
-    pub key: Pubkey,
+    pub number_of_program_accounts: usize,
     pub keyed_accounts: Vec<KeyedAccount<'a>>,
     pub keyed_accounts_range: std::ops::Range<usize>,
 }
 
 impl<'a> InvokeContextStackFrame<'a> {
-    pub fn new(key: Pubkey, keyed_accounts: Vec<KeyedAccount<'a>>) -> Self {
+    pub fn new(number_of_program_accounts: usize, keyed_accounts: Vec<KeyedAccount<'a>>) -> Self {
         let keyed_accounts_range = std::ops::Range {
             start: 0,
             end: keyed_accounts.len(),
         };
         Self {
-            key,
+            number_of_program_accounts,
             keyed_accounts,
             keyed_accounts_range,
         }
     }
+
+    pub fn program_id(&self) -> Option<&Pubkey> {
+        self.keyed_accounts
+            .get(self.number_of_program_accounts.saturating_sub(1))
+            .map(|keyed_account| keyed_account.unsigned_key())
+    }
 }
 
 /// Invocation context passed to loaders
@@ -55,7 +61,6 @@ pub trait InvokeContext {
     /// Push a stack frame onto the invocation stack
     fn push(
         &mut self,
-        key: &Pubkey,
         message: &Message,
         instruction: &CompiledInstruction,
         program_indices: &[usize],
@@ -473,9 +478,17 @@ impl<'a> MockInvokeContext<'a> {
             fee_calculator: FeeCalculator::default(),
             return_data: (Pubkey::default(), Vec::new()),
         };
+        let number_of_program_accounts = keyed_accounts
+            .iter()
+            .position(|keyed_account| keyed_account.unsigned_key() == program_id)
+            .unwrap_or(0)
+            .saturating_add(1);
         invoke_context
             .invoke_stack
-            .push(InvokeContextStackFrame::new(*program_id, keyed_accounts));
+            .push(InvokeContextStackFrame::new(
+                number_of_program_accounts,
+                keyed_accounts,
+            ));
         invoke_context
     }
 }
@@ -498,14 +511,13 @@ pub fn mock_set_sysvar<T: Sysvar>(
 impl<'a> InvokeContext for MockInvokeContext<'a> {
     fn push(
         &mut self,
-        _key: &Pubkey,
         _message: &Message,
         _instruction: &CompiledInstruction,
         _program_indices: &[usize],
         _account_indices: Option<&[usize]>,
     ) -> Result<(), InstructionError> {
         self.invoke_stack.push(InvokeContextStackFrame::new(
-            *_key,
+            0,
             create_keyed_accounts_unified(&[]),
         ));
         Ok(())
@@ -535,7 +547,7 @@ impl<'a> InvokeContext for MockInvokeContext<'a> {
     fn get_caller(&self) -> Result<&Pubkey, InstructionError> {
         self.invoke_stack
             .last()
-            .map(|frame| &frame.key)
+            .and_then(|frame| frame.program_id())
             .ok_or(InstructionError::CallDepth)
     }
     fn remove_first_keyed_account(&mut self) -> Result<(), InstructionError> {