diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs index 8dc1e728a2b94d..03547d8403b9e5 100644 --- a/programs/stake/src/stake_state.rs +++ b/programs/stake/src/stake_state.rs @@ -872,24 +872,20 @@ pub fn withdraw( return Err(StakeError::LockupInForce.into()); } - let lamports_and_reserve = checked_add(lamports, reserve)?; - // if the stake is active, we mustn't allow the account to go away - if is_staked // line coverage for branch coverage - && lamports_and_reserve > stake_account.get_lamports() - { - return Err(InstructionError::InsufficientFunds); - } - - if lamports != stake_account.get_lamports() // not a full withdrawal - && lamports_and_reserve > stake_account.get_lamports() - { - assert!(!is_staked); - return Err(InstructionError::InsufficientFunds); - } - - // Deinitialize state upon zero balance if lamports == stake_account.get_lamports() { + // if the stake is active, we mustn't allow the account to go away + if is_staked { + return Err(InstructionError::InsufficientFunds); + } + + // Deinitialize state upon zero balance stake_account.set_state(&StakeStateV2::Uninitialized)?; + } else { + // Don't allow withdrawing the reserved rent balance or active stake + let lamports_and_reserve = checked_add(lamports, reserve)?; + if lamports_and_reserve > stake_account.get_lamports() { + return Err(InstructionError::InsufficientFunds); + } } stake_account.checked_sub_lamports(lamports)?;