spl-associated-token-account: Add feature for new program (backport #22648) (#22718)

* spl-associated-token-account: Add feature for new program (#22648)

* spl-associated-token-account: Add feature for new program

* Address feedback

(cherry picked from commit fc21af4e6e)

# Conflicts:
#	runtime/src/bank.rs
#	runtime/src/lib.rs
#	sdk/src/feature_set.rs

* Fix merge conflicts

Co-authored-by: Jon Cinque <jon.cinque@gmail.com>

runtime/src/bank.rs

@@ -48,7 +48,7 @@ use {
         cost_tracker::CostTracker,
         epoch_stakes::{EpochStakes, NodeVoteAccounts},
         hashed_transaction::{HashedTransaction, HashedTransactionSlice},
-        inline_spl_token,
+        inline_spl_associated_token_account, inline_spl_token,
         instruction_recorder::InstructionRecorder,
         log_collector::LogCollector,
         message_processor::{Executors, MessageProcessor, TransactionExecutor},
@@ -5572,7 +5572,20 @@ impl Bank {
         }
 
         if new_feature_activations.contains(&feature_set::spl_token_v3_3_0_release::id()) {
-            self.apply_spl_token_v3_3_0_release();
+            self.replace_program_account(
+                &inline_spl_token::id(),
+                &inline_spl_token::new_token_program::id(),
+                "bank-apply_spl_token_v3_3_0_release",
+            );
+        }
+
+        if new_feature_activations.contains(&feature_set::spl_associated_token_account_v1_0_4::id())
+        {
+            self.replace_program_account(
+                &inline_spl_associated_token_account::id(),
+                &inline_spl_associated_token_account::program_v1_0_4::id(),
+                "bank-apply_spl_associated_token_account_v1_4_0",
+            );
         }
         if new_feature_activations.contains(&feature_set::rent_for_sysvars::id()) {
             // when this feature is activated, immediately all of existing sysvars are susceptible
@@ -5695,33 +5708,31 @@ impl Bank {
         }
     }
 
-    fn apply_spl_token_v3_3_0_release(&mut self) {
+    fn replace_program_account(
-        if let Some(old_account) = self.get_account_with_fixed_root(&inline_spl_token::id()) {
+        &mut self,
-            if let Some(new_account) =
+        old_address: &Pubkey,
-                self.get_account_with_fixed_root(&inline_spl_token::new_token_program::id())
+        new_address: &Pubkey,
-            {
+        datapoint_name: &'static str,
-                datapoint_info!(
+    ) {
-                    "bank-apply_spl_token_v3_3_0_release",
+        if let Some(old_account) = self.get_account_with_fixed_root(old_address) {
-                    ("slot", self.slot, i64),
+            if let Some(new_account) = self.get_account_with_fixed_root(new_address) {
-                );
+                datapoint_info!(datapoint_name, ("slot", self.slot, i64));
 
-                // Burn lamports in the old token account
+                // Burn lamports in the old account
                 self.capitalization
                     .fetch_sub(old_account.lamports(), Relaxed);
 
-                // Transfer new token account to old token account
+                // Transfer new account to old account
-                self.store_account(&inline_spl_token::id(), &new_account);
+                self.store_account(old_address, &new_account);
 
-                // Clear new token account
+                // Clear new account
-                self.store_account(
+                self.store_account(new_address, &AccountSharedData::default());
-                    &inline_spl_token::new_token_program::id(),
-                    &AccountSharedData::default(),
-                );
 
-                self.remove_executor(&inline_spl_token::id());
+                self.remove_executor(old_address);
             }
         }
     }
+
     fn reconfigure_token2_native_mint(&mut self) {
         let reconfigure_token2_native_mint = match self.cluster_type() {
             ClusterType::Development => true,
@@ -12272,49 +12283,39 @@ pub(crate) mod tests {
     }
 
     #[test]
-    fn test_spl_token_replacement() {
+    fn test_program_replacement() {
         let (genesis_config, _mint_keypair) = create_genesis_config(0);
         let mut bank = Bank::new(&genesis_config);
 
-        // Setup original token account
+        // Setup original program account
+        let old_address = Pubkey::new_unique();
+        let new_address = Pubkey::new_unique();
         bank.store_account_and_update_capitalization(
-            &inline_spl_token::id(),
+            &old_address,
             &AccountSharedData::from(Account {
                 lamports: 100,
                 ..Account::default()
             }),
         );
-        assert_eq!(bank.get_balance(&inline_spl_token::id()), 100);
+        assert_eq!(bank.get_balance(&old_address), 100);
 
-        // Setup new token account
+        // Setup new program account
-        let new_token_account = AccountSharedData::from(Account {
+        let new_program_account = AccountSharedData::from(Account {
             lamports: 123,
             ..Account::default()
         });
-        bank.store_account_and_update_capitalization(
+        bank.store_account_and_update_capitalization(&new_address, &new_program_account);
-            &inline_spl_token::new_token_program::id(),
+        assert_eq!(bank.get_balance(&new_address), 123);
-            &new_token_account,
-        );
-        assert_eq!(
-            bank.get_balance(&inline_spl_token::new_token_program::id()),
-            123
-        );
 
         let original_capitalization = bank.capitalization();
 
-        bank.apply_spl_token_v3_3_0_release();
+        bank.replace_program_account(&old_address, &new_address, "bank-apply_program_replacement");
 
-        // New token account is now empty
+        // New program account is now empty
-        assert_eq!(
+        assert_eq!(bank.get_balance(&new_address), 0);
-            bank.get_balance(&inline_spl_token::new_token_program::id()),
-            0
-        );
 
-        // Old token account holds the new token account
+        // Old program account holds the new program account
-        assert_eq!(
+        assert_eq!(bank.get_account(&old_address), Some(new_program_account));
-            bank.get_account(&inline_spl_token::id()),
-            Some(new_token_account)
-        );
 
         // Lamports in the old token account were burnt
         assert_eq!(bank.capitalization(), original_capitalization - 100);

runtime/src/inline_spl_associated_token_account.rs

@@ -0,0 +1,6 @@
+// Partial SPL Associated Token Account declarations inlined to avoid an external dependency on the spl-associated-token-account crate
+solana_sdk::declare_id!("ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL");
+
+pub(crate) mod program_v1_0_4 {
+    solana_sdk::declare_id!("nata4apBRD9S9256v3X8RxDQ7jmK7wLEsGNHhRNWFq3");
+}

runtime/src/lib.rs

@@ -25,6 +25,7 @@ pub mod execute_cost_table;
 pub mod genesis_utils;
 pub mod hardened_unpack;
 pub mod hashed_transaction;
+pub mod inline_spl_associated_token_account;
 pub mod inline_spl_token;
 pub mod instruction_recorder;
 pub mod loader_utils;

sdk/src/feature_set.rs

@@ -285,6 +285,10 @@ pub mod evict_invalid_stakes_cache_entries {
     solana_sdk::declare_id!("EMX9Q7TVFAmQ9V1CggAkhMzhXSg8ECp7fHrWQX2G1chf");
 }
 
+pub mod spl_associated_token_account_v1_0_4 {
+    solana_sdk::declare_id!("FaTa4SpiaSNH44PGC4z8bnGVTkSRYaWvrBs3KTu8XQQq");
+}
+
 lazy_static! {
     /// Map of feature identifiers to user-visible description
     pub static ref FEATURE_NAMES: HashMap<Pubkey, &'static str> = [
@@ -356,6 +360,7 @@ lazy_static! {
         (spl_token_v3_3_0_release::id(), "spl-token v3.3.0 release"),
         (reject_non_rent_exempt_vote_withdraws::id(), "fail vote withdraw instructions which leave the account non-rent-exempt"),
         (evict_invalid_stakes_cache_entries::id(), "evict invalid stakes cache entries on epoch boundaries"),
+        (spl_associated_token_account_v1_0_4::id(), "SPL Associated Token Account Program release version 1.0.4, tied to token 3.3.0 #22648"),
         /*************** ADD NEW FEATURES HERE ***************/
     ]
     .iter()