stake-pool: Check transient stake state explicitly during increase / decrease / redelegate

stake-pool/program/src/processor.rs

@@ -221,6 +221,42 @@ fn stake_is_inactive_without_history(stake: &stake::state::Stake, epoch: Epoch)
             && stake.delegation.deactivation_epoch == epoch)
 }
 
+/// Roughly checks if a stake account is deactivating or inactive
+fn check_if_stake_is_deactivating_or_inactive(
+    account_info: &AccountInfo,
+    vote_account_address: &Pubkey,
+) -> Result<(), ProgramError> {
+    let (_, stake) = get_stake_state(account_info)?;
+    if stake.delegation.deactivation_epoch == Epoch::MAX {
+        msg!(
+            "Existing stake {} delegated to {} is activating or active",
+            account_info.key,
+            vote_account_address
+        );
+        Err(StakePoolError::WrongStakeState.into())
+    } else {
+        Ok(())
+    }
+}
+
+/// Roughly checks if a stake account is activating or active
+fn check_if_stake_is_activating_or_active(
+    account_info: &AccountInfo,
+    vote_account_address: &Pubkey,
+) -> Result<(), ProgramError> {
+    let (_, stake) = get_stake_state(account_info)?;
+    if stake.delegation.deactivation_epoch != Epoch::MAX {
+        msg!(
+            "Existing stake {} delegated to {} is deactivating or inactive",
+            account_info.key,
+            vote_account_address
+        );
+        Err(StakePoolError::WrongStakeState.into())
+    } else {
+        Ok(())
+    }
+}
+
 /// Check that the stake state is correct: usable by the pool and delegated to
 /// the expected validator
 fn check_stake_state(
@@ -1336,8 +1372,10 @@ impl Processor {
                 );
                 return Err(ProgramError::InvalidSeeds);
             }
-            // Let the runtime check to see if the merge is valid, so there's no
-            // explicit check here that the transient stake is decreasing
+            check_if_stake_is_deactivating_or_inactive(
+                transient_stake_account_info,
+                &vote_account_address,
+            )?;
         }
 
         let stake_minimum_delegation = stake::tools::get_minimum_delegation()?;
@@ -1585,8 +1623,10 @@ impl Processor {
                 );
                 return Err(ProgramError::InvalidSeeds);
             }
-            // Let the runtime check to see if the merge is valid, so there's no
-            // explicit check here that the transient stake is increasing
+            check_if_stake_is_activating_or_active(
+                transient_stake_account_info,
+                vote_account_address,
+            )?;
         }
 
         check_validator_stake_account(
@@ -2036,6 +2076,10 @@ impl Processor {
                     destination_transient_stake_seed,
                     &stake_pool.lockup,
                 )?;
+                check_if_stake_is_activating_or_active(
+                    destination_transient_stake_account_info,
+                    vote_account_address,
+                )?;
                 Self::stake_merge(
                     stake_pool_info.key,
                     ephemeral_stake_account_info.clone(),

stake-pool/program/tests/decrease.rs

@@ -586,7 +586,7 @@ async fn fail_additional_with_increasing() {
         error,
         TransactionError::InstructionError(
             0,
-            InstructionError::Custom(stake::instruction::StakeError::MergeTransientStake as u32)
+            InstructionError::Custom(StakePoolError::WrongStakeState as u32)
         )
     );
 }

stake-pool/program/tests/increase.rs

@@ -585,7 +585,7 @@ async fn fail_additional_with_decreasing() {
         error,
         TransactionError::InstructionError(
             0,
-            InstructionError::Custom(StakeError::MergeTransientStake as u32)
+            InstructionError::Custom(StakePoolError::WrongStakeState as u32)
         )
     );
 }

stake-pool/program/tests/redelegate.rs

@@ -657,7 +657,7 @@ async fn fail_with_decreasing_stake() {
         error,
         TransactionError::InstructionError(
             0,
-            InstructionError::Custom(StakeError::MergeTransientStake as u32)
+            InstructionError::Custom(StakePoolError::WrongStakeState as u32)
         )
     );
 }