Skip to content

Commit

Permalink
refactor error handling
Browse files Browse the repository at this point in the history
  • Loading branch information
@buffalojoec
buffalojoec committed Oct 4, 2023
1 parent 3e7dc30 commit 19b35de
Show file tree
Hide file tree
Showing 2 changed files with 128 additions and 152 deletions.
244 changes: 110 additions & 134 deletions runtime/src/bank/replace_account.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -38,7 +31,7 @@ pub enum ReplaceAccountError {
/// `source`: the account to replace with
/// `destination`: the account to be replaced
fn move_account<U, V>(
bank: &mut Bank,
bank: &Bank,
source_address: &Pubkey,
source_account: &V,
destination_address: &Pubkey,
Expand Down Expand Up @@ -75,36 +68,35 @@ fn move_account<U, V>(
/// `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
Expand All @@ -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::<UpgradeableLoaderState>(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::<UpgradeableLoaderState>(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(())
}
Loading

0 comments on commit 19b35de

Please sign in to comment.