stake-pool: Update pool token supply on update, in case of burns (#2108)

Author: joncinque

stake-pool/program/src/processor.rs

@@ -1621,6 +1621,10 @@ impl Processor {
         }
         stake_pool.total_stake_lamports = total_stake_lamports;
         stake_pool.last_update_epoch = clock.epoch;
+
+        let pool_mint = Mint::unpack_from_slice(&pool_mint_info.data.borrow())?;
+        stake_pool.pool_token_supply = pool_mint.supply;
+
         stake_pool.serialize(&mut *stake_pool_info.data.borrow_mut())?;
 
         Ok(())

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

@@ -161,6 +161,33 @@ pub async fn mint_tokens(
     Ok(())
 }
 
+pub async fn burn_tokens(
+    banks_client: &mut BanksClient,
+    payer: &Keypair,
+    recent_blockhash: &Hash,
+    mint: &Pubkey,
+    account: &Pubkey,
+    authority: &Keypair,
+    amount: u64,
+) -> Result<(), TransportError> {
+    let transaction = Transaction::new_signed_with_payer(
+        &[spl_token::instruction::burn(
+            &spl_token::id(),
+            account,
+            mint,
+            &authority.pubkey(),
+            &[],
+            amount,
+        )
+        .unwrap()],
+        Some(&payer.pubkey()),
+        &[payer, authority],
+        *recent_blockhash,
+    );
+    banks_client.process_transaction(transaction).await?;
+    Ok(())
+}
+
 pub async fn get_token_balance(banks_client: &mut BanksClient, token: &Pubkey) -> u64 {
     let token_account = banks_client.get_account(*token).await.unwrap().unwrap();
     let account_info: spl_token::state::Account =
@@ -1019,7 +1046,7 @@ impl DepositStakeAccount {
     }
 
     pub async fn deposit(
-        &self,
+        &mut self,
         banks_client: &mut BanksClient,
         payer: &Keypair,
         recent_blockhash: &Hash,
@@ -1048,6 +1075,7 @@ impl DepositStakeAccount {
                 &self.authority,
             )
             .await;
+        self.pool_tokens = get_token_balance(banks_client, &self.pool_account.pubkey()).await;
         assert!(error.is_none());
     }
 }

stake-pool/program/tests/update_validator_list_balance.rs

@@ -4,14 +4,15 @@ mod helpers;
 
 use {
     helpers::*,
-    solana_program::{borsh::try_from_slice_unchecked, pubkey::Pubkey},
+    solana_program::{borsh::try_from_slice_unchecked, program_pack::Pack, pubkey::Pubkey},
     solana_program_test::*,
     solana_sdk::signature::Signer,
     spl_stake_pool::{
         stake_program,
         state::{StakePool, StakeStatus, ValidatorList},
         MAX_VALIDATORS_TO_UPDATE, MINIMUM_ACTIVE_STAKE,
     },
+    spl_token::state::Mint,
 };
 
 async fn setup(
@@ -20,6 +21,7 @@ async fn setup(
     ProgramTestContext,
     StakePoolAccounts,
     Vec<ValidatorStakeAccount>,
+    Vec<DepositStakeAccount>,
     u64,
     u64,
     u64,
@@ -99,7 +101,7 @@ async fn setup(
         assert!(error.is_none());
     }
 
-    for deposit_account in &deposit_accounts {
+    for deposit_account in &mut deposit_accounts {
         deposit_account
             .deposit(
                 &mut context.banks_client,
@@ -131,6 +133,7 @@ async fn setup(
         context,
         stake_pool_accounts,
         stake_accounts,
+        deposit_accounts,
         TEST_STAKE_AMOUNT,
         reserve_stake_amount,
         slot,
@@ -144,6 +147,7 @@ async fn success() {
         mut context,
         stake_pool_accounts,
         stake_accounts,
+        _,
         validator_lamports,
         reserve_lamports,
         mut slot,
@@ -207,7 +211,7 @@ async fn success() {
 
 #[tokio::test]
 async fn merge_into_reserve() {
-    let (mut context, stake_pool_accounts, stake_accounts, lamports, _, mut slot) =
+    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, _, mut slot) =
         setup(MAX_VALIDATORS_TO_UPDATE).await;
 
     let pre_lamports = get_validator_list_sum(
@@ -317,7 +321,7 @@ async fn merge_into_reserve() {
 
 #[tokio::test]
 async fn merge_into_validator_stake() {
-    let (mut context, stake_pool_accounts, stake_accounts, lamports, reserve_lamports, mut slot) =
+    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, reserve_lamports, mut slot) =
         setup(MAX_VALIDATORS_TO_UPDATE).await;
 
     let rent = context.banks_client.get_rent().await.unwrap();
@@ -448,7 +452,7 @@ async fn merge_into_validator_stake() {
 
 #[tokio::test]
 async fn merge_transient_stake_after_remove() {
-    let (mut context, stake_pool_accounts, stake_accounts, lamports, reserve_lamports, mut slot) =
+    let (mut context, stake_pool_accounts, stake_accounts, _, lamports, reserve_lamports, mut slot) =
         setup(1).await;
 
     let rent = context.banks_client.get_rent().await.unwrap();
@@ -590,6 +594,75 @@ async fn merge_transient_stake_after_remove() {
     assert_eq!(validator_list.validators.len(), 0);
 }
 
+#[tokio::test]
+async fn success_with_burned_tokens() {
+    let num_validators = 5;
+    let (mut context, stake_pool_accounts, stake_accounts, deposit_accounts, _, _, mut slot) =
+        setup(num_validators).await;
+
+    let mint_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.pool_mint.pubkey(),
+    )
+    .await;
+    let mint = Mint::unpack(&mint_info.data).unwrap();
+
+    let stake_pool_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
+    assert_eq!(mint.supply, stake_pool.pool_token_supply);
+
+    burn_tokens(
+        &mut context.banks_client,
+        &context.payer,
+        &context.last_blockhash,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &deposit_accounts[0].pool_account.pubkey(),
+        &deposit_accounts[0].authority,
+        deposit_accounts[0].pool_tokens,
+    )
+    .await
+    .unwrap();
+
+    let mint_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.pool_mint.pubkey(),
+    )
+    .await;
+    let mint = Mint::unpack(&mint_info.data).unwrap();
+    assert_ne!(mint.supply, stake_pool.pool_token_supply);
+
+    let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
+    slot += slots_per_epoch;
+    context.warp_to_slot(slot).unwrap();
+
+    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;
+
+    let stake_pool_info = get_account(
+        &mut context.banks_client,
+        &stake_pool_accounts.stake_pool.pubkey(),
+    )
+    .await;
+    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
+
+    assert_eq!(mint.supply, stake_pool.pool_token_supply);
+}
+
 #[tokio::test]
 async fn fail_with_uninitialized_validator_list() {} // TODO