Program instruction to withdraw un-staked lamports from stake account (#4780)

programs/stake_api/src/stake_instruction.rs

@@ -31,6 +31,17 @@ pub enum StakeInstruction {
     ///    2 - RewardsPool Stake Account from which to redeem credits
     ///    3 - Rewards syscall Account that carries points values
     RedeemVoteCredits,
+
+    /// Withdraw unstaked lamports from the stake account
+    ///
+    /// Expects 3 Accounts:
+    ///    0 - Delegate StakeAccount
+    ///    1 - System account to which the lamports will be transferred,
+    ///    2 - Syscall Account that carries epoch
+    ///
+    /// The u64 is the portion of the Stake account balance to be withdrawn,
+    ///    must be <= StakeAccount.lamports - staked lamports
+    Withdraw(u64),
 }
 
 pub fn create_stake_account(
@@ -77,6 +88,15 @@ pub fn delegate_stake(stake_pubkey: &Pubkey, vote_pubkey: &Pubkey, stake: u64) -
     Instruction::new(id(), &StakeInstruction::DelegateStake(stake), account_metas)
 }
 
+pub fn withdraw(stake_pubkey: &Pubkey, to_pubkey: &Pubkey, lamports: u64) -> Instruction {
+    let account_metas = vec![
+        AccountMeta::new(*stake_pubkey, true),
+        AccountMeta::new(*to_pubkey, false),
+        AccountMeta::new(syscall::current::id(), false),
+    ];
+    Instruction::new(id(), &StakeInstruction::Withdraw(lamports), account_metas)
+}
+
 pub fn process_instruction(
     _program_id: &Pubkey,
     keyed_accounts: &mut [KeyedAccount],
@@ -123,6 +143,19 @@ pub fn process_instruction(
                 &syscall::rewards::from_keyed_account(&rest[0])?,
             )
         }
+        StakeInstruction::Withdraw(lamports) => {
+            if rest.len() != 2 {
+                Err(InstructionError::InvalidInstructionData)?;
+            }
+            let (to, syscall) = &mut rest.split_at_mut(1);
+            let mut to = &mut to[0];
+
+            me.withdraw(
+                lamports,
+                &mut to,
+                &syscall::current::from_keyed_account(&syscall[0])?,
+            )
+        }
     }
 }
 
@@ -168,6 +201,10 @@ mod tests {
             process_instruction(&delegate_stake(&Pubkey::default(), &Pubkey::default(), 0)),
             Err(InstructionError::InvalidAccountData),
         );
+        assert_eq!(
+            process_instruction(&withdraw(&Pubkey::default(), &Pubkey::new_rand(), 100)),
+            Err(InstructionError::InvalidAccountData),
+        );
     }
 
     #[test]
@@ -250,6 +287,41 @@ mod tests {
             ),
             Err(InstructionError::InvalidAccountData),
         );
+
+        // Tests 3rd keyed account is of correct type (Current instead of rewards) in withdraw
+        assert_eq!(
+            super::process_instruction(
+                &Pubkey::default(),
+                &mut [
+                    KeyedAccount::new(&Pubkey::default(), false, &mut Account::default()),
+                    KeyedAccount::new(&Pubkey::default(), false, &mut Account::default()),
+                    KeyedAccount::new(
+                        &syscall::rewards::id(),
+                        false,
+                        &mut syscall::rewards::create_account(1, 0.0, 0.0)
+                    ),
+                ],
+                &serialize(&StakeInstruction::Withdraw(42)).unwrap(),
+            ),
+            Err(InstructionError::InvalidArgument),
+        );
+
+        // Tests correct number of accounts are provided in withdraw
+        assert_eq!(
+            super::process_instruction(
+                &Pubkey::default(),
+                &mut [
+                    KeyedAccount::new(&Pubkey::default(), false, &mut Account::default()),
+                    KeyedAccount::new(
+                        &syscall::current::id(),
+                        false,
+                        &mut syscall::rewards::create_account(1, 0.0, 0.0)
+                    ),
+                ],
+                &serialize(&StakeInstruction::Withdraw(42)).unwrap(),
+            ),
+            Err(InstructionError::InvalidInstructionData),
+        );
     }
 
 }

programs/stake_api/src/stake_state.rs

@@ -12,6 +12,7 @@ use solana_sdk::pubkey::Pubkey;
 use solana_sdk::syscall;
 use solana_sdk::timing::Epoch;
 use solana_vote_api::vote_state::VoteState;
