Program loader helpers (#18)

anza-xyz · Jun 26, 2024 · 5d2c351 · 5d2c351
1 parent 7a75350
commit 5d2c351
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/harness/src/lib.rs b/harness/src/lib.rs
@@ -45,6 +45,7 @@ use {
     },
     solana_sdk::{
         account::AccountSharedData,
+        bpf_loader_upgradeable,
         feature_set::FeatureSet,
         hash::Hash,
         instruction::Instruction,
@@ -112,8 +113,26 @@ impl Mollusk {
     /// 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: &'static str) {
         let elf = file::load_program_elf(program_name);
-        self.program_cache
-            .add_program(program_id, &elf, &self.compute_budget, &self.feature_set);
+        self.program_cache.add_program(
+            program_id,
+            &bpf_loader_upgradeable::id(),
+            &elf,
+            &self.compute_budget,
+            &self.feature_set,
+        );
+    }
+
+    /// 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, loader_key: &Pubkey, elf: &[u8]) {
+        self.program_cache.add_program(
+            program_id,
+            loader_key,
+            elf,
+            &self.compute_budget,
+            &self.feature_set,
+        );
     }
 
     /// Warp the test environment to a slot by updating sysvars.

diff --git a/harness/src/program.rs b/harness/src/program.rs
@@ -9,6 +9,7 @@ use {
     },
     solana_sdk::{
         account::{Account, AccountSharedData},
+        bpf_loader,
         bpf_loader_upgradeable::{self, UpgradeableLoaderState},
         feature_set::FeatureSet,
         native_loader,
@@ -43,6 +44,7 @@ impl ProgramCache {
     pub fn add_program(
         &mut self,
         program_id: &Pubkey,
+        loader_key: &Pubkey,
         elf: &[u8],
         compute_budget: &ComputeBudget,
         feature_set: &FeatureSet,
@@ -55,7 +57,7 @@ impl ProgramCache {
             *program_id,
             Arc::new(
                 LoadedProgram::new(
-                    &bpf_loader_upgradeable::id(),
+                    loader_key,
                     environment,
                     0,
                     0,
@@ -125,6 +127,18 @@ pub fn bpf_loader_upgradeable_program() -> (Pubkey, AccountSharedData) {
 
 /* ... */
 
+/// Create a BPF Loader 2 program account.
+pub fn program_account_loader_2(elf: &[u8]) -> AccountSharedData {
+    let lamports = Rent::default().minimum_balance(elf.len());
+    AccountSharedData::from(Account {
+        lamports,
+        data: elf.to_vec(),
+        owner: bpf_loader::id(),
+        executable: true,
+        rent_epoch: 0,
+    })
+}
+
 /// Create a BPF Loader Upgradeable program account.
 pub fn program_account(program_id: &Pubkey) -> AccountSharedData {
     let programdata_address =