Clean & check stake (#15363)

* Add failing test

* Fix test

* Clean up redendant if case

* Demonstrate withdrawal boundaries

* Update test to fail on conditions that should be acceptable

* Fix test

* Add test for larger stake source

* Mirror changes for undelegated accounts

* Extra stake checks

* Split accounts must be the right size

Co-authored-by: Stephen Akridge <sakridge@gmail.com>

programs/stake/src/stake_state.rs

@@ -1057,6 +1057,9 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
         if split.owner()? != id() {
             return Err(InstructionError::IncorrectProgramId);
         }
+        if split.data_len()? != std::mem::size_of::<StakeState>() {
+            return Err(InstructionError::InvalidAccountData);
+        }
 
         if let StakeState::Uninitialized = split.state()? {
             // verify enough account lamports
@@ -1073,17 +1076,12 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
                         split.data_len()? as u64,
                     );
 
-                    // verify enough lamports for rent in new split account
+                    // verify enough lamports for rent and more than 0 stake in new split account
-                    if lamports < split_rent_exempt_reserve.saturating_sub(split.lamports()?)
+                    if lamports <= split_rent_exempt_reserve.saturating_sub(split.lamports()?)
-                        // verify full withdrawal can cover rent in new split account
-                        || (lamports < split_rent_exempt_reserve && lamports == self.lamports()?)
                         // if not full withdrawal
                         || (lamports != self.lamports()?
                             // verify more than 0 stake left in previous stake
-                            && (lamports + meta.rent_exempt_reserve >= self.lamports()?
+                            && checked_add(lamports, meta.rent_exempt_reserve)? >= self.lamports()?)
-                                // and verify more than 0 stake in new split account
-                                || lamports
-                                    <= split_rent_exempt_reserve.saturating_sub(split.lamports()?)))
                     {
                         return Err(InstructionError::InsufficientFunds);
                     }
@@ -1106,11 +1104,7 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
                         // different sizes.
                         let remaining_stake_delta =
                             lamports.saturating_sub(meta.rent_exempt_reserve);
-                        let split_stake_amount = std::cmp::min(
+                        (remaining_stake_delta, remaining_stake_delta)
-                            lamports - split_rent_exempt_reserve,
-                            remaining_stake_delta,
-                        );
-                        (remaining_stake_delta, split_stake_amount)
                     } else {
                         // Otherwise, the new split stake should reflect the entire split
                         // requested, less any lamports needed to cover the split_rent_exempt_reserve
@@ -1134,15 +1128,12 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
                         split.data_len()? as u64,
                     );
 
-                    // enough lamports for rent in new stake
+                    // enough lamports for rent and more than 0 stake in new split account
-                    if lamports < split_rent_exempt_reserve
+                    if lamports <= split_rent_exempt_reserve.saturating_sub(split.lamports()?)
-                    // if not full withdrawal
+                        // if not full withdrawal
-                    || (lamports != self.lamports()?
+                        || (lamports != self.lamports()?
-                        // verify more than 0 stake left in previous stake
+                            // verify more than 0 stake left in previous stake
-                        && (lamports + meta.rent_exempt_reserve >= self.lamports()?
+                            && checked_add(lamports, meta.rent_exempt_reserve)? >= self.lamports()?)
-                            // and verify more than 0 stake in new split account
-                            || lamports
-                                <= split_rent_exempt_reserve.saturating_sub(split.lamports()?)))
                     {
                         return Err(InstructionError::InsufficientFunds);
                     }