+use std::cmp;
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
 pub enum StakeState {
@@ -196,6 +197,12 @@ pub trait StakeAccount {
         rewards_account: &mut KeyedAccount,
         rewards: &syscall::rewards::Rewards,
     ) -> Result<(), InstructionError>;
+    fn withdraw(
+        &mut self,
+        lamports: u64,
+        to: &mut KeyedAccount,
+        current: &syscall::current::Current,
+    ) -> Result<(), InstructionError>;
 }
 
 impl<'a> StakeAccount for KeyedAccount<'a> {
@@ -281,6 +288,44 @@ impl<'a> StakeAccount for KeyedAccount<'a> {
             Err(InstructionError::InvalidAccountData)
         }
     }
+    fn withdraw(
+        &mut self,
+        lamports: u64,
+        to: &mut KeyedAccount,
+        current: &syscall::current::Current,
+    ) -> Result<(), InstructionError> {
+        if self.signer_key().is_none() {
+            return Err(InstructionError::MissingRequiredSignature);
+        }
+
+        match self.state()? {
+            StakeState::Stake(mut stake) => {
+                let staked = if stake.stake(current.epoch) == 0 {
+                    0
+                } else {
+                    // Assume full stake if the stake is under warmup/cooldown
+                    stake.stake
+                };
+                if lamports > self.account.lamports.saturating_sub(staked) {
+                    return Err(InstructionError::InsufficientFunds);
+                }
+                self.account.lamports -= lamports;
+                // Adjust the stake (in case balance dropped below stake)
+                stake.stake = cmp::min(stake.stake, self.account.lamports);
+                to.account.lamports += lamports;
+                Ok(())
+            }
+            StakeState::Uninitialized => {
+                if lamports > self.account.lamports {
+                    return Err(InstructionError::InsufficientFunds);
+                }
+                self.account.lamports -= lamports;
+                to.account.lamports += lamports;
+                Ok(())
+            }
+            _ => Err(InstructionError::InvalidAccountData),
+        }
+    }
 }
 
 // utility function, used by Bank, tests, genesis
