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 c0f25ab543c808a1daf3474df4ef851bc994fc6c.

Cargo.lock

@@ -3901,7 +3901,6 @@ dependencies = [
  "solana-program-test",
  "solana-sdk",
  "solana-vote-program",
- "spl-math",
  "spl-token 3.1.1",
  "thiserror",
 ]

stake-pool/cli/src/main.rs

@@ -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/program/Cargo.toml

@@ -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/src/instruction.rs

@@ -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/processor.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
+                        active_stake_lamports = stake
-                            .lamports()
+                            .delegation
-                            .saturating_sub(minimum_stake_lamports(&meta));
+                            .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 (_, validator_stake) = get_stake_state(validator_stake_account_info)?;
-        let vote_account_address = stake.delegation.voter_pubkey;
+        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();
+        msg!("Stake pre merge {}", validator_stake.delegation.stake);
-        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()
-        );
 
         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!(
+        validator_list_item.active_stake_lamports = post_validator_stake
-            "lamports post merge {}",
+            .delegation
-            validator_stake_account_info.lamports()
+            .stake
-        );
+            .checked_sub(MINIMUM_ACTIVE_STAKE)
-        validator_list_item.active_stake_lamports = validator_stake_account_info
-            .lamports()
-            .checked_sub(minimum_stake_lamports(&meta))
             .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 remaining_lamports = stake.delegation.stake.saturating_sub(withdraw_lamports);
