Manage durable nonce stored value in runtime (#7684)

* Bank: Return nonce pubkey/account from `check_tx_durable_nonce`

* Forward account with HashAgeKind::DurableNonce

* Add durable nonce helper for HashAgeKind

* Add nonce util for advancing stored nonce in runtime

* Advance nonce in runtime

* Store rolled back nonce account on TX InstructionError

* nonce: Add test for replayed InstErr fee theft

core/src/transaction_status_service.rs

@@ -59,7 +59,7 @@ impl TransactionStatusService {
             .zip(balances.post_balances)
         {
             if Bank::can_commit(&status) && !transaction.signatures.is_empty() {
-                let fee_hash = if let Some(HashAgeKind::DurableNonce) = hash_age_kind {
+                let fee_hash = if let Some(HashAgeKind::DurableNonce(_, _)) = hash_age_kind {
                     bank.last_blockhash()
                 } else {
                     transaction.message().recent_blockhash

runtime/src/accounts.rs

@@ -4,6 +4,7 @@ use crate::append_vec::StoredAccount;
 use crate::bank::{HashAgeKind, TransactionProcessResult};
 use crate::blockhash_queue::BlockhashQueue;
 use crate::message_processor::has_duplicates;
+use crate::nonce_utils::prepare_if_nonce_account;
 use crate::rent_collector::RentCollector;
 use crate::system_instruction_processor::{get_system_account_kind, SystemAccountKind};
 use log::*;
@@ -12,6 +13,7 @@ use solana_metrics::inc_new_counter_error;
 use solana_sdk::account::Account;
 use solana_sdk::bank_hash::BankHash;
 use solana_sdk::clock::Slot;
+use solana_sdk::hash::Hash;
 use solana_sdk::native_loader;
 use solana_sdk::nonce_state::NonceState;
 use solana_sdk::pubkey::Pubkey;
@@ -247,7 +249,7 @@ impl Accounts {
             .zip(lock_results.into_iter())
             .map(|etx| match etx {
                 (tx, (Ok(()), hash_age_kind)) => {
-                    let fee_hash = if let Some(HashAgeKind::DurableNonce) = hash_age_kind {
+                    let fee_hash = if let Some(HashAgeKind::DurableNonce(_, _)) = hash_age_kind {
                         hash_queue.last_hash()
                     } else {
                         tx.message().recent_blockhash
@@ -559,9 +561,16 @@ impl Accounts {
         res: &[TransactionProcessResult],
         loaded: &mut [(Result<TransactionLoadResult>, Option<HashAgeKind>)],
         rent_collector: &RentCollector,
+        last_blockhash: &Hash,
     ) {
-        let accounts_to_store =
+        let accounts_to_store = self.collect_accounts_to_store(
-            self.collect_accounts_to_store(txs, txs_iteration_order, res, loaded, rent_collector);
+            txs,
+            txs_iteration_order,
+            res,
+            loaded,
+            rent_collector,
+            last_blockhash,
+        );
         self.accounts_db.store(slot, &accounts_to_store);
     }
 
@@ -582,6 +591,7 @@ impl Accounts {
         res: &'a [TransactionProcessResult],
         loaded: &'a mut [(Result<TransactionLoadResult>, Option<HashAgeKind>)],
         rent_collector: &RentCollector,
+        last_blockhash: &Hash,
     ) -> Vec<(&'a Pubkey, &'a Account)> {
         let mut accounts = Vec::with_capacity(loaded.len());
         for (i, ((raccs, _hash_age_kind), tx)) in loaded
@@ -589,10 +599,19 @@ impl Accounts {
             .zip(OrderedIterator::new(txs, txs_iteration_order))
             .enumerate()
         {
-            let (res, _hash_age_kind) = &res[i];
+            if raccs.is_err() {
-            if res.is_err() || raccs.is_err() {
                 continue;
             }
+            let (res, hash_age_kind) = &res[i];
+            let maybe_nonce = match (res, hash_age_kind) {
+                (Ok(_), Some(HashAgeKind::DurableNonce(pubkey, acc))) => Some((pubkey, acc)),
+                (
+                    Err(TransactionError::InstructionError(_, _)),
+                    Some(HashAgeKind::DurableNonce(pubkey, acc)),
+                ) => Some((pubkey, acc)),
+                (Ok(_), _hash_age_kind) => None,
+                (Err(_), _hash_age_kind) => continue,
+            };
 
             let message = &tx.message();
             let acc = raccs.as_mut().unwrap();
@@ -602,6 +621,7 @@ impl Accounts {
                 .enumerate()
                 .zip(acc.0.iter_mut())
             {
+                prepare_if_nonce_account(account, key, res, maybe_nonce, last_blockhash);
                 if message.is_writable(i) {
                     if account.rent_epoch == 0 {
                         account.rent_epoch = rent_collector.epoch;
@@ -1601,8 +1621,14 @@ mod tests {
                 },
             );
         }
-        let collected_accounts =
+        let collected_accounts = accounts.collect_accounts_to_store(
-            accounts.collect_accounts_to_store(&txs, None, &loaders, &mut loaded, &rent_collector);
+            &txs,
+            None,
+            &loaders,
+            &mut loaded,
+            &rent_collector,
+            &Hash::default(),
+        );
         assert_eq!(collected_accounts.len(), 2);
         assert!(collected_accounts
             .iter()

runtime/src/bank.rs

@@ -185,7 +185,16 @@ pub type TransactionBalances = Vec<Vec<u64>>;
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub enum HashAgeKind {
     Extant,
-    DurableNonce,
+    DurableNonce(Pubkey, Account),
+}
+
+impl HashAgeKind {
+    pub fn is_durable_nonce(&self) -> bool {
+        match self {
+            HashAgeKind::DurableNonce(_, _) => true,
+            _ => false,
+        }
+    }
 }
 
 /// Manager for the state of all accounts and programs after processing its entries.
@@ -982,8 +991,8 @@ impl Bank {
                     let message = tx.message();
                     if hash_queue.check_hash_age(&message.recent_blockhash, max_age) {
                         (Ok(()), Some(HashAgeKind::Extant))
-                    } else if self.check_tx_durable_nonce(&tx) {
+                    } else if let Some((pubkey, acc)) = self.check_tx_durable_nonce(&tx) {
-                        (Ok(()), Some(HashAgeKind::DurableNonce))
+                        (Ok(()), Some(HashAgeKind::DurableNonce(pubkey, acc)))
                     } else {
                         error_counters.reserve_blockhash += 1;
                         (Err(TransactionError::BlockhashNotFound), None)
@@ -1037,16 +1046,16 @@ impl Bank {
             .check_hash_age(hash, max_age)
     }
 
-    pub fn check_tx_durable_nonce(&self, tx: &Transaction) -> bool {
+    pub fn check_tx_durable_nonce(&self, tx: &Transaction) -> Option<(Pubkey, Account)> {
         nonce_utils::transaction_uses_durable_nonce(&tx)
             .and_then(|nonce_ix| nonce_utils::get_nonce_pubkey_from_instruction(&nonce_ix, &tx))
-            .and_then(|nonce_pubkey| self.get_account(&nonce_pubkey))
+            .and_then(|nonce_pubkey| {
-            .map_or_else(
+                self.get_account(&nonce_pubkey)
-                || false,
+                    .map(|acc| (*nonce_pubkey, acc))
-                |nonce_account| {
+            })
-                    nonce_utils::verify_nonce(&nonce_account, &tx.message().recent_blockhash)
+            .filter(|(_pubkey, nonce_account)| {
-                },
+                nonce_utils::verify_nonce(nonce_account, &tx.message().recent_blockhash)
-            )
+            })
     }
 
     pub fn check_transactions(
@@ -1236,7 +1245,11 @@ impl Bank {
         let results = OrderedIterator::new(txs, iteration_order)
             .zip(executed.iter())
             .map(|(tx, (res, hash_age_kind))| {
-                let fee_hash = if let Some(HashAgeKind::DurableNonce) = hash_age_kind {
+                let is_durable_nonce = hash_age_kind
+                    .as_ref()
+                    .map(|hash_age_kind| hash_age_kind.is_durable_nonce())
+                    .unwrap_or(false);
+                let fee_hash = if is_durable_nonce {
                     self.last_blockhash()
                 } else {
                     tx.message().recent_blockhash
@@ -1252,7 +1265,12 @@ impl Bank {
                         // credit the transaction fee even in case of InstructionError
                         // necessary to withdraw from account[0] here because previous
                         // work of doing so (in accounts.load()) is ignored by store_account()
-                        self.withdraw(&message.account_keys[0], fee)?;
+                        //
+                        // ...except nonce accounts, which will have their post-load,
+                        // pre-execute account state stored
+                        if !is_durable_nonce {
+                            self.withdraw(&message.account_keys[0], fee)?;
+                        }
                         fees += fee;
                         Ok(())
                     }
@@ -1304,6 +1322,7 @@ impl Bank {
             executed,
             loaded_accounts,
             &self.rent_collector,
+            &self.last_blockhash(),
         );
         self.collect_rent(executed, loaded_accounts);
 
@@ -1931,6 +1950,14 @@ mod tests {
     use std::{io::Cursor, result, time::Duration};
     use tempfile::TempDir;
 
+    #[test]
+    fn test_hash_age_kind_is_durable_nonce() {
+        assert!(
+            HashAgeKind::DurableNonce(Pubkey::default(), Account::default()).is_durable_nonce()
+        );
+        assert!(!HashAgeKind::Extant.is_durable_nonce());
+    }
+
     #[test]
     fn test_bank_unix_timestamp() {
         let (genesis_config, _mint_keypair) = create_genesis_config(1);
@@ -4739,7 +4766,11 @@ mod tests {
             &[&custodian_keypair, &nonce_keypair],
             nonce_hash,
         );
-        assert!(bank.check_tx_durable_nonce(&tx));
+        let nonce_account = bank.get_account(&nonce_pubkey).unwrap();
+        assert_eq!(
+            bank.check_tx_durable_nonce(&tx),
+            Some((nonce_pubkey, nonce_account))
+        );
     }
 
     #[test]
@@ -4759,7 +4790,7 @@ mod tests {
             &[&custodian_keypair, &nonce_keypair],
             nonce_hash,
         );
-        assert!(!bank.check_tx_durable_nonce(&tx));
+        assert!(bank.check_tx_durable_nonce(&tx).is_none());
     }
 
     #[test]
@@ -4780,7 +4811,7 @@ mod tests {
             nonce_hash,
         );
         tx.message.instructions[0].accounts.clear();
-        assert!(!bank.check_tx_durable_nonce(&tx));
+        assert!(bank.check_tx_durable_nonce(&tx).is_none());
     }
 
     #[test]
@@ -4802,7 +4833,7 @@ mod tests {
             &[&custodian_keypair, &nonce_keypair],
             nonce_hash,
         );
-        assert!(!bank.check_tx_durable_nonce(&tx));
+        assert!(bank.check_tx_durable_nonce(&tx).is_none());
     }
 
     #[test]
@@ -4821,7 +4852,7 @@ mod tests {
             &[&custodian_keypair, &nonce_keypair],
             Hash::default(),
         );
-        assert!(!bank.check_tx_durable_nonce(&tx));
+        assert!(bank.check_tx_durable_nonce(&tx).is_none());
     }
 
     #[test]
@@ -4934,10 +4965,11 @@ mod tests {
             bank.process_transaction(&durable_tx),
             Err(TransactionError::BlockhashNotFound)
         );
-        /* Check fee not charged */
+        /* Check fee not charged and nonce not advanced */
         assert_eq!(bank.get_balance(&custodian_pubkey), 4_640_000);
+        assert_eq!(new_nonce, get_nonce(&bank, &nonce_pubkey).unwrap());
 
-        let nonce_hash = get_nonce(&bank, &nonce_pubkey).unwrap();
+        let nonce_hash = new_nonce;
 
         /* Kick nonce hash off the blockhash_queue */
         for _ in 0..MAX_RECENT_BLOCKHASHES + 1 {
@@ -4961,8 +4993,16 @@ mod tests {
                 system_instruction::SystemError::ResultWithNegativeLamports.into()
             ))
         );
-        /* Check fee charged */
+        /* Check fee charged and nonce has advanced */
         assert_eq!(bank.get_balance(&custodian_pubkey), 4_630_000);
+        assert_ne!(nonce_hash, get_nonce(&bank, &nonce_pubkey).unwrap());
+        /* Confirm replaying a TX that failed with InstructionError::* now
+         * fails with TransactionError::BlockhashNotFound
+         */
+        assert_eq!(
+            bank.process_transaction(&durable_tx),
+            Err(TransactionError::BlockhashNotFound),
+        );
     }
 
     #[test]

runtime/src/nonce_utils.rs

@@ -1,7 +1,14 @@
 use solana_sdk::{
-    account::Account, account_utils::State, hash::Hash, instruction::CompiledInstruction,
+    account::Account,
-    instruction_processor_utils::limited_deserialize, nonce_state::NonceState, pubkey::Pubkey,
+    account_utils::State,
-    system_instruction::SystemInstruction, system_program, transaction::Transaction,
+    hash::Hash,
+    instruction::CompiledInstruction,
+    instruction_processor_utils::limited_deserialize,
+    nonce_state::NonceState,
+    pubkey::Pubkey,
+    system_instruction::SystemInstruction,
+    system_program,
+    transaction::{self, Transaction},
 };
 
 pub fn transaction_uses_durable_nonce(tx: &Transaction) -> Option<&CompiledInstruction> {
@@ -39,12 +46,38 @@ pub fn verify_nonce(acc: &Account, hash: &Hash) -> bool {
     }
 }
 
+pub fn prepare_if_nonce_account(
+    account: &mut Account,
+    account_pubkey: &Pubkey,
+    tx_result: &transaction::Result<()>,
+    maybe_nonce: Option<(&Pubkey, &Account)>,
+    last_blockhash: &Hash,
+) {
+    if let Some((nonce_key, nonce_acc)) = maybe_nonce {
+        if account_pubkey == nonce_key {
+            // Nonce TX failed with an InstructionError. Roll back
+            // its account state
+            if tx_result.is_err() {
+                *account = nonce_acc.clone()
+            }
+            // Since hash_age_kind is DurableNonce, unwrap is safe here
+            if let NonceState::Initialized(meta, _) = account.state().unwrap() {
+                account
+                    .set_state(&NonceState::Initialized(meta, *last_blockhash))
+                    .unwrap();
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use solana_sdk::{
+        account::Account,
         hash::Hash,
-        nonce_state::{with_test_keyed_account, NonceAccount},
+        instruction::InstructionError,
+        nonce_state::{with_test_keyed_account, Meta, NonceAccount},
         pubkey::Pubkey,
         signature::{Keypair, KeypairUtil},
         system_instruction,
@@ -193,4 +226,141 @@ mod tests {
             ));
         });
     }
+
+    fn create_accounts_prepare_if_nonce_account() -> (Pubkey, Account, Account, Hash) {
+        let stored_nonce = Hash::default();
+        let account = Account::new_data(
+            42,
+            &NonceState::Initialized(Meta::new(&Pubkey::default()), stored_nonce),
+            &system_program::id(),
+        )
+        .unwrap();
+        let pre_account = Account {
+            lamports: 43,
+            ..account.clone()
+        };
+        (
+            Pubkey::default(),
+            pre_account,
+            account,
+            Hash::new(&[1u8; 32]),
+        )
+    }
+
+    fn run_prepare_if_nonce_account_test(
+        account: &mut Account,
+        account_pubkey: &Pubkey,
+        tx_result: &transaction::Result<()>,
+        maybe_nonce: Option<(&Pubkey, &Account)>,
+        last_blockhash: &Hash,
+        expect_account: &Account,
+    ) -> bool {
+        // Verify expect_account's relationship
+        match maybe_nonce {
+            Some((nonce_pubkey, _nonce_account))
+                if nonce_pubkey == account_pubkey && tx_result.is_ok() =>
+            {
+                assert_ne!(expect_account, account)
+            }
+            Some((nonce_pubkey, nonce_account)) if nonce_pubkey == account_pubkey => {
+                assert_ne!(expect_account, nonce_account)
+            }
+            _ => assert_eq!(expect_account, account),
+        }
+
+        prepare_if_nonce_account(
+            account,
+            account_pubkey,
+            tx_result,
+            maybe_nonce,
+            last_blockhash,
+        );
+        expect_account == account
+    }
+
+    #[test]
+    fn test_prepare_if_nonce_account_expected() {
+        let (pre_account_pubkey, pre_account, mut post_account, last_blockhash) =
+            create_accounts_prepare_if_nonce_account();
+        let post_account_pubkey = pre_account_pubkey;
+
+        let mut expect_account = post_account.clone();
+        expect_account
+            .set_state(&NonceState::Initialized(
+                Meta::new(&Pubkey::default()),
+                last_blockhash,
+            ))
+            .unwrap();
+
+        assert!(run_prepare_if_nonce_account_test(
+            &mut post_account,
+            &post_account_pubkey,
+            &Ok(()),
+            Some((&pre_account_pubkey, &pre_account)),
+            &last_blockhash,
+            &expect_account,
+        ));
+    }
+
+    #[test]
+    fn test_prepare_if_nonce_account_not_nonce_tx() {
+        let (pre_account_pubkey, _pre_account, _post_account, last_blockhash) =
+            create_accounts_prepare_if_nonce_account();
+        let post_account_pubkey = pre_account_pubkey;
+
+        let mut post_account = Account::default();
+        let expect_account = post_account.clone();
+        assert!(run_prepare_if_nonce_account_test(
+            &mut post_account,
+            &post_account_pubkey,
+            &Ok(()),
+            None,
+            &last_blockhash,
+            &expect_account,
+        ));
+    }
+
+    #[test]
+    fn test_prepare_if_nonce_naccount_ot_nonce_pubkey() {
+        let (pre_account_pubkey, pre_account, mut post_account, last_blockhash) =
+            create_accounts_prepare_if_nonce_account();
+
+        let expect_account = post_account.clone();
+        // Wrong key
+        assert!(run_prepare_if_nonce_account_test(
+            &mut post_account,
+            &Pubkey::new(&[1u8; 32]),
+            &Ok(()),
+            Some((&pre_account_pubkey, &pre_account)),
+            &last_blockhash,
+            &expect_account,
+        ));
+    }
+
+    #[test]
+    fn test_prepare_if_nonce_account_tx_error() {
+        let (pre_account_pubkey, pre_account, mut post_account, last_blockhash) =
+            create_accounts_prepare_if_nonce_account();
+        let post_account_pubkey = pre_account_pubkey;
+
+        let mut expect_account = pre_account.clone();
+        expect_account
+            .set_state(&NonceState::Initialized(
+                Meta::new(&Pubkey::default()),
+                last_blockhash,
+            ))
+            .unwrap();
+
+        assert!(run_prepare_if_nonce_account_test(
+            &mut post_account,
+            &post_account_pubkey,
+            &Err(transaction::TransactionError::InstructionError(
+                0,
+                InstructionError::InvalidArgument.into(),
+            )),
+            Some((&pre_account_pubkey, &pre_account)),
+            &last_blockhash,
+            &expect_account,
+        ));
+    }
 }