@@ -316,6 +361,7 @@ mod tests {
     use solana_sdk::account::Account;
     use solana_sdk::pubkey::Pubkey;
     use solana_sdk::signature::{Keypair, KeypairUtil};
+    use solana_sdk::system_program;
     use solana_vote_api::vote_state;
 
     #[test]
@@ -418,6 +464,143 @@ mod tests {
         assert_eq!(stake.stake(STAKE_WARMUP_EPOCHS * 42), 0);
     }
 
+    #[test]
+    fn test_withdraw_stake() {
+        let stake_pubkey = Pubkey::new_rand();
+        let mut total_lamports = 100;
+        let stake_lamports = 42;
+        let mut stake_account =
+            Account::new(total_lamports, std::mem::size_of::<StakeState>(), &id());
+
+        let current = syscall::current::Current::default();
+
+        let to = Pubkey::new_rand();
+        let mut to_account = Account::new(1, 0, &system_program::id());
+        let mut to_keyed_account = KeyedAccount::new(&to, false, &mut to_account);
+
+        // unsigned keyed account
+        let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, false, &mut stake_account);
+        assert_eq!(
+            stake_keyed_account.withdraw(total_lamports, &mut to_keyed_account, &current),
+            Err(InstructionError::MissingRequiredSignature)
+        );
+
+        // signed keyed account but uninitialized
+        // try withdrawing more than balance
+        let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &mut stake_account);
+        assert_eq!(
+            stake_keyed_account.withdraw(total_lamports + 1, &mut to_keyed_account, &current),
+            Err(InstructionError::InsufficientFunds)
+        );
+
+        // try withdrawing some (enough for rest of the test to carry forward)
+        let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &mut stake_account);
+        assert_eq!(
+            stake_keyed_account.withdraw(5, &mut to_keyed_account, &current),
+            Ok(())
+        );
+        total_lamports -= 5;
+
+        // Stake some lamports (available lampoorts for withdrawls will reduce)
+        let vote_pubkey = Pubkey::new_rand();
+        let mut vote_account =
+            vote_state::create_account(&vote_pubkey, &Pubkey::new_rand(), 0, 100);
+        let mut vote_keyed_account = KeyedAccount::new(&vote_pubkey, false, &mut vote_account);
+        vote_keyed_account.set_state(&VoteState::default()).unwrap();
+        assert_eq!(
+            stake_keyed_account.delegate_stake(&vote_keyed_account, stake_lamports, &current),
+            Ok(())
+        );
+
+        // Try to withdraw more than what's available
+        assert_eq!(
+            stake_keyed_account.withdraw(
+                total_lamports - stake_lamports + 1,
+                &mut to_keyed_account,
+                &current
+            ),
+            Err(InstructionError::InsufficientFunds)
+        );
+
+        // Try to withdraw all unstaked lamports
+        assert_eq!(
+            stake_keyed_account.withdraw(
+                total_lamports - stake_lamports,
+                &mut to_keyed_account,
+                &current
+            ),
+            Ok(())
+        );
+    }
+
+    #[test]
+    fn test_withdraw_stake_before_warmup() {
+        let stake_pubkey = Pubkey::new_rand();
+        let total_lamports = 100;
+        let stake_lamports = 42;
+        let mut stake_account =
+            Account::new(total_lamports, std::mem::size_of::<StakeState>(), &id());
+
+        let current = syscall::current::Current::default();
+        let mut future = syscall::current::Current::default();
+        future.epoch += 16;
+
+        let to = Pubkey::new_rand();
+        let mut to_account = Account::new(1, 0, &system_program::id());
+        let mut to_keyed_account = KeyedAccount::new(&to, false, &mut to_account);
+
+        let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &mut stake_account);
+
+        // Stake some lamports (available lampoorts for withdrawls will reduce)
+        let vote_pubkey = Pubkey::new_rand();
+        let mut vote_account =
+            vote_state::create_account(&vote_pubkey, &Pubkey::new_rand(), 0, 100);
+        let mut vote_keyed_account = KeyedAccount::new(&vote_pubkey, false, &mut vote_account);
+        vote_keyed_account.set_state(&VoteState::default()).unwrap();
+        assert_eq!(
+            stake_keyed_account.delegate_stake(&vote_keyed_account, stake_lamports, &future),
+            Ok(())
+        );
+
+        // Try to withdraw including staked
+        assert_eq!(
+            stake_keyed_account.withdraw(
+                total_lamports - stake_lamports + 1,
+                &mut to_keyed_account,
+                &current
+            ),
+            Ok(())
+        );
+    }
+
+    #[test]
+    fn test_withdraw_stake_invalid_state() {
+        let stake_pubkey = Pubkey::new_rand();
+        let total_lamports = 100;
+        let mut stake_account =
+            Account::new(total_lamports, std::mem::size_of::<StakeState>(), &id());
+
+        let current = syscall::current::Current::default();
+        let mut future = syscall::current::Current::default();
+        future.epoch += 16;
+
+        let to = Pubkey::new_rand();
+        let mut to_account = Account::new(1, 0, &system_program::id());
+        let mut to_keyed_account = KeyedAccount::new(&to, false, &mut to_account);
+
+        let mut stake_keyed_account = KeyedAccount::new(&stake_pubkey, true, &mut stake_account);
+        let stake_state = StakeState::MiningPool {
+            epoch: 0,
+            point_value: 0.0,
+        };
+        stake_keyed_account.set_state(&stake_state).unwrap();
+
+        assert_eq!(
+            stake_keyed_account.withdraw(total_lamports, &mut to_keyed_account, &current),
+            Err(InstructionError::InvalidAccountData)
+        );
+    }
+
     #[test]
     fn test_stake_state_calculate_rewards() {
         let mut vote_state = VoteState::default();

wallet/src/wallet.rs

@@ -52,6 +52,7 @@ pub enum WalletCommand {
     ShowVoteAccount(Pubkey),
     CreateStakeAccount(Pubkey, u64),
     DelegateStake(Keypair, Pubkey, u64),
+    WithdrawStake(Keypair, Pubkey, u64),
     RedeemVoteCredits(Pubkey, Pubkey),
     ShowStakeAccount(Pubkey),
     CreateStorageMiningPoolAccount(Pubkey, u64),
@@ -248,6 +249,18 @@ pub fn parse_command(
                 stake,
             ))
         }
