From 2198ee068bc9b5454cd55bd7200c6a7ba40c44e7 Mon Sep 17 00:00:00 2001 From: Athar Mohammad <56029409+atharmohammad@users.noreply.github.com> Date: Tue, 23 Aug 2022 03:06:13 +0530 Subject: [PATCH] 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 | 37 ++++++++- stake-pool/cli/src/main.rs | 106 ++++++++++++++++++++---- stake-pool/js/src/index.ts | 92 +++++++++++++++++++- stake-pool/js/src/layouts.ts | 62 ++++++++++++++ stake-pool/js/test/instructions.test.ts | 80 ++++++++++++++++-- stake-pool/js/test/mocks.ts | 78 ++++++++++++++++- 6 files changed, 426 insertions(+), 29 deletions(-) diff --git a/stake-pool/cli/scripts/withdraw.sh b/stake-pool/cli/scripts/withdraw.sh index deeee7a0..733acec6 100755 --- a/stake-pool/cli/scripts/withdraw.sh +++ b/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") -keys_dir=keys +withdraw_stakes_to_stake_receiver () { + 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" diff --git a/stake-pool/cli/src/main.rs b/stake-pool/cli/src/main.rs index 299d0aac..07315aac 100644 --- a/stake-pool/cli/src/main.rs +++ b/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,19 +1528,21 @@ fn command_withdraw_stake( withdraw_account.stake_address, ); } - - // Use separate mutable variable because withdraw might create a new account - let stake_receiver = stake_receiver_param.unwrap_or_else(|| { - let stake_keypair = new_stake_account( - &config.fee_payer.pubkey(), - &mut instructions, - stake_account_rent_exemption, - ); - let stake_pubkey = stake_keypair.pubkey(); - total_rent_free_balances += stake_account_rent_exemption; - new_stake_keypairs.push(stake_keypair); - stake_pubkey - }); + let stake_receiver = + if (stake_receiver_param.is_none()) || (maybe_stake_receiver_state.is_some()) { + // Creating new account to split the stake into new account + let stake_keypair = new_stake_account( + &config.fee_payer.pubkey(), + &mut instructions, + stake_account_rent_exemption, + ); + let stake_pubkey = stake_keypair.pubkey(); + 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, diff --git a/stake-pool/js/src/index.ts b/stake-pool/js/src/index.ts index 7eac0b4d..6cd146b3 100644 --- a/stake-pool/js/src/index.ts +++ b/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 { + 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) { + if (!stakeReceiver || (stakeReceiverAccount && stakeReceiverAccount.type === 'delegated')) { 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, diff --git a/stake-pool/js/src/layouts.ts b/stake-pool/js/src/layouts.ts index ba5d9489..e7cd77b6 100644 --- a/stake-pool/js/src/layouts.ts +++ b/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; +export const StakeAccountType = enums(['uninitialized', 'initialized', 'delegated', 'rewardsPool']); + +export type StakeMeta = Infer; +export const StakeMeta = type({ + rentExemptReserve: BigNumFromString, + authorized: type({ + staker: PublicKeyFromString, + withdrawer: PublicKeyFromString, + }), + lockup: type({ + unixTimestamp: number(), + epoch: number(), + custodian: PublicKeyFromString, + }), +}); + +export type StakeAccountInfo = Infer; +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; +export const StakeAccount = type({ + type: StakeAccountType, + info: optional(StakeAccountInfo), +}); + export interface StakePool { accountType: AccountType; manager: PublicKey; diff --git a/stake-pool/js/test/instructions.test.ts b/stake-pool/js/test/instructions.test.ts index 8a9681b1..458eb0cb 100644 --- a/stake-pool/js/test/instructions.test.ts +++ b/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); + }); }); }); diff --git a/stake-pool/js/test/mocks.ts b/stake-pool/js/test/mocks.ts index df2c0c04..db93fa5a 100644 --- a/stake-pool/js/test/mocks.ts +++ b/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 >{ + executable: false, + owner: StakeProgram.programId, + lamports: 3000000000, + data, + }; +} + export function mockValidatorList() { const data = Buffer.alloc(1024); ValidatorListLayout.encode(validatorListMock, data);