fix: transaction account dedupe and roles

harness/src/keys.rs

@@ -0,0 +1,168 @@
+//! Instruction <-> Transaction key deduplication and privilege handling.
+//!
+//! Solana instructions and transactions are designed to be intentionally
+//! verbosely declarative, to provide the runtime with granular directives
+//! for manipulating chain state.
+//!
+//! As a result, when a transaction is _compiled_, many steps occur:
+//! * Ensuring there is a fee payer.
+//! * Ensuring there is a signature.
+//! * Deduplicating account keys.
+//! * Configuring the highest role awarded to each account key.
+//! * ...
+//!
+//! Since Mollusk does not use transactions or fee payers, the deduplication
+//! of account keys and handling of roles are the only two steps necessary
+//! to perform under the hood within the harness.
+//!
+//! This implementation closely follows the implementation in the Anza SDK
+//! for `Message::new_with_blockhash`. For more information, see:
+//! <https://github.com/anza-xyz/agave/blob/c6e8239843af8e6301cd198e39d0a44add427bef/sdk/program/src/message/legacy.rs#L357>.
+
+use {
+    solana_sdk::{
+        account::{AccountSharedData, WritableAccount},
+        instruction::Instruction,
+        pubkey::Pubkey,
+        transaction_context::{IndexOfAccount, InstructionAccount, TransactionAccount},
+    },
+    std::collections::HashMap,
+};
+
+struct KeyMap(HashMap<Pubkey, (bool, bool)>);
+
+impl KeyMap {
+    fn compile(instruction: &Instruction) -> Self {
+        let mut map: HashMap<Pubkey, (bool, bool)> = HashMap::new();
+        map.entry(instruction.program_id).or_default();
+        for meta in instruction.accounts.iter() {
+            let entry = map.entry(meta.pubkey).or_default();
+            entry.0 |= meta.is_signer;
+            entry.1 |= meta.is_writable;
+        }
+        Self(map)
+    }
+
+    fn is_signer(&self, key: &Pubkey) -> bool {
+        self.0.get(key).map(|(s, _)| *s).unwrap_or(false)
+    }
+
+    fn is_writable(&self, key: &Pubkey) -> bool {
+        self.0.get(key).map(|(_, w)| *w).unwrap_or(false)
+    }
+}
+
+struct Keys<'a> {
+    keys: Vec<&'a Pubkey>,
+    key_map: &'a KeyMap,
+}
+
+impl<'a> Keys<'a> {
+    fn new(key_map: &'a KeyMap) -> Self {
+        Self {
+            keys: key_map.0.keys().collect(),
+            key_map,
+        }
+    }
+
+    fn position(&self, key: &Pubkey) -> u8 {
+        self.keys.iter().position(|k| *k == key).unwrap() as u8
+    }
+
+    fn is_signer(&self, index: usize) -> bool {
+        self.key_map.is_signer(self.keys[index])
+    }
+
+    fn is_writable(&self, index: usize) -> bool {
+        self.key_map.is_writable(self.keys[index])
+    }
+}
+
+// Helper struct so Mollusk doesn't have to clone instruction data.
+struct CompiledInstructionWithoutData {
+    program_id_index: u8,
+    accounts: Vec<u8>,
+}
+
+pub struct CompiledAccounts {
+    pub program_id_index: u16,
+    pub instruction_accounts: Vec<InstructionAccount>,
+    pub transaction_accounts: Vec<TransactionAccount>,
+}
+
+pub fn compile_accounts(
+    instruction: &Instruction,
+    accounts: &[(Pubkey, AccountSharedData)],
+    loader_key: Pubkey,
+) -> CompiledAccounts {
+    let key_map = KeyMap::compile(instruction);
+    let keys = Keys::new(&key_map);
+
+    let compiled_instruction = CompiledInstructionWithoutData {
+        program_id_index: keys.position(&instruction.program_id),
+        accounts: instruction
+            .accounts
+            .iter()
+            .map(|account_meta| keys.position(&account_meta.pubkey))
+            .collect(),
+    };
+
+    let instruction_accounts: Vec<InstructionAccount> = compiled_instruction
+        .accounts
+        .iter()
+        .enumerate()
+        .map(|(ix_account_index, &index_in_transaction)| {
+            let index_in_callee = compiled_instruction
+                .accounts
+                .get(0..ix_account_index)
+                .unwrap()
+                .iter()
+                .position(|&account_index| account_index == index_in_transaction)
+                .unwrap_or(ix_account_index) as IndexOfAccount;
+            let index_in_transaction = index_in_transaction as usize;
+            InstructionAccount {
+                index_in_transaction: index_in_transaction as IndexOfAccount,
+                index_in_caller: index_in_transaction as IndexOfAccount,
+                index_in_callee,
+                is_signer: keys.is_signer(index_in_transaction),
+                is_writable: keys.is_writable(index_in_transaction),
+            }
+        })
+        .collect();
+
+    let transaction_accounts: Vec<TransactionAccount> = keys
+        .keys
+        .iter()
+        .map(|key| {
+            if *key == &instruction.program_id {
+                (**key, stub_out_program_account(loader_key))
+            } else {
+                let account = accounts
+                    .iter()
+                    .find(|(k, _)| k == *key)
+                    .map(|(_, account)| account.clone())
+                    .unwrap_or_else(|| {
+                        panic!(
+                            "    [mollusk]: An account required by the instruction was not \
+                             provided: {:?}",
+                            key,
+                        )
+                    });
+                (**key, account)
+            }
+        })
+        .collect();
+
+    CompiledAccounts {
+        program_id_index: compiled_instruction.program_id_index as u16,
+        instruction_accounts,
+        transaction_accounts,
+    }
+}
+
+fn stub_out_program_account(loader_key: Pubkey) -> AccountSharedData {
+    let mut program_account = AccountSharedData::default();
+    program_account.set_owner(loader_key);
+    program_account.set_executable(true);
+    program_account
+}

