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
This commit is contained in:
Athar Mohammad 2022-08-23 03:06:13 +05:30 committed by GitHub
parent babf51941a
commit 2198ee068b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 426 additions and 29 deletions

View File

@ -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
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"

View File

@ -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,
);
}
// Use separate mutable variable because withdraw might create a new account
let stake_receiver = stake_receiver_param.unwrap_or_else(|| {
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,
@ -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,

View File

@ -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) {
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,

View File

@ -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;

View File

@ -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);
});
});
});

View File

@ -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);