@@ -1429,16 +1420,18 @@ impl MergeKind {
             (Self::Inactive(_, _), Self::Inactive(_, _)) => None,
             (Self::Inactive(_, _), Self::ActivationEpoch(_, _)) => None,
             (Self::ActivationEpoch(meta, mut stake), Self::Inactive(_, source_lamports)) => {
-                stake.delegation.stake += source_lamports;
+                stake.delegation.stake = checked_add(stake.delegation.stake, source_lamports)?;
                 Some(StakeState::Stake(meta, stake))
             }
             (
                 Self::ActivationEpoch(meta, mut stake),
                 Self::ActivationEpoch(source_meta, source_stake),
             ) => {
-                let source_lamports =
+                let source_lamports = checked_add(
-                    source_meta.rent_exempt_reserve + source_stake.delegation.stake;
+                    source_meta.rent_exempt_reserve,
-                stake.delegation.stake += source_lamports;
+                    source_stake.delegation.stake,
+                )?;
+                stake.delegation.stake = checked_add(stake.delegation.stake, source_lamports)?;
                 Some(StakeState::Stake(meta, stake))
             }
             (Self::FullyActive(meta, mut stake), Self::FullyActive(_, source_stake)) => {
@@ -1446,7 +1439,8 @@ impl MergeKind {
                 // protect against the magic activation loophole. It will
                 // instead be moved into the destination account as extra,
                 // withdrawable `lamports`
-                stake.delegation.stake += source_stake.delegation.stake;
+                stake.delegation.stake =
+                    checked_add(stake.delegation.stake, source_stake.delegation.stake)?;
                 Some(StakeState::Stake(meta, stake))
             }
             _ => return Err(StakeError::MergeMismatch.into()),
@@ -4738,7 +4732,7 @@ mod tests {
     }
 
     #[test]
-    fn test_split_to_larger_account_edge_case() {
+    fn test_split_to_larger_account() {
         let stake_pubkey = solana_sdk::pubkey::new_rand();
         let rent = Rent::default();
         let rent_exempt_reserve = rent.minimum_balance(std::mem::size_of::<StakeState>());
@@ -4792,77 +4786,15 @@ mod tests {
             .expect("stake_account");
             let stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &stake_account);
 
-            // should return error when initial_balance < expected_rent_exempt_reserve
+            // should always return error when splitting to larger account
-            let split_attempt =
+            let split_result =
                 stake_keyed_account.split(split_amount, &split_stake_keyed_account, &signers);
-            if initial_balance < expected_rent_exempt_reserve {
+            assert_eq!(split_result, Err(InstructionError::InvalidAccountData));
-                assert_eq!(split_attempt, Err(InstructionError::InsufficientFunds));
-            } else {
-                assert_eq!(split_attempt, Ok(()));
-            }
-        }
-    }
 
-    #[test]
+            // Splitting 100% of source should not make a difference
-    fn test_split_100_percent_of_source_to_larger_account_edge_case() {
+            let split_result =
-        let stake_pubkey = solana_sdk::pubkey::new_rand();
+                stake_keyed_account.split(stake_lamports, &split_stake_keyed_account, &signers);
-        let rent = Rent::default();
+            assert_eq!(split_result, Err(InstructionError::InvalidAccountData));
-        let rent_exempt_reserve = rent.minimum_balance(std::mem::size_of::<StakeState>());
-        let stake_lamports = rent_exempt_reserve + 1;
-
-        let split_stake_pubkey = solana_sdk::pubkey::new_rand();
-        let signers = vec![stake_pubkey].into_iter().collect();
-
-        let meta = Meta {
-            authorized: Authorized::auto(&stake_pubkey),
-            rent_exempt_reserve,
-            ..Meta::default()
-        };
-
-        let state = StakeState::Stake(
-            meta,
-            Stake::just_stake(stake_lamports - rent_exempt_reserve),
-        );
-
-        let expected_rent_exempt_reserve = calculate_split_rent_exempt_reserve(
-            meta.rent_exempt_reserve,
-            std::mem::size_of::<StakeState>() as u64,
-            std::mem::size_of::<StakeState>() as u64 + 100,
-        );
-        assert!(expected_rent_exempt_reserve > stake_lamports);
-
-        let split_lamport_balances = vec![
-            0,
-            1,
-            expected_rent_exempt_reserve,
-            expected_rent_exempt_reserve + 1,
-        ];
-        for initial_balance in split_lamport_balances {
-            let split_stake_account = Account::new_ref_data_with_space(
-                initial_balance,
-                &StakeState::Uninitialized,
-                std::mem::size_of::<StakeState>() + 100,
-                &id(),
-            )
-            .expect("stake_account");
-
-            let split_stake_keyed_account =
-                KeyedAccount::new(&split_stake_pubkey, true, &split_stake_account);
-
-            let stake_account = Account::new_ref_data_with_space(
-                stake_lamports,
-                &state,
-                std::mem::size_of::<StakeState>(),
-                &id(),
-            )
-            .expect("stake_account");
-            let stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &stake_account);
-
-            // should return error
-            assert_eq!(
-                stake_keyed_account.split(stake_lamports, &split_stake_keyed_account, &signers),
-                Err(InstructionError::InsufficientFunds)
-            );
         }
     }
 
@@ -5052,8 +4984,7 @@ mod tests {
                 Stake::just_stake(stake_lamports - rent_exempt_reserve),
             ),
         ] {
-            // Test that splitting to a larger account with greater rent-exempt requirement fails
+            // Test that splitting to a larger account fails
-            // if split amount is too small
             let split_stake_account = Account::new_ref_data_with_space(
                 0,
                 &StakeState::Uninitialized,
@@ -5075,7 +5006,7 @@ mod tests {
 
             assert_eq!(
                 stake_keyed_account.split(stake_lamports, &split_stake_keyed_account, &signers),
-                Err(InstructionError::InsufficientFunds)
+                Err(InstructionError::InvalidAccountData)
             );
 
             // Test that splitting from a larger account to a smaller one works.