token-2022: Fail to close account with withheld fees (#2871)

* token-2022: Fail to close account with withheld fees

* Add comment to use `get_extension`

Author: joncinque

token/program-2022/src/error.rs

@@ -124,6 +124,10 @@ pub enum TokenError {
     /// The owner authority cannot be changed
     #[error("The owner authority cannot be changed")]
     ImmutableOwner,
+    /// An account can only be closed if its withheld fee balance is zero, harvest fees to the
+    /// mint and try again
+    #[error("An account can only be closed if its withheld fee balance is zero, harvest fees to the mint and try again")]
+    AccountHasWithheldTransferFees,
 }
 impl From<TokenError> for ProgramError {
     fn from(e: TokenError) -> Self {

token/program-2022/src/extension/transfer_fee/mod.rs

@@ -1,10 +1,11 @@
 use {
     crate::{
+        error::TokenError,
         extension::{Extension, ExtensionType},
         pod::*,
     },
     bytemuck::{Pod, Zeroable},
-    solana_program::clock::Epoch,
+    solana_program::{clock::Epoch, entrypoint::ProgramResult},
     std::{cmp, convert::TryFrom},
 };
 
@@ -90,6 +91,16 @@ pub struct TransferFeeAmount {
     /// Amount withheld during transfers, to be harvested to the mint
     pub withheld_amount: PodU64,
 }
+impl TransferFeeAmount {
+    /// Check if the extension is in a closable state
+    pub fn closable(&self) -> ProgramResult {
+        if self.withheld_amount == 0.into() {
+            Ok(())
+        } else {
+            Err(TokenError::AccountHasWithheldTransferFees.into())
+        }
+    }
+}
 impl Extension for TransferFeeAmount {
     const TYPE: ExtensionType = ExtensionType::TransferFeeAmount;
 }

token/program-2022/src/processor.rs

@@ -847,11 +847,20 @@ impl Processor {
                 account_info_iter.as_slice(),
             )?;
 
+            // TODO use get_extension when
+            // https://github.com/solana-labs/solana-program-library/pull/2822 lands
             if let Ok(confidential_transfer_state) =
                 source_account.get_extension_mut::<ConfidentialTransferAccount>()
             {
                 confidential_transfer_state.closable()?
             }
+
+            // TODO use get_extension when
+            // https://github.com/solana-labs/solana-program-library/pull/2822 lands
+            if let Ok(transfer_fee_state) = source_account.get_extension_mut::<TransferFeeAmount>()
+            {
+                transfer_fee_state.closable()?
+            }
         } else if let Ok(mut mint) =
             StateWithExtensionsMut::<Mint>::unpack(&mut source_account_data)
         {
@@ -1284,6 +1293,9 @@ impl PrintProgramError for TokenError {
             TokenError::ImmutableOwner => {
                 msg!("The owner authority cannot be changed");
             }
+            TokenError::AccountHasWithheldTransferFees => {
+                msg!("Error: An account can only be closed if its withheld fee balance is zero, harvest fees to the mint and try again");
+            }
         }
     }
 }

token/program-2022/tests/transfer_fee.rs

@@ -1539,3 +1539,62 @@ async fn withdraw_withheld_tokens_from_accounts() {
         )))
     );
 }
+
+#[tokio::test]
+async fn fail_close_with_withheld() {
+    let amount = TEST_MAXIMUM_FEE;
+    let alice_amount = amount * 100;
+    let TokenWithAccounts {
+        token,
+        transfer_fee_config,
+        alice,
+        alice_account,
+        decimals,
+        ..
+    } = create_mint_with_accounts(alice_amount).await;
+
+    // accrue withheld fees on new account
+    let account = create_and_transfer_to_account(
+        &token,
+        &alice_account,
+        &alice,
+        &alice.pubkey(),
+        amount,
+        decimals,
+    )
+    .await;
+
+    // empty the account
+    let fee = transfer_fee_config.calculate_epoch_fee(0, amount).unwrap();
+    token
+        .transfer_checked(&account, &alice_account, &alice, amount - fee, decimals)
+        .await
+        .unwrap();
+
+    // fail to close
+    let error = token
+        .close_account(&account, &Pubkey::new_unique(), &alice)
+        .await
+        .unwrap_err();
+    assert_eq!(
+        error,
+        TokenClientError::Client(Box::new(TransportError::TransactionError(
+            TransactionError::InstructionError(
+                0,
+                InstructionError::Custom(TokenError::AccountHasWithheldTransferFees as u32)
+            )
+        )))
+    );
+
+    // harvest the fees to the mint
+    token
+        .harvest_withheld_tokens_to_mint(&[&account])
+        .await
+        .unwrap();
+
+    // successfully close
+    token
+        .close_account(&account, &Pubkey::new_unique(), &alice)
+        .await
+        .unwrap();
+}