Add StakeInstruction::Merge (#10503) (#10506)

automerge

programs/stake/src/stake_instruction.rs

@@ -33,6 +33,12 @@ pub enum StakeError {
 
     #[error("split amount is more than is staked")]
     InsufficientStake,
+
+    #[error("stake account with activated stake cannot be merged")]
+    MergeActivatedStake,
+
+    #[error("stake account merge failed due to different authority or lockups")]
+    MergeMismatch,
 }
 
 impl<E> DecodeError<E> for StakeError {
@@ -124,6 +130,17 @@ pub enum StakeInstruction {
     ///    0 - initialized StakeAccount
     ///
     SetLockup(LockupArgs),
+
+    /// Merge two stake accounts. Both accounts must be deactivated and have identical lockup and
+    /// authority keys.
+    ///
+    /// # Account references
+    ///   0. [WRITE] Destination stake account for the merge
+    ///   1. [WRITE] Source stake account for to merge.  This account will be drained
+    ///   2. [] Clock sysvar
+    ///   3. [] Stake history sysvar that carries stake warmup/cooldown history
+    ///   4. [SIGNER] Stake authority
+    Merge,
 }
 
 #[derive(Default, Debug, Serialize, Deserialize, PartialEq, Clone, Copy)]
@@ -251,6 +268,26 @@ pub fn split_with_seed(
     ]
 }
 
+pub fn merge(
+    destination_stake_pubkey: &Pubkey,
+    source_stake_pubkey: &Pubkey,
+    authorized_pubkey: &Pubkey,
+) -> Vec<Instruction> {
+    let account_metas = vec![
+        AccountMeta::new(*destination_stake_pubkey, false),
+        AccountMeta::new(*source_stake_pubkey, false),
+        AccountMeta::new_readonly(sysvar::clock::id(), false),
+        AccountMeta::new_readonly(sysvar::stake_history::id(), false),
+        AccountMeta::new_readonly(*authorized_pubkey, true),
+    ];
+
+    vec![Instruction::new(
+        id(),
+        &StakeInstruction::Merge,
+        account_metas,
+    )]
+}
+
 pub fn create_account_and_delegate_stake(
     from_pubkey: &Pubkey,
     stake_pubkey: &Pubkey,
@@ -407,6 +444,15 @@ pub fn process_instruction(
             let split_stake = &next_keyed_account(keyed_accounts)?;
             me.split(lamports, split_stake, &signers)
         }
+        StakeInstruction::Merge => {
+            let source_stake = &next_keyed_account(keyed_accounts)?;
+            me.merge(
+                source_stake,
+                &Clock::from_keyed_account(next_keyed_account(keyed_accounts)?)?,
+                &StakeHistory::from_keyed_account(next_keyed_account(keyed_accounts)?)?,
+                &signers,
+            )
+        }
 
         StakeInstruction::Withdraw(lamports) => {
             let to = &next_keyed_account(keyed_accounts)?;
@@ -500,6 +546,12 @@ mod tests {
             ),
             Err(InstructionError::InvalidAccountData),
         );
+        assert_eq!(
+            process_instruction(
+                &merge(&Pubkey::default(), &Pubkey::default(), &Pubkey::default(),)[0]
+            ),
+            Err(InstructionError::InvalidAccountData),
+        );
         assert_eq!(
             process_instruction(
                 &split_with_seed(

programs/stake/src/stake_state.rs

@@ -543,6 +543,13 @@ pub trait StakeAccount {
         split_stake: &KeyedAccount,
         signers: &HashSet<Pubkey>,
     ) -> Result<(), InstructionError>;
+    fn merge(
+        &self,
+        source_stake: &KeyedAccount,
+        clock: &Clock,
+        stake_history: &StakeHistory,
+        signers: &HashSet<Pubkey>,
+    ) -> Result<(), InstructionError>;
     fn withdraw(
         &self,
         lamports: u64,
@@ -726,6 +733,51 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
         }
     }
 
+    fn merge(
+        &self,
+        source_stake: &KeyedAccount,
+        clock: &Clock,
+        stake_history: &StakeHistory,
+        signers: &HashSet<Pubkey>,
+    ) -> Result<(), InstructionError> {
+        let meta = match self.state()? {
+            StakeState::Stake(meta, stake) => {
+                // stake must be fully de-activated
+                if stake.stake(clock.epoch, Some(stake_history)) != 0 {
+                    return Err(StakeError::MergeActivatedStake.into());
+                }
+                meta
+            }
+            StakeState::Initialized(meta) => meta,
+            _ => return Err(InstructionError::InvalidAccountData),
+        };
+        // Authorized staker is allowed to split/merge accounts
+        meta.authorized.check(signers, StakeAuthorize::Staker)?;
+
+        let source_meta = match source_stake.state()? {
+            StakeState::Stake(meta, stake) => {
+                // stake must be fully de-activated
+                if stake.stake(clock.epoch, Some(stake_history)) != 0 {
+                    return Err(StakeError::MergeActivatedStake.into());
+                }
+                meta
+            }
+            StakeState::Initialized(meta) => meta,
+            _ => return Err(InstructionError::InvalidAccountData),
+        };
+
+        // Meta must match for both accounts
+        if meta != source_meta {
+            return Err(StakeError::MergeMismatch.into());
+        }
+
+        // Drain the source stake account
+        let lamports = source_stake.lamports()?;
+        source_stake.try_account_ref_mut()?.lamports -= lamports;
+        self.try_account_ref_mut()?.lamports += lamports;
+        Ok(())
+    }
+
     fn withdraw(
         &self,
         lamports: u64,
@@ -2518,6 +2570,16 @@ mod tests {
                 ..Stake::default()
             }
         }
+        fn just_bootstrap_stake(stake: u64) -> Self {
+            Self {
+                delegation: Delegation {
+                    stake,
+                    activation_epoch: std::u64::MAX,
+                    ..Delegation::default()
+                },
+                ..Stake::default()
+            }
+        }
     }
 
     #[test]
@@ -2844,6 +2906,249 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_merge() {
+        let stake_pubkey = Pubkey::new_rand();
+        let source_stake_pubkey = Pubkey::new_rand();
+        let authorized_pubkey = Pubkey::new_rand();
+        let stake_lamports = 42;
+
+        let signers = vec![authorized_pubkey].into_iter().collect();
+
+        for state in &[
+            StakeState::Initialized(Meta::auto(&authorized_pubkey)),
+            StakeState::Stake(
+                Meta::auto(&authorized_pubkey),
+                Stake::just_stake(stake_lamports),
+            ),
+        ] {
+            for source_state in &[
+                StakeState::Initialized(Meta::auto(&authorized_pubkey)),
+                StakeState::Stake(
+                    Meta::auto(&authorized_pubkey),
+                    Stake::just_stake(stake_lamports),
+                ),
+            ] {
+                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);
+
+                let source_stake_account = Account::new_ref_data_with_space(
+                    stake_lamports,
+                    source_state,
+                    std::mem::size_of::<StakeState>(),
+                    &id(),
+                )
+                .expect("source_stake_account");
+                let source_stake_keyed_account =
+                    KeyedAccount::new(&source_stake_pubkey, true, &source_stake_account);
+
+                // Authorized staker signature required...
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &HashSet::new()
+                    ),
+                    Err(InstructionError::MissingRequiredSignature)
+                );
+
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &signers
+                    ),
+                    Ok(())
+                );
+
+                // check lamports
+                assert_eq!(
+                    stake_keyed_account.account.borrow().lamports,
+                    stake_lamports * 2
+                );
+                assert_eq!(source_stake_keyed_account.account.borrow().lamports, 0);
+            }
+        }
+    }
+
+    #[test]
+    fn test_merge_incorrect_authorized_staker() {
+        let stake_pubkey = Pubkey::new_rand();
+        let source_stake_pubkey = Pubkey::new_rand();
+        let authorized_pubkey = Pubkey::new_rand();
+        let wrong_authorized_pubkey = Pubkey::new_rand();
+        let stake_lamports = 42;
+
+        let signers = vec![authorized_pubkey].into_iter().collect();
+        let wrong_signers = vec![wrong_authorized_pubkey].into_iter().collect();
+
+        for state in &[
+            StakeState::Initialized(Meta::auto(&authorized_pubkey)),
+            StakeState::Stake(
+                Meta::auto(&authorized_pubkey),
+                Stake::just_stake(stake_lamports),
+            ),
+        ] {
+            for source_state in &[
+                StakeState::Initialized(Meta::auto(&wrong_authorized_pubkey)),
+                StakeState::Stake(
+                    Meta::auto(&wrong_authorized_pubkey),
+                    Stake::just_stake(stake_lamports),
+                ),
+            ] {
+                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);
+
+                let source_stake_account = Account::new_ref_data_with_space(
+                    stake_lamports,
+                    source_state,
+                    std::mem::size_of::<StakeState>(),
+                    &id(),
+                )
+                .expect("source_stake_account");
+                let source_stake_keyed_account =
+                    KeyedAccount::new(&source_stake_pubkey, true, &source_stake_account);
+
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &wrong_signers,
+                    ),
+                    Err(InstructionError::MissingRequiredSignature)
+                );
+
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &signers,
+                    ),
+                    Err(StakeError::MergeMismatch.into())
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_merge_invalid_account_data() {
+        let stake_pubkey = Pubkey::new_rand();
+        let source_stake_pubkey = Pubkey::new_rand();
+        let authorized_pubkey = Pubkey::new_rand();
+        let stake_lamports = 42;
+        let signers = vec![authorized_pubkey].into_iter().collect();
+
+        for state in &[
+            StakeState::Uninitialized,
+            StakeState::RewardsPool,
+            StakeState::Initialized(Meta::auto(&authorized_pubkey)),
+            StakeState::Stake(
+                Meta::auto(&authorized_pubkey),
+                Stake::just_stake(stake_lamports),
+            ),
+        ] {
+            for source_state in &[StakeState::Uninitialized, StakeState::RewardsPool] {
+                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);
+
+                let source_stake_account = Account::new_ref_data_with_space(
+                    stake_lamports,
+                    source_state,
+                    std::mem::size_of::<StakeState>(),
+                    &id(),
+                )
+                .expect("source_stake_account");
+                let source_stake_keyed_account =
+                    KeyedAccount::new(&source_stake_pubkey, true, &source_stake_account);
+
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &signers,
+                    ),
+                    Err(InstructionError::InvalidAccountData)
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_merge_active_stake() {
+        let stake_pubkey = Pubkey::new_rand();
+        let source_stake_pubkey = Pubkey::new_rand();
+        let authorized_pubkey = Pubkey::new_rand();
+        let stake_lamports = 42;
+
+        let signers = vec![authorized_pubkey].into_iter().collect();
+
+        for state in &[
+            StakeState::Initialized(Meta::auto(&authorized_pubkey)),
+            StakeState::Stake(
+                Meta::auto(&authorized_pubkey),
+                Stake::just_bootstrap_stake(stake_lamports),
+            ),
+        ] {
+            for source_state in &[StakeState::Stake(
+                Meta::auto(&authorized_pubkey),
+                Stake::just_bootstrap_stake(stake_lamports),
+            )] {
+                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);
+
+                let source_stake_account = Account::new_ref_data_with_space(
+                    stake_lamports,
+                    source_state,
+                    std::mem::size_of::<StakeState>(),
+                    &id(),
+                )
+                .expect("source_stake_account");
+                let source_stake_keyed_account =
+                    KeyedAccount::new(&source_stake_pubkey, true, &source_stake_account);
+
+                // Authorized staker signature required...
+                assert_eq!(
+                    stake_keyed_account.merge(
+                        &source_stake_keyed_account,
+                        &Clock::default(),
+                        &StakeHistory::default(),
+                        &signers,
+                    ),
+                    Err(StakeError::MergeActivatedStake.into())
+                );
+            }
+        }
+    }
+
     #[test]
     fn test_lockup_is_expired() {
         let custodian = Pubkey::new_rand();

stake-monitor/src/lib.rs

@@ -21,6 +21,7 @@ pub enum AccountOperation {
     SplitDestination,
     SystemAccountEnroll,
     FailedToMaintainMinimumBalance,
+    MergeSource,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -172,6 +173,29 @@ fn process_transaction(
                             }
                         }
                     }
+                    StakeInstruction::Merge => {
+                        // Merge invalidates the source account, but does not affect the
+                        // destination account
+                        let source_merge_account_index = instruction.accounts[1] as usize;
+
+                        let source_stake_pubkey =
+                            message.account_keys[source_merge_account_index].to_string();
+
+                        if let Some(mut source_account_info) =
+                            accounts.get_mut(&source_stake_pubkey)
+                        {
+                            if source_account_info.compliant_since.is_some() {
+                                source_account_info.compliant_since = None;
+                                source_account_info
+                                    .transactions
+                                    .push(AccountTransactionInfo {
+                                        op: AccountOperation::MergeSource,
+                                        slot,
+                                        signature: signature.clone(),
+                                    });
+                            }
+                        }
+                    }
                     StakeInstruction::Withdraw(_) => {
                         // Withdrawing is not permitted