harvest fees when closing token2022 token account (#19)

* harvest fees when closing token2022 token account

* remove signer for harvesting fees

* fix harvesting fees

* add test for harvest fee during cancelation

* fix test to include transfer_fee calculation

* nit: fix comments and function name

README.md

@@ -5,21 +5,23 @@ Open source program to allow user to lock token based on a vesting plan
 ## For developer
 
 To run all the test, please run the following commands:
-```
+
+```bash
 anchor test -- --features localnet
 ```
 
 Build programs:
-```
+
+```bash
 anchor build
 ```
 
-
 ## Audit
-Jupiter Locker has been audited by Sec3 (prev.Soteria) and OtterSec. View the audit report [here](./audits).
 
+Jupiter Locker has been audited by Sec3 (prev.Soteria) and OtterSec. View the audit report [here](./audits).
 
 ## License
+
 Anchor is licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
 
 Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in Jupiter Locker by you, as defined in the Apache-2.0 license, shall be licensed as above, without any additional terms or conditions.

programs/locker/src/instructions/v2/cancel_vesting_escrow.rs

@@ -1,9 +1,11 @@
 use anchor_spl::memo::Memo;
+use anchor_spl::token_2022::spl_token_2022::extension::confidential_transfer_fee::instruction::WithdrawWithheldTokensFromAccountsData;
 use anchor_spl::token_interface::{
     close_account, CloseAccount, Mint, TokenAccount, TokenInterface,
 };
 use util::{
-    parse_remaining_accounts, AccountsType, ParsedRemainingAccounts, TRANSFER_MEMO_CANCEL_VESTING,
+    harvest_fees, parse_remaining_accounts, AccountsType, ParsedRemainingAccounts,
+    TRANSFER_MEMO_CANCEL_VESTING,
 };
 
 use crate::safe_math::SafeMath;
@@ -23,6 +25,7 @@ pub struct CancelVestingEscrow<'info> {
     pub escrow: AccountLoader<'info, VestingEscrow>,
 
     /// Mint.
+    #[account(mut)]
     pub token_mint: Box<InterfaceAccount<'info, Mint>>,
 
     /// Escrow Token Account.
@@ -146,6 +149,13 @@ pub fn handle_cancel_vesting_escrow<'c: 'info, 'info>(
         parsed_transfer_hook_accounts.transfer_hook_escrow,
     )?;
 