harness/src/lib.rs

@@ -28,6 +28,7 @@
 //!   series of checks on the result, panicking if any checks fail.
 
 pub mod file;
+mod keys;
 pub mod program;
 pub mod result;
 pub mod sysvar;
@@ -38,30 +39,23 @@ use {
         result::{Check, InstructionResult},
         sysvar::Sysvars,
     },
+    keys::CompiledAccounts,
     solana_compute_budget::compute_budget::ComputeBudget,
     solana_program_runtime::{
         invoke_context::{EnvironmentConfig, InvokeContext},
         sysvar_cache::SysvarCache,
         timings::ExecuteTimings,
     },
     solana_sdk::{
-        account::AccountSharedData,
-        bpf_loader_upgradeable,
-        feature_set::FeatureSet,
-        fee::FeeStructure,
-        hash::Hash,
-        instruction::Instruction,
-        pubkey::Pubkey,
-        transaction_context::{InstructionAccount, TransactionContext},
+        account::AccountSharedData, bpf_loader_upgradeable, feature_set::FeatureSet,
+        fee::FeeStructure, hash::Hash, instruction::Instruction, pubkey::Pubkey,
+        transaction_context::TransactionContext,
     },
-    std::{collections::HashMap, sync::Arc},
+    std::sync::Arc,
 };
 
 const DEFAULT_LOADER_KEY: Pubkey = bpf_loader_upgradeable::id();
 
-const PROGRAM_ACCOUNTS_LEN: usize = 1;
-const PROGRAM_INDICES: &[u16] = &[0];
-
 /// The Mollusk API, providing a simple interface for testing Solana programs.
 ///
 /// All fields can be manipulated through a handful of helper methods, but
