SIMD-0266: Efficient Token program (#7125)

* Add p-token feature

* Add feature activation

* Use migration v2

* Add tests

* Simplify program data address calculation

* Fix clippy

* Rename upgrade function

* Update function name on tests

* Fix review comments

* Add upgrade tests

* Move tests

* Fix lint

* Move helper

* Update effective slot

* Typo

* Add delay visibility option

* Address review comments

* Fix module name

* Fix syntax

* Add snapshot test

* Remove delay visibility flag

* Update snapshot test

* Fix snapshot test

* Add program to cache

* Fix program account size check

* Use agave_snapshots

* Fix rebase

feature-set/src/lib.rs

@@ -1187,6 +1187,18 @@ pub mod fix_alt_bn128_pairing_length_check {
     solana_pubkey::declare_id!("bnYzodLwmybj7e1HAe98yZrdJTd7we69eMMLgCXqKZm");
 }
 
+pub mod replace_spl_token_with_p_token {
+    use super::Pubkey;
+
+    solana_pubkey::declare_id!("ptokSWRqZz5u2xdqMdstkMKpFurauUpVen7TZXgDpkQ");
+
+    pub const SPL_TOKEN_PROGRAM_ID: Pubkey =
+        Pubkey::from_str_const("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");
+
+    pub const PTOKEN_PROGRAM_BUFFER: Pubkey =
+        Pubkey::from_str_const("ptokNfvuU7terQ2r2452RzVXB3o4GT33yPWo1fUkkZ2");
+}
+
 pub static FEATURE_NAMES: LazyLock<AHashMap<Pubkey, &'static str>> = LazyLock::new(|| {
     [
         (secp256k1_program_enabled::id(), "secp256k1 program"),
@@ -2134,6 +2146,10 @@ pub static FEATURE_NAMES: LazyLock<AHashMap<Pubkey, &'static str>> = LazyLock::n
             fix_alt_bn128_pairing_length_check::id(),
             "SIMD-0334: Fix alt_bn128_pairing length check",
         ),
+        (
+            replace_spl_token_with_p_token::id(),
+            "SIMD-0266: Efficient Token program",
+        ),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()

runtime/src/bank.rs

@@ -5360,6 +5360,19 @@ impl Bank {
         if new_feature_activations.contains(&feature_set::increase_cpi_account_info_limit::id()) {
             self.apply_simd_0339_invoke_cost_changes();
         }
+
+        if new_feature_activations.contains(&feature_set::replace_spl_token_with_p_token::id()) {
+            if let Err(e) = self.upgrade_loader_v2_program_with_loader_v3_program(
+                &feature_set::replace_spl_token_with_p_token::SPL_TOKEN_PROGRAM_ID,
+                &feature_set::replace_spl_token_with_p_token::PTOKEN_PROGRAM_BUFFER,
+                "replace_spl_token_with_p_token",
+            ) {
+                warn!(
+                    "Failed to replace SPL Token with p-token buffer '{}': {e}",
+                    feature_set::replace_spl_token_with_p_token::PTOKEN_PROGRAM_BUFFER,
+                );
+            }
+        }
     }
 
     fn apply_new_builtin_program_feature_transitions(

runtime/src/bank/builtins/core_bpf_migration/mod.rs

@@ -1,10 +1,11 @@
 pub(crate) mod error;
 mod source_buffer;
+mod target_bpf_v2;
 mod target_builtin;
 mod target_core_bpf;
 
 use {
-    crate::bank::Bank,
+    crate::bank::{builtins::core_bpf_migration::target_bpf_v2::TargetBpfV2, Bank},
     error::CoreBpfMigrationError,
     num_traits::{CheckedAdd, CheckedSub},
     solana_account::{AccountSharedData, ReadableAccount, WritableAccount},
@@ -44,10 +45,10 @@ impl Bank {
     /// account address.
     fn new_target_program_account(
         &self,
-        target: &TargetBuiltin,
+        program_data_address: &Pubkey,
     ) -> Result<AccountSharedData, CoreBpfMigrationError> {
         let state = UpgradeableLoaderState::Program {
-            programdata_address: target.program_data_address,
+            programdata_address: *program_data_address,
         };
         let lamports =
             self.get_minimum_balance_for_rent_exemption(UpgradeableLoaderState::size_of_program());
@@ -188,7 +189,9 @@ impl Bank {
                 program_runtime_environments.program_runtime_v1.clone(),
                 program_id,
                 &bpf_loader_upgradeable::id(),
-                data_len,
+                // The size of the program cache entry is the size of the program account
+                // + size of the program data account.
+                UpgradeableLoaderState::size_of_program().saturating_add(data_len),
                 elf,
                 self.slot,
             )?;
@@ -229,7 +232,8 @@ impl Bank {
         };
 
         // Attempt serialization first before modifying the bank.
-        let new_target_program_account = self.new_target_program_account(&target)?;
+        let new_target_program_account =
+            self.new_target_program_account(&target.program_data_address)?;
         let new_target_program_data_account =
             self.new_target_program_data_account(&source, config.upgrade_authority_address)?;
 
@@ -364,6 +368,92 @@ impl Bank {
         Ok(())
     }
 
+    /// Upgrade a Loader v2 BPF program to a Loader v3 BPF program.
+    ///
+    /// To use this function, add a feature-gated callsite to bank's
+    /// `apply_feature_activations` function, similar to below.
+    ///
+    /// ```ignore
+    /// if new_feature_activations.contains(&agave_feature_set::test_upgrade_program::id()) {
+    ///     self.upgrade_loader_v2_program_with_loader_v3_program(
+    ///        &bpf_loader_v2_program_address,
+    ///        &source_buffer_address,
+    ///        "test_upgrade_loader_v2_program_with_loader_v3_program",
+    ///     );
+    /// }
+    /// ```
+    /// The `source_buffer_address` must point to a Loader v3 buffer account
+    /// (state equal to [`UpgradeableLoaderState::Buffer`]).
+    #[allow(dead_code)] // Only used when an upgrade is configured.
+    pub(crate) fn upgrade_loader_v2_program_with_loader_v3_program(
+        &mut self,
+        loader_v2_bpf_program_address: &Pubkey,
+        source_buffer_address: &Pubkey,
+        datapoint_name: &'static str,
+    ) -> Result<(), CoreBpfMigrationError> {
+        datapoint_info!(datapoint_name, ("slot", self.slot, i64));
+
+        let target = TargetBpfV2::new_checked(self, loader_v2_bpf_program_address)?;
+        let source = SourceBuffer::new_checked(self, source_buffer_address)?;
+
+        // Attempt serialization first before modifying the bank.
+        let new_target_program_account =
+            self.new_target_program_account(&target.program_data_address)?;
+        // Loader v2 programs do not have an upgrade authority, so pass `None` when
+        // creating the new program data account.
+        let new_target_program_data_account =
+            self.new_target_program_data_account(&source, None)?;
+
+        // Gather old and new account data sizes, for updating the bank's
+        // accounts data size delta off-chain.
+        // The old data size is the total size of all original accounts
+        // involved.
+        // The new data size is the total size of all the new program accounts.
+        let old_data_size = checked_add(
+            target.program_account.data().len(),
+            source.buffer_account.data().len(),
+        )?;
+        let new_data_size = checked_add(
+            new_target_program_account.data().len(),
+            new_target_program_data_account.data().len(),
+        )?;
+
+        // Deploy the new Loader v3 program.
+        // This step will validate the program ELF against the current runtime
+        // environment, as well as update the program cache.
+        self.directly_invoke_loader_v3_deploy(
+            &target.program_address,
+            new_target_program_data_account.data(),
+        )?;
+
+        // Calculate the lamports to burn.
+        // The target program account will be replaced, so burn its lamports.
+        // The source buffer account will be cleared, so burn its lamports.
+        // The two new program accounts will need to be funded.
+        let lamports_to_burn = checked_add(
+            target.program_account.lamports(),
+            source.buffer_account.lamports(),
+        )?;
+        let lamports_to_fund = checked_add(
+            new_target_program_account.lamports(),
+            new_target_program_data_account.lamports(),
+        )?;
+        self.update_captalization(lamports_to_burn, lamports_to_fund)?;
+
+        // Store the new program accounts and clear the source buffer account.
+        self.store_account(&target.program_address, &new_target_program_account);
+        self.store_account(
+            &target.program_data_address,
+            &new_target_program_data_account,
+        );
+        self.store_account(&source.buffer_address, &AccountSharedData::default());
+
+        // Update the account data size delta.
+        self.calculate_and_update_accounts_data_size_delta_off_chain(old_data_size, new_data_size);
+
+        Ok(())
+    }
+
     fn update_captalization(
         &mut self,
         lamports_to_burn: u64,
@@ -389,16 +479,23 @@ impl Bank {
 pub(crate) mod tests {
     use {
         super::*,
-        crate::bank::{
-            test_utils::goto_end_of_slot,
-            tests::{create_genesis_config, create_simple_test_bank},
-            Bank,
+        crate::{
+            bank::{
+                test_utils::goto_end_of_slot,
+                tests::{create_genesis_config, create_simple_test_bank},
+                Bank,
+            },
+            runtime_config::RuntimeConfig,
+            snapshot_bank_utils::{bank_from_snapshot_archives, bank_to_full_snapshot_archive},
+            snapshot_utils::create_tmp_accounts_dir_for_tests,
         },
         agave_feature_set::FeatureSet,
+        agave_snapshots::snapshot_config::SnapshotConfig,
         assert_matches::assert_matches,
         solana_account::{
             state_traits::StateMut, AccountSharedData, ReadableAccount, WritableAccount,
         },
+        solana_accounts_db::accounts_db::ACCOUNTS_DB_CONFIG_FOR_TESTING,
         solana_builtins::{
             core_bpf_migration::{CoreBpfMigrationConfig, CoreBpfMigrationTargetType},
             prototype::{BuiltinPrototype, StatelessBuiltinPrototype},
@@ -408,12 +505,13 @@ pub(crate) mod tests {
         solana_epoch_schedule::EpochSchedule,
         solana_feature_gate_interface::{self as feature, Feature},
         solana_instruction::{AccountMeta, Instruction},
+        solana_keypair::Keypair,
         solana_loader_v3_interface::{get_program_data_address, state::UpgradeableLoaderState},
         solana_message::Message,
         solana_native_token::LAMPORTS_PER_SOL,
         solana_program_runtime::loaded_programs::{ProgramCacheEntry, ProgramCacheEntryType},
         solana_pubkey::Pubkey,
-        solana_sdk_ids::{bpf_loader_upgradeable, native_loader},
+        solana_sdk_ids::{bpf_loader, bpf_loader_upgradeable, native_loader},
         solana_signer::Signer,
         solana_transaction::Transaction,
         std::{fs::File, io::Read, sync::Arc},
@@ -619,7 +717,10 @@ pub(crate) mod tests {
                 .unwrap();
 
             // The target program entry should be updated.
-            assert_eq!(target_entry.account_size, program_data_account.data().len());
+            assert_eq!(
+                target_entry.account_size,
+                program_account.data().len() + program_data_account.data().len()
+            );
             assert_eq!(target_entry.deployment_slot, migration_or_upgrade_slot);
             assert_eq!(target_entry.effective_slot, migration_or_upgrade_slot + 1);
 
@@ -1162,53 +1263,17 @@ pub(crate) mod tests {
         }
     }
 
-    // This test can't be used to the `compute_budget` program, unless a valid
-    // `compute_budget` program is provided as the replacement (source).
-    // See program_runtime::compute_budget_processor::process_compute_budget_instructions`.`
-    // It also can't test the `bpf_loader_upgradeable` program, as it's used in
-    // the SVM's loader to invoke programs.
-    // See `solana_svm::account_loader::load_transaction_accounts`.
-    #[test_case(TestPrototype::Builtin(&BUILTINS[0]); "system")]
-    #[test_case(TestPrototype::Builtin(&BUILTINS[1]); "vote")]
-    #[test_case(TestPrototype::Builtin(&BUILTINS[2]); "bpf_loader_deprecated")]
-    #[test_case(TestPrototype::Builtin(&BUILTINS[3]); "bpf_loader")]
-    fn test_migrate_builtin_e2e(prototype: TestPrototype) {
-        let (mut genesis_config, mint_keypair) =
-            create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
-        let slots_per_epoch = 32;
-        genesis_config.epoch_schedule =
-            EpochSchedule::custom(slots_per_epoch, slots_per_epoch, false);
-
-        let mut root_bank = Bank::new_for_tests(&genesis_config);
-
-        // Set up the CPI mockup to test CPI'ing to the migrated program.
-        let cpi_program_id = Pubkey::new_unique();
-        let cpi_program_name = "mock_cpi_program";
-        root_bank.add_builtin(
-            cpi_program_id,
-            cpi_program_name,
-            ProgramCacheEntry::new_builtin(0, cpi_program_name.len(), cpi_mockup::Entrypoint::vm),
-        );
-
-        let (builtin_id, config) = prototype.deconstruct();
-        let feature_id = &config.feature_id;
-        let source_buffer_address = &config.source_buffer_address;
-        let upgrade_authority_address = config.upgrade_authority_address;
-
-        // Add the feature to the bank's inactive feature set.
-        // Note this will add the feature ID if it doesn't exist.
-        let mut feature_set = FeatureSet::all_enabled();
-        feature_set.deactivate(feature_id);
-        root_bank.feature_set = Arc::new(feature_set);
-
-        // Initialize the source buffer account.
-        let test_context = TestContext::new(
-            &root_bank,
-            builtin_id,
-            source_buffer_address,
-            upgrade_authority_address,
-        );
-
+    /// Activate a feature and run checks on the test context.
+    fn activate_feature_and_run_checks(
+        root_bank: Bank,
+        test_context: &TestContext,
+        program_id: &Pubkey,
+        feature_id: &Pubkey,
+        source_buffer_address: &Pubkey,
+        mint_keypair: &Keypair,
+        slots_per_epoch: u64,
+        cpi_program_id: &Pubkey,
+    ) {
         let (bank, bank_forks) = root_bank.wrap_with_bank_forks_for_tests();
 
         // Advance to the next epoch without activating the feature.
@@ -1247,32 +1312,32 @@ pub(crate) mod tests {
         assert!(bank.feature_set.is_active(feature_id));
         test_context.run_program_checks(&bank, migration_slot);
 
-        // Advance one slot so that the new BPF builtin program becomes
+        // Advance one slot so that the new BPF loader v3 program becomes
         // effective in the program cache.
         goto_end_of_slot(bank.clone());
         let next_slot = bank.slot() + 1;
         let bank =
             Bank::new_from_parent_with_bank_forks(&bank_forks, bank, &Pubkey::default(), next_slot);
 
-        // Successfully invoke the new BPF builtin program.
+        // Successfully invoke the new BPF loader v3 program.
         bank.process_transaction(&Transaction::new(
             &vec![&mint_keypair],
             Message::new(
-                &[Instruction::new_with_bytes(*builtin_id, &[], Vec::new())],
+                &[Instruction::new_with_bytes(*program_id, &[], Vec::new())],
                 Some(&mint_keypair.pubkey()),
             ),
             bank.last_blockhash(),
         ))
         .unwrap();
 
-        // Successfully invoke the new BPF builtin program via CPI.
+        // Successfully invoke the new BPF loader v3 program via CPI.
         bank.process_transaction(&Transaction::new(
             &vec![&mint_keypair],
             Message::new(
                 &[Instruction::new_with_bytes(
-                    cpi_program_id,
+                    *cpi_program_id,
                     &[],
-                    vec![AccountMeta::new_readonly(*builtin_id, false)],
+                    vec![AccountMeta::new_readonly(*program_id, false)],
                 )],
                 Some(&mint_keypair.pubkey()),
             ),
@@ -1294,25 +1359,25 @@ pub(crate) mod tests {
         assert!(bank.feature_set.is_active(feature_id));
         test_context.run_program_checks(&bank, migration_slot);
 
-        // Again, successfully invoke the new BPF builtin program.
+        // Again, successfully invoke the new BPF loader v3 program.
         bank.process_transaction(&Transaction::new(
             &vec![&mint_keypair],
             Message::new(
-                &[Instruction::new_with_bytes(*builtin_id, &[], Vec::new())],
+                &[Instruction::new_with_bytes(*program_id, &[], Vec::new())],
                 Some(&mint_keypair.pubkey()),
             ),
             bank.last_blockhash(),
         ))
         .unwrap();
 
-        // Again, successfully invoke the new BPF builtin program via CPI.
+        // Again, successfully invoke the new BPF loader v3 program via CPI.
         bank.process_transaction(&Transaction::new(
             &vec![&mint_keypair],
             Message::new(
                 &[Instruction::new_with_bytes(
-                    cpi_program_id,
+                    *cpi_program_id,
                     &[],
-                    vec![AccountMeta::new_readonly(*builtin_id, false)],
+                    vec![AccountMeta::new_readonly(*program_id, false)],
                 )],
                 Some(&mint_keypair.pubkey()),
             ),
@@ -1321,6 +1386,66 @@ pub(crate) mod tests {
         .unwrap();
     }
 
+    // This test can't be used to the `compute_budget` program, unless a valid
+    // `compute_budget` program is provided as the replacement (source).
+    // See program_runtime::compute_budget_processor::process_compute_budget_instructions`.`
+    // It also can't test the `bpf_loader_upgradeable` program, as it's used in
+    // the SVM's loader to invoke programs.
+    // See `solana_svm::account_loader::load_transaction_accounts`.
+    #[test_case(TestPrototype::Builtin(&BUILTINS[0]); "system")]
+    #[test_case(TestPrototype::Builtin(&BUILTINS[1]); "vote")]
+    #[test_case(TestPrototype::Builtin(&BUILTINS[2]); "bpf_loader_deprecated")]
+    #[test_case(TestPrototype::Builtin(&BUILTINS[3]); "bpf_loader")]
+    fn test_migrate_builtin_e2e(prototype: TestPrototype) {
+        let (mut genesis_config, mint_keypair) =
+            create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+        let slots_per_epoch = 32;
+        genesis_config.epoch_schedule =
+            EpochSchedule::custom(slots_per_epoch, slots_per_epoch, false);
+
+        let mut root_bank = Bank::new_for_tests(&genesis_config);
+
+        // Set up the CPI mockup to test CPI'ing to the migrated program.
+        let cpi_program_id = Pubkey::new_unique();
+        let cpi_program_name = "mock_cpi_program";
+        root_bank.add_builtin(
+            cpi_program_id,
+            cpi_program_name,
+            ProgramCacheEntry::new_builtin(0, cpi_program_name.len(), cpi_mockup::Entrypoint::vm),
+        );
+
+        let (builtin_id, config) = prototype.deconstruct();
+        let feature_id = &config.feature_id;
+        let source_buffer_address = &config.source_buffer_address;
+        let upgrade_authority_address = config.upgrade_authority_address;
+
+        // Add the feature to the bank's inactive feature set.
+        // Note this will add the feature ID if it doesn't exist.
+        let mut feature_set = FeatureSet::all_enabled();
+        feature_set.deactivate(feature_id);
+        root_bank.feature_set = Arc::new(feature_set);
+
+        // Initialize the source buffer account.
+        let test_context = TestContext::new(
+            &root_bank,
+            builtin_id,
+            source_buffer_address,
+            upgrade_authority_address,
+        );
+
+        // Activate the feature and run the migration checks.
+        activate_feature_and_run_checks(
+            root_bank,
+            &test_context,
+            builtin_id,
+            feature_id,
+            source_buffer_address,
+            &mint_keypair,
+            slots_per_epoch,
+            &cpi_program_id,
+        );
+    }
+
     // Simulate a failure to migrate the program.
     // Here we want to see that the bank handles the failure gracefully and
     // advances to the next epoch without issue.
@@ -1626,4 +1751,404 @@ pub(crate) mod tests {
 
         check_builtin_is_bpf(&bank);
     }
+
+    #[test]
+    fn test_upgrade_loader_v2_program_with_loader_v3_program() {
+        let mut bank = create_simple_test_bank(0);
+
+        let bpf_loader_v2_program_address = Pubkey::new_unique();
+        let source_buffer_address = Pubkey::new_unique();
+
+        {
+            let program_account = {
+                let elf = [4u8; 200]; // Mock ELF to start.
+                let space = elf.len();
+                let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+                let owner = &bpf_loader::id();
+
+                let mut account = AccountSharedData::new(lamports, space, owner);
+                account.set_executable(true);
+                account.data_as_mut_slice().copy_from_slice(&elf);
+                bank.store_account_and_update_capitalization(
+                    &bpf_loader_v2_program_address,
+                    &account,
+                );
+                account
+            };
+
+            assert_eq!(
+                &bank.get_account(&bpf_loader_v2_program_address).unwrap(),
+                &program_account
+            );
+        };
+
+        let test_context = TestContext::new(
+            &bank,
+            &bpf_loader_v2_program_address,
+            &source_buffer_address,
+            None,
+        );
+        let TestContext {
+            source_buffer_address,
+            ..
+        } = test_context;
+
+        let (
+            expected_post_upgrade_capitalization,
+            expected_post_upgrade_accounts_data_size_delta_off_chain,
+        ) = test_context
+            .calculate_post_migration_capitalization_and_accounts_data_size_delta_off_chain(&bank);
+
+        // Perform the upgrade.
+        let upgrade_slot = bank.slot();
+        bank.upgrade_loader_v2_program_with_loader_v3_program(
+            &bpf_loader_v2_program_address,
+            &source_buffer_address,
+            "test_upgrade_loader_v2_program_with_loader_v3_program",
+        )
+        .unwrap();
+
+        // Run the post-upgrade program checks.
+        test_context.run_program_checks(&bank, upgrade_slot);
+
+        // Check the bank's capitalization.
+        assert_eq!(bank.capitalization(), expected_post_upgrade_capitalization);
+
+        // Check the bank's accounts data size delta off-chain.
+        assert_eq!(
+            bank.accounts_data_size_delta_off_chain.load(Relaxed),
+            expected_post_upgrade_accounts_data_size_delta_off_chain
+        );
+
+        // Check the migrated program account is now owned by the upgradeable loader.
+        let migrated_program_account = bank.get_account(&bpf_loader_v2_program_address).unwrap();
+        assert_eq!(
+            migrated_program_account.owner(),
+            &bpf_loader_upgradeable::id()
+        );
+    }
+
+    #[test]
+    fn test_upgrade_loader_v2_program_with_loader_v3_program_fail_invalid_buffer() {
+        let mut bank = create_simple_test_bank(0);
+
+        let bpf_loader_v2_program_address = Pubkey::new_unique();
+        let source_buffer_address = Pubkey::new_unique();
+
+        {
+            let program_account = {
+                let elf = [4u8; 200]; // Mock ELF to start.
+                let space = elf.len();
+                let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+                let owner = &bpf_loader::id();
+
+                let mut account = AccountSharedData::new(lamports, space, owner);
+                account.set_executable(true);
+                account.data_as_mut_slice().copy_from_slice(&elf);
+                bank.store_account_and_update_capitalization(
+                    &bpf_loader_v2_program_address,
+                    &account,
+                );
+                account
+            };
+
+            assert_eq!(
+                &bank.get_account(&bpf_loader_v2_program_address).unwrap(),
+                &program_account
+            );
+        };
+
+        // Set up the source buffer with a valid authority, but the migration
+        // config will define the upgrade authority to be `None`.
+        {
+            let elf = test_elf();
+            let buffer_metadata_size = UpgradeableLoaderState::size_of_buffer_metadata();
+            let space = buffer_metadata_size + elf.len();
+            let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+            let owner = &bpf_loader_upgradeable::id();
+
+            let buffer_metadata = UpgradeableLoaderState::Program {
+                programdata_address: Pubkey::new_unique(),
+            };
+
+            let mut account =
+                AccountSharedData::new_data_with_space(lamports, &buffer_metadata, space, owner)
+                    .unwrap();
+            account.data_as_mut_slice()[buffer_metadata_size..].copy_from_slice(&elf);
+
+            bank.store_account_and_update_capitalization(&source_buffer_address, &account);
+        }
+
+        // Try to perform the upgrade.
+        assert_matches!(
+            bank.upgrade_loader_v2_program_with_loader_v3_program(
+                &bpf_loader_v2_program_address,
+                &source_buffer_address,
+                "test_upgrade_loader_v2_program_with_loader_v3_program",
+            )
+            .unwrap_err(),
+            CoreBpfMigrationError::InvalidBufferAccount(_)
+        )
+    }
+
+    /// Mock BPF loader v2 program for testing.
+    fn mock_bpf_loader_v2_program(bank: &Bank) -> AccountSharedData {
+        let elf = [4u8; 200]; // Mock ELF to start.
+        let space = elf.len();
+        let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+        let owner = &bpf_loader::id();
+
+        let mut account = AccountSharedData::new(lamports, space, owner);
+        account.set_executable(true);
+        account.data_as_mut_slice().copy_from_slice(&elf);
+
+        account
+    }
+
+    #[test]
+    fn test_replace_spl_token_with_p_token_e2e() {
+        let (mut genesis_config, mint_keypair) =
+            create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+        let slots_per_epoch = 32;
+        genesis_config.epoch_schedule =
+            EpochSchedule::custom(slots_per_epoch, slots_per_epoch, false);
+
+        let mut root_bank = Bank::new_for_tests(&genesis_config);
+
+        let feature_id = agave_feature_set::replace_spl_token_with_p_token::id();
+        let program_id = agave_feature_set::replace_spl_token_with_p_token::SPL_TOKEN_PROGRAM_ID;
+        let source_buffer_address =
+            agave_feature_set::replace_spl_token_with_p_token::PTOKEN_PROGRAM_BUFFER;
+
+        // Set up a mock BPF loader v2 program.
+        {
+            let program_account = mock_bpf_loader_v2_program(&root_bank);
+            root_bank.store_account_and_update_capitalization(&program_id, &program_account);
+            assert_eq!(
+                &root_bank.get_account(&program_id).unwrap(),
+                &program_account
+            );
+        };
+
+        // Set up the CPI mockup to test CPI'ing to the migrated program.
+        let cpi_program_id = Pubkey::new_unique();
+        let cpi_program_name = "mock_cpi_program";
+        root_bank.add_builtin(
+            cpi_program_id,
+            cpi_program_name,
+            ProgramCacheEntry::new_builtin(0, cpi_program_name.len(), cpi_mockup::Entrypoint::vm),
+        );
+
+        // Add the feature to the bank's inactive feature set.
+        // Note this will add the feature ID if it doesn't exist.
+        let mut feature_set = FeatureSet::all_enabled();
+        feature_set.deactivate(&feature_id);
+        root_bank.feature_set = Arc::new(feature_set);
+
+        // Initialize the source buffer account.
+        let test_context = TestContext::new(&root_bank, &program_id, &source_buffer_address, None);
+
+        // Activate the feature and run the necessary checks.
+        activate_feature_and_run_checks(
+            root_bank,
+            &test_context,
+            &program_id,
+            &feature_id,
+            &source_buffer_address,
+            &mint_keypair,
+            slots_per_epoch,
+            &cpi_program_id,
+        );
+    }
+
+    // Simulate a failure to migrate the program.
+    // Here we want to see that the bank handles the failure gracefully and
+    // advances to the next epoch without issue.
+    #[test]
+    fn test_replace_spl_token_with_p_token_e2e_failure() {
+        let (genesis_config, _mint_keypair) = create_genesis_config(0);
+        let mut root_bank = Bank::new_for_tests(&genesis_config);
+
+        let feature_id = &agave_feature_set::replace_spl_token_with_p_token::id();
+        let program_id = &agave_feature_set::replace_spl_token_with_p_token::SPL_TOKEN_PROGRAM_ID;
+        let source_buffer_address =
+            &agave_feature_set::replace_spl_token_with_p_token::PTOKEN_PROGRAM_BUFFER;
+
+        // Set up a mock BPF loader v2 program.
+        {
+            let program_account = mock_bpf_loader_v2_program(&root_bank);
+            root_bank.store_account_and_update_capitalization(program_id, &program_account);
+            assert_eq!(
+                &root_bank.get_account(program_id).unwrap(),
+                &program_account
+            );
+        };
+
+        // Add the feature to the bank's inactive feature set.
+        let mut feature_set = FeatureSet::all_enabled();
+        feature_set.inactive_mut().insert(*feature_id);
+        root_bank.feature_set = Arc::new(feature_set);
+
+        // Initialize the source buffer account.
+        let _test_context = TestContext::new(&root_bank, program_id, source_buffer_address, None);
+
+        let (bank, bank_forks) = root_bank.wrap_with_bank_forks_for_tests();
+
+        // Intentionally nuke the source buffer account to force the migration
+        // to fail.
+        bank.store_account_and_update_capitalization(
+            source_buffer_address,
+            &AccountSharedData::default(),
+        );
+
+        // Activate the feature.
+        bank.store_account_and_update_capitalization(
+            feature_id,
+            &feature::create_account(&Feature::default(), 42),
+        );
+
+        // Advance the bank to cross the epoch boundary and activate the
+        // feature.
+        goto_end_of_slot(bank.clone());
+        let bank = Bank::new_from_parent_with_bank_forks(&bank_forks, bank, &Pubkey::default(), 33);
+
+        // Assert the feature _was_ activated but the program was not migrated.
+        assert!(bank.feature_set.is_active(feature_id));
+        assert_eq!(
+            bank.get_account(program_id).unwrap().owner(),
+            &bpf_loader::id()
+        );
+
+        // Simulate crossing an epoch boundary again.
+        goto_end_of_slot(bank.clone());
+        let bank = Bank::new_from_parent_with_bank_forks(&bank_forks, bank, &Pubkey::default(), 96);
+
+        // Again, assert the feature is still active and the program still was
+        // not migrated.
+        assert!(bank.feature_set.is_active(feature_id));
+        assert_eq!(
+            bank.get_account(program_id).unwrap().owner(),
+            &bpf_loader::id()
+        );
+    }
+
+    // Simulate creating a bank from a snapshot after p-token migration feature was
+    // activated and the migration was successful.
+    #[test]
+    fn test_startup_from_snapshot_after_replace_spl_token_with_p_token() {
+        let (genesis_config, _mint_keypair) = create_genesis_config(0);
+        let mut bank = Bank::new_for_tests(&genesis_config);
+
+        let bpf_loader_v2_program_address = Pubkey::new_unique();
+        let source_buffer_address = Pubkey::new_unique();
+
+        {
+            let program_account = {
+                let elf = [4u8; 200]; // Mock ELF to start.
+                let space = elf.len();
+                let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+                let owner = &bpf_loader::id();
+
+                let mut account = AccountSharedData::new(lamports, space, owner);
+                account.set_executable(true);
+                account.data_as_mut_slice().copy_from_slice(&elf);
+                bank.store_account_and_update_capitalization(
+                    &bpf_loader_v2_program_address,
+                    &account,
+                );
+                account
+            };
+
+            assert_eq!(
+                &bank.get_account(&bpf_loader_v2_program_address).unwrap(),
+                &program_account
+            );
+        };
+
+        let test_context = TestContext::new(
+            &bank,
+            &bpf_loader_v2_program_address,
+            &source_buffer_address,
+            None,
+        );
+        let TestContext {
+            source_buffer_address,
+            ..
+        } = test_context;
+
+        // Perform the upgrade.
+        let upgrade_slot = bank.slot();
+        bank.upgrade_loader_v2_program_with_loader_v3_program(
+            &bpf_loader_v2_program_address,
+            &source_buffer_address,
+            "test_upgrade_loader_v2_program_with_loader_v3_program",
+        )
+        .unwrap();
+
+        // Run the post-upgrade program checks.
+        test_context.run_program_checks(&bank, upgrade_slot);
+
+        bank.fill_bank_with_ticks_for_tests();
+        // Force flush the bank to create the account storage entry
+        bank.squash();
+        bank.force_flush_accounts_cache();
+
+        // Create a snapshot of the bank.
+        let (_tmp_dir, accounts_dir) = create_tmp_accounts_dir_for_tests();
+        let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
+        let snapshot_archives_dir = tempfile::TempDir::new().unwrap();
+        let snapshot_archive_format = SnapshotConfig::default().archive_format;
+
+        let full_snapshot_archive_info = bank_to_full_snapshot_archive(
+            bank_snapshots_dir.path(),
+            &bank,
+            None,
+            snapshot_archives_dir.path(),
+            snapshot_archives_dir.path(),
+            snapshot_archive_format,
+        )
+        .unwrap();
+
+        // Restore the bank from the snapshot and run checks.
+        let roundtrip_bank = bank_from_snapshot_archives(
+            &[accounts_dir],
+            bank_snapshots_dir.path(),
+            &full_snapshot_archive_info,
+            None,
+            &genesis_config,
+            &RuntimeConfig::default(),
+            None,
+            None,
+            false,
+            false,
+            false,
+            ACCOUNTS_DB_CONFIG_FOR_TESTING,
+            None,
+            Arc::default(),
+        )
+        .unwrap();
+
+        // Load the migrated program to the cache and run checks.
+        let entry = roundtrip_bank
+            .load_program(&bpf_loader_v2_program_address, false, upgrade_slot)
+            .unwrap();
+
+        let mut program_cache = roundtrip_bank
+            .transaction_processor
+            .global_program_cache
+            .write()
+            .unwrap();
+
+        program_cache.assign_program(
+            &roundtrip_bank.transaction_processor.environments,
+            bpf_loader_v2_program_address,
+            entry,
+        );
+        // Release the lock on the program cache.
+        drop(program_cache);
+
+        // Run the post-upgrade program checks on the restored bank.
+        test_context.run_program_checks(&roundtrip_bank, upgrade_slot);
+        assert_eq!(bank, roundtrip_bank);
+    }
 }

runtime/src/bank/builtins/core_bpf_migration/target_bpf_v2.rs

@@ -0,0 +1,128 @@
+#![allow(dead_code)]
+use {
+    super::error::CoreBpfMigrationError,
+    crate::bank::Bank,
+    solana_account::{AccountSharedData, ReadableAccount},
+    solana_loader_v3_interface::get_program_data_address,
+    solana_pubkey::Pubkey,
+    solana_sdk_ids::bpf_loader,
+};
+
+/// The account details of a Loader v2 BPF program slated to be upgraded.
+#[derive(Debug)]
+pub(crate) struct TargetBpfV2 {
+    pub program_address: Pubkey,
+    pub program_account: AccountSharedData,
+    pub program_data_address: Pubkey,
+}
+
+impl TargetBpfV2 {
+    /// Collects the details of a Loader v2 BPF program and verifies it is properly
+    /// configured.
+    ///
+    /// The program account should exist and it should be marked as executable.
+    pub(crate) fn new_checked(
+        bank: &Bank,
+        program_address: &Pubkey,
+    ) -> Result<Self, CoreBpfMigrationError> {
+        // The program account should exist.
+        let program_account = bank
+            .get_account_with_fixed_root(program_address)
+            .ok_or(CoreBpfMigrationError::AccountNotFound(*program_address))?;
+
+        // The program account should be owned by the loader v2.
+        if program_account.owner() != &bpf_loader::id() {
+            return Err(CoreBpfMigrationError::IncorrectOwner(*program_address));
+        }
+
+        // The program account should be executable.
+        if !program_account.executable() {
+            return Err(CoreBpfMigrationError::ProgramAccountNotExecutable(
+                *program_address,
+            ));
+        }
+
+        let program_data_address = get_program_data_address(program_address);
+
+        // The program data account should not exist.
+        if bank
+            .get_account_with_fixed_root(&program_data_address)
+            .is_some()
+        {
+            return Err(CoreBpfMigrationError::ProgramHasDataAccount(
+                *program_address,
+            ));
+        }
+
+        Ok(Self {
+            program_address: *program_address,
+            program_account,
+            program_data_address,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use {
+        super::*, crate::bank::tests::create_simple_test_bank, assert_matches::assert_matches,
+        solana_account::WritableAccount, solana_sdk_ids::bpf_loader,
+    };
+
+    fn store_account(bank: &Bank, address: &Pubkey, data: &[u8], owner: &Pubkey, executable: bool) {
+        let space = data.len();
+        let lamports = bank.get_minimum_balance_for_rent_exemption(space);
+        let mut account = AccountSharedData::new(lamports, space, owner);
+        account.set_executable(executable);
+        account.data_as_mut_slice().copy_from_slice(data);
+        bank.store_account_and_update_capitalization(address, &account);
+    }
+
+    #[test]
+    fn test_target_bpf_v2() {
+        let bank = create_simple_test_bank(0);
+
+        let program_address = Pubkey::new_unique();
+        let elf = vec![4u8; 200];
+
+        // Fail if the program account does not exist.
+        assert_matches!(
+            TargetBpfV2::new_checked(&bank, &program_address).unwrap_err(),
+            CoreBpfMigrationError::AccountNotFound(..)
+        );
+
+        // Fail if the program account is not owned by the loader v2.
+        store_account(
+            &bank,
+            &program_address,
+            &elf,
+            &Pubkey::new_unique(), // Not the loader v2
+            true,
+        );
+        assert_matches!(
+            TargetBpfV2::new_checked(&bank, &program_address).unwrap_err(),
+            CoreBpfMigrationError::IncorrectOwner(..)
+        );
+
+        // Fail if the program account is not executable.
+        store_account(
+            &bank,
+            &program_address,
+            &elf,
+            &bpf_loader::id(),
+            false, // Not executable
+        );
+        assert_matches!(
+            TargetBpfV2::new_checked(&bank, &program_address).unwrap_err(),
+            CoreBpfMigrationError::ProgramAccountNotExecutable(..)
+        );
+
+        // Success
+        store_account(&bank, &program_address, &elf, &bpf_loader::id(), true);
+
+        let target_bpf_v2 = TargetBpfV2::new_checked(&bank, &program_address).unwrap();
+
+        assert_eq!(target_bpf_v2.program_address, program_address);
+        assert_eq!(target_bpf_v2.program_account.data(), elf.as_slice());
+    }
+}