token-2022: Add withdraw withheld tokens from accounts (#2857)

* token-2022: Add withdraw withheld tokens from accounts

* Add instruction field for number of signers

* Update instruction field

* Rename field and update comment

* Fix test?

Author: joncinque

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

@@ -6,6 +6,7 @@ use {
         program_option::COption,
         pubkey::Pubkey,
     },
+    std::convert::TryFrom,
 };
 
 /// Transfer Fee extension instructions
@@ -95,7 +96,10 @@ pub enum TransferFeeInstruction {
     ///   2. `[]` The mint's `withdraw_withheld_authority`'s multisignature owner/delegate.
     ///   3. ..3+M `[signer]` M signer accounts.
     ///   3+M+1. ..3+M+N `[writable]` The source accounts to withdraw from.
-    WithdrawWithheldTokensFromAccounts,
+    WithdrawWithheldTokensFromAccounts {
+        /// Number of token accounts harvested
+        num_token_accounts: u8,
+    },
     /// Permissionless instruction to transfer all withheld tokens to the mint.
     ///
     /// Succeeds for frozen accounts.
@@ -161,7 +165,11 @@ impl TransferFeeInstruction {
                 (instruction, rest)
             }
             2 => (Self::WithdrawWithheldTokensFromMint, rest),
-            3 => (Self::WithdrawWithheldTokensFromAccounts, rest),
+            3 => {
+                let (&num_token_accounts, rest) = rest.split_first().ok_or(InvalidInstruction)?;
+                let instruction = Self::WithdrawWithheldTokensFromAccounts { num_token_accounts };
+                (instruction, rest)
+            }
             4 => (Self::HarvestWithheldTokensToMint, rest),
             5 => {
                 let (transfer_fee_basis_points, rest) = TokenInstruction::unpack_u16(rest)?;
@@ -204,8 +212,9 @@ impl TransferFeeInstruction {
             Self::WithdrawWithheldTokensFromMint => {
                 buffer.push(2);
             }
-            Self::WithdrawWithheldTokensFromAccounts => {
+            Self::WithdrawWithheldTokensFromAccounts { num_token_accounts } => {
                 buffer.push(3);
+                buffer.push(num_token_accounts);
             }
             Self::HarvestWithheldTokensToMint => {
                 buffer.push(4);
@@ -326,6 +335,8 @@ pub fn withdraw_withheld_tokens_from_accounts(
     sources: &[&Pubkey],
 ) -> Result<Instruction, ProgramError> {
     check_program_account(token_program_id)?;
+    let num_token_accounts =
+        u8::try_from(sources.len()).map_err(|_| ProgramError::InvalidInstructionData)?;
     let mut accounts = Vec::with_capacity(3 + signers.len() + sources.len());
     accounts.push(AccountMeta::new_readonly(*mint, false));
     accounts.push(AccountMeta::new(*destination, false));
@@ -341,7 +352,7 @@ pub fn withdraw_withheld_tokens_from_accounts(
         program_id: *token_program_id,
         accounts,
         data: TokenInstruction::TransferFeeExtension(
-            TransferFeeInstruction::WithdrawWithheldTokensFromAccounts,
+            TransferFeeInstruction::WithdrawWithheldTokensFromAccounts { num_token_accounts },
         )
         .pack(),
     })
@@ -446,11 +457,12 @@ mod test {
         let unpacked = TokenInstruction::unpack(&expect).unwrap();
         assert_eq!(unpacked, check);
 
+        let num_token_accounts = 255;
         let check = TokenInstruction::TransferFeeExtension(
-            TransferFeeInstruction::WithdrawWithheldTokensFromAccounts,
+            TransferFeeInstruction::WithdrawWithheldTokensFromAccounts { num_token_accounts },
         );
         let packed = check.pack();
-        let expect = [23u8, 3];
+        let expect = [23u8, 3, num_token_accounts];
         assert_eq!(packed, expect);
         let unpacked = TokenInstruction::unpack(&expect).unwrap();
         assert_eq!(unpacked, check);

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

@@ -7,7 +7,7 @@ use {
                 instruction::TransferFeeInstruction, TransferFee, TransferFeeAmount,
                 TransferFeeConfig, MAX_FEE_BASIS_POINTS,
             },
-            StateWithExtensionsMut,
+            StateWithExtensions, StateWithExtensionsMut,
         },
         processor::Processor,
         state::{Account, Mint},
@@ -154,9 +154,8 @@ fn process_withdraw_withheld_tokens_from_mint(
 
 fn harvest_from_account<'a, 'b>(
     mint_key: &'b Pubkey,
-    mint_extension: &'b mut TransferFeeConfig,
     token_account_info: &'b AccountInfo<'a>,
-) -> Result<(), TokenError> {
+) -> Result<u64, TokenError> {
     let mut token_account_data = token_account_info.data.borrow_mut();
     let mut token_account = StateWithExtensionsMut::<Account>::unpack(&mut token_account_data)
         .map_err(|_| TokenError::InvalidState)?;
@@ -168,13 +167,8 @@ fn harvest_from_account<'a, 'b>(
         .get_extension_mut::<TransferFeeAmount>()
         .map_err(|_| TokenError::InvalidState)?;
     let account_withheld_amount = u64::from(token_account_extension.withheld_amount);
-    let mint_withheld_amount = u64::from(mint_extension.withheld_amount);
-    mint_extension.withheld_amount = mint_withheld_amount
-        .checked_add(account_withheld_amount)
-        .ok_or(TokenError::Overflow)?
-        .into();
     token_account_extension.withheld_amount = 0.into();
-    Ok(())
+    Ok(account_withheld_amount)
 }
 
 fn process_harvest_withheld_tokens_to_mint(accounts: &[AccountInfo]) -> ProgramResult {
@@ -187,18 +181,92 @@ fn process_harvest_withheld_tokens_to_mint(accounts: &[AccountInfo]) -> ProgramR
     let mint_extension = mint.get_extension_mut::<TransferFeeConfig>()?;
 
     for token_account_info in token_account_infos {
-        match harvest_from_account(mint_account_info.key, mint_extension, token_account_info) {
-            // Shouldn't ever happen, but if it does, we don't want to propagate any half-done changes
-            Err(TokenError::Overflow) => return Err(TokenError::Overflow.into()),
+        match harvest_from_account(mint_account_info.key, token_account_info) {
+            Ok(amount) => {
+                let mint_withheld_amount = u64::from(mint_extension.withheld_amount);
+                mint_extension.withheld_amount = mint_withheld_amount
+                    .checked_add(amount)
+                    .ok_or(TokenError::Overflow)?
+                    .into();
+            }
             Err(e) => {
                 msg!("Error harvesting from {}: {}", token_account_info.key, e);
             }
-            Ok(_) => {}
         }
     }
     Ok(())
 }
 
+fn process_withdraw_withheld_tokens_from_accounts(
+    program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    num_token_accounts: u8,
+) -> ProgramResult {
+    let account_info_iter = &mut accounts.iter();
+    let mint_account_info = next_account_info(account_info_iter)?;
+    let dest_account_info = next_account_info(account_info_iter)?;
+    let authority_info = next_account_info(account_info_iter)?;
+    let authority_info_data_len = authority_info.data_len();
+    let account_infos = account_info_iter.as_slice();
+    let num_signers = account_infos
+        .len()
+        .saturating_sub(num_token_accounts as usize);
+
+    let mint_data = mint_account_info.data.borrow();
+    let mint = StateWithExtensions::<Mint>::unpack(&mint_data)?;
+    let extension = mint.get_extension::<TransferFeeConfig>()?;
+
+    let withdraw_withheld_authority = Option::<Pubkey>::from(extension.withdraw_withheld_authority)
+        .ok_or(TokenError::NoAuthorityExists)?;
+    Processor::validate_owner(
+        program_id,
+        &withdraw_withheld_authority,
+        authority_info,
+        authority_info_data_len,
+        &account_infos[..num_signers],
+    )?;
+
+    let mut dest_account_data = dest_account_info.data.borrow_mut();
+    let mut dest_account = StateWithExtensionsMut::<Account>::unpack(&mut dest_account_data)?;
+    if dest_account.base.mint != *mint_account_info.key {
+        return Err(TokenError::MintMismatch.into());
+    }
+    if dest_account.base.is_frozen() {
+        return Err(TokenError::AccountFrozen.into());
+    }
+    for account_info in &account_infos[num_signers..] {
+        // self-harvest, can't double-borrow the underlying data
+        if account_info.key == dest_account_info.key {
+            let token_account_extension = dest_account
+                .get_extension_mut::<TransferFeeAmount>()
+                .map_err(|_| TokenError::InvalidState)?;
+            let account_withheld_amount = u64::from(token_account_extension.withheld_amount);
+            token_account_extension.withheld_amount = 0.into();
+            dest_account.base.amount = dest_account
+                .base
+                .amount
+                .checked_add(account_withheld_amount)
+                .ok_or(TokenError::Overflow)?;
+        } else {
+            match harvest_from_account(mint_account_info.key, account_info) {
+                Ok(amount) => {
+                    dest_account.base.amount = dest_account
+                        .base
+                        .amount
+                        .checked_add(amount)
+                        .ok_or(TokenError::Overflow)?;
+                }
+                Err(e) => {
+                    msg!("Error harvesting from {}: {}", account_info.key, e);
+                }
+            }
+        }
+    }
+    dest_account.pack_base();
+
+    Ok(())
+}
+
 pub(crate) fn process_instruction(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
@@ -231,8 +299,9 @@ pub(crate) fn process_instruction(
             msg!("TransferFeeInstruction: WithdrawWithheldTokensFromMint");
             process_withdraw_withheld_tokens_from_mint(program_id, accounts)
         }
-        TransferFeeInstruction::WithdrawWithheldTokensFromAccounts => {
-            unimplemented!();
+        TransferFeeInstruction::WithdrawWithheldTokensFromAccounts { num_token_accounts } => {
+            msg!("TransferFeeInstruction: WithdrawWithheldTokensFromAccounts");
+            process_withdraw_withheld_tokens_from_accounts(program_id, accounts, num_token_accounts)
         }
         TransferFeeInstruction::HarvestWithheldTokensToMint => {
             msg!("TransferFeeInstruction: HarvestWithheldTokensToMint");