Manage durable nonce stored value in runtime (#7684) (#7760)

automerge

runtime/src/nonce_utils.rs

@@ -1,7 +1,14 @@
 use solana_sdk::{
-    account::Account, account_utils::State, hash::Hash, instruction::CompiledInstruction,
-    instruction_processor_utils::limited_deserialize, nonce_state::NonceState, pubkey::Pubkey,
-    system_instruction::SystemInstruction, system_program, transaction::Transaction,
+    account::Account,
+    account_utils::State,
+    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,
+        ));
+    }
 }