Add check for existing program data account on migration (#8982)

* Add check for existing program data account

* Move bank mutation

* Update p-token feature id

* Simplify program data account lamports handling

* Improve owner check

(cherry picked from commit 3e5af27241c78e56226546360f680e4081c87bda)

feature-set/src/lib.rs

@@ -1190,7 +1190,7 @@ pub mod fix_alt_bn128_pairing_length_check {
 pub mod replace_spl_token_with_p_token {
     use super::Pubkey;
 
-    solana_pubkey::declare_id!("ptokSWRqZz5u2xdqMdstkMKpFurauUpVen7TZXgDpkQ");
+    solana_pubkey::declare_id!("ptokEXBPT9HuYdAQRysaStZNTY9bHsAcQzNscEoA6HC");
 
     pub const SPL_TOKEN_PROGRAM_ID: Pubkey =
         Pubkey::from_str_const("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");

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

@@ -261,12 +261,15 @@ impl Bank {
 
         // Calculate the lamports to burn.
         // The target program account will be replaced, so burn its lamports.
+        // The target program data account might have lamports if it existed,
+        // so burn its lamports if any.
         // 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(),
-        )?;
+        )
+        .and_then(|v| checked_add(v, target.program_data_account_lamports))?;
         let lamports_to_fund = checked_add(
             new_target_program_account.lamports(),
             new_target_program_data_account.lamports(),
@@ -428,12 +431,15 @@ impl Bank {
 
         // Calculate the lamports to burn.
         // The target program account will be replaced, so burn its lamports.
+        // The target program data account might have lamports if it existed,
+        // so burn its lamports if any.
         // 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(),
-        )?;
+        )
+        .and_then(|v| checked_add(v, target.program_data_account_lamports))?;
         let lamports_to_fund = checked_add(
             new_target_program_account.lamports(),
             new_target_program_data_account.lamports(),
@@ -511,7 +517,7 @@ pub(crate) mod tests {
         solana_native_token::LAMPORTS_PER_SOL,
         solana_program_runtime::loaded_programs::{ProgramCacheEntry, ProgramCacheEntryType},
         solana_pubkey::Pubkey,
-        solana_sdk_ids::{bpf_loader, bpf_loader_upgradeable, native_loader},
+        solana_sdk_ids::{bpf_loader, bpf_loader_upgradeable, native_loader, system_program},
         solana_signer::Signer,
         solana_transaction::Transaction,
         std::{fs::File, io::Read, sync::Arc},
@@ -2151,4 +2157,166 @@ pub(crate) mod tests {
         test_context.run_program_checks(&roundtrip_bank, upgrade_slot);
         assert_eq!(bank, roundtrip_bank);
     }
+
+    #[test]
+    fn test_replace_spl_token_with_p_token_and_funded_program_data_account_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);
+
+        // Fund the program data account so it will appear as an existing account.
+        let program_data_account = AccountSharedData::new(1_000_000_000, 0, &system_program::ID);
+        root_bank.store_account_and_update_capitalization(
+            &get_program_data_address(&program_id),
+            &program_data_account,
+        );
+
+        // 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,
+        );
+    }
+
+    #[test]
+    fn test_replace_spl_token_with_p_token_and_existing_program_data_account_failure() {
+        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);
+
+        // Create the program data account to simulate existing account owned by
+        // the upgradeable loader.
+        let program_data_account =
+            AccountSharedData::new(1_000_000_000, 0, &bpf_loader_upgradeable::ID);
+        root_bank.store_account_and_update_capitalization(
+            &get_program_data_address(&program_id),
+            &program_data_account,
+        );
+
+        // Activate the feature and run the necessary checks.
+        let (bank, bank_forks) = root_bank.wrap_with_bank_forks_for_tests();
+
+        // Advance to the next epoch without activating the feature.
+        let mut first_slot_in_next_epoch = slots_per_epoch + 1;
+        let bank = Bank::new_from_parent_with_bank_forks(
+            &bank_forks,
+            bank,
+            &Pubkey::default(),
+            first_slot_in_next_epoch,
+        );
+
+        // Assert the feature was not activated and the program was not
+        // migrated.
+        assert!(!bank.feature_set.is_active(&feature_id));
+        assert!(bank.get_account(&source_buffer_address).is_some());
+
+        // Store the account to 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());
+        first_slot_in_next_epoch += slots_per_epoch;
+        let _migration_slot = first_slot_in_next_epoch;
+        let bank = Bank::new_from_parent_with_bank_forks(
+            &bank_forks,
+            bank,
+            &Pubkey::default(),
+            first_slot_in_next_epoch,
+        );
+
+        // Check that the feature was activated.
+        assert!(bank.feature_set.is_active(&feature_id));
+        // The program should still be owned by the bpf loader v2.
+        let program_account = bank.get_account(&program_id).unwrap();
+        assert_eq!(program_account.owner(), &bpf_loader::id());
+        // The program data should have zero data and still have
+        // 1_000_000_000 lamports.
+        let program_data_account = bank
+            .get_account(&get_program_data_address(&program_id))
+            .unwrap();
+        assert_eq!(program_data_account.data().len(), 0);
+        assert_eq!(program_data_account.lamports(), 1_000_000_000);
+    }
 }

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

