Program Test: Add BPF program account helpers (#1614)

* program-test: add program account helpers

* program-test: add test for Core BPF programs

* program-test: add BPF program fixture

* wrap into `add_upgradeable_program_to_genesis`

* wrap ELF loading into `add_upgradeable_program_to_genesis`

* cleanup

program-test/src/lib.rs

@@ -614,6 +614,33 @@ impl ProgramTest {
         self.add_account(address, account.into());
     }
 
+    /// Add a BPF Upgradeable program to the test environment's genesis config.
+    ///
+    /// When testing BPF programs using the program ID of a runtime builtin
+    /// program - such as Core BPF programs - the program accounts must be
+    /// added to the genesis config in order to make them available to the new
+    /// Bank as it's being initialized.
+    ///
+    /// The presence of these program accounts will cause Bank to skip adding
+    /// the builtin version of the program, allowing the provided BPF program
+    /// to be used at the designated program ID instead.
+    ///
+    /// See https://github.com/anza-xyz/agave/blob/c038908600b8a1b0080229dea015d7fc9939c418/runtime/src/bank.rs#L5109-L5126.
+    pub fn add_upgradeable_program_to_genesis(
+        &mut self,
+        program_name: &'static str,
+        program_id: &Pubkey,
+    ) {
+        let program_file = find_file(&format!("{program_name}.so"))
+            .expect("Program file data not available for {program_name} ({program_id})");
+        let elf = read_file(program_file);
+        let program_accounts =
+            programs::bpf_loader_upgradeable_program_accounts(program_id, &elf, &Rent::default());
+        for (address, account) in program_accounts {
+            self.add_genesis_account(address, account);
+        }
+    }
+
     /// Add a SBF program to the test environment.
     ///
     /// `program_name` will also be used to locate the SBF shared object in the current or fixtures

program-test/src/programs.rs

@@ -1,6 +1,7 @@
 use solana_sdk::{
     account::{Account, AccountSharedData},
-    bpf_loader_upgradeable::UpgradeableLoaderState,
+    bpf_loader,
+    bpf_loader_upgradeable::{self, get_program_data_address, UpgradeableLoaderState},
     pubkey::Pubkey,
     rent::Rent,
 };
@@ -39,48 +40,87 @@ static SPL_PROGRAMS: &[(Pubkey, Pubkey, &[u8])] = &[
     ),
 ];
 
