From 86774834e624d54433350a68e8cd453fdba92504 Mon Sep 17 00:00:00 2001 From: billythedummy Date: Fri, 15 Mar 2024 13:46:34 -0400 Subject: [PATCH] checked ops --- libs/spl_stake_pool/src/lib.rs | 6 +++++- .../src/stakedex_traits/withdraw_stake.rs | 11 ++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/libs/spl_stake_pool/src/lib.rs b/libs/spl_stake_pool/src/lib.rs index d38fbd2..4752326 100644 --- a/libs/spl_stake_pool/src/lib.rs +++ b/libs/spl_stake_pool/src/lib.rs @@ -48,7 +48,11 @@ impl SplStakePoolStakedex { validator_index: usize, withdraw_amount: u64, ) -> Result { - let validator_list_entry = self.validator_list.validators.get(validator_index).unwrap(); + let validator_list_entry = self + .validator_list + .validators + .get(validator_index) + .ok_or(StakePoolError::ValidatorNotFound)?; // only handle withdrawal from active stake accounts for simplicity. // Likely other stake pools can't accept non active stake anyway if validator_list_entry.status != StakeStatus::Active.into() { diff --git a/libs/spl_stake_pool/src/stakedex_traits/withdraw_stake.rs b/libs/spl_stake_pool/src/stakedex_traits/withdraw_stake.rs index bfc9e51..94ac588 100644 --- a/libs/spl_stake_pool/src/stakedex_traits/withdraw_stake.rs +++ b/libs/spl_stake_pool/src/stakedex_traits/withdraw_stake.rs @@ -30,11 +30,12 @@ impl<'a> WithdrawStakeQuoteIter<'a> { .pool .get_withdraw_stake_quote_for_validator_copied(curr_index, self.withdraw_amount) .unwrap_or_default(); - let next_state = if curr_index >= self.pool.validator_list.validators.len() - 1 { - WithdrawStakeQuoteIterState::Ended - } else { - WithdrawStakeQuoteIterState::Normal(curr_index + 1) - }; + let next_state = + if curr_index >= self.pool.validator_list.validators.len().checked_sub(1)? { + WithdrawStakeQuoteIterState::Ended + } else { + WithdrawStakeQuoteIterState::Normal(curr_index.checked_add(1)?) + }; Some((wsq, next_state)) }