add mint extensions validation on fund dividends

CoMakery · Nov 19, 2024 · 25fcd83 · 25fcd83
1 parent e07de8b
commit 25fcd83
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 56 deletions.
diff --git a/programs/dividends/src/instructions/fund_dividends.rs b/programs/dividends/src/instructions/fund_dividends.rs
@@ -3,6 +3,7 @@ use anchor_spl::token_2022::spl_token_2022::onchain::invoke_transfer_checked;
 use anchor_spl::token_interface::{Mint, TokenAccount, TokenInterface};
 
 use crate::events::FundedEvent;
+use crate::utils::validate_transfer_fee_mint_extension;
 use crate::{errors::DividendsErrorCode, MerkleDistributor};
 
 /// [merkle_distributor::fund_dividends] accounts.
@@ -57,14 +58,13 @@ pub fn fund_dividends<'info>(
         DividendsErrorCode::KeysMustNotMatch
     );
     require!(
-        amount > 0 && amount <= ctx.accounts.distributor.total_claim_amount,
-        DividendsErrorCode::InvalidFundingAmount
-    );
-    require!(
+        // Ensure the funded amount exactly matches the total claim amount
         ctx.accounts.to.amount.checked_add(amount).unwrap()
-            <= ctx.accounts.distributor.total_claim_amount,
+            == ctx.accounts.distributor.total_claim_amount,
         DividendsErrorCode::InvalidFundingAmount
     );
+    let mint_data = &ctx.accounts.mint.to_account_info();
+    validate_transfer_fee_mint_extension(mint_data)?;
 
     let distributor = &mut ctx.accounts.distributor;
     let treasury_amount_before = ctx.accounts.to.amount;

diff --git a/tests/dividends/fund-dividends.ts b/tests/dividends/fund-dividends.ts
@@ -329,69 +329,39 @@ testCases.forEach(({ tokenProgramId, programName }) => {
       }
     });
 
-    it("ready to claim after fund dividends twice", async () => {
+    it("fails when fund less than required", async () => {
       const funderATADataBefore = await mintHelper.getAccount(funderATA);
       const distributorATADataBefore = await mintHelper.getAccount(
         distributorATA
       );
 
       const fundingAmountLeft = 1;
       const fundingAmount = totalClaimAmount.subn(fundingAmountLeft);
-      await dividendsProgram.methods
-        .fundDividends(fundingAmount)
-        .accountsStrict({
-          distributor,
-          mint: mintKeypair.publicKey,
-          from: funderATA,
-          to: distributorATA,
-          funder: funderKP.publicKey,
-          payer: signer.publicKey,
-          tokenProgram: tokenProgramId,
-        })
-        .signers([funderKP, signer])
-        .rpc({ commitment });
+      try {
+        await dividendsProgram.methods
+          .fundDividends(fundingAmount)
+          .accountsStrict({
+            distributor,
+            mint: mintKeypair.publicKey,
+            from: funderATA,
+            to: distributorATA,
+            funder: funderKP.publicKey,
+            payer: signer.publicKey,
+            tokenProgram: tokenProgramId,
+          })
+          .signers([funderKP, signer])
+          .rpc({ commitment });
+      } catch ({ error }) {
+        assert.equal(error.errorCode.code, "InvalidFundingAmount");
+        assert.equal(error.errorMessage, "Invalid funding amount");
+      }
+
       const distributorATADataAfterFirstFund = await mintHelper.getAccount(
         distributorATA
       );
       assert.equal(
         distributorATADataBefore.amount,
-        distributorATADataAfterFirstFund.amount -
-          BigInt(fundingAmount.toString())
-      );
-      let distributorData =
-        await dividendsProgram.account.merkleDistributor.fetch(distributor);
-      assert.isFalse(distributorData.readyToClaim);
-
-      await dividendsProgram.methods
-        .fundDividends(new BN(fundingAmountLeft))
-        .accountsStrict({
-          distributor,
-          mint: mintKeypair.publicKey,
-          from: funderATA,
-          to: distributorATA,
-          funder: funderKP.publicKey,
-          payer: signer.publicKey,
-          tokenProgram: tokenProgramId,
-        })
-        .signers([funderKP, signer])
-        .rpc({ commitment });
-      const funderATADataAfter = await mintHelper.getAccount(funderATA);
-      const distributorATADataAfterSecondFund = await mintHelper.getAccount(
-        distributorATA
-      );
-      distributorData = await dividendsProgram.account.merkleDistributor.fetch(
-        distributor
-      );
-      assert.isTrue(distributorData.readyToClaim);
-      assert.equal(
-        funderATADataBefore.amount,
-        funderATADataAfter.amount +
-          BigInt(fundingAmountLeft) +
-          BigInt(fundingAmount.toString())
-      );
-      assert.equal(
-        distributorATADataAfterFirstFund.amount + BigInt(fundingAmountLeft),
-        distributorATADataAfterSecondFund.amount
+        distributorATADataAfterFirstFund.amount
       );
     });