+        ("withdraw-stake", Some(matches)) => {
+            let staking_account_keypair =
+                keypair_of(matches, "staking_account_keypair_file").unwrap();
+            let destination_account_pubkey =
+                value_of(matches, "destination_account_pubkey").unwrap();
+            let lamports = matches.value_of("lamports").unwrap().parse()?;
+            Ok(WalletCommand::WithdrawStake(
+                staking_account_keypair,
+                destination_account_pubkey,
+                lamports,
+            ))
+        }
         ("redeem-vote-credits", Some(matches)) => {
             let staking_account_pubkey = value_of(matches, "staking_account_pubkey").unwrap();
             let voting_account_pubkey = value_of(matches, "voting_account_pubkey").unwrap();
@@ -584,6 +597,32 @@ fn process_delegate_stake(
     Ok(signature_str.to_string())
 }
 
+fn process_withdraw_stake(
+    rpc_client: &RpcClient,
+    config: &WalletConfig,
+    staking_account_keypair: &Keypair,
+    destination_account_pubkey: &Pubkey,
+    lamports: u64,
+) -> ProcessResult {
+    let (recent_blockhash, _fee_calculator) = rpc_client.get_recent_blockhash()?;
+    let ixs = vec![stake_instruction::withdraw(
+        &staking_account_keypair.pubkey(),
+        destination_account_pubkey,
+        lamports,
+    )];
+
+    let mut tx = Transaction::new_signed_with_payer(
+        ixs,
+        Some(&config.keypair.pubkey()),
+        &[&config.keypair, &staking_account_keypair],
+        recent_blockhash,
+    );
+
+    let signature_str = rpc_client
+        .send_and_confirm_transaction(&mut tx, &[&config.keypair, &staking_account_keypair])?;
+    Ok(signature_str.to_string())
+}
+
 fn process_redeem_vote_credits(
     rpc_client: &RpcClient,
     config: &WalletConfig,
@@ -1022,6 +1061,18 @@ pub fn process_command(config: &WalletConfig) -> ProcessResult {
             )
         }
 
+        WalletCommand::WithdrawStake(
+            staking_account_keypair,
+            destination_account_pubkey,
+            lamports,
+        ) => process_withdraw_stake(
+            &rpc_client,
+            config,
+            &staking_account_keypair,
+            &destination_account_pubkey,
+            *lamports,
+        ),
+
         WalletCommand::RedeemVoteCredits(staking_account_pubkey, voting_account_pubkey) => {
             process_redeem_vote_credits(
                 &rpc_client,
@@ -1401,6 +1452,35 @@ pub fn app<'ab, 'v>(name: &str, about: &'ab str, version: &'v str) -> App<'ab, '
                         .help("The number of lamports to stake, must be less than the stake account's balance."),
                 ),
         )
+        .subcommand(
+            SubCommand::with_name("withdraw-stake")
+                .about("Withdraw the unstaked lamports from the stake account")
+                .arg(
+                    Arg::with_name("staking_account_keypair_file")
+                        .index(1)
+                        .value_name("KEYPAIR_FILE")
+                        .takes_value(true)
+                        .required(true)
+                        .help("Keypair file for the staking account, for signing the withdraw transaction."),
+                )
+                .arg(
+                    Arg::with_name("destination_account_pubkey")
+                        .index(2)
+                        .value_name("PUBKEY")
+                        .takes_value(true)
+                        .required(true)
+                        .validator(is_pubkey)
+                        .help("The account where the lamports should be transfered"),
+                )
+                .arg(
+                    Arg::with_name("lamports")
+                        .index(3)
+                        .value_name("NUM")
+                        .takes_value(true)
+                        .required(true)
+                        .help("The number of lamports to to withdraw from the stake account."),
+                ),
+        )
         .subcommand(
             SubCommand::with_name("redeem-vote-credits")
                 .about("Redeem credits in the staking account")
@@ -1837,6 +1917,22 @@ mod tests {
             WalletCommand::DelegateStake(keypair, pubkey, 42)
         );
 
+        let keypair_file = make_tmp_path("keypair_file");
+        gen_keypair_file(&keypair_file).unwrap();
+        let keypair = read_keypair(&keypair_file).unwrap();
+        // Test Withdraw from Stake Account
+        let test_withdraw_stake = test_commands.clone().get_matches_from(vec![
+            "test",
+            "withdraw-stake",
+            &keypair_file,
+            &pubkey_string,
+            "42",
+        ]);
+        assert_eq!(
+            parse_command(&pubkey, &test_withdraw_stake).unwrap(),
+            WalletCommand::WithdrawStake(keypair, pubkey, 42)
+        );
+
         // Test Deploy Subcommand
         let test_deploy =
             test_commands
@@ -2006,6 +2102,12 @@ mod tests {
         let signature = process_command(&config);
         assert_eq!(signature.unwrap(), SIGNATURE.to_string());
 
+        let bob_keypair = Keypair::new();
+        let to_pubkey = Pubkey::new_rand();
+        config.command = WalletCommand::WithdrawStake(bob_keypair.into(), to_pubkey, 100);
+        let signature = process_command(&config);
+        assert_eq!(signature.unwrap(), SIGNATURE.to_string());
+
         config.command = WalletCommand::GetTransactionCount;
         assert_eq!(process_command(&config).unwrap(), "1234");