@@ -203,40 +197,22 @@ impl Mollusk {
         let mut compute_units_consumed = 0;
         let mut timings = ExecuteTimings::default();
 
-        let instruction_accounts = {
-            // For ensuring each account has the proper privilege level (dedupe).
-            //  <pubkey, (is_signer, is_writable)>
-            let mut privileges: HashMap<Pubkey, (bool, bool)> = HashMap::new();
-
-            for meta in instruction.accounts.iter() {
-                let entry = privileges.entry(meta.pubkey).or_default();
-                entry.0 |= meta.is_signer;
-                entry.1 |= meta.is_writable;
-            }
-
-            instruction
-                .accounts
-                .iter()
-                .enumerate()
-                .map(|(i, meta)| {
-                    // Guaranteed by the last iteration.
-                    let (is_signer, is_writable) = privileges.get(&meta.pubkey).unwrap();
-                    InstructionAccount {
-                        index_in_callee: i as u16,
-                        index_in_caller: i as u16,
-                        index_in_transaction: (i + PROGRAM_ACCOUNTS_LEN) as u16,
-                        is_signer: *is_signer,
-                        is_writable: *is_writable,
-                    }
-                })
-                .collect::<Vec<_>>()
-        };
+        let loader_key = self
+            .program_cache
+            .load_program(&instruction.program_id)
+            .unwrap_or_else(|| {
+                panic!(
+                    "    [mollusk]: Program targeted by instruction is missing from cache: {:?}",
+                    instruction.program_id,
+                )
+            })
+            .account_owner();
 
-        let transaction_accounts = [(self.program_id, self.program_account.clone())]
-            .iter()
-            .chain(accounts)
-            .cloned()
-            .collect::<Vec<_>>();
+        let CompiledAccounts {
+            program_id_index,
+            instruction_accounts,
+            transaction_accounts,
+        } = crate::keys::compile_accounts(instruction, accounts, loader_key);
 
         let mut transaction_context = TransactionContext::new(
             transaction_accounts,
@@ -264,20 +240,26 @@ impl Mollusk {
             .process_instruction(
                 &instruction.data,
                 &instruction_accounts,
-                PROGRAM_INDICES,
+                &[program_id_index],
                 &mut compute_units_consumed,
                 &mut timings,
             )
         };
 
-        let resulting_accounts = transaction_context
-            .deconstruct_without_keys()
-            .unwrap()
-            .into_iter()
-            .skip(PROGRAM_ACCOUNTS_LEN)
-            .zip(instruction.accounts.iter())
-            .map(|(account, meta)| (meta.pubkey, account))
-            .collect::<Vec<_>>();
+        let resulting_accounts: Vec<(Pubkey, AccountSharedData)> = (0..transaction_context
+            .get_number_of_accounts())
+            .filter_map(|index| {
+                let key = transaction_context
+                    .get_key_of_account_at_index(index)
+                    .unwrap();
+                let account = transaction_context.get_account_at_index(index).unwrap();
+                if *key != instruction.program_id {
+                    Some((*key, account.take()))
+                } else {
+                    None
+                }
+            })
+            .collect();
 
         InstructionResult {
             compute_units_consumed,

harness/src/program.rs

@@ -42,6 +42,13 @@ impl ProgramCache {
         &self.cache
     }
 
+    /// Add a builtin program to the cache.
+    pub fn add_builtin(&mut self, builtin: Builtin) {
+        let program_id = builtin.program_id;
+        let entry = builtin.program_cache_entry();
+        self.cache.write().unwrap().replenish(program_id, entry);
+    }
+
     /// Add a program to the cache.
     pub fn add_program(
         &mut self,
@@ -72,11 +79,9 @@ impl ProgramCache {
         );
     }
 
-    /// Add a builtin program to the cache.
-    pub fn add_builtin(&mut self, builtin: Builtin) {
-        let program_id = builtin.program_id;
-        let entry = builtin.program_cache_entry();
-        self.cache.write().unwrap().replenish(program_id, entry);
+    /// Load a program from the cache.
+    pub fn load_program(&self, program_id: &Pubkey) -> Option<Arc<ProgramCacheEntry>> {
+        self.cache.read().unwrap().find(program_id)
     }
 }
 

harness/tests/bpf_program.rs

@@ -254,18 +254,6 @@ fn test_cpi() {
         )
     };
 
-    // Fail CPI target program account not provided.
-    {
-        mollusk.process_and_validate_instruction(
-            &instruction,
-            &[(key, account.clone())],
-            &[
-                Check::err(ProgramError::NotEnoughAccountKeys),
-                Check::compute_units(0), // No compute units used.
-            ],
-        );
-    }
-
     // Fail CPI target program not added to test environment.
     {
         mollusk.process_and_validate_instruction(