stake-pool: skip manager fee if invalid account, check new manager fee owned by `spl_token` (#2277)

* fix

* add check if new manager fee is spl token program

* Convert manager fee info check to fail only inside `Result` box

* update also checks validity of manager fee info

* clippy

stake-pool/program/src/processor.rs

@@ -1642,9 +1642,15 @@ impl Processor {
         }
 
         let reward_lamports = total_stake_lamports.saturating_sub(previous_lamports);
-        let fee = stake_pool
+
-            .calc_epoch_fee_amount(reward_lamports)
+        // If the manager fee info is invalid, they don't deserve to receive the fee.
-            .ok_or(StakePoolError::CalculationFailure)?;
+        let fee = if stake_pool.check_manager_fee_info(manager_fee_info).is_ok() {
+            stake_pool
+                .calc_epoch_fee_amount(reward_lamports)
+                .ok_or(StakePoolError::CalculationFailure)?
+        } else {
+            0
+        };
 
         if fee > 0 {
             Self::token_mint_to(
@@ -2205,7 +2211,11 @@ impl Processor {
             return Err(StakePoolError::InvalidState.into());
         }
 
-        let pool_tokens_fee = if stake_pool.manager_fee_account == *burn_from_pool_info.key {
+        // To prevent a faulty manager fee account from preventing withdrawals
+        // if the token program does not own the account, or if the account is not initialized
+        let pool_tokens_fee = if stake_pool.manager_fee_account == *burn_from_pool_info.key
+            || stake_pool.check_manager_fee_info(manager_fee_info).is_err()
+        {
             0
         } else {
             stake_pool
@@ -2393,6 +2403,7 @@ impl Processor {
 
         check_account_owner(stake_pool_info, program_id)?;
         let mut stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data.borrow())?;
+        check_account_owner(new_manager_fee_info, &stake_pool.token_program_id)?;
         if !stake_pool.is_valid() {
             return Err(StakePoolError::InvalidState.into());
         }

stake-pool/program/src/state.rs

@@ -1,5 +1,6 @@
 //! State transition types
 
+use spl_token::state::{Account, AccountState};
 use {
     crate::{
         big_vec::BigVec, error::StakePoolError, stake_program::Lockup, MAX_WITHDRAWAL_FEE_INCREASE,
@@ -261,6 +262,23 @@ impl StakePool {
         }
     }
 
+    /// Check if the manager fee info is a valid token program account
+    /// capable of receiving tokens from the mint.
+    pub(crate) fn check_manager_fee_info(
+        &self,
+        manager_fee_info: &AccountInfo,
+    ) -> Result<(), ProgramError> {
+        let token_account = Account::unpack(&manager_fee_info.data.borrow())?;
+        if manager_fee_info.owner != &self.token_program_id
+            || token_account.state != AccountState::Initialized
+            || token_account.mint != self.pool_mint
+        {
+            msg!("Manager fee account is not owned by token program, is not initialized, or does not match stake pool's mint");
+            return Err(StakePoolError::InvalidFeeAccount.into());
+        }
+        Ok(())
+    }
+
     /// Checks that the withdraw authority is valid
     #[inline]
     pub(crate) fn check_authority_withdraw(

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

@@ -160,6 +160,54 @@ pub async fn create_token_account(
     Ok(())
 }
 
+pub async fn close_token_account(
+    banks_client: &mut BanksClient,
+    payer: &Keypair,
+    recent_blockhash: &Hash,
+    account: &Pubkey,
+    lamports_destination: &Pubkey,
+    manager: &Keypair,
+) -> Result<(), TransportError> {
+    let mut transaction = Transaction::new_with_payer(
+        &[spl_token::instruction::close_account(
+            &spl_token::id(),
+            &account,
+            &lamports_destination,
+            &manager.pubkey(),
+            &[],
+        )
+        .unwrap()],
+        Some(&payer.pubkey()),
+    );
+    transaction.sign(&[payer, manager], *recent_blockhash);
+    banks_client.process_transaction(transaction).await?;
+    Ok(())
+}
+
+pub async fn freeze_token_account(
+    banks_client: &mut BanksClient,
+    payer: &Keypair,
+    recent_blockhash: &Hash,
+    account: &Pubkey,
+    pool_mint: &Pubkey,
+    manager: &Keypair,
+) -> Result<(), TransportError> {
+    let mut transaction = Transaction::new_with_payer(
+        &[spl_token::instruction::freeze_account(
+            &spl_token::id(),
+            &account,
+            pool_mint,
+            &manager.pubkey(),
+            &[],
+        )
+        .unwrap()],
+        Some(&payer.pubkey()),
+    );
+    transaction.sign(&[payer, manager], *recent_blockhash);
+    banks_client.process_transaction(transaction).await?;
+    Ok(())
+}
+
 pub async fn mint_tokens(
     banks_client: &mut BanksClient,
     payer: &Keypair,

stake-pool/program/tests/withdraw.rs

@@ -102,6 +102,20 @@ async fn setup() -> (
 
 #[tokio::test]
 async fn success() {
+    _success(SuccessTestType::Success).await;
+}
+
+#[tokio::test]
+async fn success_with_closed_manager_fee_account() {
+    _success(SuccessTestType::UninitializedManagerFee).await;
+}
+
+enum SuccessTestType {
+    Success,
+    UninitializedManagerFee,
+}
+
+async fn _success(test_type: SuccessTestType) {
     let (
         mut banks_client,
         payer,
@@ -148,6 +162,60 @@ async fn success() {
     )
     .await;
 
+    let destination_keypair = Keypair::new();
+    create_token_account(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &destination_keypair,
+        &stake_pool_accounts.pool_mint.pubkey(),
+        &Keypair::new().pubkey(),
+    )
+    .await
+    .unwrap();
+
+    if let SuccessTestType::UninitializedManagerFee = test_type {
+        transfer_spl_tokens(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &destination_keypair.pubkey(),
+            &stake_pool_accounts.manager,
+            pool_fee_balance_before,
+        )
+        .await;
+        // Check that the account cannot be frozen due to lack of
+        // freeze authority.
+        let transaction_error = freeze_token_account(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &stake_pool_accounts.pool_mint.pubkey(),
+            &stake_pool_accounts.manager,
+        )
+        .await
+        .unwrap_err();
+
+        match transaction_error {
+            TransportError::TransactionError(TransactionError::InstructionError(_, error)) => {
+                assert_eq!(error, InstructionError::Custom(0x10));
+            }
+            _ => panic!("Wrong error occurs while try to withdraw with wrong stake program ID"),
+        }
+        close_token_account(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &destination_keypair.pubkey(),
+            &stake_pool_accounts.manager,
+        )
+        .await
+        .unwrap();
+    }
+
     let new_authority = Pubkey::new_unique();
     let error = stake_pool_accounts
         .withdraw_stake(
@@ -169,7 +237,12 @@ async fn success() {
     let stake_pool =
         try_from_slice_unchecked::<state::StakePool>(&stake_pool.data.as_slice()).unwrap();
     // first and only deposit, lamports:pool 1:1
-    let tokens_withdrawal_fee = stake_pool_accounts.calculate_withdrawal_fee(tokens_to_withdraw);
+    let tokens_withdrawal_fee = match test_type {
+        SuccessTestType::Success => {
+            stake_pool_accounts.calculate_withdrawal_fee(tokens_to_withdraw)
+        }
+        _ => 0,
+    };
     let tokens_burnt = tokens_to_withdraw - tokens_withdrawal_fee;
     assert_eq!(
         stake_pool.total_stake_lamports,
@@ -180,16 +253,18 @@ async fn success() {
         stake_pool_before.pool_token_supply - tokens_burnt
     );
 
-    // Check manager received withdrawal fee
+    if let SuccessTestType::Success = test_type {
-    let pool_fee_balance = get_token_balance(
+        // Check manager received withdrawal fee
-        &mut banks_client,
+        let pool_fee_balance = get_token_balance(
-        &stake_pool_accounts.pool_fee_account.pubkey(),
+            &mut banks_client,
-    )
+            &stake_pool_accounts.pool_fee_account.pubkey(),
-    .await;
+        )
-    assert_eq!(
+        .await;
-        pool_fee_balance,
+        assert_eq!(
-        pool_fee_balance_before + tokens_withdrawal_fee,
+            pool_fee_balance,
-    );
+            pool_fee_balance_before + tokens_withdrawal_fee,
+        );
+    }
 
     // Check validator stake list storage
     let validator_list = get_account(