From 24f02c1404fe40b62283a8a4f1d2ce79a09ccd6e Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Fri, 17 May 2024 23:12:53 +0000 Subject: [PATCH 1/3] Refactor nonce rollback to remove overly complex code --- accounts-db/src/accounts.rs | 113 ++++++++++++++++++------------------ 1 file changed, 56 insertions(+), 57 deletions(-) diff --git a/accounts-db/src/accounts.rs b/accounts-db/src/accounts.rs index e0ed2639bf53b6..31fbd285f21214 100644 --- a/accounts-db/src/accounts.rs +++ b/accounts-db/src/accounts.rs @@ -719,11 +719,14 @@ impl Accounts { TransactionExecutionResult::NotExecuted(_) => continue, }; - let maybe_nonce = match (execution_status, &loaded_transaction.nonce) { - (Ok(_), _) => None, // Success, don't do any additional nonce processing - (Err(_), Some(nonce)) => { - Some((nonce, true /* rollback */)) - } + enum TransactionStatus<'a> { + Successful, + FailedWithNonce { nonce: &'a NonceFull }, + } + + let tx_status = match (execution_status, &loaded_transaction.nonce) { + (Ok(_), _) => TransactionStatus::Successful, // Success, don't do any additional nonce processing + (Err(_), Some(nonce)) => TransactionStatus::FailedWithNonce { nonce }, (Err(_), None) => { // Fees for failed transactions which don't use durable nonces are // deducted in Bank::filter_program_errors_and_collect_fee @@ -753,18 +756,24 @@ impl Accounts { .filter(|(i, _)| is_storable_account(message, *i)) { if message.is_writable(i) { - let is_fee_payer = i == 0; - let is_nonce_account = prepare_if_nonce_account( - address, - account, - execution_status, - is_fee_payer, - maybe_nonce, - durable_nonce, - lamports_per_signature, - ); - - if execution_status.is_ok() || is_nonce_account || is_fee_payer { + let should_store_account = match tx_status { + TransactionStatus::Successful => true, + TransactionStatus::FailedWithNonce { nonce } => { + let is_fee_payer = i == 0; + let is_nonce_account = prepare_if_nonce_account( + address, + account, + is_fee_payer, + nonce, + durable_nonce, + lamports_per_signature, + ); + + is_fee_payer || is_nonce_account + } + }; + + if should_store_account { // Add to the accounts to store accounts.push((&*address, &*account)); transactions.push(Some(tx)); @@ -779,52 +788,42 @@ impl Accounts { fn prepare_if_nonce_account( address: &Pubkey, account: &mut AccountSharedData, - execution_result: &Result<()>, is_fee_payer: bool, - maybe_nonce: Option<(&NonceFull, bool)>, + nonce: &NonceFull, &durable_nonce: &DurableNonce, lamports_per_signature: u64, ) -> bool { - if let Some((nonce, rollback)) = maybe_nonce { - if address == nonce.address() { - if rollback { - // The transaction failed which would normally drop the account - // processing changes, since this account is now being included - // in the accounts written back to the db, roll it back to - // pre-processing state. - *account = nonce.account().clone(); - } - - // Advance the stored blockhash to prevent fee theft by someone - // replaying nonce transactions that have failed with an - // `InstructionError`. - // - // Since we know we are dealing with a valid nonce account, - // unwrap is safe here - let nonce_versions = StateMut::::state(nonce.account()).unwrap(); - if let NonceState::Initialized(ref data) = nonce_versions.state() { - let nonce_state = NonceState::new_initialized( - &data.authority, - durable_nonce, - lamports_per_signature, - ); - let nonce_versions = NonceVersions::new(nonce_state); - account.set_state(&nonce_versions).unwrap(); - } - true - } else { - if execution_result.is_err() && is_fee_payer { - if let Some(fee_payer_account) = nonce.fee_payer_account() { - // Instruction error and fee-payer for this nonce tx is not - // the nonce account itself, rollback the fee payer to the - // fee-paid original state. - *account = fee_payer_account.clone(); - } - } - - false + if address == nonce.address() { + // The transaction failed which would normally drop the account + // processing changes, since this account is now being included + // in the accounts written back to the db, roll it back to + // pre-processing state. + *account = nonce.account().clone(); + + // Advance the stored blockhash to prevent fee theft by someone + // replaying nonce transactions that have failed with an + // `InstructionError`. + // + // Since we know we are dealing with a valid nonce account, + // unwrap is safe here + let nonce_versions = StateMut::::state(nonce.account()).unwrap(); + if let NonceState::Initialized(ref data) = nonce_versions.state() { + let nonce_state = + NonceState::new_initialized(&data.authority, durable_nonce, lamports_per_signature); + let nonce_versions = NonceVersions::new(nonce_state); + account.set_state(&nonce_versions).unwrap(); } + true } else { + if is_fee_payer { + if let Some(fee_payer_account) = nonce.fee_payer_account() { + // Instruction error and fee-payer for this nonce tx is not + // the nonce account itself, rollback the fee payer to the + // fee-paid original state. + *account = fee_payer_account.clone(); + } + } + false } } From 54a710499b668555f5f21a32b2521a3cc3bc91b9 Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Sat, 18 May 2024 14:32:45 +0000 Subject: [PATCH 2/3] fix tests --- accounts-db/src/accounts.rs | 122 ++++-------------------------------- 1 file changed, 11 insertions(+), 111 deletions(-) diff --git a/accounts-db/src/accounts.rs b/accounts-db/src/accounts.rs index 31fbd285f21214..a12b4d6efb7ad8 100644 --- a/accounts-db/src/accounts.rs +++ b/accounts-db/src/accounts.rs @@ -1684,29 +1684,26 @@ mod tests { fn run_prepare_if_nonce_account_test( account_address: &Pubkey, account: &mut AccountSharedData, - tx_result: &Result<()>, is_fee_payer: bool, - maybe_nonce: Option<(&NonceFull, bool)>, + nonce: &NonceFull, durable_nonce: &DurableNonce, lamports_per_signature: u64, expect_account: &AccountSharedData, ) -> bool { // Verify expect_account's relationship if !is_fee_payer { - match maybe_nonce { - Some((nonce, _)) if nonce.address() == account_address => { - assert_ne!(expect_account, nonce.account()) - } - _ => assert_eq!(expect_account, account), + if nonce.address() == account_address { + assert_ne!(expect_account, nonce.account()); + } else { + assert_eq!(expect_account, account); } } prepare_if_nonce_account( account_address, account, - tx_result, is_fee_payer, - maybe_nonce, + nonce, durable_nonce, lamports_per_signature, ); @@ -1741,35 +1738,8 @@ mod tests { assert!(run_prepare_if_nonce_account_test( &post_account_address, &mut post_account, - &Ok(()), false, - Some((&nonce, true)), - &blockhash, - lamports_per_signature, - &expect_account, - )); - } - - #[test] - fn test_prepare_if_nonce_account_not_nonce_tx() { - let ( - pre_account_address, - _pre_account, - _post_account, - blockhash, - lamports_per_signature, - _maybe_fee_payer_account, - ) = create_accounts_prepare_if_nonce_account(); - let post_account_address = pre_account_address; - - let mut post_account = AccountSharedData::default(); - let expect_account = post_account.clone(); - assert!(run_prepare_if_nonce_account_test( - &post_account_address, - &mut post_account, - &Ok(()), - false, - None, + &nonce, &blockhash, lamports_per_signature, &expect_account, @@ -1794,45 +1764,8 @@ mod tests { assert!(run_prepare_if_nonce_account_test( &Pubkey::from([1u8; 32]), &mut post_account, - &Ok(()), false, - Some((&nonce, true)), - &blockhash, - lamports_per_signature, - &expect_account, - )); - } - - #[test] - fn test_prepare_if_nonce_account_tx_error() { - let ( - pre_account_address, - pre_account, - mut post_account, - blockhash, - lamports_per_signature, - maybe_fee_payer_account, - ) = create_accounts_prepare_if_nonce_account(); - let post_account_address = pre_account_address; - let mut expect_account = pre_account.clone(); - - let nonce = NonceFull::new(pre_account_address, pre_account, maybe_fee_payer_account); - - expect_account - .set_state(&NonceVersions::new(NonceState::Initialized( - nonce::state::Data::new(Pubkey::default(), blockhash, lamports_per_signature), - ))) - .unwrap(); - - assert!(run_prepare_if_nonce_account_test( - &post_account_address, - &mut post_account, - &Err(TransactionError::InstructionError( - 0, - InstructionError::InvalidArgument, - )), - false, - Some((&nonce, true)), + &nonce, &blockhash, lamports_per_signature, &expect_account, @@ -1844,7 +1777,7 @@ mod tests { let nonce_account = AccountSharedData::new_data(1, &(), &system_program::id()).unwrap(); let pre_fee_payer_account = AccountSharedData::new_data(42, &(), &system_program::id()).unwrap(); - let mut post_fee_payer_account = + let post_fee_payer_account = AccountSharedData::new_data(84, &[1, 2, 3, 4], &system_program::id()).unwrap(); let nonce = NonceFull::new( Pubkey::new_unique(), @@ -1855,23 +1788,8 @@ mod tests { assert!(run_prepare_if_nonce_account_test( &Pubkey::new_unique(), &mut post_fee_payer_account.clone(), - &Err(TransactionError::InstructionError( - 0, - InstructionError::InvalidArgument, - )), false, - Some((&nonce, true)), - &DurableNonce::default(), - 1, - &post_fee_payer_account, - )); - - assert!(run_prepare_if_nonce_account_test( - &Pubkey::new_unique(), - &mut post_fee_payer_account.clone(), - &Ok(()), - true, - Some((&nonce, true)), + &nonce, &DurableNonce::default(), 1, &post_fee_payer_account, @@ -1880,26 +1798,8 @@ mod tests { assert!(run_prepare_if_nonce_account_test( &Pubkey::new_unique(), &mut post_fee_payer_account.clone(), - &Err(TransactionError::InstructionError( - 0, - InstructionError::InvalidArgument, - )), - true, - None, - &DurableNonce::default(), - 1, - &post_fee_payer_account, - )); - - assert!(run_prepare_if_nonce_account_test( - &Pubkey::new_unique(), - &mut post_fee_payer_account, - &Err(TransactionError::InstructionError( - 0, - InstructionError::InvalidArgument, - )), true, - Some((&nonce, true)), + &nonce, &DurableNonce::default(), 1, &pre_fee_payer_account, From 67852938dbdadca21f0974c12472d40470ac9a1a Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Mon, 3 Jun 2024 19:28:47 +0000 Subject: [PATCH 3/3] feedback --- accounts-db/src/accounts.rs | 96 +++++++++++++++++-------------------- 1 file changed, 45 insertions(+), 51 deletions(-) diff --git a/accounts-db/src/accounts.rs b/accounts-db/src/accounts.rs index a12b4d6efb7ad8..1cb88c817862b8 100644 --- a/accounts-db/src/accounts.rs +++ b/accounts-db/src/accounts.rs @@ -719,14 +719,14 @@ impl Accounts { TransactionExecutionResult::NotExecuted(_) => continue, }; - enum TransactionStatus<'a> { - Successful, + enum AccountCollectionMode<'a> { + Normal, FailedWithNonce { nonce: &'a NonceFull }, } - let tx_status = match (execution_status, &loaded_transaction.nonce) { - (Ok(_), _) => TransactionStatus::Successful, // Success, don't do any additional nonce processing - (Err(_), Some(nonce)) => TransactionStatus::FailedWithNonce { nonce }, + let collection_mode = match (execution_status, &loaded_transaction.nonce) { + (Ok(_), _) => AccountCollectionMode::Normal, + (Err(_), Some(nonce)) => AccountCollectionMode::FailedWithNonce { nonce }, (Err(_), None) => { // Fees for failed transactions which don't use durable nonces are // deducted in Bank::filter_program_errors_and_collect_fee @@ -756,14 +756,15 @@ impl Accounts { .filter(|(i, _)| is_storable_account(message, *i)) { if message.is_writable(i) { - let should_store_account = match tx_status { - TransactionStatus::Successful => true, - TransactionStatus::FailedWithNonce { nonce } => { + let should_collect_account = match collection_mode { + AccountCollectionMode::Normal => true, + AccountCollectionMode::FailedWithNonce { nonce } => { let is_fee_payer = i == 0; - let is_nonce_account = prepare_if_nonce_account( - address, + let is_nonce_account = address == nonce.address(); + post_process_failed_nonce( account, is_fee_payer, + is_nonce_account, nonce, durable_nonce, lamports_per_signature, @@ -773,7 +774,7 @@ impl Accounts { } }; - if should_store_account { + if should_collect_account { // Add to the accounts to store accounts.push((&*address, &*account)); transactions.push(Some(tx)); @@ -785,15 +786,15 @@ impl Accounts { } } -fn prepare_if_nonce_account( - address: &Pubkey, +fn post_process_failed_nonce( account: &mut AccountSharedData, is_fee_payer: bool, + is_nonce_account: bool, nonce: &NonceFull, &durable_nonce: &DurableNonce, lamports_per_signature: u64, -) -> bool { - if address == nonce.address() { +) { + if is_nonce_account { // The transaction failed which would normally drop the account // processing changes, since this account is now being included // in the accounts written back to the db, roll it back to @@ -806,25 +807,20 @@ fn prepare_if_nonce_account( // // Since we know we are dealing with a valid nonce account, // unwrap is safe here - let nonce_versions = StateMut::::state(nonce.account()).unwrap(); + let nonce_versions = StateMut::::state(account).unwrap(); if let NonceState::Initialized(ref data) = nonce_versions.state() { let nonce_state = NonceState::new_initialized(&data.authority, durable_nonce, lamports_per_signature); let nonce_versions = NonceVersions::new(nonce_state); account.set_state(&nonce_versions).unwrap(); } - true - } else { - if is_fee_payer { - if let Some(fee_payer_account) = nonce.fee_payer_account() { - // Instruction error and fee-payer for this nonce tx is not - // the nonce account itself, rollback the fee payer to the - // fee-paid original state. - *account = fee_payer_account.clone(); - } + } else if is_fee_payer { + if let Some(fee_payer_account) = nonce.fee_payer_account() { + // Instruction error and fee-payer for this nonce tx is not + // the nonce account itself, rollback the fee payer to the + // fee-paid original state. + *account = fee_payer_account.clone(); } - - false } } @@ -1658,7 +1654,7 @@ mod tests { accounts.accounts_db.clean_accounts_for_tests(); } - fn create_accounts_prepare_if_nonce_account() -> ( + fn create_accounts_post_process_failed_nonce() -> ( Pubkey, AccountSharedData, AccountSharedData, @@ -1681,10 +1677,10 @@ mod tests { ) } - fn run_prepare_if_nonce_account_test( - account_address: &Pubkey, + fn run_post_process_failed_nonce_test( account: &mut AccountSharedData, is_fee_payer: bool, + is_nonce_account: bool, nonce: &NonceFull, durable_nonce: &DurableNonce, lamports_per_signature: u64, @@ -1692,17 +1688,17 @@ mod tests { ) -> bool { // Verify expect_account's relationship if !is_fee_payer { - if nonce.address() == account_address { + if is_nonce_account { assert_ne!(expect_account, nonce.account()); } else { assert_eq!(expect_account, account); } } - prepare_if_nonce_account( - account_address, + post_process_failed_nonce( account, is_fee_payer, + is_nonce_account, nonce, durable_nonce, lamports_per_signature, @@ -1712,7 +1708,7 @@ mod tests { } #[test] - fn test_prepare_if_nonce_account_expected() { + fn test_post_process_failed_nonce_expected() { let ( pre_account_address, pre_account, @@ -1720,8 +1716,7 @@ mod tests { blockhash, lamports_per_signature, maybe_fee_payer_account, - ) = create_accounts_prepare_if_nonce_account(); - let post_account_address = pre_account_address; + ) = create_accounts_post_process_failed_nonce(); let nonce = NonceFull::new( pre_account_address, pre_account.clone(), @@ -1735,10 +1730,10 @@ mod tests { ))) .unwrap(); - assert!(run_prepare_if_nonce_account_test( - &post_account_address, + assert!(run_post_process_failed_nonce_test( &mut post_account, - false, + false, // is_fee_payer + true, // is_nonce_account &nonce, &blockhash, lamports_per_signature, @@ -1747,7 +1742,7 @@ mod tests { } #[test] - fn test_prepare_if_nonce_account_not_nonce_address() { + fn test_post_process_failed_nonce_not_nonce_address() { let ( pre_account_address, pre_account, @@ -1755,16 +1750,15 @@ mod tests { blockhash, lamports_per_signature, maybe_fee_payer_account, - ) = create_accounts_prepare_if_nonce_account(); + ) = create_accounts_post_process_failed_nonce(); let nonce = NonceFull::new(pre_account_address, pre_account, maybe_fee_payer_account); let expect_account = post_account.clone(); - // Wrong key - assert!(run_prepare_if_nonce_account_test( - &Pubkey::from([1u8; 32]), + assert!(run_post_process_failed_nonce_test( &mut post_account, - false, + false, // is_fee_payer + false, // is_nonce_account &nonce, &blockhash, lamports_per_signature, @@ -1785,20 +1779,20 @@ mod tests { Some(pre_fee_payer_account.clone()), ); - assert!(run_prepare_if_nonce_account_test( - &Pubkey::new_unique(), + assert!(run_post_process_failed_nonce_test( &mut post_fee_payer_account.clone(), - false, + false, // is_fee_payer + false, // is_nonce_account &nonce, &DurableNonce::default(), 1, &post_fee_payer_account, )); - assert!(run_prepare_if_nonce_account_test( - &Pubkey::new_unique(), + assert!(run_post_process_failed_nonce_test( &mut post_fee_payer_account.clone(), - true, + true, // is_fee_payer + false, // is_nonce_account &nonce, &DurableNonce::default(), 1,