From 0a05bbca2face21c2d6bfd35bae85dc4296e0345 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 6 Nov 2020 20:27:06 +0000
Subject: [PATCH] Fix stake redelegate (bp #13358) (#13449)

* stake: Add redelegation failing test

(cherry picked from commit 491ad59d2ec11859587f8e73e20a09183dc471f6)

* stake: Consider withdraws we redelegating

(cherry picked from commit fe1e08b9ad42969b15a00411cc69a099cc9348cd)

Co-authored-by: Trent Nelson <trent@solana.com>
---
 programs/stake/src/stake_state.rs | 109 +++++++++++++++++++++++++++++-
 1 file changed, 108 insertions(+), 1 deletion(-)

diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs
index 9167e5de05..92e5a7f0df 100644
--- a/programs/stake/src/stake_state.rs
+++ b/programs/stake/src/stake_state.rs
@@ -492,6 +492,7 @@ impl Stake {
 
     fn redelegate(
         &mut self,
+        stake_lamports: u64,
         voter_pubkey: &Pubkey,
         vote_state: &VoteState,
         clock: &Clock,
@@ -504,6 +505,7 @@ impl Stake {
         if self.stake(clock.epoch, Some(stake_history)) != 0 {
             return Err(StakeError::TooSoonToRedelegate);
         }
+        self.delegation.stake = stake_lamports;
         self.delegation.activation_epoch = clock.epoch;
         self.delegation.deactivation_epoch = std::u64::MAX;
         self.delegation.voter_pubkey = *voter_pubkey;
@@ -702,6 +704,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
             StakeState::Stake(meta, mut stake) => {
                 meta.authorized.check(signers, StakeAuthorize::Staker)?;
                 stake.redelegate(
+                    self.lamports()?.saturating_sub(meta.rent_exempt_reserve), // can't stake the rent ;)
                     vote_account.unsigned_key(),
                     &State::<VoteStateVersions>::state(vote_account)?.convert_to_current(),
                     clock,
@@ -1098,7 +1101,7 @@ mod tests {
     use crate::id;
     use solana_sdk::{account::Account, native_token, pubkey::Pubkey, system_program};
     use solana_vote_program::vote_state;
-    use std::cell::RefCell;
+    use std::{cell::RefCell, iter::FromIterator};
 
     impl Meta {
         pub fn auto(authorized: &Pubkey) -> Self {
@@ -3688,4 +3691,108 @@ mod tests {
         // Test another staking action
         assert_eq!(stake_keyed_account.deactivate(&clock, &new_signers), Ok(()));
     }
+
+    #[test]
+    fn test_redelegate_consider_balance_changes() {
+        let initial_lamports = 4242424242;
+        let rent = Rent::default();
+        let rent_exempt_reserve = rent.minimum_balance(std::mem::size_of::<StakeState>());
+        let withdrawer_pubkey = Pubkey::new_unique();
+        let stake_lamports = rent_exempt_reserve + initial_lamports;
+
+        let meta = Meta {
+            rent_exempt_reserve,
+            ..Meta::auto(&withdrawer_pubkey)
+        };
+        let stake_account = Account::new_ref_data_with_space(
+            stake_lamports,
+            &StakeState::Initialized(meta),
+            std::mem::size_of::<StakeState>(),
+            &id(),
+        )
+        .expect("stake_account");
+        let stake_keyed_account = KeyedAccount::new(&withdrawer_pubkey, true, &stake_account);
+
+        let vote_pubkey = Pubkey::new_unique();
+        let vote_account = RefCell::new(vote_state::create_account(
+            &vote_pubkey,
+            &Pubkey::new_unique(),
+            0,
+            100,
+        ));
+        let vote_keyed_account = KeyedAccount::new(&vote_pubkey, false, &vote_account);
+
+        let signers = HashSet::from_iter(vec![withdrawer_pubkey]);
+        let config = Config::default();
+        let stake_history = StakeHistory::default();
+        let mut clock = Clock::default();
+        stake_keyed_account
+            .delegate(
+                &vote_keyed_account,
+                &clock,
+                &stake_history,
+                &config,
+                &signers,
+            )
+            .unwrap();
+
+        clock.epoch += 1;
+        stake_keyed_account.deactivate(&clock, &signers).unwrap();
+
+        clock.epoch += 1;
+        let to = Pubkey::new_unique();
+        let to_account = Account::new_ref(1, 0, &system_program::id());
+        let to_keyed_account = KeyedAccount::new(&to, false, &to_account);
+        let withdraw_lamports = initial_lamports / 2;
+        stake_keyed_account
+            .withdraw(
+                withdraw_lamports,
+                &to_keyed_account,
+                &clock,
+                &stake_history,
+                &stake_keyed_account,
+                None,
+            )
+            .unwrap();
+        let expected_balance = rent_exempt_reserve + initial_lamports - withdraw_lamports;
+        assert_eq!(stake_keyed_account.lamports().unwrap(), expected_balance);
+
+        clock.epoch += 1;
+        stake_keyed_account
+            .delegate(
+                &vote_keyed_account,
+                &clock,
+                &stake_history,
+                &config,
+                &signers,
+            )
+            .unwrap();
+        let stake = StakeState::stake_from(&stake_account.borrow()).unwrap();
+        assert_eq!(
+            stake.delegation.stake,
+            stake_keyed_account.lamports().unwrap() - rent_exempt_reserve,
+        );
+
+        clock.epoch += 1;
+        stake_keyed_account.deactivate(&clock, &signers).unwrap();
+
+        // Out of band deposit
+        stake_keyed_account.try_account_ref_mut().unwrap().lamports += withdraw_lamports;
+
+        clock.epoch += 1;
+        stake_keyed_account
+            .delegate(
+                &vote_keyed_account,
+                &clock,
+                &stake_history,
+                &config,
+                &signers,
+            )
+            .unwrap();
+        let stake = StakeState::stake_from(&stake_account.borrow()).unwrap();
+        assert_eq!(
+            stake.delegation.stake,
+            stake_keyed_account.lamports().unwrap() - rent_exempt_reserve,
+        );
+    }
 }