@@ -5,7 +5,7 @@ use {
     solana_account::{AccountSharedData, ReadableAccount},
     solana_loader_v3_interface::get_program_data_address,
     solana_pubkey::Pubkey,
-    solana_sdk_ids::bpf_loader,
+    solana_sdk_ids::{bpf_loader, system_program},
 };
 
 /// The account details of a Loader v2 BPF program slated to be upgraded.
@@ -14,6 +14,7 @@ pub(crate) struct TargetBpfV2 {
     pub program_address: Pubkey,
     pub program_account: AccountSharedData,
     pub program_data_address: Pubkey,
+    pub program_data_account_lamports: u64,
 }
 
 impl TargetBpfV2 {
@@ -44,20 +45,26 @@ impl TargetBpfV2 {
 
         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,
-            ));
-        }
+        // The program data account is expected not to exist.
+        let program_data_account_lamports =
+            if let Some(account) = bank.get_account_with_fixed_root(&program_data_address) {
+                // The program data account should not exist, but a system account with funded
+                // lamports is acceptable.
+                if account.owner() != &system_program::id() {
+                    return Err(CoreBpfMigrationError::ProgramHasDataAccount(
+                        *program_address,
+                    ));
+                }
+                account.lamports()
+            } else {
+                0
+            };
 
         Ok(Self {
             program_address: *program_address,
             program_account,
             program_data_address,
+            program_data_account_lamports,
         })
     }
 }

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

@@ -5,7 +5,9 @@ use {
     solana_builtins::core_bpf_migration::CoreBpfMigrationTargetType,
     solana_loader_v3_interface::get_program_data_address,
     solana_pubkey::Pubkey,
-    solana_sdk_ids::native_loader::ID as NATIVE_LOADER_ID,
+    solana_sdk_ids::{
+        native_loader::ID as NATIVE_LOADER_ID, system_program::ID as SYSTEM_PROGRAM_ID,
+    },
 };
 
 /// The account details of a built-in program to be migrated to Core BPF.
@@ -14,6 +16,7 @@ pub(crate) struct TargetBuiltin {
     pub program_address: Pubkey,
     pub program_account: AccountSharedData,
     pub program_data_address: Pubkey,
+    pub program_data_account_lamports: u64,
 }
 
 impl TargetBuiltin {
@@ -50,20 +53,26 @@ impl TargetBuiltin {
 
         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,
-            ));
-        }
+        // The program data account is expected not to exist.
+        let program_data_account_lamports =
+            if let Some(account) = bank.get_account_with_fixed_root(&program_data_address) {
+                // The program data account should not exist, but a system account with funded
+                // lamports is acceptable.
+                if account.owner() != &SYSTEM_PROGRAM_ID {
+                    return Err(CoreBpfMigrationError::ProgramHasDataAccount(
+                        *program_address,
+                    ));
+                }
+                account.lamports()
+            } else {
+                0
+            };
 
         Ok(Self {
             program_address: *program_address,
             program_account,
             program_data_address,
+            program_data_account_lamports,
         })
     }
 }