From 50710473a8ee89ab6ca848f0b5afa430148a44b5 Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Mon, 21 Dec 2020 19:54:49 -0700 Subject: [PATCH] Deinitialize stake data upon zero balance --- programs/stake/src/stake_state.rs | 122 +++++++++++++++++++----------- 1 file changed, 77 insertions(+), 45 deletions(-) diff --git a/programs/stake/src/stake_state.rs b/programs/stake/src/stake_state.rs index 6da143faf9..9a0f15e8f3 100644 --- a/programs/stake/src/stake_state.rs +++ b/programs/stake/src/stake_state.rs @@ -99,9 +99,17 @@ impl StakeState { } pub fn lockup(&self) -> Option { + self.meta().map(|meta| meta.lockup) + } + + pub fn meta_from(account: &Account) -> Option { + Self::from(account).and_then(|state: Self| state.meta()) + } + + pub fn meta(&self) -> Option { match self { - StakeState::Stake(meta, _stake) => Some(meta.lockup), - StakeState::Initialized(meta) => Some(meta.lockup), + StakeState::Stake(meta, _stake) => Some(*meta), + StakeState::Initialized(meta) => Some(*meta), _ => None, } } @@ -1054,6 +1062,11 @@ impl<'a> StakeAccount for KeyedAccount<'a> { _ => return Err(InstructionError::InvalidAccountData), } + // Deinitialize state upon zero balance + if lamports == self.lamports()? { + self.set_state(&StakeState::Uninitialized)?; + } + split.try_account_ref_mut()?.lamports += lamports; self.try_account_ref_mut()?.lamports -= lamports; Ok(()) @@ -1090,6 +1103,9 @@ impl<'a> StakeAccount for KeyedAccount<'a> { self.set_state(&merged_state)?; } + // Source is about to be drained, deinitialize its state + source_account.set_state(&StakeState::Uninitialized)?; + // Drain the source stake account let lamports = source_account.lamports()?; source_account.try_account_ref_mut()?.lamports -= lamports; @@ -1166,6 +1182,11 @@ impl<'a> StakeAccount for KeyedAccount<'a> { return Err(InstructionError::InsufficientFunds); } + // Deinitialize state upon zero balance + if lamports == self.lamports()? { + self.set_state(&StakeState::Uninitialized)?; + } + self.try_account_ref_mut()?.lamports -= lamports; to.try_account_ref_mut()?.lamports += lamports; Ok(()) @@ -2828,6 +2849,7 @@ mod tests { Ok(()) ); assert_eq!(stake_account.borrow().lamports, 0); + assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized)); // reset balance stake_account.borrow_mut().lamports = stake_lamports; @@ -2953,6 +2975,7 @@ mod tests { Ok(()) ); assert_eq!(stake_account.borrow().lamports, 0); + assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized)); } #[test] @@ -3128,6 +3151,7 @@ mod tests { ), Ok(()) ); + assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized)); } #[test] @@ -3184,6 +3208,7 @@ mod tests { ), Ok(()) ); + assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized)); } } @@ -3600,6 +3625,7 @@ mod tests { ), Ok(()) ); + assert_eq!(stake_keyed_account.state(), Ok(StakeState::Uninitialized)); } #[test] @@ -4498,7 +4524,7 @@ mod tests { match state { StakeState::Initialized(_) => { assert_eq!(Ok(*state), split_stake_keyed_account.state()); - assert_eq!(Ok(*state), stake_keyed_account.state()); + assert_eq!(Ok(StakeState::Uninitialized), stake_keyed_account.state()); } StakeState::Stake(meta, stake) => { assert_eq!( @@ -4514,19 +4540,7 @@ mod tests { )), split_stake_keyed_account.state() ); - assert_eq!( - Ok(StakeState::Stake( - *meta, - Stake { - delegation: Delegation { - stake: 0, - ..stake.delegation - }, - ..*stake - } - )), - stake_keyed_account.state() - ); + assert_eq!(Ok(StakeState::Uninitialized), stake_keyed_account.state()); } _ => unreachable!(), } @@ -4608,19 +4622,7 @@ mod tests { )), split_stake_keyed_account.state() ); - assert_eq!( - Ok(StakeState::Stake( - meta, - Stake { - delegation: Delegation { - stake: 0, - ..stake.delegation - }, - ..stake - } - )), - stake_keyed_account.state() - ); + assert_eq!(Ok(StakeState::Uninitialized), stake_keyed_account.state()); } } } @@ -4722,9 +4724,9 @@ mod tests { Ok(StakeState::Initialized(expected_split_meta)), split_stake_keyed_account.state() ); - assert_eq!(Ok(*state), stake_keyed_account.state()); + assert_eq!(Ok(StakeState::Uninitialized), stake_keyed_account.state()); } - StakeState::Stake(meta, stake) => { + StakeState::Stake(_meta, stake) => { // Expected stake should reflect original stake amount so that extra lamports // from the rent_exempt_reserve inequality do not magically activate let expected_stake = stake_lamports - rent_exempt_reserve; @@ -4748,19 +4750,7 @@ mod tests { + expected_rent_exempt_reserve + (rent_exempt_reserve - expected_rent_exempt_reserve) ); - assert_eq!( - Ok(StakeState::Stake( - *meta, - Stake { - delegation: Delegation { - stake: 0, - ..stake.delegation - }, - ..*stake - } - )), - stake_keyed_account.state() - ); + assert_eq!(Ok(StakeState::Uninitialized), stake_keyed_account.state()); } _ => unreachable!(), } @@ -4836,6 +4826,44 @@ mod tests { stake_lamports * 2 ); assert_eq!(source_stake_keyed_account.account.borrow().lamports, 0); + + // check state + match state { + StakeState::Initialized(meta) => { + assert_eq!( + stake_keyed_account.state(), + Ok(StakeState::Initialized(*meta)), + ); + } + StakeState::Stake(meta, stake) => { + let expected_stake = stake.delegation.stake + + source_state + .stake() + .map(|stake| stake.delegation.stake) + .unwrap_or_else(|| { + stake_lamports + - source_state.meta().unwrap().rent_exempt_reserve + }); + assert_eq!( + stake_keyed_account.state(), + Ok(StakeState::Stake( + *meta, + Stake { + delegation: Delegation { + stake: expected_stake, + ..stake.delegation + }, + ..*stake + } + )), + ); + } + _ => unreachable!(), + } + assert_eq!( + source_stake_keyed_account.state(), + Ok(StakeState::Uninitialized) + ); } } } @@ -5127,7 +5155,11 @@ mod tests { let test_source_keyed = KeyedAccount::new(source_account.unsigned_key(), true, &test_source_account); - test_stake_keyed.merge(&test_source_keyed, clock, stake_history, signers) + let result = test_stake_keyed.merge(&test_source_keyed, clock, stake_history, signers); + if result.is_ok() { + assert_eq!(test_source_keyed.state(), Ok(StakeState::Uninitialized),); + } + result } // stake activation epoch, source initialized succeeds