stake-pool: Add ability to withdraw from transient stakes (#1890)

* stake-pool: Add ability to withdraw from transient stakes

It's possible for a very malicious pool staker to constantly increase /
decrease the stake on validators, making it impossible for people to get
their SOL out.

Update the accounting to know how much of the stake is active and how
much is transient and allow users to withdraw from transient accounts,
but only if there's no more active stake.

* Remove mut ickiness

stake-pool/cli/src/main.rs

@@ -641,7 +641,7 @@ fn command_list(config: &Config, stake_pool_address: &Pubkey) -> CommandResult {
         println!(
             "Validator Vote Account: {}\tBalance: {}\tLast Update Epoch: {}{}",
             validator.vote_account_address,
-            Sol(validator.stake_lamports),
+            Sol(validator.stake_lamports()),
             validator.last_update_epoch,
             if validator.last_update_epoch != epoch_info.epoch {
                 " [UPDATE REQUIRED]"

stake-pool/program/src/instruction.rs

@@ -129,8 +129,8 @@ pub enum StakePoolInstruction {
     ///  0. `[]` Stake pool
     ///  1. `[s]` Stake pool staker
     ///  2. `[]` Stake pool withdraw authority
-    ///  3. `[]` Validator list
+    ///  3. `[w]` Validator list
-    ///  5. `[w]` Canonical stake account to split from
+    ///  4. `[w]` Canonical stake account to split from
     ///  5. `[w]` Transient stake account to receive split
     ///  6. `[]` Clock sysvar
     ///  7. `[]` Rent sysvar
@@ -444,7 +444,7 @@ pub fn decrease_validator_stake(
         AccountMeta::new_readonly(*stake_pool, false),
         AccountMeta::new_readonly(*staker, true),
         AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
-        AccountMeta::new_readonly(*validator_list, false),
+        AccountMeta::new(*validator_list, false),
         AccountMeta::new(*validator_stake, false),
         AccountMeta::new(*transient_stake, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),

stake-pool/program/src/processor.rs

@@ -736,7 +736,8 @@ impl Processor {
         validator_list.validators.push(ValidatorStakeInfo {
             status: StakeStatus::Active,
             vote_account_address,
-            stake_lamports: stake_lamports.saturating_sub(minimum_lamport_amount),
+            active_stake_lamports: stake_lamports.saturating_sub(minimum_lamport_amount),
+            transient_stake_lamports: 0,
             last_update_epoch: clock.epoch,
         });
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
@@ -916,7 +917,7 @@ impl Processor {
 
         stake_pool.check_validator_list(validator_list_info)?;
         check_account_owner(validator_list_info, program_id)?;
-        let validator_list =
+        let mut validator_list =
             try_from_slice_unchecked::<ValidatorList>(&validator_list_info.data.borrow())?;
         if !validator_list.is_valid() {
             return Err(StakePoolError::InvalidState.into());
@@ -944,13 +945,15 @@ impl Processor {
             &[transient_stake_bump_seed],
         ];
 
-        if !validator_list.contains(&vote_account_address) {
+        let maybe_validator_list_entry = validator_list.find_mut(&vote_account_address);
+        if maybe_validator_list_entry.is_none() {
             msg!(
                 "Vote account {} not found in stake pool",
                 vote_account_address
             );
             return Err(StakePoolError::ValidatorNotFound.into());
         }
+        let mut validator_list_entry = maybe_validator_list_entry.unwrap();
 
         let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
         if lamports <= stake_rent {
@@ -996,6 +999,13 @@ impl Processor {
             stake_pool.withdraw_bump_seed,
         )?;
 
+        validator_list_entry.active_stake_lamports = validator_list_entry
+            .active_stake_lamports
+            .checked_sub(lamports)
+            .ok_or(StakePoolError::CalculationFailure)?;
+        validator_list_entry.transient_stake_lamports = lamports;
+        validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
+
         Ok(())
     }
 
@@ -1146,10 +1156,7 @@ impl Processor {
             stake_pool.withdraw_bump_seed,
         )?;
 
-        validator_list_entry.stake_lamports = validator_list_entry
+        validator_list_entry.transient_stake_lamports = lamports;
-            .stake_lamports
-            .checked_add(lamports)
-            .ok_or(StakePoolError::CalculationFailure)?;
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
 
         Ok(())
@@ -1274,7 +1281,8 @@ impl Processor {
                 continue;
             };
 
-            let mut stake_lamports = 0;
+            let mut active_stake_lamports = 0;
+            let mut transient_stake_lamports = 0;
             let validator_stake_state = try_from_slice_unchecked::<stake_program::StakeState>(
                 &validator_stake_info.data.borrow(),
             )
@@ -1293,7 +1301,7 @@ impl Processor {
             match transient_stake_state {
                 Some(stake_program::StakeState::Initialized(_meta)) => {
                     if no_merge {
-                        stake_lamports += transient_stake_info.lamports();
+                        transient_stake_lamports = transient_stake_info.lamports();
                     } else {
                         // merge into reserve
                         Self::stake_merge(
@@ -1316,7 +1324,7 @@ impl Processor {
                 }
                 Some(stake_program::StakeState::Stake(_, stake)) => {
                     if no_merge {
-                        stake_lamports += transient_stake_info.lamports();
+                        transient_stake_lamports = transient_stake_info.lamports();
                     } else if stake.delegation.deactivation_epoch < clock.epoch {
                         // deactivated, merge into reserve
                         Self::stake_merge(
@@ -1355,15 +1363,15 @@ impl Processor {
                                 )?;
                             } else {
                                 msg!("Stake activating or just active, not ready to merge");
-                                stake_lamports += transient_stake_info.lamports();
+                                transient_stake_lamports = transient_stake_info.lamports();
                             }
                         } else {
                             msg!("Transient stake is activating or active, but validator stake is not, need to add the validator stake account on {} back into the stake pool", stake.delegation.voter_pubkey);
-                            stake_lamports += transient_stake_info.lamports();
+                            transient_stake_lamports = transient_stake_info.lamports();
                         }
                     } else {
                         msg!("Transient stake not ready to be merged anywhere");
-                        stake_lamports += transient_stake_info.lamports();
+                        transient_stake_lamports = transient_stake_info.lamports();
                     }
                 }
                 None
@@ -1377,7 +1385,7 @@ impl Processor {
             match validator_stake_state {
                 Some(stake_program::StakeState::Stake(meta, _)) => {
                     if validator_stake_record.status == StakeStatus::Active {
-                        stake_lamports += validator_stake_info
+                        active_stake_lamports = validator_stake_info
                             .lamports()
                             .saturating_sub(minimum_stake_lamports(&meta));
                     } else {
@@ -1393,7 +1401,8 @@ impl Processor {
             }
 
             validator_stake_record.last_update_epoch = clock.epoch;
-            validator_stake_record.stake_lamports = stake_lamports;
+            validator_stake_record.active_stake_lamports = active_stake_lamports;
+            validator_stake_record.transient_stake_lamports = transient_stake_lamports;
             changes = true;
         }
 
@@ -1465,7 +1474,7 @@ impl Processor {
                 return Err(StakePoolError::StakeListOutOfDate.into());
             }
             total_stake_lamports = total_stake_lamports
-                .checked_add(validator_stake_record.stake_lamports)
+                .checked_add(validator_stake_record.stake_lamports())
                 .ok_or(StakePoolError::CalculationFailure)?;
         }
 
@@ -1674,7 +1683,7 @@ impl Processor {
             "lamports post merge {}",
             validator_stake_account_info.lamports()
         );
-        validator_list_item.stake_lamports = validator_stake_account_info
+        validator_list_item.active_stake_lamports = validator_stake_account_info
             .lamports()
             .checked_sub(minimum_stake_lamports(&meta))
             .ok_or(StakePoolError::CalculationFailure)?;
@@ -1738,19 +1747,19 @@ impl Processor {
             .calc_lamports_withdraw_amount(pool_tokens)
             .ok_or(StakePoolError::CalculationFailure)?;
 
-        let validator_list_item = if *stake_split_from.key == stake_pool.reserve_stake {
+        let validator_list_item_info = if *stake_split_from.key == stake_pool.reserve_stake {
             // check that the validator stake accounts have no withdrawable stake
             if let Some(withdrawable_entry) = validator_list
                 .validators
                 .iter()
-                .find(|&&x| x.stake_lamports != 0)
+                .find(|&&x| x.stake_lamports() != 0)
             {
                 let (validator_stake_address, _) = crate::find_stake_program_address(
                     &program_id,
                     &withdrawable_entry.vote_account_address,
                     stake_pool_info.key,
                 );
-                msg!("Error withdrawing from reserve: validator stake account {} has {} lamports available, please use that first.", validator_stake_address, withdrawable_entry.stake_lamports);
+                msg!("Error withdrawing from reserve: validator stake account {} has {} lamports available, please use that first.", validator_stake_address, withdrawable_entry.stake_lamports());
                 return Err(StakePoolError::StakeLamportsNotEqualToMinimum.into());
             }
 
@@ -1767,12 +1776,6 @@ impl Processor {
         } else {
             let (meta, stake) = get_stake_state(stake_split_from)?;
             let vote_account_address = stake.delegation.voter_pubkey;
-            check_validator_stake_address(
-                program_id,
-                stake_pool_info.key,
-                stake_split_from.key,
-                &vote_account_address,
-            )?;
 
             if let Some(preferred_withdraw_validator) =
                 validator_list.preferred_withdraw_validator_vote_address
@@ -1781,13 +1784,33 @@ impl Processor {
                     .find(&preferred_withdraw_validator)
                     .ok_or(StakePoolError::ValidatorNotFound)?;
                 if preferred_withdraw_validator != vote_account_address
-                    && preferred_validator_info.stake_lamports > 0
+                    && preferred_validator_info.active_stake_lamports > 0
                 {
-                    msg!("Validator vote address {} is preferred for withdrawals, it currently has {} lamports available. Please withdraw those before using other validator stake accounts.", preferred_withdraw_validator, preferred_validator_info.stake_lamports);
+                    msg!("Validator vote address {} is preferred for withdrawals, it currently has {} lamports available. Please withdraw those before using other validator stake accounts.", preferred_withdraw_validator, preferred_validator_info.active_stake_lamports);
                     return Err(StakePoolError::IncorrectWithdrawVoteAddress.into());
                 }
             }
 
+            // if there's any active stake, we must withdraw from an active
+            // stake account
+            let withdrawing_from_transient_stake = if validator_list.has_active_stake() {
+                check_validator_stake_address(
+                    program_id,
+                    stake_pool_info.key,
+                    stake_split_from.key,
+                    &vote_account_address,
+                )?;
+                false
+            } else {
+                check_transient_stake_address(
+                    program_id,
+                    stake_pool_info.key,
+                    stake_split_from.key,
+                    &vote_account_address,
+                )?;
+                true
+            };
+
             let validator_list_item = validator_list
                 .find_mut(&vote_account_address)
                 .ok_or(StakePoolError::ValidatorNotFound)?;
@@ -1804,7 +1827,7 @@ impl Processor {
                 msg!("Attempting to withdraw {} lamports from validator account with {} lamports, {} must remain", withdraw_lamports, current_lamports, required_lamports);
                 return Err(StakePoolError::StakeLamportsNotEqualToMinimum.into());
             }
-            Some(validator_list_item)
+            Some((validator_list_item, withdrawing_from_transient_stake))
         };
 
         Self::token_burn(
@@ -1846,11 +1869,20 @@ impl Processor {
             .ok_or(StakePoolError::CalculationFailure)?;
         stake_pool.serialize(&mut *stake_pool_info.data.borrow_mut())?;
 
-        if let Some(validator_list_item) = validator_list_item {
+        if let Some((validator_list_item, withdrawing_from_transient_stake_account)) =
-            validator_list_item.stake_lamports = validator_list_item
+            validator_list_item_info
-                .stake_lamports
+        {
-                .checked_sub(withdraw_lamports)
+            if withdrawing_from_transient_stake_account {
-                .ok_or(StakePoolError::CalculationFailure)?;
+                validator_list_item.transient_stake_lamports = validator_list_item
+                    .transient_stake_lamports
+                    .checked_sub(withdraw_lamports)
+                    .ok_or(StakePoolError::CalculationFailure)?;
+            } else {
+                validator_list_item.active_stake_lamports = validator_list_item
+                    .active_stake_lamports
+                    .checked_sub(withdraw_lamports)
+                    .ok_or(StakePoolError::CalculationFailure)?;
+            }
             validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
         }
 

stake-pool/program/src/state.rs

@@ -328,15 +328,29 @@ pub struct ValidatorStakeInfo {
     /// Validator vote account address
     pub vote_account_address: Pubkey,
 
-    /// Amount of stake delegated to this validator
+    /// Amount of active stake delegated to this validator
-    /// Note that if `last_update_epoch` does not match the current epoch then this field may not
+    /// Note that if `last_update_epoch` does not match the current epoch then
-    /// be accurate
+    /// this field may not be accurate
-    pub stake_lamports: u64,
+    pub active_stake_lamports: u64,
 
-    /// Last epoch the `stake_lamports` field was updated
+    /// Amount of transient stake delegated to this validator
+    /// Note that if `last_update_epoch` does not match the current epoch then
+    /// this field may not be accurate
+    pub transient_stake_lamports: u64,
+
+    /// Last epoch the active and transient stake lamports fields were updated
     pub last_update_epoch: u64,
 }
 
+impl ValidatorStakeInfo {
+    /// Get the total lamports delegated to this validator (active and transient)
+    pub fn stake_lamports(&self) -> u64 {
+        self.active_stake_lamports
+            .checked_add(self.transient_stake_lamports)
+            .unwrap()
+    }
+}
+
 impl ValidatorList {
     /// Create an empty instance containing space for `max_validators` and preferred validator keys
     pub fn new(max_validators: u32) -> Self {
@@ -352,7 +366,7 @@ impl ValidatorList {
     /// Calculate the number of validator entries that fit in the provided length
     pub fn calculate_max_validators(buffer_length: usize) -> usize {
         let header_size = 1 + 4 + 4 + 33 + 33;
-        buffer_length.saturating_sub(header_size) / 49
+        buffer_length.saturating_sub(header_size) / 57
     }
 
     /// Check if contains validator with particular pubkey
@@ -384,6 +398,11 @@ impl ValidatorList {
     pub fn is_uninitialized(&self) -> bool {
         self.account_type == AccountType::Uninitialized
     }
+
+    /// Check if the list has any active stake
+    pub fn has_active_stake(&self) -> bool {
+        self.validators.iter().any(|x| x.active_stake_lamports > 0)
+    }
 }
 
 /// Fee rate as a ratio, minted on `UpdateStakePoolBalance` as a proportion of
@@ -407,18 +426,53 @@ mod test {
         solana_program::native_token::LAMPORTS_PER_SOL,
     };
 
-    #[test]
+    fn uninitialized_validator_list() -> ValidatorList {
-    fn test_state_packing() {
+        ValidatorList {
-        let max_validators = 10_000;
-        let size = get_instance_packed_len(&ValidatorList::new(max_validators)).unwrap();
-        // Not initialized
-        let stake_list = ValidatorList {
             account_type: AccountType::Uninitialized,
             preferred_deposit_validator_vote_address: None,
             preferred_withdraw_validator_vote_address: None,
             max_validators: 0,
             validators: vec![],
-        };
+        }
+    }
+
+    fn test_validator_list(max_validators: u32) -> ValidatorList {
+        ValidatorList {
+            account_type: AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: Some(Pubkey::new_unique()),
+            preferred_withdraw_validator_vote_address: Some(Pubkey::new_unique()),
+            max_validators,
+            validators: vec![
+                ValidatorStakeInfo {
+                    status: StakeStatus::Active,
+                    vote_account_address: Pubkey::new_from_array([1; 32]),
+                    active_stake_lamports: 123456789,
+                    transient_stake_lamports: 1111111,
+                    last_update_epoch: 987654321,
+                },
+                ValidatorStakeInfo {
+                    status: StakeStatus::DeactivatingTransient,
+                    vote_account_address: Pubkey::new_from_array([2; 32]),
+                    active_stake_lamports: 998877665544,
+                    transient_stake_lamports: 222222222,
+                    last_update_epoch: 11223445566,
+                },
+                ValidatorStakeInfo {
+                    status: StakeStatus::ReadyForRemoval,
+                    vote_account_address: Pubkey::new_from_array([3; 32]),
+                    active_stake_lamports: 0,
+                    transient_stake_lamports: 0,
+                    last_update_epoch: 999999999999999,
+                },
+            ],
+        }
+    }
+
+    #[test]
+    fn state_packing() {
+        let max_validators = 10_000;
+        let size = get_instance_packed_len(&ValidatorList::new(max_validators)).unwrap();
+        let stake_list = uninitialized_validator_list();
         let mut byte_vec = vec![0u8; size];
         let mut bytes = byte_vec.as_mut_slice();
         stake_list.serialize(&mut bytes).unwrap();
@@ -440,32 +494,7 @@ mod test {
         assert_eq!(stake_list_unpacked, stake_list);
 
         // With several accounts
-        let stake_list = ValidatorList {
+        let stake_list = test_validator_list(max_validators);
-            account_type: AccountType::ValidatorList,
-            preferred_deposit_validator_vote_address: Some(Pubkey::new_unique()),
-            preferred_withdraw_validator_vote_address: Some(Pubkey::new_unique()),
-            max_validators,
-            validators: vec![
-                ValidatorStakeInfo {
-                    status: StakeStatus::Active,
-                    vote_account_address: Pubkey::new_from_array([1; 32]),
-                    stake_lamports: 123456789,
-                    last_update_epoch: 987654321,
-                },
-                ValidatorStakeInfo {
-                    status: StakeStatus::DeactivatingTransient,
-                    vote_account_address: Pubkey::new_from_array([2; 32]),
-                    stake_lamports: 998877665544,
-                    last_update_epoch: 11223445566,
-                },
-                ValidatorStakeInfo {
-                    status: StakeStatus::ReadyForRemoval,
-                    vote_account_address: Pubkey::new_from_array([3; 32]),
-                    stake_lamports: 0,
-                    last_update_epoch: 999999999999999,
-                },
-            ],
-        };
         let mut byte_vec = vec![0u8; size];
         let mut bytes = byte_vec.as_mut_slice();
         stake_list.serialize(&mut bytes).unwrap();
@@ -473,6 +502,17 @@ mod test {
         assert_eq!(stake_list_unpacked, stake_list);
     }
 
+    #[test]
+    fn validator_list_active_stake() {
+        let max_validators = 10_000;
+        let mut validator_list = test_validator_list(max_validators);
+        assert!(validator_list.has_active_stake());
+        for validator in validator_list.validators.iter_mut() {
+            validator.active_stake_lamports = 0;
+        }
+        assert!(!validator_list.has_active_stake());
+    }
+
     proptest! {
         #[test]
         fn stake_list_size_calculation(test_amount in 0..=100_000_u32) {

stake-pool/program/tests/deposit.rs

@@ -191,8 +191,8 @@ async fn success() {
         .find(&validator_stake_account.vote.pubkey())
         .unwrap();
     assert_eq!(
-        validator_stake_item.stake_lamports,
+        validator_stake_item.stake_lamports(),
-        validator_stake_item_before.stake_lamports + stake_lamports
+        validator_stake_item_before.stake_lamports() + stake_lamports
     );
 
     // Check validator stake account actual SOL balance
@@ -203,8 +203,9 @@ async fn success() {
     let meta = stake_state.meta().unwrap();
     assert_eq!(
         validator_stake_account.lamports - minimum_stake_lamports(&meta),
-        validator_stake_item.stake_lamports
+        validator_stake_item.stake_lamports()
     );
+    assert_eq!(validator_stake_item.transient_stake_lamports, 0);
 }
 
 #[tokio::test]

stake-pool/program/tests/helpers/mod.rs

@@ -1124,7 +1124,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())
         .sum();
     let rent = banks_client.get_rent().await.unwrap();
     let rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());

stake-pool/program/tests/update_validator_list_balance.rs

@@ -536,8 +536,9 @@ async fn merge_transient_stake_after_remove() {
         validator_list.validators[0].status,
         StakeStatus::DeactivatingTransient
     );
+    assert_eq!(validator_list.validators[0].active_stake_lamports, 0);
     assert_eq!(
-        validator_list.validators[0].stake_lamports,
+        validator_list.validators[0].transient_stake_lamports,
         deactivated_lamports
     );
 
@@ -569,7 +570,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(), 0);
 
     let reserve_stake = context
         .banks_client

stake-pool/program/tests/vsa_add.rs

@@ -91,7 +91,8 @@ async fn success() {
                 status: state::StakeStatus::Active,
                 vote_account_address: user_stake.vote.pubkey(),
                 last_update_epoch: 0,
-                stake_lamports: 0,
+                active_stake_lamports: 0,
+                transient_stake_lamports: 0,
             }]
         }
     );

stake-pool/program/tests/vsa_remove.rs

@@ -529,7 +529,8 @@ async fn success_with_deactivating_transient_stake() {
             status: state::StakeStatus::DeactivatingTransient,
             vote_account_address: validator_stake.vote.pubkey(),
             last_update_epoch: 0,
-            stake_lamports: TEST_STAKE_AMOUNT + stake_rent,
+            active_stake_lamports: 0,
+            transient_stake_lamports: TEST_STAKE_AMOUNT + stake_rent,
         }],
     };
     assert_eq!(validator_list, expected_list);

stake-pool/program/tests/withdraw.rs

@@ -181,8 +181,12 @@ async fn success() {
         .find(&validator_stake_account.vote.pubkey())
         .unwrap();
     assert_eq!(
-        validator_stake_item.stake_lamports,
+        validator_stake_item.stake_lamports(),
-        validator_stake_item_before.stake_lamports - tokens_to_burn
+        validator_stake_item_before.stake_lamports() - tokens_to_burn
+    );
+    assert_eq!(
+        validator_stake_item.active_stake_lamports,
+        validator_stake_item.stake_lamports(),
     );
 
     // Check tokens burned
@@ -201,7 +205,7 @@ async fn success() {
     let meta = stake_state.meta().unwrap();
     assert_eq!(
         validator_stake_account.lamports - minimum_stake_lamports(&meta),
-        validator_stake_item.stake_lamports
+        validator_stake_item.active_stake_lamports
     );
 
     // Check user recipient stake account balance
@@ -1131,3 +1135,172 @@ async fn fail_with_wrong_preferred_withdraw() {
         _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
     }
 }
+
+#[tokio::test]
+async fn success_withdraw_from_transient() {
+    let mut context = program_test().start_with_context().await;
+    let stake_pool_accounts = StakePoolAccounts::new();
+    let initial_reserve_lamports = 1;
+    stake_pool_accounts
+        .initialize_stake_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            initial_reserve_lamports,
+        )
+        .await
+        .unwrap();
+
+    // add a preferred withdraw validator, keep it empty, to be sure that this works
+    let preferred_validator = simple_add_validator_to_pool(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts,
+    )
+    .await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            instruction::PreferredValidatorType::Withdraw,
+            Some(preferred_validator.vote.pubkey()),
+        )
+        .await;
+
+    let validator_stake = simple_add_validator_to_pool(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts,
+    )
+    .await;
+
+    let deposit_lamports = TEST_STAKE_AMOUNT;
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+
+    let deposit_info = simple_deposit(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts,
+        &validator_stake,
+        deposit_lamports,
+    )
+    .await
+    .unwrap();
+
+    // Delegate tokens for burning during withdraw
+    let user_transfer_authority = Keypair::new();
+    delegate_tokens(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &deposit_info.pool_account.pubkey(),
+        &deposit_info.authority,
+        &user_transfer_authority.pubkey(),
+        deposit_info.pool_tokens,
+    )
+    .await;
+
+    // decrease minimum stake
+    let error = stake_pool_accounts
+        .decrease_validator_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &validator_stake.stake_account,
+            &validator_stake.transient_stake_account,
+            stake_rent + 1,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let withdraw_destination = Keypair::new();
+    let withdraw_destination_authority = Pubkey::new_unique();
+    let _initial_stake_lamports = create_blank_stake_account(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &withdraw_destination,
+    )
+    .await;
+
+    // fail withdrawing from transient, still a lamport in the validator stake account
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &withdraw_destination.pubkey(),
+            &user_transfer_authority,
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.transient_stake_account,
+            &withdraw_destination_authority,
+            deposit_info.pool_tokens / 2,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(
+        error,
+        TransactionError::InstructionError(
+            0,
+            InstructionError::Custom(StakePoolError::InvalidStakeAccountAddress as u32)
+        )
+    );
+
+    // warp forward to deactivation
+    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    context
+        .warp_to_slot(first_normal_slot + slots_per_epoch)
+        .unwrap();
+
+    // update to merge deactivated stake into reserve
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &[
+                preferred_validator.vote.pubkey(),
+                validator_stake.vote.pubkey(),
+            ],
+            false,
+        )
+        .await;
+
+    // decrease rest of stake
+    let error = stake_pool_accounts
+        .decrease_validator_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &validator_stake.stake_account,
+            &validator_stake.transient_stake_account,
+            deposit_lamports - 1,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // nothing left in the validator stake account (or any others), so withdrawing
+    // from the transient account is ok!
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &withdraw_destination.pubkey(),
+            &user_transfer_authority,
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.transient_stake_account,
+            &withdraw_destination_authority,
+            deposit_info.pool_tokens / 4,
+        )
+        .await;
+    assert!(error.is_none());
+}