From eccb68d1c2b7373037fbac2ae4286e6109949bc5 Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Mon, 4 Nov 2024 18:41:15 +0000 Subject: [PATCH 1/3] refactor: step 1 --- programs/stake/src/stake_state.rs | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs index 8dc1e728a2b94d..1dc0633753a1a9 100644 --- a/programs/stake/src/stake_state.rs +++ b/programs/stake/src/stake_state.rs @@ -873,18 +873,16 @@ pub fn withdraw( } 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_and_reserve > 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); + } - if lamports != stake_account.get_lamports() // not a full withdrawal - && lamports_and_reserve > stake_account.get_lamports() - { - assert!(!is_staked); - return Err(InstructionError::InsufficientFunds); + // fail if not a full withdrawal + if lamports != stake_account.get_lamports() { + return Err(InstructionError::InsufficientFunds); + } } // Deinitialize state upon zero balance From cebea2f892bb0985d76966c590684fdef4912f87 Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Mon, 4 Nov 2024 18:41:31 +0000 Subject: [PATCH 2/3] refactor: step 2 --- programs/stake/src/stake_state.rs | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs index 1dc0633753a1a9..d8b3225ce87c63 100644 --- a/programs/stake/src/stake_state.rs +++ b/programs/stake/src/stake_state.rs @@ -872,22 +872,28 @@ pub fn withdraw( return Err(StakeError::LockupInForce.into()); } - let lamports_and_reserve = checked_add(lamports, reserve)?; - if lamports_and_reserve > stake_account.get_lamports() { + 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); } - // fail if not a full withdrawal - if lamports != stake_account.get_lamports() { - return Err(InstructionError::InsufficientFunds); - } - } - - // Deinitialize state upon zero balance - if lamports == stake_account.get_lamports() { + // 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() { + // if the stake is active, we mustn't allow the account to go away + if is_staked { + return Err(InstructionError::InsufficientFunds); + } + + // fail if not a full withdrawal + if lamports != stake_account.get_lamports() { + return Err(InstructionError::InsufficientFunds); + } + } } stake_account.checked_sub_lamports(lamports)?; From 84a5994c3f638eca6f224f48ea6f20b32222889c Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Mon, 4 Nov 2024 20:34:39 +0000 Subject: [PATCH 3/3] refactor: step 3 --- programs/stake/src/stake_state.rs | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs index d8b3225ce87c63..03547d8403b9e5 100644 --- a/programs/stake/src/stake_state.rs +++ b/programs/stake/src/stake_state.rs @@ -884,15 +884,7 @@ pub fn withdraw( // 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() { - // if the stake is active, we mustn't allow the account to go away - if is_staked { - return Err(InstructionError::InsufficientFunds); - } - - // fail if not a full withdrawal - if lamports != stake_account.get_lamports() { - return Err(InstructionError::InsufficientFunds); - } + return Err(InstructionError::InsufficientFunds); } }