stake-pool: Split from stake account during removal (#2367)

If the stake pool gives over the validator stake account to the staker,
they may keep it and make it impossible to re-add that validator in the
future.

Split the whole amount into a new stake account on removal.

Author: joncinque

stake-pool/cli/src/main.rs

@@ -152,6 +152,33 @@ fn checked_transaction_with_signers<T: Signers>(
     Ok(transaction)
 }
 
+fn new_stake_account(
+    fee_payer: &Pubkey,
+    instructions: &mut Vec<Instruction>,
+    lamports: u64,
+) -> Keypair {
+    // Account for tokens not specified, creating one
+    let stake_receiver_keypair = Keypair::new();
+    let stake_receiver_pubkey = stake_receiver_keypair.pubkey();
+    println!(
+        "Creating account to receive stake {}",
+        stake_receiver_pubkey
+    );
+
+    instructions.push(
+        // Creating new account
+        system_instruction::create_account(
+            fee_payer,
+            &stake_receiver_pubkey,
+            lamports,
+            STAKE_STATE_LEN as u64,
+            &stake_program::id(),
+        ),
+    );
+
+    stake_receiver_keypair
+}
+
 #[allow(clippy::too_many_arguments)]
 fn command_create_pool(
     config: &Config,
@@ -423,6 +450,7 @@ fn command_vsa_remove(
     stake_pool_address: &Pubkey,
     vote_account: &Pubkey,
     new_authority: &Option<Pubkey>,
+    stake_receiver: &Option<Pubkey>,
 ) -> CommandResult {
     if !config.no_update {
         command_update(config, stake_pool_address, false, false)?;
@@ -437,6 +465,20 @@ fn command_vsa_remove(
 
     let stake_pool = get_stake_pool(&config.rpc_client, stake_pool_address)?;
 
+    let mut instructions = vec![];
+    let mut stake_keypair = None;
+
+    let stake_receiver = stake_receiver.unwrap_or_else(|| {
+        let new_stake_keypair = new_stake_account(
+            &config.fee_payer.pubkey(),
+            &mut instructions,
+            /* stake_receiver_account_balance = */ 0,
+        );
+        let stake_pubkey = new_stake_keypair.pubkey();
+        stake_keypair = Some(new_stake_keypair);
+        stake_pubkey
+    });
+
     let staker_pubkey = config.staker.pubkey();
     let new_authority = new_authority.as_ref().unwrap_or(&staker_pubkey);
 
@@ -446,6 +488,9 @@ fn command_vsa_remove(
         .ok_or("Vote account not found in validator list")?;
 
     let mut signers = vec![config.fee_payer.as_ref(), config.staker.as_ref()];
+    if let Some(stake_keypair) = stake_keypair.as_ref() {
+        signers.push(stake_keypair);
+    }
     unique_signers!(signers);
     let transaction = checked_transaction_with_signers(
         config,
@@ -458,6 +503,7 @@ fn command_vsa_remove(
                 vote_account,
                 new_authority,
                 validator_stake_info.transient_seed_suffix_start,
+                &stake_receiver,
             ),
         ],
         &signers,
@@ -1270,33 +1316,19 @@ fn command_withdraw(
 
         // Use separate mutable variable because withdraw might create a new account
         let stake_receiver = stake_receiver_param.unwrap_or_else(|| {
-            // Account for tokens not specified, creating one
-            let stake_receiver_account = Keypair::new(); // Will be added to signers if creating new account
-            let stake_receiver_pubkey = stake_receiver_account.pubkey();
-            println!(
-                "Creating account to receive stake {}",
-                stake_receiver_pubkey
-            );
-
             let stake_receiver_account_balance = config
                 .rpc_client
                 .get_minimum_balance_for_rent_exemption(STAKE_STATE_LEN)
                 .unwrap();
-
-            instructions.push(
-                // Creating new account
-                system_instruction::create_account(
-                    &config.fee_payer.pubkey(),
-                    &stake_receiver_pubkey,
-                    stake_receiver_account_balance,
-                    STAKE_STATE_LEN as u64,
-                    &stake_program::id(),
-                ),
+            let stake_keypair = new_stake_account(
+                &config.fee_payer.pubkey(),
+                &mut instructions,
+                stake_receiver_account_balance,
             );
-
+            let stake_pubkey = stake_keypair.pubkey();
             total_rent_free_balances += stake_receiver_account_balance;
-            new_stake_keypairs.push(stake_receiver_account);
-            stake_receiver_pubkey
+            new_stake_keypairs.push(stake_keypair);
+            stake_pubkey
         });
 
         instructions.push(spl_stake_pool::instruction::withdraw_stake(
@@ -1753,6 +1785,14 @@ fn main() {
                     .help("New authority to set as Staker and Withdrawer in the stake account removed from the pool.
                           Defaults to the client keypair."),
             )
+            .arg(
+                Arg::with_name("stake_receiver")
+                    .long("stake-receiver")
+                    .validator(is_pubkey)
+                    .value_name("ADDRESS")
+                    .takes_value(true)
+                    .help("Stake account to receive SOL from the stake pool. Defaults to a new stake account."),
+            )
         )
         .subcommand(SubCommand::with_name("increase-validator-stake")
             .about("Increase stake to a validator, drawing from the stake pool reserve. Must be signed by the pool staker.")
@@ -2318,8 +2358,15 @@ fn main() {
         ("remove-validator", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
             let vote_account = pubkey_of(arg_matches, "vote_account").unwrap();
-            let new_authority: Option<Pubkey> = pubkey_of(arg_matches, "new_authority");
-            command_vsa_remove(&config, &stake_pool_address, &vote_account, &new_authority)
+            let new_authority = pubkey_of(arg_matches, "new_authority");
+            let stake_receiver = pubkey_of(arg_matches, "stake_receiver");
+            command_vsa_remove(
+                &config,
+                &stake_pool_address,
+                &vote_account,
+                &new_authority,
+                &stake_receiver,
+            )
         }
         ("increase-validator-stake", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();

stake-pool/program/src/instruction.rs

@@ -127,8 +127,9 @@ pub enum StakePoolInstruction {
     ///   4. `[w]` Validator stake list storage account
     ///   5. `[w]` Stake account to remove from the pool
     ///   6. `[]` Transient stake account, to check that that we're not trying to activate
-    ///   7. '[]' Sysvar clock
-    ///   8. `[]` Stake program id,
+    ///   7. `[w]` Destination stake account, to receive the minimum SOL from the validator stake account. Must be
+    ///   8. `[]` Sysvar clock
+    ///   9. `[]` Stake program id,
     RemoveValidatorFromPool,
 
     /// (Staker only) Decrease active stake on a validator, eventually moving it to the reserve
@@ -483,6 +484,7 @@ pub fn remove_validator_from_pool(
     validator_list: &Pubkey,
     stake_account: &Pubkey,
     transient_stake_account: &Pubkey,
+    destination_stake_account: &Pubkey,
 ) -> Instruction {
     let accounts = vec![
         AccountMeta::new(*stake_pool, false),
@@ -492,6 +494,7 @@ pub fn remove_validator_from_pool(
         AccountMeta::new(*validator_list, false),
         AccountMeta::new(*stake_account, false),
         AccountMeta::new_readonly(*transient_stake_account, false),
+        AccountMeta::new(*destination_stake_account, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(stake_program::id(), false),
     ];
@@ -658,6 +661,7 @@ pub fn remove_validator_from_pool_with_vote(
     vote_account_address: &Pubkey,
     new_stake_account_authority: &Pubkey,
     transient_stake_seed: u64,
+    destination_stake_address: &Pubkey,
 ) -> Instruction {
     let pool_withdraw_authority =
         find_withdraw_authority_program_address(program_id, stake_pool_address).0;
@@ -678,6 +682,7 @@ pub fn remove_validator_from_pool_with_vote(
         &stake_pool.validator_list,
         &stake_account_address,
         &transient_stake_account,
+        destination_stake_address,
     )
 }
 

stake-pool/program/src/processor.rs

@@ -886,6 +886,7 @@ impl Processor {
         let validator_list_info = next_account_info(account_info_iter)?;
         let stake_account_info = next_account_info(account_info_iter)?;
         let transient_stake_account_info = next_account_info(account_info_iter)?;
+        let destination_stake_account_info = next_account_info(account_info_iter)?;
         let clock_info = next_account_info(account_info_iter)?;
         let clock = &Clock::from_account_info(clock_info)?;
         let stake_program_info = next_account_info(account_info_iter)?;
@@ -994,12 +995,23 @@ impl Processor {
             StakeStatus::ReadyForRemoval
         };
 
-        Self::stake_authorize_signed(
+        // split whole thing into destination stake account
+        Self::stake_split(
             stake_pool_info.key,
             stake_account_info.clone(),
             withdraw_authority_info.clone(),
             AUTHORITY_WITHDRAW,
             stake_pool.stake_withdraw_bump_seed,
+            stake_account_info.lamports(),
+            destination_stake_account_info.clone(),
+        )?;
+
+        Self::stake_authorize_signed(
+            stake_pool_info.key,
+            destination_stake_account_info.clone(),
+            withdraw_authority_info.clone(),
+            AUTHORITY_WITHDRAW,
+            stake_pool.stake_withdraw_bump_seed,
             new_stake_authority_info.key,
             clock_info.clone(),
             stake_program_info.clone(),

stake-pool/program/tests/helpers/mod.rs

@@ -1093,20 +1093,31 @@ impl StakePoolAccounts {
         new_authority: &Pubkey,
         validator_stake: &Pubkey,
         transient_stake: &Pubkey,
+        destination_stake: &Keypair,
     ) -> Option<TransportError> {
         let transaction = Transaction::new_signed_with_payer(
-            &[instruction::remove_validator_from_pool(
-                &id(),
-                &self.stake_pool.pubkey(),
-                &self.staker.pubkey(),
-                &self.withdraw_authority,
-                &new_authority,
-                &self.validator_list.pubkey(),
-                validator_stake,
-                transient_stake,
-            )],
+            &[
+                system_instruction::create_account(
+                    &payer.pubkey(),
+                    &destination_stake.pubkey(),
+                    0,
+                    std::mem::size_of::<stake_program::StakeState>() as u64,
+                    &stake_program::id(),
+                ),
+                instruction::remove_validator_from_pool(
+                    &id(),
+                    &self.stake_pool.pubkey(),
+                    &self.staker.pubkey(),
+                    &self.withdraw_authority,
+                    &new_authority,
+                    &self.validator_list.pubkey(),
+                    validator_stake,
+                    transient_stake,
+                    &destination_stake.pubkey(),
+                ),
+            ],
             Some(&payer.pubkey()),
-            &[payer, &self.staker],
+            &[payer, &self.staker, destination_stake],
             *recent_blockhash,
         );
         banks_client.process_transaction(transaction).await.err()

stake-pool/program/tests/huge_pool.rs

@@ -404,6 +404,7 @@ async fn remove_validator_from_pool() {
     );
 
     let new_authority = Pubkey::new_unique();
+    let destination_stake = Keypair::new();
     let error = stake_pool_accounts
         .remove_validator_from_pool(
             &mut context.banks_client,
@@ -412,6 +413,7 @@ async fn remove_validator_from_pool() {
             &new_authority,
             &stake_address,
             &transient_stake_address,
+            &destination_stake,
         )
         .await;
     assert!(error.is_none());
@@ -431,6 +433,7 @@ async fn remove_validator_from_pool() {
     );
 
     let new_authority = Pubkey::new_unique();
+    let destination_stake = Keypair::new();
     let error = stake_pool_accounts
         .remove_validator_from_pool(
             &mut context.banks_client,
@@ -439,6 +442,7 @@ async fn remove_validator_from_pool() {
             &new_authority,
             &stake_address,
             &transient_stake_address,
+            &destination_stake,
         )
         .await;
     assert!(error.is_none());
@@ -455,6 +459,7 @@ async fn remove_validator_from_pool() {
     );
 
     let new_authority = Pubkey::new_unique();
+    let destination_stake = Keypair::new();
     let error = stake_pool_accounts
         .remove_validator_from_pool(
             &mut context.banks_client,
@@ -463,6 +468,7 @@ async fn remove_validator_from_pool() {
             &new_authority,
             &stake_address,
             &transient_stake_address,
+            &destination_stake,
         )
         .await;
     assert!(error.is_none());

stake-pool/program/tests/set_preferred.rs

@@ -223,6 +223,7 @@ async fn fail_ready_for_removal() {
         transient_stake_seed,
     );
     let new_authority = Pubkey::new_unique();
+    let destination_stake = Keypair::new();
     let remove_err = stake_pool_accounts
         .remove_validator_from_pool(
             &mut banks_client,
@@ -231,6 +232,7 @@ async fn fail_ready_for_removal() {
             &new_authority,
             &validator_stake_account.stake_account,
             &transient_stake_address,
+            &destination_stake,
         )
         .await;
     assert!(remove_err.is_none());

stake-pool/program/tests/update_validator_list_balance.rs

@@ -6,7 +6,11 @@ use {
     helpers::*,
     solana_program::{borsh::try_from_slice_unchecked, program_pack::Pack, pubkey::Pubkey},
     solana_program_test::*,
-    solana_sdk::{signature::Signer, system_instruction, transaction::Transaction},
+    solana_sdk::{
+        signature::{Keypair, Signer},
+        system_instruction,
+        transaction::Transaction,
+    },
     spl_stake_pool::{
         find_transient_stake_program_address, id, instruction, stake_program,
         state::{StakePool, StakeStatus, ValidatorList},
@@ -462,6 +466,7 @@ async fn merge_transient_stake_after_remove() {
     let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
     let deactivated_lamports = lamports;
     let new_authority = Pubkey::new_unique();
+    let destination_stake = Keypair::new();
     // Decrease and remove all validators
     for stake_account in &stake_accounts {
         let error = stake_pool_accounts
@@ -484,6 +489,7 @@ async fn merge_transient_stake_after_remove() {
                 &new_authority,
                 &stake_account.stake_account,
                 &stake_account.transient_stake_account,
+                &destination_stake,
             )
             .await;
         assert!(error.is_none());