stake-pool: Use delegation.stake for accounting (#1934)

* stake-pool: Use `delegation.stake` for accounting

* Clippy

* Improve test and remove spl-math dependency

* Address smaller review comments

* Fix calculation, improve tests

* Clippy

* Remove esm from mocha

* Revert "Remove esm from mocha"

This reverts commit c0f25ab.

Author: joncinque

Cargo.lock

@@ -597,6 +597,7 @@ fn command_deposit(
             &stake,
             &config.staker.pubkey(),
             &validator_stake_account,
+            &stake_pool.reserve_stake,
             &token_receiver,
             &stake_pool.pool_mint,
             &spl_token::id(),
@@ -610,6 +611,7 @@ fn command_deposit(
             &stake,
             &config.staker.pubkey(),
             &validator_stake_account,
+            &stake_pool.reserve_stake,
             &token_receiver,
             &stake_pool.pool_mint,
             &spl_token::id(),

stake-pool/cli/src/main.rs

@@ -20,7 +20,6 @@ num_enum = "0.5.1"
 serde = "1.0.121"
 serde_derive = "1.0.103"
 solana-program = "1.6.11"
-spl-math = { version = "0.1", path = "../../libraries/math", features = [ "no-entrypoint" ] }
 spl-token = { version = "3.1", path = "../../token/program", features = [ "no-entrypoint" ] }
 thiserror = "1.0"
 bincode = "1.3.1"

stake-pool/program/Cargo.toml

@@ -236,9 +236,10 @@ pub enum StakePoolInstruction {
     ///   3. `[]` Stake pool withdraw authority
     ///   4. `[w]` Stake account to join the pool (withdraw authority for the stake account should be first set to the stake pool deposit authority)
     ///   5. `[w]` Validator stake account for the stake account to be merged with
-    ///   6. `[w]` User account to receive pool tokens
+    ///   6. `[w]` Reserve stake account, to withdraw rent exempt reserve
+    ///   7. `[w]` User account to receive pool tokens
     ///   8. `[w]` Pool token mint account
-    ///   9. '[]' Sysvar clock account (required)
+    ///   9. '[]' Sysvar clock account
     ///   10. '[]' Sysvar stake history account
     ///   11. `[]` Pool token program id,
     ///   12. `[]` Stake program id,
@@ -781,6 +782,7 @@ pub fn deposit(
     deposit_stake_address: &Pubkey,
     deposit_stake_withdraw_authority: &Pubkey,
     validator_stake_account: &Pubkey,
+    reserve_stake_account: &Pubkey,
     pool_tokens_to: &Pubkey,
     pool_mint: &Pubkey,
     token_program_id: &Pubkey,
@@ -794,6 +796,7 @@ pub fn deposit(
         AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
         AccountMeta::new(*deposit_stake_address, false),
         AccountMeta::new(*validator_stake_account, false),
+        AccountMeta::new(*reserve_stake_account, false),
         AccountMeta::new(*pool_tokens_to, false),
         AccountMeta::new(*pool_mint, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
@@ -834,6 +837,7 @@ pub fn deposit_with_authority(
     deposit_stake_address: &Pubkey,
     deposit_stake_withdraw_authority: &Pubkey,
     validator_stake_account: &Pubkey,
+    reserve_stake_account: &Pubkey,
     pool_tokens_to: &Pubkey,
     pool_mint: &Pubkey,
     token_program_id: &Pubkey,
@@ -845,6 +849,7 @@ pub fn deposit_with_authority(
         AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
         AccountMeta::new(*deposit_stake_address, false),
         AccountMeta::new(*validator_stake_account, false),
+        AccountMeta::new(*reserve_stake_account, false),
         AccountMeta::new(*pool_tokens_to, false),
         AccountMeta::new(*pool_mint, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),

stake-pool/program/src/instruction.rs

@@ -357,6 +357,47 @@ impl Processor {
         )
     }
 
+    /// Issue stake_program::withdraw instruction to move additional lamports
+    #[allow(clippy::too_many_arguments)]
+    fn stake_withdraw<'a>(
+        stake_pool: &Pubkey,
+        source_account: AccountInfo<'a>,
+        authority: AccountInfo<'a>,
+        authority_type: &[u8],
+        bump_seed: u8,
+        destination_account: AccountInfo<'a>,
+        clock: AccountInfo<'a>,
+        stake_history: AccountInfo<'a>,
+        stake_program_info: AccountInfo<'a>,
+        lamports: u64,
+    ) -> Result<(), ProgramError> {
+        let me_bytes = stake_pool.to_bytes();
+        let authority_signature_seeds = [&me_bytes[..32], authority_type, &[bump_seed]];
+        let signers = &[&authority_signature_seeds[..]];
+        let custodian_pubkey = None;
+
+        let withdraw_instruction = stake_program::withdraw(
+            source_account.key,
+            authority.key,
+            destination_account.key,
+            lamports,
+            custodian_pubkey,
+        );
+
+        invoke_signed(
+            &withdraw_instruction,
+            &[
+                source_account,
+                destination_account,
+                clock,
+                stake_history,
+                authority,
+                stake_program_info,
+            ],
+            signers,
+        )
+    }
+
     /// Issue a spl_token `Burn` instruction.
     #[allow(clippy::too_many_arguments)]
     fn token_burn<'a>(
@@ -1337,9 +1378,13 @@ impl Processor {
                         }
                     }
                 }
-                Some(stake_program::StakeState::Stake(_, stake)) => {
+                Some(stake_program::StakeState::Stake(meta, stake)) => {
+                    let account_stake = meta
+                        .rent_exempt_reserve
+                        .checked_add(stake.delegation.stake)
+                        .ok_or(StakePoolError::CalculationFailure)?;
                     if no_merge {
-                        transient_stake_lamports = transient_stake_info.lamports();
+                        transient_stake_lamports = account_stake;
                     } else if stake.delegation.deactivation_epoch < clock.epoch {
                         // deactivated, merge into reserve
                         Self::stake_merge(
@@ -1378,15 +1423,15 @@ impl Processor {
                                 )?;
                             } else {
                                 msg!("Stake activating or just active, not ready to merge");
-                                transient_stake_lamports = transient_stake_info.lamports();
+                                transient_stake_lamports = account_stake;
                             }
                         } else {
                             msg!("Transient stake is activating or active, but validator stake is not, need to add the validator stake account on {} back into the stake pool", stake.delegation.voter_pubkey);
-                            transient_stake_lamports = transient_stake_info.lamports();
+                            transient_stake_lamports = account_stake;
                         }
                     } else {
                         msg!("Transient stake not ready to be merged anywhere");
-                        transient_stake_lamports = transient_stake_info.lamports();
+                        transient_stake_lamports = account_stake;
                     }
                 }
                 None
@@ -1398,11 +1443,13 @@ impl Processor {
             //  * active -> do everything
             //  * any other state / not a stake -> error state, but account for transient stake
             match validator_stake_state {
-                Some(stake_program::StakeState::Stake(meta, _)) => {
+                Some(stake_program::StakeState::Stake(_, stake)) => {
                     if validator_stake_record.status == StakeStatus::Active {
-                        active_stake_lamports = validator_stake_info
-                            .lamports()
-                            .saturating_sub(minimum_stake_lamports(&meta));
+                        active_stake_lamports = stake
+                            .delegation
+                            .stake
+                            .checked_sub(MINIMUM_ACTIVE_STAKE)
+                            .ok_or(StakePoolError::CalculationFailure)?;
                     } else {
                         msg!("Validator stake account no longer part of the pool, ignoring");
                     }
@@ -1566,6 +1613,7 @@ impl Processor {
         let withdraw_authority_info = next_account_info(account_info_iter)?;
         let stake_info = next_account_info(account_info_iter)?;
         let validator_stake_account_info = next_account_info(account_info_iter)?;
+        let reserve_stake_account_info = next_account_info(account_info_iter)?;
         let dest_user_info = next_account_info(account_info_iter)?;
         let pool_mint_info = next_account_info(account_info_iter)?;
         let clock_info = next_account_info(account_info_iter)?;
@@ -1594,6 +1642,7 @@ impl Processor {
         stake_pool.check_deposit_authority(deposit_authority_info.key)?;
         stake_pool.check_mint(pool_mint_info)?;
         stake_pool.check_validator_list(validator_list_info)?;
+        stake_pool.check_reserve_stake(reserve_stake_account_info)?;
 
         if stake_pool.token_program_id != *token_program_info.key {
             return Err(ProgramError::IncorrectProgramId);
@@ -1610,8 +1659,9 @@ impl Processor {
             return Err(StakePoolError::InvalidState.into());
         }
 
-        let (meta, stake) = get_stake_state(validator_stake_account_info)?;
-        let vote_account_address = stake.delegation.voter_pubkey;
+        let (_, validator_stake) = get_stake_state(validator_stake_account_info)?;
+        let pre_all_validator_lamports = validator_stake_account_info.lamports();
+        let vote_account_address = validator_stake.delegation.voter_pubkey;
         check_validator_stake_address(
             program_id,
             stake_pool_info.key,
@@ -1633,15 +1683,7 @@ impl Processor {
             return Err(StakePoolError::ValidatorNotFound.into());
         }
 
-        let stake_lamports = **stake_info.lamports.borrow();
-        let new_pool_tokens = stake_pool
-            .calc_pool_tokens_for_deposit(stake_lamports)
-            .ok_or(StakePoolError::CalculationFailure)?;
-
-        msg!(
-            "lamports pre merge {}",
-            validator_stake_account_info.lamports()
-        );
+        msg!("Stake pre merge {}", validator_stake.delegation.stake);
 
         let (deposit_authority_program_address, deposit_bump_seed) =
             find_deposit_authority_program_address(program_id, stake_pool_info.key);
@@ -1678,6 +1720,21 @@ impl Processor {
             stake_program_info.clone(),
         )?;
 
+        let (_, post_validator_stake) = get_stake_state(validator_stake_account_info)?;
+        let post_all_validator_lamports = validator_stake_account_info.lamports();
+        msg!("Stake post merge {}", post_validator_stake.delegation.stake);
+        let all_deposit_lamports = post_all_validator_lamports
+            .checked_sub(pre_all_validator_lamports)
+            .ok_or(StakePoolError::CalculationFailure)?;
+        let stake_deposit_lamports = post_validator_stake
+            .delegation
+            .stake
+            .checked_sub(validator_stake.delegation.stake)
+            .ok_or(StakePoolError::CalculationFailure)?;
+        let new_pool_tokens = stake_pool
+            .calc_pool_tokens_for_deposit(all_deposit_lamports)
+            .ok_or(StakePoolError::CalculationFailure)?;
+
         Self::token_mint_to(
             stake_pool_info.key,
             token_program_info.clone(),
@@ -1689,23 +1746,39 @@ impl Processor {
             new_pool_tokens,
         )?;
 
+        // withdraw additional lamports to the reserve
+        let additional_lamports = all_deposit_lamports
+            .checked_sub(stake_deposit_lamports)
+            .ok_or(StakePoolError::CalculationFailure)?;
+        if additional_lamports > 0 {
+            Self::stake_withdraw(
+                stake_pool_info.key,
+                validator_stake_account_info.clone(),
+                withdraw_authority_info.clone(),
+                AUTHORITY_WITHDRAW,
+                stake_pool.withdraw_bump_seed,
+                reserve_stake_account_info.clone(),
+                clock_info.clone(),
+                stake_history_info.clone(),
+                stake_program_info.clone(),
+                additional_lamports,
+            )?;
+        }
+
         stake_pool.pool_token_supply = stake_pool
             .pool_token_supply
             .checked_add(new_pool_tokens)
             .ok_or(StakePoolError::CalculationFailure)?;
         stake_pool.total_stake_lamports = stake_pool
             .total_stake_lamports
-            .checked_add(stake_lamports)
+            .checked_add(all_deposit_lamports)
             .ok_or(StakePoolError::CalculationFailure)?;
         stake_pool.serialize(&mut *stake_pool_info.data.borrow_mut())?;
 
-        msg!(
-            "lamports post merge {}",
-            validator_stake_account_info.lamports()
-        );
-        validator_list_item.active_stake_lamports = validator_stake_account_info
-            .lamports()
-            .checked_sub(minimum_stake_lamports(&meta))
+        validator_list_item.active_stake_lamports = post_validator_stake
+            .delegation
+            .stake
+            .checked_sub(MINIMUM_ACTIVE_STAKE)
             .ok_or(StakePoolError::CalculationFailure)?;
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
 
@@ -1794,7 +1867,7 @@ impl Processor {
                 .ok_or(StakePoolError::StakeLamportsNotEqualToMinimum)?;
             None
         } else {
-            let (meta, stake) = get_stake_state(stake_split_from)?;
+            let (_, stake) = get_stake_state(stake_split_from)?;
             let vote_account_address = stake.delegation.voter_pubkey;
 
             if let Some(preferred_withdraw_validator) =
@@ -1840,11 +1913,9 @@ impl Processor {
                 return Err(StakePoolError::ValidatorNotFound.into());
             }
 
-            let required_lamports = minimum_stake_lamports(&meta);
-            let current_lamports = stake_split_from.lamports();
-            let remaining_lamports = current_lamports.saturating_sub(withdraw_lamports);
-            if remaining_lamports < required_lamports {
-                msg!("Attempting to withdraw {} lamports from validator account with {} lamports, {} must remain", withdraw_lamports, current_lamports, required_lamports);
+            let remaining_lamports = stake.delegation.stake.saturating_sub(withdraw_lamports);
+            if remaining_lamports < MINIMUM_ACTIVE_STAKE {
+                msg!("Attempting to withdraw {} lamports from validator account with {} stake lamports, {} must remain", withdraw_lamports, stake.delegation.stake, MINIMUM_ACTIVE_STAKE);
                 return Err(StakePoolError::StakeLamportsNotEqualToMinimum.into());
             }
             Some((validator_list_item, withdrawing_from_transient_stake))

stake-pool/program/src/processor.rs

@@ -645,6 +645,29 @@ pub fn deactivate_stake(stake_pubkey: &Pubkey, authorized_pubkey: &Pubkey) -> In
     Instruction::new_with_bincode(id(), &StakeInstruction::Deactivate, account_metas)
 }
 
+/// FIXME copied from the stake program
+pub fn withdraw(
+    stake_pubkey: &Pubkey,
+    withdrawer_pubkey: &Pubkey,
+    to_pubkey: &Pubkey,
+    lamports: u64,
+    custodian_pubkey: Option<&Pubkey>,
+) -> Instruction {
+    let mut account_metas = vec![
+        AccountMeta::new(*stake_pubkey, false),
+        AccountMeta::new(*to_pubkey, false),
+        AccountMeta::new_readonly(sysvar::clock::id(), false),
+        AccountMeta::new_readonly(sysvar::stake_history::id(), false),
+        AccountMeta::new_readonly(*withdrawer_pubkey, true),
+    ];
+
+    if let Some(custodian_pubkey) = custodian_pubkey {
+        account_metas.push(AccountMeta::new_readonly(*custodian_pubkey, true));
+    }
+
+    Instruction::new_with_bincode(id(), &StakeInstruction::Withdraw(lamports), account_metas)
+}
+
 #[cfg(test)]
 mod test {
     use {super::*, bincode::serialize, solana_program::borsh::try_from_slice_unchecked};

stake-pool/program/src/stake_program.rs

@@ -4,7 +4,6 @@ use {
     crate::error::StakePoolError,
     borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
     solana_program::{account_info::AccountInfo, msg, program_error::ProgramError, pubkey::Pubkey},
-    spl_math::checked_ceil_div::CheckedCeilDiv,
     std::convert::TryFrom,
 };
 
@@ -99,13 +98,6 @@ impl StakePool {
         )
         .ok()
     }
-    /// calculate the pool tokens that should be burned for a withdrawal of `stake_lamports`
-    pub fn calc_pool_tokens_for_withdraw(&self, stake_lamports: u64) -> Option<u64> {
-        let (quotient, _) = (stake_lamports as u128)
-            .checked_mul(self.pool_token_supply as u128)?
-            .checked_ceil_div(self.total_stake_lamports as u128)?;
-        u64::try_from(quotient).ok()
-    }
 
     /// calculate lamports amount on withdrawal
     pub fn calc_lamports_withdraw_amount(&self, pool_tokens: u64) -> Option<u64> {