+    // Do fee harvesting
+    harvest_fees(
+        &ctx.accounts.token_program,
+        &ctx.accounts.escrow_token,
+        &ctx.accounts.token_mint,
+    )?;
+
     ctx.accounts.close_escrow_token()?;
 
     emit_cpi!(EventCancelVestingEscrow {

programs/locker/src/util/token2022.rs

@@ -2,15 +2,18 @@ use anchor_lang::prelude::*;
 use anchor_spl::memo;
 use anchor_spl::memo::{BuildMemo, Memo};
 use anchor_spl::token::Token;
+use anchor_spl::token_2022::spl_token_2022::extension::transfer_fee::{
+    TransferFee, TransferFeeConfig, MAX_FEE_BASIS_POINTS,
+};
 use anchor_spl::token_2022::spl_token_2022::{
     self,
     extension::{self, StateWithExtensions},
 };
-use anchor_spl::token_2022::spl_token_2022::extension::transfer_fee::{
-    MAX_FEE_BASIS_POINTS, TransferFee, TransferFeeConfig,
-};
-use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface};
 use anchor_spl::token_interface::spl_token_2022::extension::BaseStateWithExtensions;
+use anchor_spl::token_interface::{
+    harvest_withheld_tokens_to_mint, HarvestWithheldTokensToMint, Mint, TokenAccount,
+    TokenInterface,
+};
 
 use crate::{LockerError, VestingEscrow};
 
@@ -219,6 +222,42 @@ pub fn calculate_transfer_fee_included_amount(
     Ok(actual_amount)
 }
 
+pub fn harvest_fees<'c: 'info, 'info>(
+    token_program_id: &Interface<'info, TokenInterface>,
+    token_account: &InterfaceAccount<'info, TokenAccount>,
+    mint: &InterfaceAccount<'info, Mint>,
+) -> Result<()> {
+    let mint_info = mint.to_account_info();
+    if mint_info.owner.key() == Token::id() {
+        return Result::Ok(());
+    }
+
+    let token_mint_data = mint_info.try_borrow_data()?;
+    let token_mint_unpacked =
+        StateWithExtensions::<spl_token_2022::state::Mint>::unpack(&token_mint_data)?;
+    let mut is_harvestable = false;
+    if let Ok(_transfer_fee_config) = token_mint_unpacked.get_extension::<TransferFeeConfig>() {
+        is_harvestable = true;
+    }
+    // need to do this because Rust says we are still borrowing the data
+    drop(token_mint_data);
+
+    if is_harvestable {
+        harvest_withheld_tokens_to_mint(
+            CpiContext::new(
+                token_program_id.to_account_info(),
+                HarvestWithheldTokensToMint {
+                    token_program_id: token_program_id.to_account_info(),
+                    mint: mint.to_account_info(),
+                },
+            ),
+            vec![token_account.to_account_info()],
+        )?;
+    }
+
+    Ok(())
+}
+
 // Memo Extension support
 pub fn is_transfer_memo_required(token_account: &InterfaceAccount<TokenAccount>) -> Result<bool> {
     let token_account_info = token_account.to_account_info();

tests/common/index.ts

@@ -117,3 +117,7 @@ export const getCurrentBlockTime = async (connection: web3.Connection) => {
   const currentBlockTime = await connection.getBlockTime(currentSlot);
   return currentBlockTime;
 };
+
+export const getCurrentEpoch = async (connection: web3.Connection) => {
+  return (await connection.getEpochInfo()).epoch;
+};

tests/locker_utils/index.ts

@@ -9,8 +9,12 @@ import {
 import { Locker } from "../../target/types/locker";
 import {
   ASSOCIATED_TOKEN_PROGRAM_ID,
+  calculateEpochFee,
   createAssociatedTokenAccountInstruction,
   getAssociatedTokenAddressSync,
+  getEpochFee,
+  getMint,
+  getTransferFeeConfig,
   TOKEN_2022_PROGRAM_ID,
   TOKEN_PROGRAM_ID,
 } from "@solana/spl-token";
@@ -21,6 +25,7 @@ import {
   RemainingAccountsType,
 } from "./token_2022/remaining-accounts";
 import { AccountMeta, ComputeBudgetProgram } from "@solana/web3.js";
+import { getCurrentEpoch } from "../common";
 
 export const LOCKER_PROGRAM_ID = new web3.PublicKey(
   "2r5VekMNiWPzi1pWwvJczrdPaZnJG59u91unSrTunwJg"
@@ -582,6 +587,32 @@ export async function cancelVestingPlan(
     .signers([signer])
     .rpc();
 
+  let creator_fee = 0;
+  let claimer_fee = 0;
+  if (tokenProgram == TOKEN_2022_PROGRAM_ID) {
+    const feeConfig = getTransferFeeConfig(
+      await getMint(
+        program.provider.connection,
+        escrowState.tokenMint,
+        undefined,
+        TOKEN_2022_PROGRAM_ID
+      )
+    );
+    const epoch = BigInt(await getCurrentEpoch(program.provider.connection));
+    creator_fee = feeConfig
+      ? Number(
+          calculateEpochFee(
+            feeConfig,
+            epoch,
+            BigInt(total_amount - claimable_amount)
+          )
+        )
+      : 0;
+    claimer_fee = feeConfig
+      ? Number(calculateEpochFee(feeConfig, epoch, BigInt(claimable_amount)))
+      : 0;
+  }
+
   if (isAssertion) {
     const escrowState = await program.account.vestingEscrow.fetch(escrow);
     expect(escrowState.cancelledAt.toNumber()).greaterThan(0);
@@ -595,14 +626,17 @@ export async function cancelVestingPlan(
       await program.provider.connection.getTokenAccountBalance(creatorToken)
     ).value.amount;
     expect(
-      parseInt(creator_token_balance_before) + total_amount - claimable_amount
+      parseInt(creator_token_balance_before) +
+        total_amount -
+        claimable_amount -
+        creator_fee
     ).eq(parseInt(creator_token_balance));
 
     const recipient_token_balance = (
       await program.provider.connection.getTokenAccountBalance(recipientToken)
     ).value.amount;
-    expect(parseInt(recipient_token_balance_before) + claimable_amount).eq(
-      parseInt(recipient_token_balance)
-    );
+    expect(
+      parseInt(recipient_token_balance_before) + claimable_amount - claimer_fee
+    ).eq(parseInt(recipient_token_balance));
   }
 }