From 32ec9147bbff10e56223259083cf376c12a3dc4f Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 12 Feb 2021 23:59:08 +0000
Subject: [PATCH] Rework spl_token_v2_self_transfer_fix to avoid any SPL Token
 downtime (#15306)

(cherry picked from commit 2e7aebf0bbcfec19b959ba55797ae5fb01f96322)

Co-authored-by: Michael Vines <mvines@gmail.com>
---
 runtime/src/bank.rs                  | 62 ++++++++++++++++++++++++----
 runtime/src/inline_spl_token_v2_0.rs |  4 ++
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index f41ab6cc8c..ab84ad1062 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4785,11 +4785,29 @@ impl Bank {
     }
 
     fn apply_spl_token_v2_self_transfer_fix(&mut self) {
-        if let Some(mut account) = self.get_account(&inline_spl_token_v2_0::id()) {
-            self.capitalization.fetch_sub(account.lamports, Relaxed);
-            account.lamports = 0;
-            self.store_account(&inline_spl_token_v2_0::id(), &account);
-            self.remove_executor(&inline_spl_token_v2_0::id());
+        if let Some(old_account) = self.get_account(&inline_spl_token_v2_0::id()) {
+            if let Some(new_account) =
+                self.get_account(&inline_spl_token_v2_0::new_token_program::id())
+            {
+                datapoint_info!(
+                    "bank-apply_spl_token_v2_self_transfer_fix",
+                    ("slot", self.slot, i64),
+                );
+
+                // Burn lamports in the old token account
+                self.capitalization.fetch_sub(old_account.lamports, Relaxed);
+
+                // Transfer new token account to old token account
+                self.store_account(&inline_spl_token_v2_0::id(), &new_account);
+
+                // Clear new token account
+                self.store_account(
+                    &inline_spl_token_v2_0::new_token_program::id(),
+                    &Account::default(),
+                );
+
+                self.remove_executor(&inline_spl_token_v2_0::id());
+            }
         }
     }
 
@@ -11080,7 +11098,7 @@ pub(crate) mod tests {
         let (genesis_config, _mint_keypair) = create_genesis_config(0);
         let mut bank = Bank::new(&genesis_config);
 
-        // Setup a simulated account
+        // Setup original token account
         bank.store_account_and_update_capitalization(
             &inline_spl_token_v2_0::id(),
             &Account {
@@ -11089,12 +11107,40 @@ pub(crate) mod tests {
             },
         );
         assert_eq!(bank.get_balance(&inline_spl_token_v2_0::id()), 100);
+
+        // Setup new token account
+        let new_token_account = Account {
+            lamports: 123,
+            data: vec![1, 2, 3],
+            executable: true,
+            ..Account::default()
+        };
+        bank.store_account_and_update_capitalization(
+            &inline_spl_token_v2_0::new_token_program::id(),
+            &new_token_account,
+        );
+        assert_eq!(
+            bank.get_balance(&inline_spl_token_v2_0::new_token_program::id()),
+            123
+        );
+
         let original_capitalization = bank.capitalization();
 
         bank.apply_spl_token_v2_self_transfer_fix();
 
-        // Account is now empty, and the account lamports were burnt
-        assert_eq!(bank.get_balance(&inline_spl_token_v2_0::id()), 0);
+        // New token account is now empty
+        assert_eq!(
+            bank.get_balance(&inline_spl_token_v2_0::new_token_program::id()),
+            0
+        );
+
+        // Old token account holds the new token account
+        assert_eq!(
+            bank.get_account(&inline_spl_token_v2_0::id()),
+            Some(new_token_account)
+        );
+
+        // Lamports in the old token account were burnt
         assert_eq!(bank.capitalization(), original_capitalization - 100);
     }
 
diff --git a/runtime/src/inline_spl_token_v2_0.rs b/runtime/src/inline_spl_token_v2_0.rs
index 27f746f257..0f09bf259d 100644
--- a/runtime/src/inline_spl_token_v2_0.rs
+++ b/runtime/src/inline_spl_token_v2_0.rs
@@ -1,6 +1,10 @@
 // Partial SPL Token v2.0.x declarations inlined to avoid an external dependency on the spl-token crate
 solana_sdk::declare_id!("TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA");
 
+pub(crate) mod new_token_program {
+    solana_sdk::declare_id!("t31zsgDmRntje65uXV3LrnWaJtJJpMd4LyJxq2R2VrU");
+}
+
 /*
     spl_token::state::Account {
         mint: Pubkey,