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

@@ -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,

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 validator_stake_info = validator_stakes
-            let transient_stake_info = validator_stakes.last().unwrap();
+                .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)?;
         }
 

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<u64, StakePoolError> {
         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<Self, ProgramError> {

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(),
+        post_validator_stake_item.stake_lamports().unwrap(),
-        pre_validator_stake_item.stake_lamports() + stake_lamports - stake_rent,
+        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(),
+        post_validator_stake_item.stake_lamports().unwrap(),
-        pre_validator_stake_item.stake_lamports() + stake_lamports - stake_rent,
+        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);
 

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::<stake::state::StakeState>());

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

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.stake_lamports().unwrap(),
-        validator_stake_item_before.stake_lamports() - tokens_burnt
+        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