stake-pool: Add preferred deposit / withdraw validator accounts (#1831)

* stake-pool: Add set preferred validator instruction

* Add functionality for using the preferred deposit / withdraw

* Add CLI support

stake-pool/cli/src/main.rs

@@ -32,6 +32,7 @@ use {
         self,
         borsh::get_instance_packed_len,
         find_stake_program_address, find_withdraw_authority_program_address,
+        instruction::PreferredValidatorType,
         stake_program::{self, StakeState},
         state::{Fee, StakePool, ValidatorList},
     },
@@ -466,6 +467,34 @@ fn command_decrease_validator_stake(
     Ok(())
 }
 
+fn command_set_preferred_validator(
+    config: &Config,
+    stake_pool_address: &Pubkey,
+    preferred_type: PreferredValidatorType,
+    vote_address: Option<Pubkey>,
+) -> CommandResult {
+    let stake_pool = get_stake_pool(&config.rpc_client, stake_pool_address)?;
+    let mut transaction = Transaction::new_with_payer(
+        &[spl_stake_pool::instruction::set_preferred_validator(
+            &spl_stake_pool::id(),
+            &stake_pool_address,
+            &config.staker.pubkey(),
+            &stake_pool.validator_list,
+            preferred_type,
+            vote_address,
+        )],
+        Some(&config.fee_payer.pubkey()),
+    );
+
+    let (recent_blockhash, fee_calculator) = config.rpc_client.get_recent_blockhash()?;
+    check_fee_payer_balance(config, fee_calculator.calculate_fee(&transaction.message()))?;
+    let mut signers = vec![config.fee_payer.as_ref(), config.staker.as_ref()];
+    unique_signers!(signers);
+    transaction.sign(&signers, recent_blockhash);
+    send_transaction(&config, transaction)?;
+    Ok(())
+}
+
 fn add_associated_token_account(
     config: &Config,
     mint: &Pubkey,
@@ -1340,6 +1369,46 @@ fn main() {
                     .help("Amount in lamports to remove from the validator stake account. Must be at least the rent-exempt amount for a stake."),
             )
         )
+        .subcommand(SubCommand::with_name("set-preferred-validator")
+            .about("Set the preferred validator for deposits or withdrawals. Must be signed by the pool staker.")
+            .arg(
+                Arg::with_name("pool")
+                    .index(1)
+                    .validator(is_pubkey)
+                    .value_name("POOL_ADDRESS")
+                    .takes_value(true)
+                    .required(true)
+                    .help("Stake pool address"),
+            )
+            .arg(
+                Arg::with_name("preferred_type")
+                    .index(2)
+                    .value_name("OPERATION")
+                    .possible_values(&["deposit", "withdraw"]) // PreferredValidatorType enum
+                    .takes_value(true)
+                    .required(true)
+                    .help("Operation for which to restrict the validator"),
+            )
+            .arg(
+                Arg::with_name("vote_account")
+                    .long("vote-account")
+                    .validator(is_pubkey)
+                    .value_name("VOTE_ACCOUNT_ADDRESS")
+                    .takes_value(true)
+                    .help("Vote account for the validator that users must deposit into."),
+            )
+            .arg(
+                Arg::with_name("unset")
+                    .long("unset")
+                    .takes_value(false)
+                    .help("Unset the preferred validator."),
+            )
+            .group(ArgGroup::with_name("validator")
+                .arg("vote_account")
+                .arg("unset")
+                .required(true)
+            )
+        )
         .subcommand(SubCommand::with_name("deposit")
             .about("Add stake account to the stake pool")
             .arg(
@@ -1678,6 +1747,24 @@ fn main() {
             let amount = value_t_or_exit!(arg_matches, "amount", f64);
             command_decrease_validator_stake(&config, &stake_pool_address, &vote_account, amount)
         }
+        ("set-preferred-validator", Some(arg_matches)) => {
+            let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
+            let preferred_type = match arg_matches.value_of("preferred_type").unwrap() {
+                "deposit" => PreferredValidatorType::Deposit,
+                "withdraw" => PreferredValidatorType::Withdraw,
+                _ => unreachable!(),
+            };
+            let vote_account = pubkey_of(arg_matches, "vote_account");
+            let _unset = arg_matches.is_present("unset");
+            // since unset and vote_account can't both be set, if unset is set
+            // then vote_account will be None, which is valid for the program
+            command_set_preferred_validator(
+                &config,
+                &stake_pool_address,
+                preferred_type,
+                vote_account,
+            )
+        }
         ("deposit", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
             let stake_account = pubkey_of(arg_matches, "stake_account").unwrap();

stake-pool/program/src/error.rs

@@ -88,6 +88,14 @@ pub enum StakePoolError {
     /// The lamports in the validator stake account is not equal to the minimum
     #[error("StakeLamportsNotEqualToMinimum")]
     StakeLamportsNotEqualToMinimum,
+    /// The provided deposit stake account is not delegated to the preferred deposit vote account
+    #[error("IncorrectDepositVoteAddress")]
+    IncorrectDepositVoteAddress,
+
+    // 25.
+    /// The provided withdraw stake account is not the preferred deposit vote account
+    #[error("IncorrectWithdrawVoteAddress")]
+    IncorrectWithdrawVoteAddress,
 }
 impl From<StakePoolError> for ProgramError {
     fn from(e: StakePoolError) -> Self {

stake-pool/program/src/instruction.rs

@@ -18,6 +18,17 @@ use {
     },
 };
 
+/// Defines which validator vote account is set during the
+/// `SetPreferredValidator` instruction
+#[repr(C)]
+#[derive(Clone, Debug, PartialEq, BorshSerialize, BorshDeserialize, BorshSchema)]
+pub enum PreferredValidatorType {
+    /// Set preferred validator for deposits
+    Deposit,
+    /// Set preferred validator for withdraws
+    Withdraw,
+}
+
 /// Instructions supported by the StakePool program.
 #[repr(C)]
 #[derive(Clone, Debug, PartialEq, BorshSerialize, BorshDeserialize, BorshSchema)]
@@ -154,6 +165,28 @@ pub enum StakePoolInstruction {
     ///  userdata: amount of lamports to split into the transient stake account
     IncreaseValidatorStake(u64),
 
+    /// (Staker only) Set the preferred deposit or withdraw stake account for the
+    /// stake pool
+    ///
+    /// In order to avoid users abusing the stake pool as a free conversion
+    /// between SOL staked on different validators, the staker can force all
+    /// deposits and/or withdraws to go to one chosen account, or unset that account.
+    ///
+    /// 0. `[]` Stake pool
+    /// 1. `[s]` Stake pool staker
+    /// 2. `[w]` Validator list
+    ///
+    /// Fails if the validator is not part of the stake pool.
+    SetPreferredValidator {
+        /// Affected operation (deposit or withdraw)
+        #[allow(dead_code)] // but it's not
+        validator_type: PreferredValidatorType,
+        /// Validator vote account that deposits or withdraws must go through,
+        /// unset with None
+        #[allow(dead_code)] // but it's not
+        validator_vote_address: Option<Pubkey>,
+    },
+
     ///  Updates balances of validator and transient stake accounts in the pool
     ///
     ///  While going through the pairs of validator and transient stake accounts,
@@ -465,6 +498,31 @@ pub fn increase_validator_stake(
     }
 }
 
+/// Creates `SetPreferredDepositValidator` instruction
+pub fn set_preferred_validator(
+    program_id: &Pubkey,
+    stake_pool_address: &Pubkey,
+    staker: &Pubkey,
+    validator_list_address: &Pubkey,
+    validator_type: PreferredValidatorType,
+    validator_vote_address: Option<Pubkey>,
+) -> Instruction {
+    Instruction {
+        program_id: *program_id,
+        accounts: vec![
+            AccountMeta::new_readonly(*stake_pool_address, false),
+            AccountMeta::new_readonly(*staker, true),
+            AccountMeta::new(*validator_list_address, false),
+        ],
+        data: StakePoolInstruction::SetPreferredValidator {
+            validator_type,
+            validator_vote_address,
+        }
+        .try_to_vec()
+        .unwrap(),
+    }
+}
+
 /// Creates `CreateValidatorStakeAccount` instruction with a vote account
 pub fn create_validator_stake_account_with_vote(
     program_id: &Pubkey,

stake-pool/program/src/processor.rs

@@ -5,7 +5,7 @@ use {
         borsh::try_from_slice_unchecked,
         error::StakePoolError,
         find_deposit_authority_program_address,
-        instruction::StakePoolInstruction,
+        instruction::{PreferredValidatorType, StakePoolInstruction},
         minimum_reserve_lamports, minimum_stake_lamports, stake_program,
         state::{AccountType, Fee, StakePool, StakeStatus, ValidatorList, ValidatorStakeInfo},
         AUTHORITY_DEPOSIT, AUTHORITY_WITHDRAW, MINIMUM_ACTIVE_STAKE, TRANSIENT_STAKE_SEED,
@@ -432,6 +432,8 @@ impl Processor {
             return Err(StakePoolError::UnexpectedValidatorListAccountSize.into());
         }
         validator_list.account_type = AccountType::ValidatorList;
+        validator_list.preferred_deposit_validator_vote_address = None;
+        validator_list.preferred_withdraw_validator_vote_address = None;
         validator_list.validators.clear();
         validator_list.max_validators = max_validators;
 
@@ -862,6 +864,13 @@ impl Processor {
                 .retain(|item| item.vote_account_address != vote_account_address),
             _ => unreachable!(),
         }
+
+        if validator_list.preferred_deposit_validator_vote_address == Some(vote_account_address) {
+            validator_list.preferred_deposit_validator_vote_address = None;
+        }
+        if validator_list.preferred_withdraw_validator_vote_address == Some(vote_account_address) {
+            validator_list.preferred_withdraw_validator_vote_address = None;
+        }
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
 
         Ok(())
@@ -1150,6 +1159,55 @@ impl Processor {
         Ok(())
     }
 
+    /// Process `SetPreferredValidator` instruction
+    fn process_set_preferred_validator(
+        program_id: &Pubkey,
+        accounts: &[AccountInfo],
+        validator_type: PreferredValidatorType,
+        vote_account_address: Option<Pubkey>,
+    ) -> ProgramResult {
+        let account_info_iter = &mut accounts.iter();
+        let stake_pool_info = next_account_info(account_info_iter)?;
+        let staker_info = next_account_info(account_info_iter)?;
+        let validator_list_info = next_account_info(account_info_iter)?;
+
+        check_account_owner(stake_pool_info, program_id)?;
+        check_account_owner(validator_list_info, program_id)?;
+
+        let stake_pool = StakePool::try_from_slice(&stake_pool_info.data.borrow())?;
+        if !stake_pool.is_valid() {
+            msg!("Expected valid stake pool");
+            return Err(StakePoolError::InvalidState.into());
+        }
+
+        stake_pool.check_staker(staker_info)?;
+        stake_pool.check_validator_list(validator_list_info)?;
+
+        let mut validator_list =
+            try_from_slice_unchecked::<ValidatorList>(&validator_list_info.data.borrow())?;
+        if !validator_list.is_valid() {
+            return Err(StakePoolError::InvalidState.into());
+        }
+
+        if let Some(vote_account_address) = vote_account_address {
+            if !validator_list.contains(&vote_account_address) {
+                msg!("Validator for {} not present in the stake pool, cannot set as preferred deposit account");
+                return Err(StakePoolError::ValidatorNotFound.into());
+            }
+        }
+
+        match validator_type {
+            PreferredValidatorType::Deposit => {
+                validator_list.preferred_deposit_validator_vote_address = vote_account_address
+            }
+            PreferredValidatorType::Withdraw => {
+                validator_list.preferred_withdraw_validator_vote_address = vote_account_address
+            }
+        };
+        validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
+        Ok(())
+    }
+
     /// Processes `UpdateValidatorListBalance` instruction.
     fn process_update_validator_list_balance(
         program_id: &Pubkey,
@@ -1484,7 +1542,6 @@ impl Processor {
         let clock_info = next_account_info(account_info_iter)?;
         let clock = &Clock::from_account_info(clock_info)?;
         let stake_history_info = next_account_info(account_info_iter)?;
-        //let stake_history = &StakeHistory::from_account_info(stake_history_info)?;
         let token_program_info = next_account_info(account_info_iter)?;
         let stake_program_info = next_account_info(account_info_iter)?;
 
@@ -1533,6 +1590,11 @@ impl Processor {
             validator_stake_account_info.key,
             &vote_account_address,
         )?;
+        if let Some(preferred_deposit) = validator_list.preferred_deposit_validator_vote_address {
+            if preferred_deposit != vote_account_address {
+                return Err(StakePoolError::IncorrectDepositVoteAddress.into());
+            }
+        }
 
         let validator_list_item = validator_list
             .find_mut(&vote_account_address)
@@ -1717,6 +1779,20 @@ impl Processor {
                 &vote_account_address,
             )?;
 
+            if let Some(preferred_withdraw_validator) =
+                validator_list.preferred_withdraw_validator_vote_address
+            {
+                let preferred_validator_info = validator_list
+                    .find(&preferred_withdraw_validator)
+                    .ok_or(StakePoolError::ValidatorNotFound)?;
+                if preferred_withdraw_validator != vote_account_address
+                    && preferred_validator_info.stake_lamports > 0
+                {
+                    msg!("Validator vote address {} is preferred for withdrawals, it currently has {} lamports available. Please withdraw those before using other validator stake accounts.", preferred_withdraw_validator, preferred_validator_info.stake_lamports);
+                    return Err(StakePoolError::IncorrectWithdrawVoteAddress.into());
+                }
+            }
+
             let validator_list_item = validator_list
                 .find_mut(&vote_account_address)
                 .ok_or(StakePoolError::ValidatorNotFound)?;
@@ -1901,6 +1977,18 @@ impl Processor {
                 msg!("Instruction: IncreaseValidatorStake");
                 Self::process_increase_validator_stake(program_id, accounts, amount)
             }
+            StakePoolInstruction::SetPreferredValidator {
+                validator_type,
+                validator_vote_address,
+            } => {
+                msg!("Instruction: SetPreferredValidator");
+                Self::process_set_preferred_validator(
+                    program_id,
+                    accounts,
+                    validator_type,
+                    validator_vote_address,
+                )
+            }
             StakePoolInstruction::UpdateValidatorListBalance {
                 start_index,
                 no_merge,
@@ -1973,6 +2061,8 @@ impl PrintProgramError for StakePoolError {
             StakePoolError::WrongStaker=> msg!("Error: Wrong pool staker account"),
             StakePoolError::NonZeroPoolTokenSupply => msg!("Error: Pool token supply is not zero on initialization"),
             StakePoolError::StakeLamportsNotEqualToMinimum => msg!("Error: The lamports in the validator stake account is not equal to the minimum"),
+            StakePoolError::IncorrectDepositVoteAddress => msg!("Error: The provided deposit stake account is not delegated to the preferred deposit vote account"),
+            StakePoolError::IncorrectWithdrawVoteAddress => msg!("Error: The provided withdraw stake account is not the preferred deposit vote account"),
         }
     }
 }

stake-pool/program/src/state.rs

@@ -286,6 +286,12 @@ pub struct ValidatorList {
     /// Account type, must be ValidatorList currently
     pub account_type: AccountType,
 
+    /// Preferred deposit validator vote account pubkey
+    pub preferred_deposit_validator_vote_address: Option<Pubkey>,
+
+    /// Preferred withdraw validator vote account pubkey
+    pub preferred_withdraw_validator_vote_address: Option<Pubkey>,
+
     /// Maximum allowable number of validators
     pub max_validators: u32,
 
@@ -332,10 +338,12 @@ pub struct ValidatorStakeInfo {
 }
 
 impl ValidatorList {
-    /// Create an empty instance containing space for `max_validators`
+    /// Create an empty instance containing space for `max_validators` and preferred validator keys
     pub fn new(max_validators: u32) -> Self {
         Self {
             account_type: AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: Some(Pubkey::default()),
+            preferred_withdraw_validator_vote_address: Some(Pubkey::default()),
             max_validators,
             validators: vec![ValidatorStakeInfo::default(); max_validators as usize],
         }
@@ -343,7 +351,7 @@ impl ValidatorList {
 
     /// Calculate the number of validator entries that fit in the provided length
     pub fn calculate_max_validators(buffer_length: usize) -> usize {
-        let header_size = 1 + 4 + 4;
+        let header_size = 1 + 4 + 4 + 33 + 33;
         buffer_length.saturating_sub(header_size) / 49
     }
 
@@ -406,6 +414,8 @@ mod test {
         // Not initialized
         let stake_list = ValidatorList {
             account_type: AccountType::Uninitialized,
+            preferred_deposit_validator_vote_address: None,
+            preferred_withdraw_validator_vote_address: None,
             max_validators: 0,
             validators: vec![],
         };
@@ -415,9 +425,11 @@ mod test {
         let stake_list_unpacked = try_from_slice_unchecked::<ValidatorList>(&byte_vec).unwrap();
         assert_eq!(stake_list_unpacked, stake_list);
 
-        // Empty
+        // Empty, one preferred key
         let stake_list = ValidatorList {
             account_type: AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: Some(Pubkey::new_unique()),
+            preferred_withdraw_validator_vote_address: None,
             max_validators: 0,
             validators: vec![],
         };
@@ -430,6 +442,8 @@ mod test {
         // With several accounts
         let stake_list = ValidatorList {
             account_type: AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: Some(Pubkey::new_unique()),
+            preferred_withdraw_validator_vote_address: Some(Pubkey::new_unique()),
             max_validators,
             validators: vec![
                 ValidatorStakeInfo {

stake-pool/program/tests/deposit.rs

@@ -32,6 +32,10 @@ async fn setup() -> (
     Hash,
     StakePoolAccounts,
     ValidatorStakeAccount,
+    Keypair,
+    Pubkey,
+    Pubkey,
+    u64,
 ) {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
@@ -40,7 +44,7 @@ async fn setup() -> (
         .await
         .unwrap();
 
-    let validator_stake_account: ValidatorStakeAccount = simple_add_validator_to_pool(
+    let validator_stake_account = simple_add_validator_to_pool(
         &mut banks_client,
         &payer,
         &recent_blockhash,
@@ -48,73 +52,77 @@ async fn setup() -> (
     )
     .await;
 
-    (
-        banks_client,
-        payer,
-        recent_blockhash,
-        stake_pool_accounts,
-        validator_stake_account,
-    )
-}
-
-#[tokio::test]
-async fn test_stake_pool_deposit() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
-        setup().await;
-
     let user = Keypair::new();
     // make stake account
-    let user_stake = Keypair::new();
+    let deposit_stake = Keypair::new();
     let lockup = stake_program::Lockup::default();
 
-    let stake_authority = Keypair::new();
     let authorized = stake_program::Authorized {
-        staker: stake_authority.pubkey(),
-        withdrawer: stake_authority.pubkey(),
+        staker: user.pubkey(),
+        withdrawer: user.pubkey(),
     };
 
     let stake_lamports = create_independent_stake_account(
         &mut banks_client,
         &payer,
         &recent_blockhash,
-        &user_stake,
+        &deposit_stake,
         &authorized,
         &lockup,
         TEST_STAKE_AMOUNT,
     )
     .await;
 
-    create_vote(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &validator_stake_account.validator,
-        &validator_stake_account.vote,
-    )
-    .await;
     delegate_stake_account(
         &mut banks_client,
         &payer,
         &recent_blockhash,
-        &user_stake.pubkey(),
-        &stake_authority,
+        &deposit_stake.pubkey(),
+        &user,
         &validator_stake_account.vote.pubkey(),
     )
     .await;
 
     // make pool token account
-    let user_pool_account = Keypair::new();
+    let pool_token_account = Keypair::new();
     create_token_account(
         &mut banks_client,
         &payer,
         &recent_blockhash,
-        &user_pool_account,
+        &pool_token_account,
         &stake_pool_accounts.pool_mint.pubkey(),
         &user.pubkey(),
     )
     .await
     .unwrap();
 
+    (
+        banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake.pubkey(),
+        pool_token_account.pubkey(),
+        stake_lamports,
+    )
+}
+
+#[tokio::test]
+async fn success() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        stake_lamports,
+    ) = setup().await;
+
     // Save stake pool state before depositing
     let stake_pool_before =
         get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await;
@@ -133,22 +141,22 @@ async fn test_stake_pool_deposit() {
         .find(&validator_stake_account.vote.pubkey())
         .unwrap();
 
-    stake_pool_accounts
+    let error = stake_pool_accounts
         .deposit_stake(
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.pubkey(),
-            &user_pool_account.pubkey(),
+            &deposit_stake,
+            &pool_token_account,
             &validator_stake_account.stake_account,
-            &stake_authority,
+            &user,
         )
-        .await
-        .unwrap();
+        .await;
+    assert!(error.is_none());
 
     // Original stake account should be drained
     assert!(banks_client
-        .get_account(user_stake.pubkey())
+        .get_account(deposit_stake)
         .await
         .expect("get_account")
         .is_none());
@@ -168,8 +176,7 @@ async fn test_stake_pool_deposit() {
     );
 
     // Check minted tokens
-    let user_token_balance =
-        get_token_balance(&mut banks_client, &user_pool_account.pubkey()).await;
+    let user_token_balance = get_token_balance(&mut banks_client, &pool_token_account).await;
     assert_eq!(user_token_balance, tokens_issued);
 
     // Check balances in validator stake account list storage
@@ -201,26 +208,18 @@ async fn test_stake_pool_deposit() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_with_wrong_stake_program_id() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
-        setup().await;
-
-    let user = Keypair::new();
-    // make stake account
-    let user_stake = Keypair::new();
-
-    // make pool token account
-    let user_pool_account = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
+async fn fail_with_wrong_stake_program_id() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        _user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup().await;
 
     let wrong_stake_program = Pubkey::new_unique();
 
@@ -229,9 +228,9 @@ async fn test_stake_pool_deposit_with_wrong_stake_program_id() {
         AccountMeta::new(stake_pool_accounts.validator_list.pubkey(), false),
         AccountMeta::new_readonly(stake_pool_accounts.deposit_authority, false),
         AccountMeta::new_readonly(stake_pool_accounts.withdraw_authority, false),
-        AccountMeta::new(user_stake.pubkey(), false),
+        AccountMeta::new(deposit_stake, false),
         AccountMeta::new(validator_stake_account.stake_account, false),
-        AccountMeta::new(user_pool_account.pubkey(), false),
+        AccountMeta::new(pool_token_account, false),
         AccountMeta::new(stake_pool_accounts.pool_mint.pubkey(), false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(sysvar::stake_history::id(), false),
@@ -263,41 +262,18 @@ async fn test_stake_pool_deposit_with_wrong_stake_program_id() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_with_wrong_token_program_id() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
-        setup().await;
-
-    let user = Keypair::new();
-    // make stake account
-    let user_stake = Keypair::new();
-    let lockup = stake_program::Lockup::default();
-    let authorized = stake_program::Authorized {
-        staker: user.pubkey(),
-        withdrawer: user.pubkey(),
-    };
-    create_independent_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake,
-        &authorized,
-        &lockup,
-        TEST_STAKE_AMOUNT,
-    )
-    .await;
-
-    // make pool token account
-    let user_pool_account = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
+async fn fail_with_wrong_token_program_id() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup().await;
 
     let wrong_token_program = Keypair::new();
 
@@ -307,10 +283,10 @@ async fn test_stake_pool_deposit_with_wrong_token_program_id() {
             &stake_pool_accounts.stake_pool.pubkey(),
             &stake_pool_accounts.validator_list.pubkey(),
             &stake_pool_accounts.withdraw_authority,
-            &user_stake.pubkey(),
+            &deposit_stake,
             &user.pubkey(),
             &validator_stake_account.stake_account,
-            &user_pool_account.pubkey(),
+            &pool_token_account,
             &stake_pool_accounts.pool_mint.pubkey(),
             &wrong_token_program.pubkey(),
         ),
@@ -332,47 +308,19 @@ async fn test_stake_pool_deposit_with_wrong_token_program_id() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_with_wrong_validator_list_account() {
+async fn fail_with_wrong_validator_list_account() {
     let (
         mut banks_client,
         payer,
         recent_blockhash,
         mut stake_pool_accounts,
         validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
     ) = setup().await;
 
-    let user = Keypair::new();
-    // make stake account
-    let user_stake = Keypair::new();
-    let lockup = stake_program::Lockup::default();
-    let authorized = stake_program::Authorized {
-        staker: user.pubkey(),
-        withdrawer: user.pubkey(),
-    };
-    create_independent_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake,
-        &authorized,
-        &lockup,
-        TEST_STAKE_AMOUNT,
-    )
-    .await;
-
-    // make pool token account
-    let user_pool_account = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
-
     let wrong_validator_list = Keypair::new();
     stake_pool_accounts.validator_list = wrong_validator_list;
 
@@ -381,20 +329,20 @@ async fn test_stake_pool_deposit_with_wrong_validator_list_account() {
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.pubkey(),
-            &user_pool_account.pubkey(),
+            &deposit_stake,
+            &pool_token_account,
             &validator_stake_account.stake_account,
             &user,
         )
         .await
-        .err()
+        .unwrap()
         .unwrap();
 
     match transaction_error {
-        TransportError::TransactionError(TransactionError::InstructionError(
+        TransactionError::InstructionError(
             _,
             InstructionError::Custom(error_index),
-        )) => {
+        ) => {
             let program_error = error::StakePoolError::InvalidValidatorStakeList as u32;
             assert_eq!(error_index, program_error);
         }
@@ -403,7 +351,7 @@ async fn test_stake_pool_deposit_with_wrong_validator_list_account() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_to_unknown_validator() {
+async fn fail_with_unknown_validator() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
     let stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts
@@ -464,14 +412,11 @@ async fn test_stake_pool_deposit_to_unknown_validator() {
             &user,
         )
         .await
-        .err()
+        .unwrap()
         .unwrap();
 
     match transaction_error {
-        TransportError::TransactionError(TransactionError::InstructionError(
-            _,
-            InstructionError::Custom(error_index),
-        )) => {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
             let program_error = error::StakePoolError::ValidatorNotFound as u32;
             assert_eq!(error_index, program_error);
         }
@@ -482,68 +427,37 @@ async fn test_stake_pool_deposit_to_unknown_validator() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_with_wrong_withdraw_authority() {
+async fn fail_with_wrong_withdraw_authority() {
     let (
         mut banks_client,
         payer,
         recent_blockhash,
         mut stake_pool_accounts,
         validator_stake_account,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
     ) = setup().await;
 
-    let user = Keypair::new();
-    // make stake account
-    let user_stake = Keypair::new();
-    let lockup = stake_program::Lockup::default();
-    let authorized = stake_program::Authorized {
-        staker: user.pubkey(),
-        withdrawer: user.pubkey(),
-    };
-    create_independent_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake,
-        &authorized,
-        &lockup,
-        TEST_STAKE_AMOUNT,
-    )
-    .await;
-
-    // make pool token account
-    let user_pool_account = Keypair::new();
-    create_token_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_pool_account,
-        &stake_pool_accounts.pool_mint.pubkey(),
-        &user.pubkey(),
-    )
-    .await
-    .unwrap();
-
-    stake_pool_accounts.withdraw_authority = Keypair::new().pubkey();
+    stake_pool_accounts.withdraw_authority = Pubkey::new_unique();
 
     let transaction_error = stake_pool_accounts
         .deposit_stake(
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.pubkey(),
-            &user_pool_account.pubkey(),
+            &deposit_stake,
+            &pool_token_account,
             &validator_stake_account.stake_account,
             &user,
         )
         .await
-        .err()
+        .unwrap()
         .unwrap();
 
     match transaction_error {
-        TransportError::TransactionError(TransactionError::InstructionError(
-            _,
-            InstructionError::Custom(error_index),
-        )) => {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
             let program_error = error::StakePoolError::InvalidProgramAddress as u32;
             assert_eq!(error_index, program_error);
         }
@@ -552,28 +466,18 @@ async fn test_stake_pool_deposit_with_wrong_withdraw_authority() {
 }
 
 #[tokio::test]
-async fn test_stake_pool_deposit_with_wrong_mint_for_receiver_acc() {
-    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
-        setup().await;
-
-    // make stake account
-    let user = Keypair::new();
-    let user_stake = Keypair::new();
-    let lockup = stake_program::Lockup::default();
-    let authorized = stake_program::Authorized {
-        staker: user.pubkey(),
-        withdrawer: user.pubkey(),
-    };
-    create_independent_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake,
-        &authorized,
-        &lockup,
-        TEST_STAKE_AMOUNT,
-    )
-    .await;
+async fn fail_with_wrong_mint_for_receiver_acc() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        user,
+        deposit_stake,
+        _pool_token_account,
+        _stake_lamports,
+    ) = setup().await;
 
     let outside_mint = Keypair::new();
     let outside_withdraw_auth = Keypair::new();
@@ -606,20 +510,17 @@ async fn test_stake_pool_deposit_with_wrong_mint_for_receiver_acc() {
             &mut banks_client,
             &payer,
             &recent_blockhash,
-            &user_stake.pubkey(),
+            &deposit_stake,
             &outside_pool_fee_acc.pubkey(),
             &validator_stake_account.stake_account,
             &user,
         )
         .await
-        .err()
+        .unwrap()
         .unwrap();
 
     match transaction_error {
-        TransportError::TransactionError(TransactionError::InstructionError(
-            _,
-            InstructionError::Custom(error_index),
-        )) => {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
             let program_error = token_error::TokenError::MintMismatch as u32;
             assert_eq!(error_index, program_error);
         }
@@ -628,10 +529,10 @@ async fn test_stake_pool_deposit_with_wrong_mint_for_receiver_acc() {
 }
 
 #[tokio::test]
-async fn test_deposit_with_uninitialized_validator_list() {} // TODO
+async fn fail_with_uninitialized_validator_list() {} // TODO
 
 #[tokio::test]
-async fn test_deposit_with_out_of_dated_pool_balances() {} // TODO
+async fn fail_with_out_of_dated_pool_balances() {} // TODO
 
 #[tokio::test]
 async fn success_with_deposit_authority() {
@@ -700,7 +601,7 @@ async fn success_with_deposit_authority() {
     .await
     .unwrap();
 
-    stake_pool_accounts
+    let error = stake_pool_accounts
         .deposit_stake(
             &mut banks_client,
             &payer,
@@ -710,8 +611,8 @@ async fn success_with_deposit_authority() {
             &validator_stake_account.stake_account,
             &user,
         )
-        .await
-        .unwrap();
+        .await;
+    assert!(error.is_none());
 }
 
 #[tokio::test]
@@ -796,7 +697,7 @@ async fn fail_without_deposit_authority_signature() {
             &user,
         )
         .await
-        .unwrap_err()
+        .unwrap()
         .unwrap();
 
     match error {
@@ -809,3 +710,97 @@ async fn fail_without_deposit_authority_signature() {
         _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
     }
 }
+
+#[tokio::test]
+async fn success_with_preferred_deposit() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup().await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Deposit,
+            Some(validator_stake.vote.pubkey()),
+        )
+        .await;
+
+    let error = stake_pool_accounts
+        .deposit_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &deposit_stake,
+            &pool_token_account,
+            &validator_stake.stake_account,
+            &user,
+        )
+        .await;
+    assert!(error.is_none());
+}
+
+#[tokio::test]
+async fn fail_with_wrong_preferred_deposit() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake,
+        user,
+        deposit_stake,
+        pool_token_account,
+        _stake_lamports,
+    ) = setup().await;
+
+    let preferred_validator = simple_add_validator_to_pool(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+    )
+    .await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Deposit,
+            Some(preferred_validator.vote.pubkey()),
+        )
+        .await;
+
+    let error = stake_pool_accounts
+        .deposit_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &deposit_stake,
+            &pool_token_account,
+            &validator_stake.stake_account,
+            &user,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    match error {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
+            assert_eq!(
+                error_index,
+                error::StakePoolError::IncorrectDepositVoteAddress as u32
+            );
+        }
+        _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
+    }
+}

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

@@ -606,7 +606,7 @@ impl StakePoolAccounts {
         pool_account: &Pubkey,
         validator_stake_account: &Pubkey,
         current_staker: &Keypair,
-    ) -> Result<(), TransportError> {
+    ) -> Option<TransportError> {
         let mut signers = vec![payer, current_staker];
         let instructions = if let Some(deposit_authority) = self.deposit_authority_keypair.as_ref()
         {
@@ -644,8 +644,7 @@ impl StakePoolAccounts {
             &signers,
             *recent_blockhash,
         );
-        banks_client.process_transaction(transaction).await?;
-        Ok(())
+        banks_client.process_transaction(transaction).await.err()
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -873,6 +872,30 @@ impl StakePoolAccounts {
         );
         banks_client.process_transaction(transaction).await.err()
     }
+
+    pub async fn set_preferred_validator(
+        &self,
+        banks_client: &mut BanksClient,
+        payer: &Keypair,
+        recent_blockhash: &Hash,
+        validator_type: instruction::PreferredValidatorType,
+        validator: Option<Pubkey>,
+    ) -> Option<TransportError> {
+        let transaction = Transaction::new_signed_with_payer(
+            &[instruction::set_preferred_validator(
+                &id(),
+                &self.stake_pool.pubkey(),
+                &self.staker.pubkey(),
+                &self.validator_list.pubkey(),
+                validator_type,
+                validator,
+            )],
+            Some(&payer.pubkey()),
+            &[payer, &self.staker],
+            *recent_blockhash,
+        );
+        banks_client.process_transaction(transaction).await.err()
+    }
 }
 
 pub async fn simple_add_validator_to_pool(
@@ -986,7 +1009,7 @@ impl DepositStakeAccount {
         .await
         .unwrap();
 
-        stake_pool_accounts
+        let error = stake_pool_accounts
             .deposit_stake(
                 banks_client,
                 payer,
@@ -996,8 +1019,8 @@ impl DepositStakeAccount {
                 &self.validator_stake_account,
                 &self.authority,
             )
-            .await
-            .unwrap();
+            .await;
+        assert!(error.is_none());
     }
 }
 
@@ -1051,7 +1074,7 @@ pub async fn simple_deposit(
     .unwrap();
 
     let validator_stake_account = validator_stake_account.stake_account;
-    stake_pool_accounts
+    let error = stake_pool_accounts
         .deposit_stake(
             banks_client,
             payer,
@@ -1061,8 +1084,11 @@ pub async fn simple_deposit(
             &validator_stake_account,
             &authority,
         )
-        .await
-        .ok()?;
+        .await;
+    // backwards, but oh well!
+    if error.is_some() {
+        return None;
+    }
 
     let pool_tokens = get_token_balance(banks_client, &pool_account.pubkey()).await;
 

stake-pool/program/tests/set_fee.rs

@@ -8,8 +8,9 @@ use {
     solana_program::hash::Hash,
     solana_program_test::*,
     solana_sdk::{
-        instruction::InstructionError, signature::Keypair, signature::Signer,
-        transaction::Transaction, transaction::TransactionError,
+        instruction::InstructionError,
+        signature::{Keypair, Signer},
+        transaction::{Transaction, TransactionError},
     },
     spl_stake_pool::{
         error, id, instruction,

stake-pool/program/tests/set_preferred.rs

@@ -0,0 +1,238 @@
+#![cfg(feature = "test-bpf")]
+
+mod helpers;
+
+use {
+    helpers::*,
+    solana_program::hash::Hash,
+    solana_program_test::*,
+    solana_sdk::{
+        instruction::InstructionError,
+        pubkey::Pubkey,
+        signature::{Keypair, Signer},
+        transaction::{Transaction, TransactionError},
+    },
+    spl_stake_pool::{
+        borsh::try_from_slice_unchecked,
+        error, id,
+        instruction::{self, PreferredValidatorType},
+        state::ValidatorList,
+    },
+};
+
+async fn setup() -> (
+    BanksClient,
+    Keypair,
+    Hash,
+    StakePoolAccounts,
+    ValidatorStakeAccount,
+) {
+    let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
+    let stake_pool_accounts = StakePoolAccounts::new();
+    stake_pool_accounts
+        .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1)
+        .await
+        .unwrap();
+
+    let validator_stake_account = simple_add_validator_to_pool(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+    )
+    .await;
+
+    (
+        banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+    )
+}
+
+#[tokio::test]
+async fn success_deposit() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
+        setup().await;
+
+    let vote_account_address = validator_stake_account.vote.pubkey();
+    let error = stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            PreferredValidatorType::Deposit,
+            Some(vote_account_address),
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(&validator_list.data.as_slice()).unwrap();
+
+    assert_eq!(
+        validator_list.preferred_deposit_validator_vote_address,
+        Some(vote_account_address)
+    );
+    assert_eq!(
+        validator_list.preferred_withdraw_validator_vote_address,
+        None
+    );
+}
+
+#[tokio::test]
+async fn success_withdraw() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
+        setup().await;
+
+    let vote_account_address = validator_stake_account.vote.pubkey();
+
+    let error = stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            PreferredValidatorType::Withdraw,
+            Some(vote_account_address),
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(&validator_list.data.as_slice()).unwrap();
+
+    assert_eq!(
+        validator_list.preferred_deposit_validator_vote_address,
+        None
+    );
+    assert_eq!(
+        validator_list.preferred_withdraw_validator_vote_address,
+        Some(vote_account_address)
+    );
+}
+
+#[tokio::test]
+async fn success_unset() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake_account) =
+        setup().await;
+
+    let vote_account_address = validator_stake_account.vote.pubkey();
+    let error = stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            PreferredValidatorType::Withdraw,
+            Some(vote_account_address),
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(&validator_list.data.as_slice()).unwrap();
+
+    assert_eq!(
+        validator_list.preferred_withdraw_validator_vote_address,
+        Some(vote_account_address)
+    );
+
+    let error = stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            PreferredValidatorType::Withdraw,
+            None,
+        )
+        .await;
+    assert!(error.is_none());
+
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<ValidatorList>(&validator_list.data.as_slice()).unwrap();
+
+    assert_eq!(
+        validator_list.preferred_withdraw_validator_vote_address,
+        None
+    );
+}
+
+#[tokio::test]
+async fn fail_wrong_staker() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, _) = setup().await;
+
+    let wrong_staker = Keypair::new();
+    let transaction = Transaction::new_signed_with_payer(
+        &[instruction::set_preferred_validator(
+            &id(),
+            &stake_pool_accounts.stake_pool.pubkey(),
+            &wrong_staker.pubkey(),
+            &stake_pool_accounts.validator_list.pubkey(),
+            PreferredValidatorType::Withdraw,
+            None,
+        )],
+        Some(&payer.pubkey()),
+        &[&payer, &wrong_staker],
+        recent_blockhash,
+    );
+    let error = banks_client
+        .process_transaction(transaction)
+        .await
+        .err()
+        .unwrap()
+        .unwrap();
+
+    match error {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
+            let program_error = error::StakePoolError::WrongStaker as u32;
+            assert_eq!(error_index, program_error);
+        }
+        _ => panic!("Wrong error occurs while malicious try to set manager"),
+    }
+}
+
+#[tokio::test]
+async fn fail_not_present_validator() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, _) = setup().await;
+
+    let validator_vote_address = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            PreferredValidatorType::Withdraw,
+            Some(validator_vote_address),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+
+    match error {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
+            let program_error = error::StakePoolError::ValidatorNotFound as u32;
+            assert_eq!(error_index, program_error);
+        }
+        _ => panic!("Wrong error occurs while malicious try to set manager"),
+    }
+}

stake-pool/program/tests/vsa_add.rs

@@ -84,6 +84,8 @@ async fn success() {
         validator_list,
         state::ValidatorList {
             account_type: state::AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: None,
+            preferred_withdraw_validator_vote_address: None,
             max_validators: stake_pool_accounts.max_validators,
             validators: vec![state::ValidatorStakeInfo {
                 status: state::StakeStatus::Active,

stake-pool/program/tests/vsa_remove.rs

@@ -97,6 +97,8 @@ async fn success() {
         validator_list,
         state::ValidatorList {
             account_type: state::AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: None,
+            preferred_withdraw_validator_vote_address: None,
             max_validators: stake_pool_accounts.max_validators,
             validators: vec![]
         }
@@ -520,6 +522,8 @@ async fn success_with_deactivating_transient_stake() {
         try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
     let expected_list = state::ValidatorList {
         account_type: state::AccountType::ValidatorList,
+        preferred_deposit_validator_vote_address: None,
+        preferred_withdraw_validator_vote_address: None,
         max_validators: stake_pool_accounts.max_validators,
         validators: vec![state::ValidatorStakeInfo {
             status: state::StakeStatus::DeactivatingTransient,
@@ -552,6 +556,74 @@ async fn success_with_deactivating_transient_stake() {
     assert_eq!(validator_list, expected_list);
 }
 
+#[tokio::test]
+async fn success_resets_preferred_validator() {
+    let (mut banks_client, payer, recent_blockhash, stake_pool_accounts, validator_stake) =
+        setup().await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Deposit,
+            Some(validator_stake.vote.pubkey()),
+        )
+        .await;
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Withdraw,
+            Some(validator_stake.vote.pubkey()),
+        )
+        .await;
+
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .remove_validator_from_pool(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &new_authority,
+            &validator_stake.stake_account,
+            &validator_stake.transient_stake_account,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // Check if account was removed from the list of stake accounts
+    let validator_list = get_account(
+        &mut banks_client,
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let validator_list =
+        try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
+    assert_eq!(
+        validator_list,
+        state::ValidatorList {
+            account_type: state::AccountType::ValidatorList,
+            preferred_deposit_validator_vote_address: None,
+            preferred_withdraw_validator_vote_address: None,
+            max_validators: stake_pool_accounts.max_validators,
+            validators: vec![]
+        }
+    );
+
+    // Check of stake account authority has changed
+    let stake = get_account(&mut banks_client, &validator_stake.stake_account).await;
+    let stake_state = deserialize::<stake_program::StakeState>(&stake.data).unwrap();
+    match stake_state {
+        stake_program::StakeState::Stake(meta, _) => {
+            assert_eq!(&meta.authorized.staker, &new_authority);
+            assert_eq!(&meta.authorized.withdrawer, &new_authority);
+        }
+        _ => panic!(),
+    }
+}
+
 #[tokio::test]
 async fn fail_not_updated_stake_pool() {} // TODO
 

stake-pool/program/tests/withdraw.rs

@@ -33,6 +33,7 @@ async fn setup() -> (
     ValidatorStakeAccount,
     DepositStakeAccount,
     Keypair,
+    Keypair,
     u64,
 ) {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
@@ -76,6 +77,16 @@ async fn setup() -> (
     )
     .await;
 
+    // Create stake account to withdraw to
+    let user_stake_recipient = Keypair::new();
+    create_blank_stake_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &user_stake_recipient,
+    )
+    .await;
+
     (
         banks_client,
         payer,
@@ -84,6 +95,7 @@ async fn setup() -> (
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     )
 }
@@ -98,25 +110,21 @@ async fn success() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-    let initial_stake_lamports = create_blank_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake_recipient,
-    )
-    .await;
-
     // Save stake pool state before withdrawal
     let stake_pool_before =
         get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await;
     let stake_pool_before =
         state::StakePool::try_from_slice(&stake_pool_before.data.as_slice()).unwrap();
 
+    // Check user recipient stake account balance
+    let initial_stake_lamports = get_account(&mut banks_client, &user_stake_recipient.pubkey())
+        .await
+        .lamports;
+
     // Save validator stake account record before withdrawal
     let validator_list = get_account(
         &mut banks_client,
@@ -215,12 +223,10 @@ async fn fail_with_wrong_stake_program() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-
     let new_authority = Pubkey::new_unique();
     let wrong_stake_program = Pubkey::new_unique();
 
@@ -276,12 +282,10 @@ async fn fail_with_wrong_withdraw_authority() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-
     let new_authority = Pubkey::new_unique();
     stake_pool_accounts.withdraw_authority = Keypair::new().pubkey();
 
@@ -322,12 +326,10 @@ async fn fail_with_wrong_token_program_id() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-
     let new_authority = Pubkey::new_unique();
     let wrong_token_program = Keypair::new();
 
@@ -374,12 +376,10 @@ async fn fail_with_wrong_validator_list() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-
     let new_authority = Pubkey::new_unique();
     stake_pool_accounts.validator_list = Keypair::new();
 
@@ -422,6 +422,7 @@ async fn fail_with_unknown_validator() {
         _,
         _,
         user_transfer_authority,
+        user_stake_recipient,
         _,
     ) = setup().await;
 
@@ -507,9 +508,6 @@ async fn fail_with_unknown_validator() {
     )
     .await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-
     let new_authority = Pubkey::new_unique();
 
     let transaction_error = stake_pool_accounts
@@ -547,19 +545,10 @@ async fn fail_double_withdraw_to_the_same_account() {
         validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-    create_blank_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake_recipient,
-    )
-    .await;
-
     let new_authority = Pubkey::new_unique();
     let error = stake_pool_accounts
         .withdraw_stake(
@@ -777,19 +766,10 @@ async fn fail_overdraw_validator() {
         _validator_stake_account,
         deposit_info,
         user_transfer_authority,
+        user_stake_recipient,
         tokens_to_burn,
     ) = setup().await;
 
-    // Create stake account to withdraw to
-    let user_stake_recipient = Keypair::new();
-    let _initial_stake_lamports = create_blank_stake_account(
-        &mut banks_client,
-        &payer,
-        &recent_blockhash,
-        &user_stake_recipient,
-    )
-    .await;
-
     let validator_stake_account = simple_add_validator_to_pool(
         &mut banks_client,
         &payer,
@@ -1013,3 +993,141 @@ async fn success_with_reserve() {
         initial_stake_lamports + deposit_info.stake_lamports + stake_rent
     );
 }
+
+#[tokio::test]
+async fn success_with_preferred_validator() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake,
+        deposit_info,
+        user_transfer_authority,
+        user_stake_recipient,
+        tokens_to_burn,
+    ) = setup().await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Withdraw,
+            Some(validator_stake.vote.pubkey()),
+        )
+        .await;
+
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &user_stake_recipient.pubkey(),
+            &user_transfer_authority,
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.stake_account,
+            &new_authority,
+            tokens_to_burn,
+        )
+        .await;
+    assert!(error.is_none());
+}
+
+#[tokio::test]
+async fn fail_with_wrong_preferred_withdraw() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake,
+        deposit_info,
+        user_transfer_authority,
+        user_stake_recipient,
+        tokens_to_burn,
+    ) = setup().await;
+
+    let preferred_validator = simple_add_validator_to_pool(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+    )
+    .await;
+
+    stake_pool_accounts
+        .set_preferred_validator(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            instruction::PreferredValidatorType::Withdraw,
+            Some(preferred_validator.vote.pubkey()),
+        )
+        .await;
+
+    // preferred is empty, this works
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &user_stake_recipient.pubkey(),
+            &user_transfer_authority,
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.stake_account,
+            &new_authority,
+            tokens_to_burn,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // deposit into preferred, then fail
+    let _preferred_deposit = simple_deposit(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts,
+        &preferred_validator,
+        TEST_STAKE_AMOUNT,
+    )
+    .await
+    .unwrap();
+
+    // Create stake account to withdraw to
+    let user_stake_recipient = Keypair::new();
+    create_blank_stake_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &user_stake_recipient,
+    )
+    .await;
+
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &user_stake_recipient.pubkey(),
+            &user_transfer_authority,
+            &deposit_info.pool_account.pubkey(),
+            &validator_stake.stake_account,
+            &new_authority,
+            tokens_to_burn,
+        )
+        .await
+        .unwrap()
+        .unwrap();
+    match error {
+        TransactionError::InstructionError(_, InstructionError::Custom(error_index)) => {
+            assert_eq!(
+                error_index,
+                StakePoolError::IncorrectWithdrawVoteAddress as u32
+            );
+        }
+        _ => panic!("Wrong error occurs while try to make a deposit with wrong stake program ID"),
+    }
+}