feat: remove program targeting (#30)

buffalojoec · Oct 18, 2024 · 2b973db · 2b973db
1 parent dfad1af
commit 2b973db
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 70 deletions.
diff --git a/harness/src/lib.rs b/harness/src/lib.rs
@@ -64,9 +64,7 @@ pub struct Mollusk {
     pub compute_budget: ComputeBudget,
     pub feature_set: FeatureSet,
     pub fee_structure: FeeStructure,
-    pub program_account: AccountSharedData,
     pub program_cache: ProgramCache,
-    pub program_id: Pubkey,
     pub sysvars: Sysvars,
 }
 
@@ -78,63 +76,36 @@ impl Default for Mollusk {
              solana_runtime::message_processor=debug,\
              solana_runtime::system_instruction_processor=trace",
         );
-        let (program_id, program_account) = program::keyed_account_for_system_program();
         Self {
             compute_budget: ComputeBudget::default(),
             feature_set: FeatureSet::all_enabled(),
             fee_structure: FeeStructure::default(),
-            program_account,
             program_cache: ProgramCache::default(),
-            program_id,
             sysvars: Sysvars::default(),
         }
     }
 }
 
 impl Mollusk {
-    /// Create a new Mollusk instance for the provided program.
+    /// Create a new Mollusk instance containing the provided program.
     ///
     /// Attempts the load the program's ELF file from the default search paths.
-    /// Once loaded, adds the program to the program cache and updates the
-    /// Mollusk instance with the program's ID and account.
+    /// Once loaded, adds the program to the program cache and returns the
+    /// newly created Mollusk instance.
     pub fn new(program_id: &Pubkey, program_name: &str) -> Self {
-        let mut mollusk = Self {
-            program_id: *program_id,
-            program_account: program::create_program_account_loader_v3(program_id),
-            ..Default::default()
-        };
-        mollusk.add_program(program_id, program_name);
+        let mut mollusk = Self::default();
+        mollusk.add_program(program_id, program_name, &DEFAULT_LOADER_KEY);
         mollusk
     }
 
     /// Add a program to the test environment.
     ///
     /// If you intend to CPI to a program, this is likely what you want to use.
-    pub fn add_program(&mut self, program_id: &Pubkey, program_name: &str) {
-        let elf = file::load_program_elf(program_name);
-        self.add_program_with_elf(program_id, &elf);
-    }
-
-    /// Add a program to the test environment under a specific loader.
-    ///
-    /// If you intend to CPI to a program, this is likely what you want to use.
-    pub fn add_program_with_loader(
-        &mut self,
-        program_id: &Pubkey,
-        program_name: &str,
-        loader_key: &Pubkey,
-    ) {
+    pub fn add_program(&mut self, program_id: &Pubkey, program_name: &str, loader_key: &Pubkey) {
         let elf = file::load_program_elf(program_name);
         self.add_program_with_elf_and_loader(program_id, &elf, loader_key);
     }
 
-    /// Add a program to the test environment using a provided ELF.
-    ///
-    /// If you intend to CPI to a program, this is likely what you want to use.
-    pub fn add_program_with_elf(&mut self, program_id: &Pubkey, elf: &[u8]) {
-        self.add_program_with_elf_and_loader(program_id, elf, &DEFAULT_LOADER_KEY);
-    }
-
     /// Add a program to the test environment using a provided ELF under a
     /// specific loader.
     ///
@@ -154,32 +125,6 @@ impl Mollusk {
         );
     }
 
-    /// Switch the target program to a different program.
-    ///
-    /// Note: The program must already be contained in the program cache.
-    pub fn switch_target_program(&mut self, program_id: &Pubkey) {
-        let loader_key: Pubkey = self
-            .program_cache
-            .cache()
-            .read()
-            .unwrap()
-            .find(program_id)
-            .expect("Program not found in cache")
-            .account_owner
-            .into();
-        if loader_key != DEFAULT_LOADER_KEY {
-            panic!("Loader not supported for target program: {:?}", loader_key);
-        }
-        self.program_id = *program_id;
-        self.program_account = program::create_program_account_loader_v3(program_id);
-    }
-
-    /// Add a program to the cache and make it the target program.
-    pub fn add_and_switch_target_program(&mut self, program_id: &Pubkey, program_name: &str) {
-        self.add_program(program_id, program_name);
-        self.switch_target_program(program_id);
-    }
-
     /// Warp the test environment to a slot by updating sysvars.
     pub fn warp_to_slot(&mut self, slot: u64) {
         self.sysvars.warp_to_slot(slot)

diff --git a/harness/src/program.rs b/harness/src/program.rs
@@ -9,8 +9,7 @@ use {
     },
     solana_sdk::{
         account::{Account, AccountSharedData},
-        bpf_loader,
-        bpf_loader_upgradeable::{self, UpgradeableLoaderState},
+        bpf_loader_upgradeable::UpgradeableLoaderState,
         feature_set::FeatureSet,
         native_loader,
         pubkey::Pubkey,
@@ -19,6 +18,14 @@ use {
     std::sync::{Arc, RwLock},
 };
 
+/// Loader keys, re-exported from `solana_sdk` for convenience.
+pub mod loader_keys {
+    pub use solana_sdk::{
+        bpf_loader::ID as LOADER_V2, bpf_loader_upgradeable::ID as LOADER_V3,
+        loader_v4::ID as LOADER_V4, native_loader::ID as NATIVE_LOADER,
+    };
+}
+
 pub struct ProgramCache {
     cache: RwLock<ProgramCacheForTxBatch>,
 }
@@ -108,12 +115,12 @@ static BUILTINS: &[Builtin] = &[
         entrypoint: solana_system_program::system_processor::Entrypoint::vm,
     },
     Builtin {
-        program_id: bpf_loader::id(),
+        program_id: loader_keys::LOADER_V2,
         name: "solana_bpf_loader_program",
         entrypoint: solana_bpf_loader_program::Entrypoint::vm,
     },
     Builtin {
-        program_id: bpf_loader_upgradeable::id(),
+        program_id: loader_keys::LOADER_V3,
         name: "solana_bpf_loader_upgradeable_program",
         entrypoint: solana_bpf_loader_program::Entrypoint::vm,
     },
@@ -160,7 +167,7 @@ pub fn create_program_account_loader_v2(elf: &[u8]) -> AccountSharedData {
     AccountSharedData::from(Account {
         lamports,
         data: elf.to_vec(),
-        owner: bpf_loader::id(),
+        owner: loader_keys::LOADER_V2,
         executable: true,
         rent_epoch: 0,
     })
@@ -169,7 +176,7 @@ pub fn create_program_account_loader_v2(elf: &[u8]) -> AccountSharedData {
 /// Create a BPF Loader v3 (Upgradeable) program account.
 pub fn create_program_account_loader_v3(program_id: &Pubkey) -> AccountSharedData {
     let programdata_address =
-        Pubkey::find_program_address(&[program_id.as_ref()], &bpf_loader_upgradeable::id()).0;
+        Pubkey::find_program_address(&[program_id.as_ref()], &loader_keys::LOADER_V3).0;
     let data = bincode::serialize(&UpgradeableLoaderState::Program {
         programdata_address,
     })
@@ -178,7 +185,7 @@ pub fn create_program_account_loader_v3(program_id: &Pubkey) -> AccountSharedDat
     AccountSharedData::from(Account {
         lamports,
         data,
-        owner: bpf_loader_upgradeable::id(),
+        owner: loader_keys::LOADER_V3,
         executable: true,
         rent_epoch: 0,
     })
@@ -205,7 +212,7 @@ pub fn create_program_data_account_loader_v3(elf: &[u8]) -> AccountSharedData {
     AccountSharedData::from(Account {
         lamports,
         data,
-        owner: bpf_loader_upgradeable::id(),
+        owner: loader_keys::LOADER_V3,
         executable: false,
         rent_epoch: 0,
     })

diff --git a/harness/tests/bpf_program.rs b/harness/tests/bpf_program.rs
@@ -273,7 +273,11 @@ fn test_cpi() {
         );
     }
 
-    mollusk.add_program(&cpi_target_program_id, "test_program_cpi_target");
+    mollusk.add_program(
+        &cpi_target_program_id,
+        "test_program_cpi_target",
+        &mollusk_svm::program::loader_keys::LOADER_V3,
+    );
 
     // Fail account not signer.
     {