+/// Returns a tuple `(Pubkey, Account)` for a BPF program, where the key is the
+/// provided program ID and the account is a valid BPF Loader program account
+/// containing the ELF.
+fn bpf_loader_program_account(program_id: &Pubkey, elf: &[u8], rent: &Rent) -> (Pubkey, Account) {
+    (
+        *program_id,
+        Account {
+            lamports: rent.minimum_balance(elf.len()).max(1),
+            data: elf.to_vec(),
+            owner: bpf_loader::id(),
+            executable: true,
+            rent_epoch: 0,
+        },
+    )
+}
+
+/// Returns two tuples `(Pubkey, Account)` for a BPF upgradeable program.
+/// The first tuple is the program account. It contains the provided program ID
+/// and an account with a pointer to its program data account.
+/// The second tuple is the program data account. It contains the program data
+/// address and an account with the program data - a valid BPF Loader Upgradeable
+/// program data account containing the ELF.
+pub(crate) fn bpf_loader_upgradeable_program_accounts(
+    program_id: &Pubkey,
+    elf: &[u8],
+    rent: &Rent,
+) -> [(Pubkey, Account); 2] {
+    let programdata_address = get_program_data_address(program_id);
+    let program_account = {
+        let space = UpgradeableLoaderState::size_of_program();
+        let lamports = rent.minimum_balance(space);
+        let data = bincode::serialize(&UpgradeableLoaderState::Program {
+            programdata_address,
+        })
+        .unwrap();
+        Account {
+            lamports,
+            data,
+            owner: bpf_loader_upgradeable::id(),
+            executable: true,
+            rent_epoch: 0,
+        }
+    };
+    let programdata_account = {
+        let space = UpgradeableLoaderState::size_of_programdata_metadata() + elf.len();
+        let lamports = rent.minimum_balance(space);
+        let mut data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
+            slot: 0,
+            upgrade_authority_address: Some(Pubkey::default()),
+        })
+        .unwrap();
+        data.extend_from_slice(elf);
+        Account {
+            lamports,
+            data,
+            owner: bpf_loader_upgradeable::id(),
+            executable: false,
+            rent_epoch: 0,
+        }
+    };
+    [
+        (*program_id, program_account),
+        (programdata_address, programdata_account),
+    ]
+}
+
 pub fn spl_programs(rent: &Rent) -> Vec<(Pubkey, AccountSharedData)> {
     SPL_PROGRAMS
         .iter()
         .flat_map(|(program_id, loader_id, elf)| {
             let mut accounts = vec![];
-            let data = if *loader_id == solana_sdk::bpf_loader_upgradeable::ID {
-                let (programdata_address, _) =
-                    Pubkey::find_program_address(&[program_id.as_ref()], loader_id);
-                let mut program_data = bincode::serialize(&UpgradeableLoaderState::ProgramData {
-                    slot: 0,
-                    upgrade_authority_address: Some(Pubkey::default()),
-                })
-                .unwrap();
-                program_data.extend_from_slice(elf);
-                accounts.push((
-                    programdata_address,
-                    AccountSharedData::from(Account {
-                        lamports: rent.minimum_balance(program_data.len()).max(1),
-                        data: program_data,
-                        owner: *loader_id,
-                        executable: false,
-                        rent_epoch: 0,
-                    }),
-                ));
-                bincode::serialize(&UpgradeableLoaderState::Program {
-                    programdata_address,
-                })
-                .unwrap()
+            if loader_id.eq(&solana_sdk::bpf_loader_upgradeable::ID) {
+                for (key, account) in bpf_loader_upgradeable_program_accounts(program_id, elf, rent)
+                {
+                    accounts.push((key, AccountSharedData::from(account)));
+                }
             } else {
-                elf.to_vec()
-            };
-            accounts.push((
-                *program_id,
-                AccountSharedData::from(Account {
-                    lamports: rent.minimum_balance(data.len()).max(1),
-                    data,
-                    owner: *loader_id,
-                    executable: true,
-                    rent_epoch: 0,
-                }),
-            ));
-            accounts.into_iter()
+                let (key, account) = bpf_loader_program_account(program_id, elf, rent);
+                accounts.push((key, AccountSharedData::from(account)));
+            }
+            accounts
         })
         .collect()
 }

program-test/tests/bpf.rs

@@ -10,11 +10,9 @@ use {
 async fn test_add_bpf_program() {
     let program_id = Pubkey::new_unique();
 
-    std::env::set_var("SBF_OUT_DIR", "../programs/bpf_loader/test_elfs/out");
-
     let mut program_test = ProgramTest::default();
     program_test.prefer_bpf(true);
-    program_test.add_program("noop_aligned", program_id, None);
+    program_test.add_program("noop_program", program_id, None);
 
     let mut context = program_test.start_with_context().await;
 

program-test/tests/core_bpf.rs

@@ -0,0 +1,41 @@
+use {
+    solana_program_test::ProgramTest,
+    solana_sdk::{
+        bpf_loader_upgradeable, instruction::Instruction, signature::Signer,
+        transaction::Transaction,
+    },
+};
+
+#[tokio::test]
+async fn test_add_bpf_program() {
+    // Core BPF program: Address Lookup Lable.
+    let program_id = solana_sdk::address_lookup_table::program::id();
+
+    let mut program_test = ProgramTest::default();
+    program_test.add_upgradeable_program_to_genesis("noop_program", &program_id);
+
+    let mut context = program_test.start_with_context().await;
+
+    // Assert the program is a BPF Loader Upgradeable program.
+    let program_account = context
+        .banks_client
+        .get_account(program_id)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(program_account.owner, bpf_loader_upgradeable::id());
+
+    // Invoke the program.
+    let instruction = Instruction::new_with_bytes(program_id, &[], Vec::new());
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction],
+        Some(&context.payer.pubkey()),
+        &[&context.payer],
+        context.last_blockhash,
+    );
+    context
+        .banks_client
+        .process_transaction(transaction)
+        .await
+        .unwrap();
+}