From 19b35de595901fd414d7eaf39a02f31ad85e4bcd Mon Sep 17 00:00:00 2001 From: Joe Date: Wed, 4 Oct 2023 14:57:13 +0200 Subject: [PATCH] refactor error handling --- runtime/src/bank/replace_account.rs | 244 +++++++++++++--------------- runtime/src/bank/tests.rs | 36 ++-- 2 files changed, 128 insertions(+), 152 deletions(-) diff --git a/runtime/src/bank/replace_account.rs b/runtime/src/bank/replace_account.rs index 1835e769413686..9fe91c3b0bddf0 100644 --- a/runtime/src/bank/replace_account.rs +++ b/runtime/src/bank/replace_account.rs @@ -12,23 +12,16 @@ use { }; /// Errors returned by `replace_account` methods -#[derive(Debug, Error, PartialEq)] +#[derive(Debug, Error)] pub enum ReplaceAccountError { - /// Source account not found - #[error("Source account not found")] - SourceAccountNotFound, - /// Source data account not found - #[error("Source data account not found")] - SourceDataAccountNotFound, - /// Destination account not found - #[error("Destination account not found")] - DestinationAccountNotFound, - /// Destination account exists - #[error("Destination account exists")] - DestinationAccountExists, - /// Unable to serialize program account state - #[error("Unable to serialize program account state")] - UnableToSerializeProgramAccountState, + /// Account not found + #[error("Account not found: {0:?}")] + AccountNotFound(Pubkey), + /// Account exists + #[error("Account exists: {0:?}")] + AccountExists(Pubkey), + #[error("Bincode Error: {0}")] + BincodeError(#[from] bincode::Error), /// Not an upgradeable program #[error("Not an upgradeable program")] NotAnUpgradeableProgram, @@ -38,7 +31,7 @@ pub enum ReplaceAccountError { /// `source`: the account to replace with /// `destination`: the account to be replaced fn move_account( - bank: &mut Bank, + bank: &Bank, source_address: &Pubkey, source_account: &V, destination_address: &Pubkey, @@ -75,36 +68,35 @@ fn move_account( /// `destination`: the non-upgradeable program account to be replaced #[allow(dead_code)] pub(crate) fn replace_non_upgradeable_program_account( - bank: &mut Bank, + bank: &Bank, source_address: &Pubkey, destination_address: &Pubkey, datapoint_name: &'static str, ) -> Result<(), ReplaceAccountError> { - if let Some(destination_account) = bank.get_account_with_fixed_root(destination_address) { - if let Some(source_account) = bank.get_account_with_fixed_root(source_address) { - datapoint_info!(datapoint_name, ("slot", bank.slot, i64)); - - move_account( - bank, - source_address, - &source_account, - destination_address, - Some(&destination_account), - ); - - // Unload a program from the bank's cache - bank.loaded_programs_cache - .write() - .unwrap() - .remove_programs([*destination_address].into_iter()); - - Ok(()) - } else { - Err(ReplaceAccountError::SourceAccountNotFound) - } - } else { - Err(ReplaceAccountError::DestinationAccountNotFound) - } + let destination_account = bank + .get_account_with_fixed_root(destination_address) + .ok_or(ReplaceAccountError::AccountNotFound(*destination_address))?; + let source_account = bank + .get_account_with_fixed_root(source_address) + .ok_or(ReplaceAccountError::AccountNotFound(*source_address))?; + + datapoint_info!(datapoint_name, ("slot", bank.slot, i64)); + + move_account( + bank, + source_address, + &source_account, + destination_address, + Some(&destination_account), + ); + + // Unload a program from the bank's cache + bank.loaded_programs_cache + .write() + .unwrap() + .remove_programs([*destination_address].into_iter()); + + Ok(()) } /// Use to replace an empty account with a program by feature activation @@ -114,102 +106,86 @@ pub(crate) fn replace_non_upgradeable_program_account( /// `source`: the upgradeable program account to replace with /// `destination`: the empty account to be replaced pub(crate) fn replace_empty_account_with_upgradeable_program( - bank: &mut Bank, + bank: &Bank, source_address: &Pubkey, destination_address: &Pubkey, datapoint_name: &'static str, ) -> Result<(), ReplaceAccountError> { // Must be attempting to replace an empty account with a program // account _and_ data account - if let Some(source_account) = bank.get_account_with_fixed_root(source_address) { - let (destination_data_address, _) = Pubkey::find_program_address( - &[destination_address.as_ref()], - &bpf_loader_upgradeable::id(), - ); - let (source_data_address, _) = - Pubkey::find_program_address(&[source_address.as_ref()], &bpf_loader_upgradeable::id()); - - // Make sure the data within the source account is the PDA of its - // data account. This also means it has at least the necessary - // lamports for rent. - if let Ok(source_state) = - bincode::deserialize::(source_account.data()) - { - match source_state { - UpgradeableLoaderState::Program { - programdata_address: _, - } => { - if let Some(source_data_account) = - bank.get_account_with_fixed_root(&source_data_address) - { - // Make sure the destination account is empty - // We aren't going to check that there isn't a data account at - // the known program-derived address (ie. `destination_data_address`), - // because if it exists, it will be overwritten - if bank - .get_account_with_fixed_root(destination_address) - .is_none() - { - let state = UpgradeableLoaderState::Program { - programdata_address: destination_data_address, - }; - if let Ok(data) = bincode::serialize(&state) { - let lamports = - bank.get_minimum_balance_for_rent_exemption(data.len()); - let created_program_account = Account { - lamports, - data, - owner: bpf_loader_upgradeable::id(), - executable: true, - rent_epoch: source_account.rent_epoch(), - }; - - datapoint_info!(datapoint_name, ("slot", bank.slot, i64)); - let change_in_capitalization = - source_account.lamports().saturating_sub(lamports); - - // Replace the destination data account with the source one - // If the destination data account does not exist, it will be created - // If it does exist, it will be overwritten - move_account( - bank, - &source_data_address, - &source_data_account, - &destination_data_address, - bank.get_account_with_fixed_root(&destination_data_address) - .as_ref(), - ); - - // Write the source data account's PDA into the destination program account - move_account( - bank, - source_address, - &created_program_account, - destination_address, - None::<&AccountSharedData>, - ); - - // Any remaining lamports in the source program account are burnt - bank.capitalization - .fetch_sub(change_in_capitalization, Relaxed); - - Ok(()) - } else { - Err(ReplaceAccountError::UnableToSerializeProgramAccountState) - } - } else { - Err(ReplaceAccountError::DestinationAccountExists) - } - } else { - Err(ReplaceAccountError::SourceDataAccountNotFound) - } - } - _ => Err(ReplaceAccountError::NotAnUpgradeableProgram), - } - } else { - Err(ReplaceAccountError::NotAnUpgradeableProgram) - } - } else { - Err(ReplaceAccountError::SourceAccountNotFound) + let source_account = bank + .get_account_with_fixed_root(source_address) + .ok_or(ReplaceAccountError::AccountNotFound(*source_address))?; + + let (destination_data_address, _) = Pubkey::find_program_address( + &[destination_address.as_ref()], + &bpf_loader_upgradeable::id(), + ); + let (source_data_address, _) = + Pubkey::find_program_address(&[source_address.as_ref()], &bpf_loader_upgradeable::id()); + + // Make sure the data within the source account is the PDA of its + // data account. This also means it has at least the necessary + // lamports for rent. + let source_state = bincode::deserialize::(source_account.data())?; + if !matches!(source_state, UpgradeableLoaderState::Program { .. }) { + return Err(ReplaceAccountError::NotAnUpgradeableProgram); + } + + let source_data_account = bank + .get_account_with_fixed_root(&source_data_address) + .ok_or(ReplaceAccountError::AccountNotFound(source_data_address))?; + + // Make sure the destination account is empty + // We aren't going to check that there isn't a data account at + // the known program-derived address (ie. `destination_data_address`), + // because if it exists, it will be overwritten + if bank + .get_account_with_fixed_root(destination_address) + .is_some() + { + return Err(ReplaceAccountError::AccountExists(*destination_address)); } + let state = UpgradeableLoaderState::Program { + programdata_address: destination_data_address, + }; + let data = bincode::serialize(&state)?; + let lamports = bank.get_minimum_balance_for_rent_exemption(data.len()); + let created_program_account = Account { + lamports, + data, + owner: bpf_loader_upgradeable::id(), + executable: true, + rent_epoch: source_account.rent_epoch(), + }; + + datapoint_info!(datapoint_name, ("slot", bank.slot, i64)); + let change_in_capitalization = source_account.lamports().saturating_sub(lamports); + + // Replace the destination data account with the source one + // If the destination data account does not exist, it will be created + // If it does exist, it will be overwritten + move_account( + bank, + &source_data_address, + &source_data_account, + &destination_data_address, + bank.get_account_with_fixed_root(&destination_data_address) + .as_ref(), + ); + + // Write the source data account's PDA into the destination program account + move_account( + bank, + source_address, + &created_program_account, + destination_address, + None::<&AccountSharedData>, + ); + + // Any remaining lamports in the source program account are burnt + bank.capitalization + .fetch_sub(change_in_capitalization, Relaxed); + + Ok(()) } diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs index f4c37780fd3508..9732b79231f278 100644 --- a/runtime/src/bank/tests.rs +++ b/runtime/src/bank/tests.rs @@ -8008,7 +8008,7 @@ fn test_compute_active_feature_set() { } fn test_program_replace_set_up_account( - bank: &mut Bank, + bank: &Bank, pubkey: &Pubkey, lamports: u64, state: &T, @@ -8038,13 +8038,13 @@ fn test_replace_non_upgradeable_program_account() { // Should replace the destination program account with the source program account: // - Destination: [*Source program data] let bpf_id = bpf_loader::id(); - let mut bank = create_simple_test_bank(0); + let bank = create_simple_test_bank(0); let destination = Pubkey::new_unique(); let destination_state = vec![0u8; 4]; let destination_lamports = bank.get_minimum_balance_for_rent_exemption(destination_state.len()); test_program_replace_set_up_account( - &mut bank, + &bank, &destination, destination_lamports, &destination_state, @@ -8056,7 +8056,7 @@ fn test_replace_non_upgradeable_program_account() { let source_state = vec![6; 30]; let source_lamports = bank.get_minimum_balance_for_rent_exemption(source_state.len()); let check_source_account = test_program_replace_set_up_account( - &mut bank, + &bank, &source, source_lamports, &source_state, @@ -8068,7 +8068,7 @@ fn test_replace_non_upgradeable_program_account() { let original_capitalization = bank.capitalization(); replace_non_upgradeable_program_account( - &mut bank, + &bank, &source, &destination, "bank-apply_program_replacement", @@ -8128,7 +8128,7 @@ fn test_replace_empty_account_with_upgradeable_program_success( // // If the destination data account exists, it will be overwritten let bpf_upgradeable_id = bpf_loader_upgradeable::id(); - let mut bank = create_simple_test_bank(0); + let bank = create_simple_test_bank(0); // Create the test source accounts, one for program and one for data let source = Pubkey::new_unique(); @@ -8141,7 +8141,7 @@ fn test_replace_empty_account_with_upgradeable_program_success( let source_data_state = vec![6; 30]; let source_data_lamports = bank.get_minimum_balance_for_rent_exemption(source_data_state.len()); test_program_replace_set_up_account( - &mut bank, + &bank, &source, source_lamports, &source_state, @@ -8149,7 +8149,7 @@ fn test_replace_empty_account_with_upgradeable_program_success( true, ); let check_source_data_account = test_program_replace_set_up_account( - &mut bank, + &bank, &source_data, source_data_lamports, &source_data_state, @@ -8168,7 +8168,7 @@ fn test_replace_empty_account_with_upgradeable_program_success( let destination_data_lamports = bank.get_minimum_balance_for_rent_exemption(destination_data_state.len()); test_program_replace_set_up_account( - &mut bank, + &bank, &destination_data, destination_data_lamports, &destination_data_state, @@ -8186,7 +8186,7 @@ fn test_replace_empty_account_with_upgradeable_program_success( // Do the replacement replace_empty_account_with_upgradeable_program( - &mut bank, + &bank, &source, &destination, "bank-apply_empty_account_replacement_for_program", @@ -8250,14 +8250,14 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists( ) { // Should not be allowed to execute replacement let bpf_upgradeable_id = bpf_loader_upgradeable::id(); - let mut bank = create_simple_test_bank(0); + let bank = create_simple_test_bank(0); // Create the test destination account with some arbitrary data and lamports balance let destination = Pubkey::new_unique(); let destination_state = vec![0, 0, 0, 0]; // Arbitrary bytes, doesn't matter let destination_lamports = bank.get_minimum_balance_for_rent_exemption(destination_state.len()); let destination_account = test_program_replace_set_up_account( - &mut bank, + &bank, &destination, destination_lamports, &destination_state, @@ -8276,7 +8276,7 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists( let source_data_state = vec![6; 30]; let source_data_lamports = bank.get_minimum_balance_for_rent_exemption(source_data_state.len()); let source_account = test_program_replace_set_up_account( - &mut bank, + &bank, &source, source_lamports, &source_state, @@ -8284,7 +8284,7 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists( true, ); let source_data_account = test_program_replace_set_up_account( - &mut bank, + &bank, &source_data, source_data_lamports, &source_data_state, @@ -8302,7 +8302,7 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists( let destination_data_lamports = bank.get_minimum_balance_for_rent_exemption(destination_data_state.len()); let destination_data_account = test_program_replace_set_up_account( - &mut bank, + &bank, &destination_data, destination_data_lamports, &destination_data_state, @@ -8317,15 +8317,15 @@ fn test_replace_empty_account_with_upgradeable_program_fail_when_account_exists( let original_capitalization = bank.capitalization(); // Attempt the replacement - assert_eq!( + assert_matches!( replace_empty_account_with_upgradeable_program( - &mut bank, + &bank, &source, &destination, "bank-apply_empty_account_replacement_for_program", ) .unwrap_err(), - ReplaceAccountError::DestinationAccountExists + ReplaceAccountError::AccountExists(..) ); // Everything should be unchanged