From 2cdef76cc9ea4ff9844ac65c32b11691f596ebb0 Mon Sep 17 00:00:00 2001 From: Jon Cinque Date: Sat, 28 Jan 2023 02:13:46 +0100 Subject: [PATCH] stake-pool: Remove some unwraps (#4003) * stake-pool: Remove a couple of unwraps * Remove one more unwrap * Update CLI too --- stake-pool/cli/src/main.rs | 2 +- stake-pool/program/src/processor.rs | 14 ++++++++++---- stake-pool/program/src/state.rs | 8 ++++++-- stake-pool/program/tests/deposit.rs | 12 ++++++------ stake-pool/program/tests/helpers/mod.rs | 2 +- .../program/tests/update_validator_list_balance.rs | 2 +- stake-pool/program/tests/withdraw.rs | 6 +++--- 7 files changed, 28 insertions(+), 18 deletions(-) diff --git a/stake-pool/cli/src/main.rs b/stake-pool/cli/src/main.rs index cba12254..2319c6e3 100644 --- a/stake-pool/cli/src/main.rs +++ b/stake-pool/cli/src/main.rs @@ -1103,7 +1103,7 @@ fn command_list(config: &Config, stake_pool_address: &Pubkey) -> CommandResult { stake_account_address: stake_account_address.to_string(), validator_active_stake_lamports: validator.active_stake_lamports, validator_last_update_epoch: validator.last_update_epoch, - validator_lamports: validator.stake_lamports(), + validator_lamports: validator.stake_lamports().unwrap(), validator_transient_stake_account_address: transient_stake_account_address .to_string(), validator_transient_stake_lamports: validator.transient_stake_lamports, diff --git a/stake-pool/program/src/processor.rs b/stake-pool/program/src/processor.rs index 581aa941..83ad845b 100644 --- a/stake-pool/program/src/processor.rs +++ b/stake-pool/program/src/processor.rs @@ -446,7 +446,9 @@ impl Processor { stake::instruction::split(stake_account.key, authority.key, amount, split_stake.key); invoke_signed( - split_instruction.last().unwrap(), + split_instruction + .last() + .ok_or(ProgramError::InvalidInstructionData)?, &[stake_account, split_stake, authority], signers, ) @@ -2237,8 +2239,12 @@ impl Processor { .zip(validator_stake_accounts.chunks_exact(2)); for (validator_stake_record, validator_stakes) in validator_iter { // chunks_exact means that we always get 2 elements, making this safe - let validator_stake_info = validator_stakes.first().unwrap(); - let transient_stake_info = validator_stakes.last().unwrap(); + let validator_stake_info = validator_stakes + .first() + .ok_or(ProgramError::InvalidInstructionData)?; + let transient_stake_info = validator_stakes + .last() + .ok_or(ProgramError::InvalidInstructionData)?; if check_validator_stake_address( program_id, stake_pool_info.key, @@ -2533,7 +2539,7 @@ impl Processor { return Err(StakePoolError::StakeListOutOfDate.into()); } total_lamports = total_lamports - .checked_add(validator_stake_record.stake_lamports()) + .checked_add(validator_stake_record.stake_lamports()?) .ok_or(StakePoolError::CalculationFailure)?; } diff --git a/stake-pool/program/src/state.rs b/stake-pool/program/src/state.rs index 30b2ef43..cede91c1 100644 --- a/stake-pool/program/src/state.rs +++ b/stake-pool/program/src/state.rs @@ -661,10 +661,10 @@ pub struct ValidatorStakeInfo { impl ValidatorStakeInfo { /// Get the total lamports on this validator (active and transient) - pub fn stake_lamports(&self) -> u64 { + pub fn stake_lamports(&self) -> Result { self.active_stake_lamports .checked_add(self.transient_stake_lamports) - .unwrap() + .ok_or(StakePoolError::CalculationFailure) } /// Performs a very cheap comparison, for checking if this validator stake @@ -680,12 +680,14 @@ impl ValidatorStakeInfo { /// Performs a comparison, used to check if this validator stake /// info has more active lamports than some limit pub fn active_lamports_greater_than(data: &[u8], lamports: &u64) -> bool { + // without this unwrap, compute usage goes up significantly u64::try_from_slice(&data[0..8]).unwrap() > *lamports } /// Performs a comparison, used to check if this validator stake /// info has more transient lamports than some limit pub fn transient_lamports_greater_than(data: &[u8], lamports: &u64) -> bool { + // without this unwrap, compute usage goes up significantly u64::try_from_slice(&data[8..16]).unwrap() > *lamports } @@ -701,6 +703,8 @@ impl Pack for ValidatorStakeInfo { const LEN: usize = 73; fn pack_into_slice(&self, data: &mut [u8]) { let mut data = data; + // Removing this unwrap would require changing from `Pack` to some other + // trait or `bytemuck`, so it stays in for now self.serialize(&mut data).unwrap(); } fn unpack_from_slice(src: &[u8]) -> Result { diff --git a/stake-pool/program/tests/deposit.rs b/stake-pool/program/tests/deposit.rs index 585f43ff..07955b9e 100644 --- a/stake-pool/program/tests/deposit.rs +++ b/stake-pool/program/tests/deposit.rs @@ -235,8 +235,8 @@ async fn success(token_program_id: Pubkey) { .find(&validator_stake_account.vote.pubkey()) .unwrap(); assert_eq!( - post_validator_stake_item.stake_lamports(), - pre_validator_stake_item.stake_lamports() + stake_lamports - stake_rent, + post_validator_stake_item.stake_lamports().unwrap(), + pre_validator_stake_item.stake_lamports().unwrap() + stake_lamports - stake_rent, ); // Check validator stake account actual SOL balance @@ -247,7 +247,7 @@ async fn success(token_program_id: Pubkey) { .await; assert_eq!( validator_stake_account.lamports, - post_validator_stake_item.stake_lamports() + post_validator_stake_item.stake_lamports().unwrap() ); assert_eq!(post_validator_stake_item.transient_stake_lamports, 0); @@ -429,8 +429,8 @@ async fn success_with_extra_stake_lamports() { .find(&validator_stake_account.vote.pubkey()) .unwrap(); assert_eq!( - post_validator_stake_item.stake_lamports(), - pre_validator_stake_item.stake_lamports() + stake_lamports - stake_rent, + post_validator_stake_item.stake_lamports().unwrap(), + pre_validator_stake_item.stake_lamports().unwrap() + stake_lamports - stake_rent, ); // Check validator stake account actual SOL balance @@ -441,7 +441,7 @@ async fn success_with_extra_stake_lamports() { .await; assert_eq!( validator_stake_account.lamports, - post_validator_stake_item.stake_lamports() + post_validator_stake_item.stake_lamports().unwrap() ); assert_eq!(post_validator_stake_item.transient_stake_lamports, 0); diff --git a/stake-pool/program/tests/helpers/mod.rs b/stake-pool/program/tests/helpers/mod.rs index a7691f59..f8a1a9e1 100644 --- a/stake-pool/program/tests/helpers/mod.rs +++ b/stake-pool/program/tests/helpers/mod.rs @@ -2127,7 +2127,7 @@ pub async fn get_validator_list_sum( let validator_sum: u64 = validator_list .validators .iter() - .map(|info| info.stake_lamports()) + .map(|info| info.stake_lamports().unwrap()) .sum(); let rent = banks_client.get_rent().await.unwrap(); let rent = rent.minimum_balance(std::mem::size_of::()); diff --git a/stake-pool/program/tests/update_validator_list_balance.rs b/stake-pool/program/tests/update_validator_list_balance.rs index 827189a2..841b8b88 100644 --- a/stake-pool/program/tests/update_validator_list_balance.rs +++ b/stake-pool/program/tests/update_validator_list_balance.rs @@ -662,7 +662,7 @@ async fn merge_transient_stake_after_remove() { validator_list.validators[0].status, StakeStatus::ReadyForRemoval ); - assert_eq!(validator_list.validators[0].stake_lamports(), 0); + assert_eq!(validator_list.validators[0].stake_lamports().unwrap(), 0); let reserve_stake = context .banks_client diff --git a/stake-pool/program/tests/withdraw.rs b/stake-pool/program/tests/withdraw.rs index 63d2c76e..85c60e71 100644 --- a/stake-pool/program/tests/withdraw.rs +++ b/stake-pool/program/tests/withdraw.rs @@ -219,12 +219,12 @@ async fn _success(token_program_id: Pubkey, test_type: SuccessTestType) { .find(&validator_stake_account.vote.pubkey()) .unwrap(); assert_eq!( - validator_stake_item.stake_lamports(), - validator_stake_item_before.stake_lamports() - tokens_burnt + validator_stake_item.stake_lamports().unwrap(), + validator_stake_item_before.stake_lamports().unwrap() - tokens_burnt ); assert_eq!( validator_stake_item.active_stake_lamports, - validator_stake_item.stake_lamports(), + validator_stake_item.stake_lamports().unwrap(), ); // Check tokens used