-            let current_lamports = stake_split_from.lamports();
+            if remaining_lamports < MINIMUM_ACTIVE_STAKE {
-            let remaining_lamports = current_lamports.saturating_sub(withdraw_lamports);
+                msg!("Attempting to withdraw {} lamports from validator account with {} stake lamports, {} must remain", withdraw_lamports, stake.delegation.stake, MINIMUM_ACTIVE_STAKE);
-            if remaining_lamports < required_lamports {
-                msg!("Attempting to withdraw {} lamports from validator account with {} lamports, {} must remain", withdraw_lamports, current_lamports, required_lamports);
                 return Err(StakePoolError::StakeLamportsNotEqualToMinimum.into());
             }
             Some((validator_list_item, withdrawing_from_transient_stake))

stake-pool/program/src/stake_program.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/state.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> {

stake-pool/program/tests/deposit.rs

@@ -8,7 +8,6 @@ use {
     helpers::*,
     solana_program::{
         borsh::try_from_slice_unchecked,
-        hash::Hash,
         instruction::{AccountMeta, Instruction, InstructionError},
         pubkey::Pubkey,
         sysvar,
@@ -25,9 +24,7 @@ use {
 };
 
 async fn setup() -> (
-    BanksClient,
+    ProgramTestContext,
-    Keypair,
-    Hash,
     StakePoolAccounts,
     ValidatorStakeAccount,
     Keypair,
@@ -35,21 +32,32 @@ async fn setup() -> (
     Pubkey,
     u64,
 ) {
-    let (mut banks_client, payer, recent_blockhash) = program_test().start().await;
+    let mut context = program_test().start_with_context().await;
+
     let stake_pool_accounts = StakePoolAccounts::new();
     stake_pool_accounts
-        .initialize_stake_pool(&mut banks_client, &payer, &recent_blockhash, 1)
+        .initialize_stake_pool(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            1,
+        )
         .await
         .unwrap();
 
     let validator_stake_account = simple_add_validator_to_pool(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &stake_pool_accounts,
     )
     .await;
 
+    let first_normal_slot = context.genesis_config().epoch_schedule.first_normal_slot;
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    let mut slot = first_normal_slot;
+    context.warp_to_slot(slot).unwrap();
+
     let user = Keypair::new();
     // make stake account
     let deposit_stake = Keypair::new();
@@ -61,9 +69,9 @@ async fn setup() -> (
     };
 
     let stake_lamports = create_independent_stake_account(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &deposit_stake,
         &authorized,
         &lockup,
@@ -72,21 +80,33 @@ async fn setup() -> (
     .await;
 
     delegate_stake_account(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &deposit_stake.pubkey(),
         &user,
         &validator_stake_account.vote.pubkey(),
     )
     .await;
 
+    slot += 2 * slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+    stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &[validator_stake_account.vote.pubkey()],
+            false,
+        )
+        .await;
+
     // make pool token account
     let pool_token_account = Keypair::new();
     create_token_account(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &pool_token_account,
         &stake_pool_accounts.pool_mint.pubkey(),
         &user.pubkey(),
@@ -95,9 +115,7 @@ async fn setup() -> (
     .unwrap();
 
     (
-        banks_client,
+        context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake_account,
         user,
@@ -110,9 +128,7 @@ async fn setup() -> (
 #[tokio::test]
 async fn success() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake_account,
         user,
@@ -121,29 +137,43 @@ async fn success() {
         stake_lamports,
     ) = setup().await;
 
+    let rent = context.banks_client.get_rent().await.unwrap();
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+
     // Save stake pool state before depositing
-    let stake_pool_before =
+    let pre_stake_pool = get_account(
-        get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await;
+        &mut context.banks_client,
-    let stake_pool_before =
+        &stake_pool_accounts.stake_pool.pubkey(),
-        try_from_slice_unchecked::<state::StakePool>(&stake_pool_before.data.as_slice()).unwrap();
+    )
+    .await;
+    let pre_stake_pool =
+        try_from_slice_unchecked::<state::StakePool>(&pre_stake_pool.data.as_slice()).unwrap();
 
     // Save validator stake account record before depositing
     let validator_list = get_account(
-        &mut banks_client,
+        &mut context.banks_client,
         &stake_pool_accounts.validator_list.pubkey(),
     )
     .await;
     let validator_list =
         try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
-    let validator_stake_item_before = validator_list
+    let pre_validator_stake_item = validator_list
         .find(&validator_stake_account.vote.pubkey())
         .unwrap();
 
+    // Save reserve state before depositing
+    let pre_reserve_lamports = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+    )
+    .await
+    .lamports;
+
     let error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &pool_token_account,
             &validator_stake_account.stake_account,
@@ -153,7 +183,8 @@ async fn success() {
     assert!(error.is_none());
 
     // Original stake account should be drained
-    assert!(banks_client
+    assert!(context
+        .banks_client
         .get_account(deposit_stake)
         .await
         .expect("get_account")
@@ -162,57 +193,72 @@ async fn success() {
     let tokens_issued = stake_lamports; // For now tokens are 1:1 to stake
 
     // Stake pool should add its balance to the pool balance
-    let stake_pool = get_account(&mut banks_client, &stake_pool_accounts.stake_pool.pubkey()).await;
+    let post_stake_pool = get_account(
-    let stake_pool =
+        &mut context.banks_client,
-        try_from_slice_unchecked::<state::StakePool>(&stake_pool.data.as_slice()).unwrap();
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let post_stake_pool =
+        try_from_slice_unchecked::<state::StakePool>(&post_stake_pool.data.as_slice()).unwrap();
     assert_eq!(
-        stake_pool.total_stake_lamports,
+        post_stake_pool.total_stake_lamports,
-        stake_pool_before.total_stake_lamports + stake_lamports
+        pre_stake_pool.total_stake_lamports + stake_lamports
     );
     assert_eq!(
-        stake_pool.pool_token_supply,
+        post_stake_pool.pool_token_supply,
-        stake_pool_before.pool_token_supply + tokens_issued
+        pre_stake_pool.pool_token_supply + tokens_issued
     );
 
     // Check minted tokens
-    let user_token_balance = get_token_balance(&mut banks_client, &pool_token_account).await;
+    let user_token_balance =
+        get_token_balance(&mut context.banks_client, &pool_token_account).await;
     assert_eq!(user_token_balance, tokens_issued);
 
     // Check balances in validator stake account list storage
     let validator_list = get_account(
-        &mut banks_client,
+        &mut context.banks_client,
         &stake_pool_accounts.validator_list.pubkey(),
     )
     .await;
     let validator_list =
         try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
-    let validator_stake_item = validator_list
+    let post_validator_stake_item = validator_list
         .find(&validator_stake_account.vote.pubkey())
         .unwrap();
     assert_eq!(
-        validator_stake_item.stake_lamports(),
+        post_validator_stake_item.stake_lamports(),
-        validator_stake_item_before.stake_lamports() + stake_lamports
+        pre_validator_stake_item.stake_lamports() + stake_lamports - stake_rent,
     );
 
     // Check validator stake account actual SOL balance
-    let validator_stake_account =
+    let validator_stake_account = get_account(
-        get_account(&mut banks_client, &validator_stake_account.stake_account).await;
+        &mut context.banks_client,
+        &validator_stake_account.stake_account,
+    )
+    .await;
     let stake_state =
         deserialize::<stake_program::StakeState>(&validator_stake_account.data).unwrap();
     let meta = stake_state.meta().unwrap();
     assert_eq!(
         validator_stake_account.lamports - minimum_stake_lamports(&meta),
-        validator_stake_item.stake_lamports()
+        post_validator_stake_item.stake_lamports()
     );
-    assert_eq!(validator_stake_item.transient_stake_lamports, 0);
+    assert_eq!(post_validator_stake_item.transient_stake_lamports, 0);
+
+    // Check reserve
+    let post_reserve_lamports = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+    )
+    .await
+    .lamports;
+    assert_eq!(post_reserve_lamports, pre_reserve_lamports + stake_rent);
 }
 
 #[tokio::test]
 async fn fail_with_wrong_stake_program_id() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake_account,
         _user,
@@ -230,6 +276,7 @@ async fn fail_with_wrong_stake_program_id() {
         AccountMeta::new_readonly(stake_pool_accounts.withdraw_authority, false),
         AccountMeta::new(deposit_stake, false),
         AccountMeta::new(validator_stake_account.stake_account, false),
+        AccountMeta::new(stake_pool_accounts.reserve_stake.pubkey(), false),
         AccountMeta::new(pool_token_account, false),
         AccountMeta::new(stake_pool_accounts.pool_mint.pubkey(), false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
@@ -245,9 +292,11 @@ async fn fail_with_wrong_stake_program_id() {
             .unwrap(),
     };
 
-    let mut transaction = Transaction::new_with_payer(&[instruction], Some(&payer.pubkey()));
+    let mut transaction =
-    transaction.sign(&[&payer], recent_blockhash);
+        Transaction::new_with_payer(&[instruction], Some(&context.payer.pubkey()));
-    let transaction_error = banks_client
+    transaction.sign(&[&context.payer], context.last_blockhash);
+    let transaction_error = context
+        .banks_client
         .process_transaction(transaction)
         .await
         .err()
@@ -264,9 +313,7 @@ async fn fail_with_wrong_stake_program_id() {
 #[tokio::test]
 async fn fail_with_wrong_token_program_id() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake_account,
         user,
@@ -286,14 +333,16 @@ async fn fail_with_wrong_token_program_id() {
             &deposit_stake,
             &user.pubkey(),
             &validator_stake_account.stake_account,
+            &stake_pool_accounts.reserve_stake.pubkey(),
             &pool_token_account,
             &stake_pool_accounts.pool_mint.pubkey(),
             &wrong_token_program.pubkey(),
         ),
-        Some(&payer.pubkey()),
+        Some(&context.payer.pubkey()),
     );
-    transaction.sign(&[&payer, &user], recent_blockhash);
+    transaction.sign(&[&context.payer, &user], context.last_blockhash);
-    let transaction_error = banks_client
+    let transaction_error = context
+        .banks_client
         .process_transaction(transaction)
         .await
         .err()
@@ -310,9 +359,7 @@ async fn fail_with_wrong_token_program_id() {
 #[tokio::test]
 async fn fail_with_wrong_validator_list_account() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         mut stake_pool_accounts,
         validator_stake_account,
         user,
@@ -326,9 +373,9 @@ async fn fail_with_wrong_validator_list_account() {
 
     let transaction_error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &pool_token_account,
             &validator_stake_account.stake_account,
@@ -429,9 +476,7 @@ async fn fail_with_unknown_validator() {
 #[tokio::test]
 async fn fail_with_wrong_withdraw_authority() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         mut stake_pool_accounts,
         validator_stake_account,
         user,
@@ -444,9 +489,9 @@ async fn fail_with_wrong_withdraw_authority() {
 
     let transaction_error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &pool_token_account,
             &validator_stake_account.stake_account,
@@ -468,9 +513,7 @@ async fn fail_with_wrong_withdraw_authority() {
 #[tokio::test]
 async fn fail_with_wrong_mint_for_receiver_acc() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake_account,
         user,
@@ -485,9 +528,9 @@ async fn fail_with_wrong_mint_for_receiver_acc() {
     let outside_pool_fee_acc = Keypair::new();
 
     create_mint(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &outside_mint,
         &outside_withdraw_auth.pubkey(),
     )
@@ -495,9 +538,9 @@ async fn fail_with_wrong_mint_for_receiver_acc() {
     .unwrap();
 
     create_token_account(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &outside_pool_fee_acc,
         &outside_mint.pubkey(),
         &outside_manager.pubkey(),
@@ -507,9 +550,9 @@ async fn fail_with_wrong_mint_for_receiver_acc() {
 
     let transaction_error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &outside_pool_fee_acc.pubkey(),
             &validator_stake_account.stake_account,
@@ -714,9 +757,7 @@ async fn fail_without_deposit_authority_signature() {
 #[tokio::test]
 async fn success_with_preferred_deposit() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake,
         user,
@@ -727,9 +768,9 @@ async fn success_with_preferred_deposit() {
 
     stake_pool_accounts
         .set_preferred_validator(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             instruction::PreferredValidatorType::Deposit,
             Some(validator_stake.vote.pubkey()),
         )
@@ -737,9 +778,9 @@ async fn success_with_preferred_deposit() {
 
     let error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &pool_token_account,
             &validator_stake.stake_account,
@@ -752,9 +793,7 @@ async fn success_with_preferred_deposit() {
 #[tokio::test]
 async fn fail_with_wrong_preferred_deposit() {
     let (
-        mut banks_client,
+        mut context,
-        payer,
-        recent_blockhash,
         stake_pool_accounts,
         validator_stake,
         user,
@@ -764,18 +803,18 @@ async fn fail_with_wrong_preferred_deposit() {
     ) = setup().await;
 
     let preferred_validator = simple_add_validator_to_pool(
-        &mut banks_client,
+        &mut context.banks_client,
-        &payer,
+        &context.payer,
-        &recent_blockhash,
+        &context.last_blockhash,
         &stake_pool_accounts,
     )
     .await;
 
     stake_pool_accounts
         .set_preferred_validator(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             instruction::PreferredValidatorType::Deposit,
             Some(preferred_validator.vote.pubkey()),
         )
@@ -783,9 +822,9 @@ async fn fail_with_wrong_preferred_deposit() {
 
     let error = stake_pool_accounts
         .deposit_stake(
-            &mut banks_client,
+            &mut context.banks_client,
-            &payer,
+            &context.payer,
-            &recent_blockhash,
+            &context.last_blockhash,
             &deposit_stake,
             &pool_token_account,
             &validator_stake.stake_account,

stake-pool/program/tests/helpers/mod.rs

@@ -622,6 +622,7 @@ impl StakePoolAccounts {
                 stake,
                 &current_staker.pubkey(),
                 validator_stake_account,
+                &self.reserve_stake.pubkey(),
                 pool_account,
                 &self.pool_mint.pubkey(),
                 &spl_token::id(),
@@ -635,6 +636,7 @@ impl StakePoolAccounts {
                 stake,
                 &current_staker.pubkey(),
                 validator_stake_account,
+                &self.reserve_stake.pubkey(),
                 pool_account,
                 &self.pool_mint.pubkey(),
                 &spl_token::id(),

stake-pool/program/tests/initialize.rs

@@ -85,7 +85,7 @@ async fn success() {
     .await;
     let validator_list =
         try_from_slice_unchecked::<state::ValidatorList>(validator_list.data.as_slice()).unwrap();
-    assert_eq!(validator_list.is_valid(), true);
+    assert!(validator_list.is_valid());
 }
 
 #[tokio::test]

stake-pool/program/tests/update_stake_pool_balance.rs

@@ -94,17 +94,10 @@ async fn success() {
         .await;
     assert!(error.is_none());
 
-    // Add extra funds, simulating rewards
+    // Increment vote credits to earn rewards
-    const EXTRA_STAKE_AMOUNT: u64 = 1_000_000;
+    const VOTE_CREDITS: u64 = 1_000;
     for stake_account in &stake_accounts {
-        transfer(
+        context.increment_vote_account_credits(&stake_account.vote.pubkey(), VOTE_CREDITS);
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-            &stake_account.stake_account,
-            EXTRA_STAKE_AMOUNT,
-        )
-        .await;
     }
 
     // Update epoch
@@ -155,10 +148,8 @@ async fn success() {
 
     let expected_fee_lamports =
         (post_balance - pre_balance) * stake_pool.fee.numerator / stake_pool.fee.denominator;
-    let actual_fee_lamports = stake_pool
+    let actual_fee_lamports = stake_pool.calc_pool_tokens_for_deposit(actual_fee).unwrap();
-        .calc_pool_tokens_for_withdraw(actual_fee)
+    assert_eq!(actual_fee_lamports, expected_fee_lamports);
-        .unwrap();
-    assert!(actual_fee_lamports <= expected_fee_lamports);
 
     let expected_fee = expected_fee_lamports * pool_token_supply / post_balance;
     assert_eq!(expected_fee, actual_fee);
@@ -166,6 +157,87 @@ async fn success() {
     assert_eq!(pool_token_supply, stake_pool.pool_token_supply);
 }
 
+#[tokio::test]
+async fn success_ignoring_extra_lamports() {
+    let (mut context, stake_pool_accounts, stake_accounts) = setup().await;
+
+    let pre_balance = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    let stake_pool = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool.data.as_slice()).unwrap();
+    assert_eq!(pre_balance, stake_pool.total_stake_lamports);
+
+    let pre_token_supply = get_token_supply(
+        &mut context.banks_client,
+        &stake_pool_accounts.pool_mint.pubkey(),
+    )
+    .await;
+
+    let error = stake_pool_accounts
+        .update_stake_pool_balance(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // Transfer extra funds, should not be taken into account
+    const EXTRA_STAKE_AMOUNT: u64 = 1_000_000;
+    for stake_account in &stake_accounts {
+        transfer(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            &stake_account.stake_account,
+            EXTRA_STAKE_AMOUNT,
+        )
+        .await;
+    }
+
+    // Update epoch
+    context.warp_to_slot(50_000).unwrap();
+
+    // Update list and pool
+    let error = stake_pool_accounts
+        .update_all(
+            &mut context.banks_client,
+            &context.payer,
+            &context.last_blockhash,
+            stake_accounts
+                .iter()
+                .map(|v| v.vote.pubkey())
+                .collect::<Vec<Pubkey>>()
+                .as_slice(),
+            false,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // Check fee
+    let post_balance = get_validator_list_sum(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+        &stake_pool_accounts.validator_list.pubkey(),
+    )
+    .await;
+    assert_eq!(post_balance, pre_balance);
+    let pool_token_supply = get_token_supply(
+        &mut context.banks_client,
+        &stake_pool_accounts.pool_mint.pubkey(),
+    )
+    .await;
+    assert_eq!(pool_token_supply, pre_token_supply);
+}
+
 #[tokio::test]
 async fn fail_with_wrong_validator_list() {
     let (mut banks_client, payer, recent_blockhash) = program_test().start().await;

stake-pool/program/tests/update_validator_list_balance.rs

@@ -6,7 +6,7 @@ use {
     helpers::*,
     solana_program::{borsh::try_from_slice_unchecked, pubkey::Pubkey},
     solana_program_test::*,
-    solana_sdk::signature::{Keypair, Signer},
+    solana_sdk::signature::Signer,
     spl_stake_pool::{
         stake_program,
         state::{StakePool, StakeStatus, ValidatorList},
@@ -42,27 +42,6 @@ async fn setup(
         .await
         .unwrap();
 
-    // so warmups / cooldowns go faster
-    let validator = Keypair::new();
-    let vote = Keypair::new();
-    create_vote(
-        &mut context.banks_client,
-        &context.payer,
-        &context.last_blockhash,
-        &validator,
-        &vote,
-    )
-    .await;
-    let deposit_account =
-        DepositStakeAccount::new_with_vote(vote.pubkey(), validator.pubkey(), 100_000_000_000);
-    deposit_account
-        .create_and_delegate(
-            &mut context.banks_client,
-            &context.payer,
-            &context.last_blockhash,
-        )
-        .await;
-
     // Add several accounts with some stake
     let mut stake_accounts: Vec<ValidatorStakeAccount> = vec![];
     let mut deposit_accounts: Vec<DepositStakeAccount> = vec![];
@@ -454,16 +433,29 @@ async fn merge_into_validator_stake() {
     }
 
     // Check validator stake accounts have the expected balance now:
-    // validator stake account minimum + deposited lamports + 2 rents + increased lamports
+    // validator stake account minimum + deposited lamports + rents + increased lamports
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
     let expected_lamports = MINIMUM_ACTIVE_STAKE
         + lamports
         + reserve_lamports / stake_accounts.len() as u64
-        + 2 * rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+        + stake_rent;
     for stake_account in &stake_accounts {
         let validator_stake =
             get_account(&mut context.banks_client, &stake_account.stake_account).await;
         assert_eq!(validator_stake.lamports, expected_lamports);
     }
+
+    // Check reserve stake accounts for expected balance:
+    // own rent, other account rents, and 1 extra lamport
+    let reserve_stake = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.reserve_stake.pubkey(),
+    )
+    .await;
+    assert_eq!(
+        reserve_stake.lamports,
+        1 + stake_rent * (1 + stake_accounts.len() as u64)
+    );
 }
 
 #[tokio::test]
@@ -473,7 +465,7 @@ async fn merge_transient_stake_after_remove() {
 
     let rent = context.banks_client.get_rent().await.unwrap();
     let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
-    let deactivated_lamports = lamports + stake_rent;
+    let deactivated_lamports = lamports;
     let new_authority = Pubkey::new_unique();
     // Decrease and remove all validators
     for stake_account in &stake_accounts {
@@ -578,7 +570,7 @@ async fn merge_transient_stake_after_remove() {
         .unwrap();
     assert_eq!(
         reserve_stake.lamports,
-        reserve_lamports + deactivated_lamports + stake_rent + 1
+        reserve_lamports + deactivated_lamports + 2 * stake_rent + 1
     );
 
     // Update stake pool balance, should be gone