Extends withdraw functionality in stake pool (#3445)

* extends withdraw to merge if stake account is provided

* check if account is stake account and delegated to same validator

* fixing tests and defining withdraw accounts for delegated stake reciever

* implementation improvements in cli

* fix js tests

* added test for uninitialized stake account

stake-pool/cli/scripts/withdraw.sh

@@ -15,6 +15,16 @@ create_keypair () {
   fi
 }
 
+create_stake_account () {
+  authority=$1
+  while read -r validator
+  do
+    solana-keygen new --no-passphrase -o "$keys_dir/stake_account_$validator.json"
+    solana create-stake-account "$keys_dir/stake_account_$validator.json" 2 
+    solana delegate-stake --force "$keys_dir/stake_account_$validator.json"  "$validator" 
+  done < "$validator_list"
+}
+
 withdraw_stakes () {
   stake_pool_pubkey=$1
   validator_list=$2
@@ -25,20 +35,41 @@ withdraw_stakes () {
   done < "$validator_list"
 }
 
-stake_pool_pubkey=$(solana-keygen pubkey "$stake_pool_keyfile")
+withdraw_stakes_to_stake_receiver () {
-keys_dir=keys
+  stake_pool_pubkey=$1
+  validator_list=$2
+  pool_amount=$3
+  while read -r validator
+  do
+    stake_receiver=$(solana-keygen pubkey "$keys_dir/stake_account_$validator.json")
+    $spl_stake_pool withdraw-stake "$stake_pool_pubkey" "$pool_amount" --vote-account "$validator" --stake-receiver "$stake_receiver"
+  done < "$validator_list"
+}
 
 spl_stake_pool=spl-stake-pool
 # Uncomment to use a locally build CLI
-#spl_stake_pool=../../../target/debug/spl-stake-pool
+# spl_stake_pool=../../../target/debug/spl-stake-pool
+
+stake_pool_pubkey=$(solana-keygen pubkey "$stake_pool_keyfile")
+keys_dir=keys
 
 echo "Setting up keys directory $keys_dir"
 mkdir -p $keys_dir
 authority=$keys_dir/authority.json
+
+create_stake_account $authority
+echo "Waiting for stakes to activate, this may take awhile depending on the network!"
+echo "If you are running on localnet with 32 slots per epoch, wait 12 seconds..."
+sleep 12
+
 echo "Setting up authority for withdrawn stake accounts at $authority"
 create_keypair $authority
 
 echo "Withdrawing stakes from stake pool"
 withdraw_stakes "$stake_pool_pubkey" "$validator_list" "$withdraw_sol_amount"
+
+echo "Withdrawing stakes from stake pool to recieve it in stake receiver account"
+withdraw_stakes_to_stake_receiver "$stake_pool_pubkey" "$validator_list" "$withdraw_sol_amount"
+
 echo "Withdrawing SOL from stake pool to authority"
 $spl_stake_pool withdraw-sol "$stake_pool_pubkey" $authority "$withdraw_sol_amount"

stake-pool/cli/src/main.rs

@@ -6,6 +6,7 @@ use {
         client::*,
         output::{CliStakePool, CliStakePoolDetails, CliStakePoolStakeAccountInfo, CliStakePools},
     },
+    bincode::deserialize,
     clap::{
         crate_description, crate_name, crate_version, value_t, value_t_or_exit, App, AppSettings,
         Arg, ArgGroup, ArgMatches, SubCommand,
@@ -1370,12 +1371,77 @@ fn command_withdraw_stake(
         .into());
     }
 
+    // Check for the delegated stake receiver
+    let maybe_stake_receiver_state = stake_receiver_param
+        .map(|stake_receiver_pubkey| {
+            let stake_account = config.rpc_client.get_account(&stake_receiver_pubkey).ok()?;
+            let stake_state: stake::state::StakeState = deserialize(stake_account.data.as_slice())
+                .map_err(|err| format!("Invalid stake account {}: {}", stake_receiver_pubkey, err))
+                .ok()?;
+            if stake_state.delegation().is_some() && stake_account.owner == stake::program::id() {
+                Some(stake_state)
+            } else {
+                None
+            }
+        })
+        .flatten();
+
     let withdraw_accounts = if use_reserve {
         vec![WithdrawAccount {
             stake_address: stake_pool.reserve_stake,
             vote_address: None,
             pool_amount,
         }]
+    } else if maybe_stake_receiver_state.is_some() {
+        let vote_account = maybe_stake_receiver_state
+            .unwrap()
+            .delegation()
+            .unwrap()
+            .voter_pubkey;
+        if let Some(vote_account_address) = vote_account_address {
+            if *vote_account_address != vote_account {
+                return Err(format!("Provided withdrawal vote account {} does not match delegation on stake receiver account {}, 
+                remove this flag or provide a different stake account delegated to {}", vote_account_address, vote_account, vote_account_address).into());
+            }
+        }
+        // Check if the vote account exists in the stake pool
+        let validator_list = get_validator_list(&config.rpc_client, &stake_pool.validator_list)?;
+        if validator_list
+            .validators
+            .into_iter()
+            .any(|x| x.vote_account_address == vote_account)
+        {
+            let (stake_account_address, _) = find_stake_program_address(
+                &spl_stake_pool::id(),
+                &vote_account,
+                stake_pool_address,
+            );
+            let stake_account = config.rpc_client.get_account(&stake_account_address)?;
+
+            let available_for_withdrawal = stake_pool
+                .calc_lamports_withdraw_amount(
+                    stake_account
+                        .lamports
+                        .saturating_sub(MINIMUM_ACTIVE_STAKE)
+                        .saturating_sub(stake_account_rent_exemption),
+                )
+                .unwrap();
+
+            if available_for_withdrawal < pool_amount {
+                return Err(format!(
+                    "Not enough lamports available for withdrawal from {}, {} asked, {} available",
+                    stake_account_address, pool_amount, available_for_withdrawal
+                )
+                .into());
+            }
+            vec![WithdrawAccount {
+                stake_address: stake_account_address,
+                vote_address: Some(vote_account),
+                pool_amount,
+            }]
+        } else {
+            return Err(format!("Provided stake account is delegated to a vote account {} which does not exist in the stake pool", vote_account).into());
+        }
     } else if let Some(vote_account_address) = vote_account_address {
         let (stake_account_address, _) = find_stake_program_address(
             &spl_stake_pool::id(),
@@ -1439,7 +1505,6 @@ fn command_withdraw_stake(
     );
 
     let mut total_rent_free_balances = 0;
-
     // Go through prepared accounts and withdraw/claim them
     for withdraw_account in withdraw_accounts {
         // Convert pool tokens amount to lamports
@@ -1463,9 +1528,9 @@ fn command_withdraw_stake(
                 withdraw_account.stake_address,
             );
         }
-
+        let stake_receiver =
-        // Use separate mutable variable because withdraw might create a new account
+            if (stake_receiver_param.is_none()) || (maybe_stake_receiver_state.is_some()) {
-        let stake_receiver = stake_receiver_param.unwrap_or_else(|| {
+                // Creating new account to split the stake into new account
                 let stake_keypair = new_stake_account(
                     &config.fee_payer.pubkey(),
                     &mut instructions,
@@ -1475,7 +1540,9 @@ fn command_withdraw_stake(
                 total_rent_free_balances += stake_account_rent_exemption;
                 new_stake_keypairs.push(stake_keypair);
                 stake_pubkey
-        });
+            } else {
+                stake_receiver_param.unwrap()
+            };
 
         instructions.push(spl_stake_pool::instruction::withdraw_stake(
             &spl_stake_pool::id(),
@@ -1494,6 +1561,17 @@ fn command_withdraw_stake(
         ));
     }
 
+    // Merging the stake with account provided by user
+    if maybe_stake_receiver_state.is_some() {
+        for new_stake_keypair in &new_stake_keypairs {
+            instructions.extend(stake::instruction::merge(
+                &stake_receiver_param.unwrap(),
+                &new_stake_keypair.pubkey(),
+                &config.fee_payer.pubkey(),
+            ));
+        }
+    }
+
     let recent_blockhash = get_latest_blockhash(&config.rpc_client)?;
     let message = Message::new_with_blockhash(
         &instructions,

stake-pool/js/src/index.ts

@@ -27,6 +27,7 @@ import {
 } from './utils';
 import { StakePoolInstruction } from './instructions';
 import {
+  StakeAccount,
   StakePool,
   StakePoolLayout,
   ValidatorList,
@@ -34,6 +35,7 @@ import {
   ValidatorStakeInfo,
 } from './layouts';
 import { MAX_VALIDATORS_TO_UPDATE, MINIMUM_ACTIVE_STAKE, STAKE_POOL_PROGRAM_ID } from './constants';
+import { create } from 'superstruct';
 
 export type { StakePool, AccountType, ValidatorList, ValidatorStakeInfo } from './layouts';
 export { STAKE_POOL_PROGRAM_ID } from './constants';
@@ -90,6 +92,28 @@ export async function getStakePoolAccount(
   };
 }
 
+/**
+ * Retrieves and deserializes a Stake account using a web3js connection and the stake address.
+ * @param connection: An active web3js connection.
+ * @param stakeAccount: The public key (address) of the stake account.
+ */
+export async function getStakeAccount(
+  connection: Connection,
+  stakeAccount: PublicKey,
+): Promise<StakeAccount> {
+  const result = (await connection.getParsedAccountInfo(stakeAccount)).value;
+  if (!result || !('parsed' in result.data)) {
+    throw new Error('Invalid stake account');
+  }
+  const program = result.data.program;
+  if (program != 'stake') {
+    throw new Error('Not a stake account');
+  }
+  const parsed = create(result.data.parsed, StakeAccount);
+
+  return parsed;
+}
+
 /**
  * Retrieves all StakePool and ValidatorList accounts that are running a particular StakePool program.
  * @param connection: An active web3js connection.
@@ -331,6 +355,7 @@ export async function withdrawStake(
     poolTokenAccount,
     stakePool.account.data.poolMint,
   );
+
   if (!tokenAccount) {
     throw new Error('Invalid token account');
   }
@@ -352,6 +377,11 @@ export async function withdrawStake(
     stakePoolAddress,
   );
 
+  let stakeReceiverAccount = null;
+  if (stakeReceiver) {
+    stakeReceiverAccount = await getStakeAccount(connection, stakeReceiver);
+  }
+
   const withdrawAccounts: WithdrawAccount[] = [];
 
   if (useReserve) {
@@ -360,6 +390,53 @@ export async function withdrawStake(
       voteAddress: undefined,
       poolAmount,
     });
+  } else if (stakeReceiverAccount && stakeReceiverAccount?.type == 'delegated') {
+    const voteAccount = stakeReceiverAccount.info?.stake?.delegation.voter;
+    if (!voteAccount) throw new Error(`Invalid stake reciever ${stakeReceiver} delegation`);
+    const validatorListAccount = await connection.getAccountInfo(
+      stakePool.account.data.validatorList,
+    );
+    const validatorList = ValidatorListLayout.decode(validatorListAccount?.data) as ValidatorList;
+    const isValidVoter = validatorList.validators.find((val) =>
+      val.voteAccountAddress.equals(voteAccount),
+    );
+    if (voteAccountAddress && voteAccountAddress !== voteAccount) {
+      throw new Error(`Provided withdrawal vote account ${voteAccountAddress} does not match delegation on stake receiver account ${voteAccount}, 
+      remove this flag or provide a different stake account delegated to ${voteAccountAddress}`);
+    }
+    if (isValidVoter) {
+      const stakeAccountAddress = await findStakeProgramAddress(
+        STAKE_POOL_PROGRAM_ID,
+        voteAccount,
+        stakePoolAddress,
+      );
+
+      const stakeAccount = await connection.getAccountInfo(stakeAccountAddress);
+      if (!stakeAccount) {
+        throw new Error(`Preferred withdraw valdator's stake account is invalid`);
+      }
+
+      const availableForWithdrawal = calcLamportsWithdrawAmount(
+        stakePool.account.data,
+        stakeAccount.lamports - MINIMUM_ACTIVE_STAKE - stakeAccountRentExemption,
+      );
+
+      if (availableForWithdrawal < poolAmount) {
+        throw new Error(
+          `Not enough lamports available for withdrawal from ${stakeAccountAddress},
+            ${poolAmount} asked, ${availableForWithdrawal} available.`,
+        );
+      }
+      withdrawAccounts.push({
+        stakeAddress: stakeAccountAddress,
+        voteAddress: voteAccount,
+        poolAmount,
+      });
+    } else {
+      throw new Error(
+        `Provided stake account is delegated to a vote account ${voteAccount} which does not exist in the stake pool`,
+      );
+    }
   } else if (voteAccountAddress) {
     const stakeAccountAddress = await findStakeProgramAddress(
       STAKE_POOL_PROGRAM_ID,
@@ -443,11 +520,9 @@ export async function withdrawStake(
     }
 
     console.info(infoMsg);
-
     let stakeToReceive;
 
-    // Use separate mutable variable because withdraw might create a new account
+    if (!stakeReceiver || (stakeReceiverAccount && stakeReceiverAccount.type === 'delegated')) {
-    if (!stakeReceiver) {
       const stakeKeypair = newStakeAccount(tokenOwner, instructions, stakeAccountRentExemption);
       signers.push(stakeKeypair);
       totalRentFreeBalances += stakeAccountRentExemption;
@@ -473,6 +548,17 @@ export async function withdrawStake(
     );
     i++;
   }
+  if (stakeReceiver && stakeReceiverAccount && stakeReceiverAccount.type === 'delegated') {
+    signers.forEach((newStakeKeypair) => {
+      instructions.concat(
+        StakeProgram.merge({
+          stakePubkey: stakeReceiver,
+          sourceStakePubKey: newStakeKeypair.publicKey,
+          authorizedPubkey: tokenOwner,
+        }).instructions,
+      );
+    });
+  }
 
   return {
     instructions,

stake-pool/js/src/layouts.ts

@@ -2,6 +2,17 @@ import { publicKey, struct, u32, u64, u8, option, vec } from '@project-serum/bor
 import { Lockup, PublicKey } from '@solana/web3.js';
 import { AccountInfo } from '@solana/spl-token';
 import BN from 'bn.js';
+import {
+  Infer,
+  number,
+  nullable,
+  enums,
+  type,
+  coerce,
+  instance,
+  string,
+  optional,
+} from 'superstruct';
 
 export interface Fee {
   denominator: BN;
@@ -33,6 +44,57 @@ export enum AccountType {
   ValidatorList,
 }
 
+export const BigNumFromString = coerce(instance(BN), string(), (value) => {
+  if (typeof value === 'string') return new BN(value, 10);
+  throw new Error('invalid big num');
+});
+
+export const PublicKeyFromString = coerce(
+  instance(PublicKey),
+  string(),
+  (value) => new PublicKey(value),
+);
+
+export type StakeAccountType = Infer<typeof StakeAccountType>;
+export const StakeAccountType = enums(['uninitialized', 'initialized', 'delegated', 'rewardsPool']);
+
+export type StakeMeta = Infer<typeof StakeMeta>;
+export const StakeMeta = type({
+  rentExemptReserve: BigNumFromString,
+  authorized: type({
+    staker: PublicKeyFromString,
+    withdrawer: PublicKeyFromString,
+  }),
+  lockup: type({
+    unixTimestamp: number(),
+    epoch: number(),
+    custodian: PublicKeyFromString,
+  }),
+});
+
+export type StakeAccountInfo = Infer<typeof StakeAccountInfo>;
+export const StakeAccountInfo = type({
+  meta: StakeMeta,
+  stake: nullable(
+    type({
+      delegation: type({
+        voter: PublicKeyFromString,
+        stake: BigNumFromString,
+        activationEpoch: BigNumFromString,
+        deactivationEpoch: BigNumFromString,
+        warmupCooldownRate: number(),
+      }),
+      creditsObserved: number(),
+    }),
+  ),
+});
+
+export type StakeAccount = Infer<typeof StakeAccount>;
+export const StakeAccount = type({
+  type: StakeAccountType,
+  info: optional(StakeAccountInfo),
+});
+
 export interface StakePool {
   accountType: AccountType;
   manager: PublicKey;

stake-pool/js/test/instructions.test.ts

@@ -15,11 +15,21 @@ import {
   depositSol,
   withdrawSol,
   withdrawStake,
+  getStakeAccount,
 } from '../src';
 
 import { decodeData } from '../src/utils';
 
-import { mockTokenAccount, mockValidatorList, stakePoolMock } from './mocks';
+import {
+  mockRpc,
+  mockTokenAccount,
+  mockValidatorList,
+  mockValidatorsStakeAccount,
+  stakePoolMock,
+  CONSTANTS,
+  stakeAccountData,
+  uninitializedStakeAccount,
+} from './mocks';
 
 describe('StakePoolProgram', () => {
   const connection = new Connection('http://127.0.0.1:8899');
@@ -146,7 +156,7 @@ describe('StakePoolProgram', () => {
         if (pubKey == stakePoolAddress) {
           return stakePoolAccount;
         }
-        if (pubKey.toBase58() == '9q2rZU5RujvyD9dmYKhzJAZfG4aGBbvQ8rWY52jCNBai') {
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
           return null;
         }
         return null;
@@ -162,7 +172,7 @@ describe('StakePoolProgram', () => {
         if (pubKey == stakePoolAddress) {
           return stakePoolAccount;
         }
-        if (pubKey.toBase58() == 'GQkqTamwqjaNDfsbNm7r3aXPJ4oTSqKC3d5t2PF9Smqd') {
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
           return mockTokenAccount(0);
         }
         return null;
@@ -182,7 +192,7 @@ describe('StakePoolProgram', () => {
         if (pubKey == stakePoolAddress) {
           return stakePoolAccount;
         }
-        if (pubKey.toBase58() == 'GQkqTamwqjaNDfsbNm7r3aXPJ4oTSqKC3d5t2PF9Smqd') {
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
           return mockTokenAccount(LAMPORTS_PER_SOL);
         }
         return null;
@@ -216,7 +226,7 @@ describe('StakePoolProgram', () => {
         if (pubKey == stakePoolAddress) {
           return stakePoolAccount;
         }
-        if (pubKey.toBase58() == 'GQkqTamwqjaNDfsbNm7r3aXPJ4oTSqKC3d5t2PF9Smqd') {
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
           return mockTokenAccount(0);
         }
         return null;
@@ -235,15 +245,14 @@ describe('StakePoolProgram', () => {
         if (pubKey == stakePoolAddress) {
           return stakePoolAccount;
         }
-        if (pubKey.toBase58() == 'GQkqTamwqjaNDfsbNm7r3aXPJ4oTSqKC3d5t2PF9Smqd') {
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
           return mockTokenAccount(LAMPORTS_PER_SOL * 2);
         }
-        if (pubKey.toBase58() == stakePoolMock.validatorList.toBase58()) {
+        if (pubKey.equals(stakePoolMock.validatorList)) {
           return mockValidatorList();
         }
         return null;
       });
-
       const res = await withdrawStake(connection, stakePoolAddress, tokenOwner, 1);
 
       expect((connection.getAccountInfo as jest.Mock).mock.calls.length).toBe(4);
@@ -252,5 +261,60 @@ describe('StakePoolProgram', () => {
       expect(res.stakeReceiver).toEqual(undefined);
       expect(res.totalRentFreeBalances).toEqual(10000);
     });
+
+    it.only('withdraw to a stake account provided', async () => {
+      const stakeReceiver = new PublicKey(20);
+      connection.getAccountInfo = jest.fn(async (pubKey: PublicKey) => {
+        if (pubKey == stakePoolAddress) {
+          return stakePoolAccount;
+        }
+        if (pubKey.equals(CONSTANTS.poolTokenAccount)) {
+          return mockTokenAccount(LAMPORTS_PER_SOL * 2);
+        }
+        if (pubKey.equals(stakePoolMock.validatorList)) {
+          return mockValidatorList();
+        }
+        if (pubKey.equals(CONSTANTS.validatorStakeAccountAddress))
+          return mockValidatorsStakeAccount();
+        return null;
+      });
+      connection.getParsedAccountInfo = jest.fn(async (pubKey: PublicKey) => {
+        if (pubKey.equals(stakeReceiver)) {
+          return mockRpc(stakeAccountData);
+        }
+        return null;
+      });
+
+      const res = await withdrawStake(
+        connection,
+        stakePoolAddress,
+        tokenOwner,
+        1,
+        undefined,
+        undefined,
+        stakeReceiver,
+      );
+
+      expect((connection.getAccountInfo as jest.Mock).mock.calls.length).toBe(4);
+      expect((connection.getParsedAccountInfo as jest.Mock).mock.calls.length).toBe(1);
+      expect(res.instructions).toHaveLength(3);
+      expect(res.signers).toHaveLength(2);
+      expect(res.stakeReceiver).toEqual(stakeReceiver);
+      expect(res.totalRentFreeBalances).toEqual(10000);
+    });
+  });
+  describe('getStakeAccount', () => {
+    it.only('returns an uninitialized parsed stake account', async () => {
+      const stakeAccount = new PublicKey(20);
+      connection.getParsedAccountInfo = jest.fn(async (pubKey: PublicKey) => {
+        if (pubKey.equals(stakeAccount)) {
+          return mockRpc(uninitializedStakeAccount);
+        }
+        return null;
+      });
+      const parsedStakeAccount = await getStakeAccount(connection, stakeAccount);
+      expect((connection.getParsedAccountInfo as jest.Mock).mock.calls.length).toBe(1);
+      expect(parsedStakeAccount).toEqual(uninitializedStakeAccount.parsed);
+    });
   });
 });

stake-pool/js/test/mocks.ts

@@ -1,8 +1,17 @@
-import { AccountInfo, LAMPORTS_PER_SOL, PublicKey } from '@solana/web3.js';
+import { AccountInfo, LAMPORTS_PER_SOL, PublicKey, StakeProgram } from '@solana/web3.js';
 import BN from 'bn.js';
 import { ValidatorStakeInfo } from '../src';
 import { ValidatorStakeInfoStatus, AccountLayout, ValidatorListLayout } from '../src/layouts';
 
+export const CONSTANTS = {
+  poolTokenAccount: new PublicKey(
+    new BN('e4f53a3a11521b9171c942ff91183ec8db4e6f347bb9aa7d4a814b7874bfd15c', 'hex'),
+  ),
+  validatorStakeAccountAddress: new PublicKey(
+    new BN('69184b7f1bc836271c4ac0e29e53eb38a38ea0e7bcde693c45b30d1592a5a678', 'hex'),
+  ),
+};
+
 export const stakePoolMock = {
   accountType: 1,
   manager: new PublicKey(11),
@@ -132,6 +141,73 @@ export function mockTokenAccount(amount = 0) {
   };
 }
 
+export const mockRpc = (data: any): any => {
+  const value = {
+    owner: StakeProgram.programId,
+    lamports: LAMPORTS_PER_SOL,
+    data: data,
+    executable: false,
+    rentEpoch: 0,
+  };
+  const result = {
+    context: {
+      slot: 11,
+    },
+    value: value,
+  };
+  return result;
+};
+
+export const stakeAccountData = {
+  program: 'stake',
+  parsed: {
+    type: 'delegated',
+    info: {
+      meta: {
+        rentExemptReserve: new BN(1),
+        lockup: {
+          epoch: 32,
+          unixTimestamp: 2,
+          custodian: new PublicKey(12),
+        },
+        authorized: {
+          staker: new PublicKey(12),
+          withdrawer: new PublicKey(12),
+        },
+      },
+      stake: {
+        delegation: {
+          voter: new PublicKey(
+            new BN('e4e37d6f2e80c0bb0f3da8a06304e57be5cda6efa2825b86780aa320d9784cf8', 'hex'),
+          ),
+          stake: new BN(0),
+          activationEpoch: new BN(1),
+          deactivationEpoch: new BN(1),
+          warmupCooldownRate: 1.2,
+        },
+        creditsObserved: 1,
+      },
+    },
+  },
+};
+
+export const uninitializedStakeAccount = {
+  program: 'stake',
+  parsed: {
+    type: 'uninitialized',
+  },
+};
+
+export function mockValidatorsStakeAccount() {
+  const data = Buffer.alloc(1024);
+  return <AccountInfo<any>>{
+    executable: false,
+    owner: StakeProgram.programId,
+    lamports: 3000000000,
+    data,
+  };
+}
+
 export function mockValidatorList() {
   const data = Buffer.alloc(1024);
   ValidatorListLayout.encode(validatorListMock, data);