address review comments

solana-labs · May 1, 2023 · 5828711 · 5828711
1 parent 8d0d713
commit 5828711
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 31 deletions.
diff --git a/program-runtime/src/invoke_context.rs b/program-runtime/src/invoke_context.rs
@@ -5,6 +5,7 @@ use {
         compute_budget::ComputeBudget,
         executor_cache::TransactionExecutorCache,
         ic_logger_msg, ic_msg,
+        loaded_programs::LoadedProgram,
         log_collector::LogCollector,
         pre_account::PreAccount,
         stable_log,
@@ -882,7 +883,8 @@ macro_rules! with_mock_invoke_context_and_builtin_programs {
         $invoke_context:ident,
         $transaction_context:ident,
         $transaction_accounts:expr,
-        $builtin_programs:expr
+        $builtin_programs:expr,
+        $loaded_programs:expr
     ) => {
         use {
             solana_sdk::{
@@ -922,14 +924,18 @@ macro_rules! with_mock_invoke_context_and_builtin_programs {
                 }
             }
         });
+        let mut executor_cache = TransactionExecutorCache::default();
+        $loaded_programs.iter().for_each(|(pubkey, program)| {
+            executor_cache.set(*pubkey, program.clone(), true, false, 0)
+        });
         let mut $invoke_context = InvokeContext::new(
             &mut $transaction_context,
             Rent::default(),
             $builtin_programs,
             &sysvar_cache,
             Some(LogCollector::new_ref()),
             compute_budget,
-            Rc::new(RefCell::new(TransactionExecutorCache::default())),
+            Rc::new(RefCell::new(executor_cache)),
             Arc::new(FeatureSet::all_enabled()),
             Hash::default(),
             0,
@@ -944,32 +950,59 @@ macro_rules! with_mock_invoke_context {
         $invoke_context:ident,
         $transaction_context:ident,
         $transaction_accounts:expr,
-        $builtin_programs:expr
+        $builtin_programs:expr,
+        $loaded_programs:expr
     ) => {
         use $crate::with_mock_invoke_context_and_builtin_programs;
         with_mock_invoke_context_and_builtin_programs!(
             $invoke_context,
             $transaction_context,
             $transaction_accounts,
-            $builtin_programs
+            $builtin_programs,
+            $loaded_programs
+        );
+    };
+
+    (
+        $invoke_context:ident,
+        $transaction_context:ident,
+        $transaction_accounts:expr,
+        $builtin_programs:expr
+    ) => {
+        use $crate::{
+            loaded_programs::LoadedProgram, with_mock_invoke_context_and_builtin_programs,
+        };
+        let loaded_programs: Vec<(Pubkey, Arc<LoadedProgram>)> = vec![];
+        with_mock_invoke_context_and_builtin_programs!(
+            $invoke_context,
+            $transaction_context,
+            $transaction_accounts,
+            $builtin_programs,
+            &loaded_programs
         );
     };
 
     ($invoke_context:ident, $transaction_context:ident, $transaction_accounts:expr) => {
         use $crate::{
-            builtin_program::BuiltinPrograms, with_mock_invoke_context_and_builtin_programs,
+            builtin_program::BuiltinPrograms, loaded_programs::LoadedProgram,
+            with_mock_invoke_context_and_builtin_programs,
         };
         let builtin_programs = BuiltinPrograms::default();
+        let loaded_programs: Vec<(Pubkey, Arc<LoadedProgram>)> = vec![];
         with_mock_invoke_context_and_builtin_programs!(
             $invoke_context,
             $transaction_context,
             $transaction_accounts,
-            &builtin_programs
+            &builtin_programs,
+            &loaded_programs
         );
     };
 }
 
-pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut InvokeContext)>(
+pub fn mock_process_instruction_with_loaded_programs<
+    F: FnMut(&mut InvokeContext),
+    G: FnMut(&mut InvokeContext),
+>(
     loader_id: &Pubkey,
     mut program_indices: Vec<IndexOfAccount>,
     instruction_data: &[u8],
@@ -979,6 +1012,7 @@ pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut Invo
     process_instruction: ProcessInstructionWithContext,
     mut pre_adjustments: F,
     mut post_adjustments: G,
+    loaded_programs: &[(Pubkey, Arc<LoadedProgram>)],
 ) -> Vec<AccountSharedData> {
     let mut instruction_accounts: Vec<InstructionAccount> =
         Vec::with_capacity(instruction_account_metas.len());
@@ -1012,7 +1046,8 @@ pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut Invo
         invoke_context,
         transaction_context,
         transaction_accounts,
-        &builtin_programs
+        &builtin_programs,
+        loaded_programs
     );
     pre_adjustments(&mut invoke_context);
     let result = invoke_context.process_instruction(
@@ -1029,6 +1064,30 @@ pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut Invo
     transaction_accounts
 }
 
+pub fn mock_process_instruction<F: FnMut(&mut InvokeContext), G: FnMut(&mut InvokeContext)>(
+    loader_id: &Pubkey,
+    program_indices: Vec<IndexOfAccount>,
+    instruction_data: &[u8],
+    transaction_accounts: Vec<TransactionAccount>,
+    instruction_account_metas: Vec<AccountMeta>,
+    expected_result: Result<(), InstructionError>,
+    process_instruction: ProcessInstructionWithContext,
+    pre_adjustments: F,
+    post_adjustments: G,
+) -> Vec<AccountSharedData> {
+    mock_process_instruction_with_loaded_programs(
+        loader_id,
+        program_indices,
+        instruction_data,
+        transaction_accounts,
+        instruction_account_metas,
+        expected_result,
+        process_instruction,
+        pre_adjustments,
+        post_adjustments,
+        &[],
+    )
+}
 #[cfg(test)]
 mod tests {
     use {

diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
@@ -1666,7 +1666,9 @@ mod tests {
         super::*,
         rand::Rng,
         solana_program_runtime::{
-            invoke_context::mock_process_instruction, with_mock_invoke_context,
+            invoke_context::mock_process_instruction,
+            invoke_context::mock_process_instruction_with_loaded_programs,
+            with_mock_invoke_context,
         },
         solana_rbpf::{
             verifier::{Verifier, VerifierError},
@@ -1730,21 +1732,17 @@ mod tests {
         expected_result: Result<(), InstructionError>,
         loaded_programs: &[(Pubkey, Arc<LoadedProgram>)],
     ) -> Vec<AccountSharedData> {
-        mock_process_instruction(
+        mock_process_instruction_with_loaded_programs(
             loader_id,
             program_indices.to_vec(),
             instruction_data,
             transaction_accounts,
             instruction_accounts,
             expected_result,
             super::process_instruction,
-            |invoke_context| {
-                let mut cache = invoke_context.tx_executor_cache.borrow_mut();
-                loaded_programs.iter().for_each(|(pubkey, program)| {
-                    cache.set(*pubkey, program.clone(), true, false, 0)
-                });
-            },
             |_invoke_context| {},
+            |_invoke_context| {},
+            loaded_programs,
         )
     }
 
@@ -1996,7 +1994,7 @@ mod tests {
         );
 
         // Case: limited budget
-        mock_process_instruction(
+        mock_process_instruction_with_loaded_programs(
             &loader_id,
             vec![0],
             &[],
@@ -2006,10 +2004,9 @@ mod tests {
             super::process_instruction,
             |invoke_context| {
                 invoke_context.mock_set_remaining(0);
-                let mut cache = invoke_context.tx_executor_cache.borrow_mut();
-                cache.set(program_id, loaded_program.clone(), true, false, 0);
             },
             |_invoke_context| {},
+            &[(program_id, loaded_program.clone())],
         );
 
         // Case: Account not a program

diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
@@ -36,7 +36,7 @@ use {
         compute_budget::{self, ComputeBudget, MAX_COMPUTE_UNIT_LIMIT},
         declare_process_instruction,
         executor_cache::TransactionExecutorCache,
-        invoke_context::mock_process_instruction,
+        invoke_context::mock_process_instruction_with_loaded_programs,
         loaded_programs::{LoadedProgram, LoadedProgramType},
         prioritization_fee::{PrioritizationFeeDetails, PrioritizationFeeType},
         timings::ExecuteTimings,
@@ -7717,7 +7717,7 @@ fn test_bpf_loader_upgradeable_deploy_with_max_len() {
         .expect("Failed to load the program");
 
     // Invoke deployed program
-    mock_process_instruction(
+    mock_process_instruction_with_loaded_programs(
         &bpf_loader_upgradeable::id(),
         vec![0, 1],
         &[],
@@ -7728,17 +7728,9 @@ fn test_bpf_loader_upgradeable_deploy_with_max_len() {
         Vec::new(),
         Ok(()),
         solana_bpf_loader_program::process_instruction,
-        |invoke_context| {
-            let mut cache = invoke_context.tx_executor_cache.borrow_mut();
-            cache.set(
-                program_keypair.pubkey(),
-                loaded_program.clone(),
-                true,
-                false,
-                0,
-            );
-        },
         |_invoke_context| {},
+        |_invoke_context| {},
+        &[(program_keypair.pubkey(), loaded_program)],
     );
 
     // Test initialized program account