stake-pool: Don't assess withdrawal fee from fee account (#2180)

stake-pool/program/src/processor.rs

@@ -1963,9 +1963,13 @@ impl Processor {
             return Err(StakePoolError::InvalidState.into());
         }
 
-        let pool_tokens_fee = stake_pool
-            .calc_pool_tokens_withdrawal_fee(pool_tokens)
-            .ok_or(StakePoolError::CalculationFailure)?;
+        let pool_tokens_fee = if stake_pool.manager_fee_account == *burn_from_info.key {
+            0
+        } else {
+            stake_pool
+                .calc_pool_tokens_withdrawal_fee(pool_tokens)
+                .ok_or(StakePoolError::CalculationFailure)?
+        };
         let pool_tokens_burnt = pool_tokens
             .checked_sub(pool_tokens_fee)
             .ok_or(StakePoolError::CalculationFailure)?;

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

@@ -99,6 +99,32 @@ pub async fn transfer(
     banks_client.process_transaction(transaction).await.unwrap();
 }
 
+pub async fn transfer_spl_tokens(
+    banks_client: &mut BanksClient,
+    payer: &Keypair,
+    recent_blockhash: &Hash,
+    source: &Pubkey,
+    destination: &Pubkey,
+    authority: &Keypair,
+    amount: u64,
+) {
+    let transaction = Transaction::new_signed_with_payer(
+        &[spl_token::instruction::transfer(
+            &spl_token::id(),
+            source,
+            destination,
+            &authority.pubkey(),
+            &[],
+            amount,
+        )
+        .unwrap()],
+        Some(&payer.pubkey()),
+        &[payer, authority],
+        *recent_blockhash,
+    );
+    banks_client.process_transaction(transaction).await.unwrap();
+}
+
 pub async fn create_token_account(
     banks_client: &mut BanksClient,
     payer: &Keypair,

stake-pool/program/tests/withdraw.rs

@@ -1336,3 +1336,72 @@ async fn success_withdraw_from_transient() {
         .await;
     assert!(error.is_none());
 }
+
+#[tokio::test]
+async fn success_withdraw_all_fee_tokens() {
+    let (
+        mut banks_client,
+        payer,
+        recent_blockhash,
+        stake_pool_accounts,
+        validator_stake_account,
+        deposit_info,
+        user_transfer_authority,
+        user_stake_recipient,
+        tokens_to_withdraw,
+    ) = setup().await;
+
+    // move tokens to fee account
+    transfer_spl_tokens(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &deposit_info.pool_account.pubkey(),
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+        &user_transfer_authority,
+        tokens_to_withdraw,
+    )
+    .await;
+
+    let fee_tokens = get_token_balance(
+        &mut banks_client,
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+    )
+    .await;
+
+    let user_transfer_authority = Keypair::new();
+    delegate_tokens(
+        &mut banks_client,
+        &payer,
+        &recent_blockhash,
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+        &stake_pool_accounts.manager,
+        &user_transfer_authority.pubkey(),
+        fee_tokens,
+    )
+    .await;
+
+    let new_authority = Pubkey::new_unique();
+    let error = stake_pool_accounts
+        .withdraw_stake(
+            &mut banks_client,
+            &payer,
+            &recent_blockhash,
+            &user_stake_recipient.pubkey(),
+            &user_transfer_authority,
+            &stake_pool_accounts.pool_fee_account.pubkey(),
+            &validator_stake_account.stake_account,
+            &new_authority,
+            fee_tokens,
+        )
+        .await;
+    assert!(error.is_none());
+
+    // Check balance is 0
+    let fee_tokens = get_token_balance(
+        &mut banks_client,
+        &stake_pool_accounts.pool_fee_account.pubkey(),
+    )
+    .await;
+    assert_eq!(fee_tokens, 0);
+}