From d646a076547b123f43c347a8152939f8cdea641e Mon Sep 17 00:00:00 2001 From: Joe Date: Wed, 16 Aug 2023 00:56:20 -0600 Subject: [PATCH] modify for all cases --- runtime/src/bank.rs | 141 +++++++++++++++++++++------- runtime/src/bank/tests.rs | 193 +++++++++++++++++++++++++++++++++----- 2 files changed, 274 insertions(+), 60 deletions(-) diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index c5de09e82d0980..34efaf577f476c 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -8062,6 +8062,26 @@ impl Bank { } } + fn replace_account( + &mut self, + old_address: &Pubkey, + new_address: &Pubkey, + old_account: Option<&AccountSharedData>, + new_account: &AccountSharedData, + ) { + let (old_lamports, old_len) = match old_account { + Some(old_account) => (old_account.lamports(), old_account.data().len()), + None => (0, 0), + }; + self.capitalization.fetch_sub(old_lamports, Relaxed); + self.store_account(old_address, new_account); + self.store_account(new_address, &AccountSharedData::default()); + self.calculate_and_update_accounts_data_size_delta_off_chain( + old_len, + new_account.data().len(), + ); + } + /// Use to replace programs by feature activation #[allow(dead_code)] fn replace_program_account( @@ -8070,46 +8090,95 @@ impl Bank { new_address: &Pubkey, datapoint_name: &'static str, ) { - let program_data_address = |address: &Pubkey| { - Pubkey::find_program_address(&[address.as_ref()], &bpf_loader_upgradeable::id()).0 - }; - let maybe_replace_program_account = - |old_address: &Pubkey, new_address: &Pubkey, is_program: bool| { - if let Some(old_account) = self.get_account_with_fixed_root(old_address) { - if let Some(new_account) = self.get_account_with_fixed_root(new_address) { - datapoint_info!(datapoint_name, ("slot", self.slot, i64)); - - // Burn lamports in the old account - self.capitalization - .fetch_sub(old_account.lamports(), Relaxed); - - // Transfer new account to old account - self.store_account(old_address, &new_account); - - // Clear new account - self.store_account(new_address, &AccountSharedData::default()); - - // Unload a program from the bank's cache - if is_program { - self.loaded_programs_cache - .write() - .unwrap() - .remove_programs([*old_address].into_iter()); - } + // Both program accounts must exist in order to attempt a replacement + if let Some(old_account) = self.get_account_with_fixed_root(old_address) { + if let Some(new_account) = self.get_account_with_fixed_root(new_address) { + // Both exist, so we can proceed with the replacement + datapoint_info!(datapoint_name, ("slot", self.slot, i64)); + + // Derive each program's data account address (PDA) + let (old_data_address, _) = Pubkey::find_program_address( + &[old_address.as_ref()], + &bpf_loader_upgradeable::id(), + ); + let (new_data_address, _) = Pubkey::find_program_address( + &[new_address.as_ref()], + &bpf_loader_upgradeable::id(), + ); - self.calculate_and_update_accounts_data_size_delta_off_chain( - old_account.data().len(), - new_account.data().len(), + // If a data account is also provided with this new program + // account, then we want to update the existing data account + if let Some(new_data_account) = self.get_account_with_fixed_root(&new_data_address) + { + // A data account exists for the new program + // Check if the old program account has a data account + if let Some(old_data_account) = + self.get_account_with_fixed_root(&old_data_address) + { + // It does. Replace it with the new data account + self.replace_account( + &old_data_address, + &new_data_address, + Some(&old_data_account), + &new_data_account, + ); + // The old program account will already house the PDA + // of the data account + } else { + // It does _not_. Create it with the new data account + self.replace_account( + &old_data_address, + &new_data_address, + None, + &new_data_account, ); + // Update the old program account to house the PDA of + // the data account + { + let mut account = Account::from(old_account.clone()); + account.data = old_data_address.as_ref().to_vec(); + let account_shared_data = AccountSharedData::from(account); + self.store_account(old_address, &account_shared_data); + } } + // We only swapped the data accounts, so now we need to + // clear the new program account + self.capitalization + .fetch_sub(new_account.lamports(), Relaxed); + self.store_account(new_address, &AccountSharedData::default()); + } else if let Some(old_data_account) = + self.get_account_with_fixed_root(&old_data_address) + { + // A data account exists for the old program, but not the + // new program + // Swap program accounts and delete the old data account + self.replace_account( + old_address, + new_address, + Some(&old_account), + &new_account, + ); + self.capitalization + .fetch_sub(old_data_account.lamports(), Relaxed); + self.store_account(&old_data_address, &AccountSharedData::default()); + } else { + // A data account does not exist for the new program + // Swap program accounts only + self.replace_account( + old_address, + new_address, + Some(&old_account), + &new_account, + ); } - }; - maybe_replace_program_account( - &program_data_address(old_address), - &program_data_address(new_address), - false, /* is_program */ - ); - maybe_replace_program_account(old_address, new_address, true /* is_program */); + + // Unload a program from the bank's cache + self.loaded_programs_cache + .write() + .unwrap() + .remove_programs([*old_address].into_iter()); + } + } } /// Get all the accounts for this bank and calculate stats diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs index 0bc68b80f35887..9ec80c46054beb 100644 --- a/runtime/src/bank/tests.rs +++ b/runtime/src/bank/tests.rs @@ -8013,9 +8013,15 @@ fn test_compute_active_feature_set() { assert!(bank.feature_set.is_active(&test_feature)); } -fn set_up_account_with_bank(bank: &mut Bank, pubkey: &Pubkey, lamports: u64) -> AccountSharedData { +fn set_up_account_with_bank( + bank: &mut Bank, + pubkey: &Pubkey, + lamports: u64, + data: Vec, +) -> AccountSharedData { let new_account = AccountSharedData::from(Account { lamports, + data, ..Account::default() }); bank.store_account_and_update_capitalization(pubkey, &new_account); @@ -8024,41 +8030,180 @@ fn set_up_account_with_bank(bank: &mut Bank, pubkey: &Pubkey, lamports: u64) -> } #[test] -fn test_program_replacement() { +fn test_non_upgradable_program_replace() { + // Non-Upgradable program + // - Old: [Old program data] + // - New: [*New program data] + // + // Should replace the old program account with the new program account: + // - Old: [*New program data] let mut bank = create_simple_test_bank(0); - // Setup original program accounts - let old_address = Pubkey::new_unique(); - set_up_account_with_bank(&mut bank, &old_address, 100); + let old = Pubkey::new_unique(); + set_up_account_with_bank(&mut bank, &old, 100, vec![0, 0, 0, 0]); - let (old_data_address, _) = - Pubkey::find_program_address(&[old_address.as_ref()], &bpf_loader_upgradeable::id()); - set_up_account_with_bank(&mut bank, &old_data_address, 102); + let new = Pubkey::new_unique(); + let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0]; + set_up_account_with_bank(&mut bank, &new, 100, new_program_bytes.clone()); - // Setup new program accounts - let new_address = Pubkey::new_unique(); - let new_program_account = set_up_account_with_bank(&mut bank, &new_address, 123); + let original_capitalization = bank.capitalization(); + + bank.replace_program_account(&old, &new, "bank-apply_program_replacement"); + + // Old program account balance is unchanged + assert_eq!(bank.get_balance(&old), 100); + + // New program account is now empty + assert_eq!(bank.get_balance(&new), 0); + + // Old program account now holds the new program data, ie: + // - Old: [*New program data] + let old_account = bank.get_account(&old).unwrap(); + assert_eq!(old_account.data(), &new_program_bytes,); + + // Lamports in the old token account was burnt + assert_eq!(bank.capitalization(), original_capitalization - 100); +} + +#[test] +fn test_non_upgradable_program_replace_with_data() { + // Non-Upgradable program + // - Old: [Old program data] + // - New: PDA(NewData) + // - NewData: [*New program data] + // + // Should replace the program account with the PDA of the data account, + // and create the data account: + // - Old: PDA(OldData) + // - OldData: [*New program data] + let bpf_id = bpf_loader_upgradeable::id(); + let mut bank = create_simple_test_bank(0); + + let old = Pubkey::new_unique(); + let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id); + set_up_account_with_bank(&mut bank, &old, 100, vec![0, 0, 0, 0]); - let (new_data_address, _) = - Pubkey::find_program_address(&[new_address.as_ref()], &bpf_loader_upgradeable::id()); - let new_program_data_account = set_up_account_with_bank(&mut bank, &new_data_address, 125); + let new = Pubkey::new_unique(); + let (new_data, _) = Pubkey::find_program_address(&[new.as_ref()], &bpf_id); + let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0]; + set_up_account_with_bank(&mut bank, &new, 100, new_data.to_bytes().to_vec()); + set_up_account_with_bank(&mut bank, &new_data, 102, new_program_bytes.clone()); let original_capitalization = bank.capitalization(); - bank.replace_program_account(&old_address, &new_address, "bank-apply_program_replacement"); + bank.replace_program_account(&old, &new, "bank-apply_program_replacement"); + + // Old program account balances are unchanged + assert_eq!(bank.get_balance(&old), 100); + assert_eq!(bank.get_balance(&old_data), 102); // New program accounts are now empty - assert_eq!(bank.get_balance(&new_address), 0); - assert_eq!(bank.get_balance(&new_data_address), 0); + assert_eq!(bank.get_balance(&new), 0); + assert_eq!(bank.get_balance(&new_data), 0); - // Old program account holds the new program account - assert_eq!(bank.get_account(&old_address), Some(new_program_account)); + // Old program account now holds the PDA, ie: + // - Old: PDA(OldData) + let old_account = bank.get_account(&old).unwrap(); + assert_eq!(old_account.data(), &old_data.to_bytes().to_vec(),); - // Old program data account holds the new program data account - assert_eq!( - bank.get_account(&old_data_address), - Some(new_program_data_account) - ); + // Old program data account has been created & now holds the new data, ie: + // - OldData: [*New program data] + let old_data_account = bank.get_account(&old_data).unwrap(); + assert_eq!(old_data_account.data(), &new_program_bytes,); + + // Lamports in the old token accounts were burnt + assert_eq!(bank.capitalization(), original_capitalization - 100); +} + +#[test] +fn test_upgradable_program_replace() { + // Upgradable program + // - Old: PDA(OldData) + // - OldData: [Old program data] + // - New: PDA(NewData) + // - NewData: [*New program data] + // + // Should _only_ replace the data account, not the program account: + // - Old: PDA(OldData) + // - OldData: [*New program data] + let bpf_id = bpf_loader_upgradeable::id(); + let mut bank = create_simple_test_bank(0); + + let old = Pubkey::new_unique(); + let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id); + set_up_account_with_bank(&mut bank, &old, 100, old_data.to_bytes().to_vec()); + set_up_account_with_bank(&mut bank, &old_data, 102, vec![0, 1, 2, 3, 4, 5, 6]); + + let new = Pubkey::new_unique(); + let (new_data, _) = Pubkey::find_program_address(&[new.as_ref()], &bpf_id); + let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0]; + set_up_account_with_bank(&mut bank, &new, 100, new_data.to_bytes().to_vec()); + set_up_account_with_bank(&mut bank, &new_data, 102, new_program_bytes.clone()); + + let original_capitalization = bank.capitalization(); + + bank.replace_program_account(&old, &new, "bank-apply_program_replacement"); + + // Old program account balances are unchanged + assert_eq!(bank.get_balance(&old), 100); + assert_eq!(bank.get_balance(&old_data), 102); + + // New program accounts are now empty + assert_eq!(bank.get_balance(&new), 0); + assert_eq!(bank.get_balance(&new_data), 0); + + // Old program account still holds the same PDA, ie: + // - Old: PDA(OldData) + let old_account = bank.get_account(&old).unwrap(); + assert_eq!(old_account.data(), &old_data.to_bytes().to_vec(),); + + // Old program data account now holds the new data, ie: + // - OldData: [*New program data] + let old_data_account = bank.get_account(&old_data).unwrap(); + assert_eq!(old_data_account.data(), &new_program_bytes,); + + // Lamports in the old token accounts were burnt + assert_eq!(bank.capitalization(), original_capitalization - 100 - 102); +} + +#[test] +fn test_upgradable_program_replace_with_no_data() { + // Upgradable program + // - Old: PDA(OldData) + // - OldData: [Old program data] + // - New: [*New program data] + // + // Should replace the program account and delete the data account: + // - Old: [*New program data] + let bpf_id = bpf_loader_upgradeable::id(); + let mut bank = create_simple_test_bank(0); + + let old = Pubkey::new_unique(); + let (old_data, _) = Pubkey::find_program_address(&[old.as_ref()], &bpf_id); + set_up_account_with_bank(&mut bank, &old, 100, old_data.to_bytes().to_vec()); + set_up_account_with_bank(&mut bank, &old_data, 102, vec![0, 1, 2, 3, 4, 5, 6]); + + let new = Pubkey::new_unique(); + let new_program_bytes = vec![6, 5, 4, 3, 2, 1, 0]; + set_up_account_with_bank(&mut bank, &new, 100, new_program_bytes.clone()); + + let original_capitalization = bank.capitalization(); + + bank.replace_program_account(&old, &new, "bank-apply_program_replacement"); + + // Old program account balance is unchanged + assert_eq!(bank.get_balance(&old), 100); + + // Old data account is now empty + assert_eq!(bank.get_balance(&old_data), 0); + + // New program account is now empty + assert_eq!(bank.get_balance(&new), 0); + + // Old program account now holds the new program bytes + // - Old: [*New program data] + let old_account = bank.get_account(&old).unwrap(); + assert_eq!(old_account.data(), &new_program_bytes,); // Lamports in the old token accounts were burnt assert_eq!(bank.capitalization(), original_capitalization - 100 - 102);