From ee0c0c4a596a8de11d944af6efc33e6a0a00298c Mon Sep 17 00:00:00 2001 From: Michael Vines Date: Mon, 26 Apr 2021 08:33:56 -0700 Subject: [PATCH] Add ALL support to withdraw-stake subcommand (cherry picked from commit cf779c63c5d5c0e2c348a5ceb0fd0b8dc4f98916) --- cli/src/cli.rs | 12 ++++---- cli/src/stake.rs | 73 ++++++++++++++++++++++++++++------------------ cli/tests/stake.rs | 4 +-- 3 files changed, 52 insertions(+), 37 deletions(-) diff --git a/cli/src/cli.rs b/cli/src/cli.rs index ba12c585ec..cc29f792c4 100644 --- a/cli/src/cli.rs +++ b/cli/src/cli.rs @@ -292,7 +292,7 @@ pub enum CliCommand { WithdrawStake { stake_account_pubkey: Pubkey, destination_account_pubkey: Pubkey, - lamports: u64, + amount: SpendAmount, withdraw_authority: SignerIndex, custodian: Option, sign_only: bool, @@ -710,7 +710,7 @@ pub fn parse_command( } // Stake Commands ("create-stake-account", Some(matches)) => { - parse_stake_create_account(matches, default_signer, wallet_manager) + parse_create_stake_account(matches, default_signer, wallet_manager) } ("delegate-stake", Some(matches)) => { parse_stake_delegate_stake(matches, default_signer, wallet_manager) @@ -1721,7 +1721,7 @@ pub fn process_command(config: &CliConfig) -> ProcessResult { CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey, - lamports, + amount, withdraw_authority, custodian, sign_only, @@ -1737,7 +1737,7 @@ pub fn process_command(config: &CliConfig) -> ProcessResult { config, &stake_account_pubkey, &destination_account_pubkey, - *lamports, + *amount, *withdraw_authority, *custodian, *sign_only, @@ -2683,7 +2683,7 @@ mod tests { config.command = CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey: to_pubkey, - lamports: 100, + amount: SpendAmount::All, withdraw_authority: 0, custodian: None, sign_only: false, @@ -2752,7 +2752,7 @@ mod tests { }; config.signers = vec![&keypair, &merge_stake_account]; let result = process_command(&config); - assert!(dbg!(result).is_ok()); + assert!(result.is_ok()); config.command = CliCommand::GetSlot; assert_eq!(process_command(&config).unwrap(), "0"); diff --git a/cli/src/stake.rs b/cli/src/stake.rs index a849e6661b..d039f4f974 100644 --- a/cli/src/stake.rs +++ b/cli/src/stake.rs @@ -348,9 +348,9 @@ impl StakeSubCommands for App<'_, '_> { .index(3) .value_name("AMOUNT") .takes_value(true) - .validator(is_amount) + .validator(is_amount_or_all) .required(true) - .help("The amount to withdraw from the stake account, in SOL") + .help("The amount to withdraw from the stake account, in SOL; accepts keyword ALL") ) .arg( Arg::with_name("seed") @@ -474,7 +474,7 @@ impl StakeSubCommands for App<'_, '_> { } } -pub fn parse_stake_create_account( +pub fn parse_create_stake_account( matches: &ArgMatches<'_>, default_signer: &DefaultSigner, wallet_manager: &mut Option>, @@ -814,7 +814,7 @@ pub fn parse_stake_withdraw_stake( pubkey_of_signer(matches, "stake_account_pubkey", wallet_manager)?.unwrap(); let destination_account_pubkey = pubkey_of_signer(matches, "destination_account_pubkey", wallet_manager)?.unwrap(); - let lamports = lamports_of_sol(matches, "amount").unwrap(); + let amount = SpendAmount::new_from_matches(matches, "amount"); let sign_only = matches.is_present(SIGN_ONLY_ARG.name); let dump_transaction_message = matches.is_present(DUMP_TRANSACTION_MESSAGE.name); let blockhash_query = BlockhashQuery::new_from_matches(matches); @@ -842,7 +842,7 @@ pub fn parse_stake_withdraw_stake( command: CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey, - lamports, + amount, withdraw_authority: signer_info.index_of(withdraw_authority_pubkey).unwrap(), sign_only, dump_transaction_message, @@ -1244,7 +1244,7 @@ pub fn process_withdraw_stake( config: &CliConfig, stake_account_pubkey: &Pubkey, destination_account_pubkey: &Pubkey, - lamports: u64, + amount: SpendAmount, withdraw_authority: SignerIndex, custodian: Option, sign_only: bool, @@ -1256,8 +1256,6 @@ pub fn process_withdraw_stake( seed: Option<&String>, fee_payer: SignerIndex, ) -> ProcessResult { - let (recent_blockhash, fee_calculator) = - blockhash_query.get_blockhash_and_fee_calculator(rpc_client, config.commitment)?; let withdraw_authority = config.signers[withdraw_authority]; let custodian = custodian.map(|index| config.signers[index]); @@ -1267,28 +1265,45 @@ pub fn process_withdraw_stake( *stake_account_pubkey }; - let ixs = vec![stake_instruction::withdraw( - &stake_account_address, - &withdraw_authority.pubkey(), - destination_account_pubkey, - lamports, - custodian.map(|signer| signer.pubkey()).as_ref(), - )] - .with_memo(memo); + let (recent_blockhash, fee_calculator) = + blockhash_query.get_blockhash_and_fee_calculator(rpc_client, config.commitment)?; let fee_payer = config.signers[fee_payer]; let nonce_authority = config.signers[nonce_authority]; - let message = if let Some(nonce_account) = &nonce_account { - Message::new_with_nonce( - ixs, - Some(&fee_payer.pubkey()), - nonce_account, - &nonce_authority.pubkey(), - ) - } else { - Message::new(&ixs, Some(&fee_payer.pubkey())) + let build_message = |lamports| { + let ixs = vec![stake_instruction::withdraw( + &stake_account_address, + &withdraw_authority.pubkey(), + destination_account_pubkey, + lamports, + custodian.map(|signer| signer.pubkey()).as_ref(), + )] + .with_memo(memo); + + if let Some(nonce_account) = &nonce_account { + Message::new_with_nonce( + ixs, + Some(&fee_payer.pubkey()), + nonce_account, + &nonce_authority.pubkey(), + ) + } else { + Message::new(&ixs, Some(&fee_payer.pubkey())) + } }; + + let (message, _) = resolve_spend_tx_and_check_account_balances( + rpc_client, + sign_only, + amount, + &fee_calculator, + &stake_account_address, + &fee_payer.pubkey(), + build_message, + config.commitment, + )?; + let mut tx = Transaction::new_unsigned(message); if sign_only { @@ -3096,7 +3111,7 @@ mod tests { command: CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey: stake_account_pubkey, - lamports: 42_000_000_000, + amount: SpendAmount::Some(42_000_000_000), withdraw_authority: 0, custodian: None, sign_only: false, @@ -3129,7 +3144,7 @@ mod tests { command: CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey: stake_account_pubkey, - lamports: 42_000_000_000, + amount: SpendAmount::Some(42_000_000_000), withdraw_authority: 1, custodian: None, sign_only: false, @@ -3167,7 +3182,7 @@ mod tests { command: CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey: stake_account_pubkey, - lamports: 42_000_000_000, + amount: SpendAmount::Some(42_000_000_000), withdraw_authority: 0, custodian: Some(1), sign_only: false, @@ -3213,7 +3228,7 @@ mod tests { command: CliCommand::WithdrawStake { stake_account_pubkey, destination_account_pubkey: stake_account_pubkey, - lamports: 42_000_000_000, + amount: SpendAmount::Some(42_000_000_000), withdraw_authority: 0, custodian: None, sign_only: false, diff --git a/cli/tests/stake.rs b/cli/tests/stake.rs index 59cd4f1dd9..2f7ed02d64 100644 --- a/cli/tests/stake.rs +++ b/cli/tests/stake.rs @@ -1455,7 +1455,7 @@ fn test_offline_nonced_create_stake_account_and_withdraw() { config_offline.command = CliCommand::WithdrawStake { stake_account_pubkey: stake_pubkey, destination_account_pubkey: recipient_pubkey, - lamports: 42, + amount: SpendAmount::Some(42), withdraw_authority: 0, custodian: None, sign_only: true, @@ -1474,7 +1474,7 @@ fn test_offline_nonced_create_stake_account_and_withdraw() { config.command = CliCommand::WithdrawStake { stake_account_pubkey: stake_pubkey, destination_account_pubkey: recipient_pubkey, - lamports: 42, + amount: SpendAmount::Some(42), withdraw_authority: 0, custodian: None, sign_only: false,