feat: track whitelist tokens, do not track all tokens, only iterate t… (#98)

bengobeil · web-flow · commit 099bcfc7691b · 2025-08-29T15:41:35.000-03:00
feat: track whitelist tokens, do not track all tokens, only iterate through whitelisted tokens
diff --git a/src/BullaFrendLendV2.sol b/src/BullaFrendLendV2.sol
@@ -43,9 +43,8 @@ contract BullaFrendLendV2 is BullaClaimControllerBase, ERC165, Ownable, IBullaFr
     uint256 public loanOfferCount;
     uint16 public protocolFeeBPS;
 
-    address[] public protocolFeeTokens;
+    address[] public whitelistedProtocolFeeTokens;
     mapping(address => uint256) public protocolFeesByToken;
-    mapping(address => bool) private _tokenExists;
 
     // Whitelist for protocol fee token withdrawals
     mapping(address => bool) public protocolFeeTokenWhitelist;
@@ -373,12 +372,8 @@ contract BullaFrendLendV2 is BullaClaimControllerBase, ERC165, Ownable, IBullaFr
 
         // Transfer the total amount from sender to this contract, to avoid double approval
         if (paymentAmount > 0) {
-            // Track protocol fee for this token if any interest was paid
+            // Track protocol fee for this token if any interest was paid, even if token is not whitelisted
             if (protocolFee > 0) {
-                if (!_tokenExists[loan.token]) {
-                    protocolFeeTokens.push(loan.token);
-                    _tokenExists[loan.token] = true;
-                }
                 protocolFeesByToken[loan.token] += protocolFee;
             }
 
@@ -423,12 +418,12 @@ contract BullaFrendLendV2 is BullaClaimControllerBase, ERC165, Ownable, IBullaFr
      * @notice Allows owner to withdraw accumulated protocol fees
      */
     function withdrawAllFees() external onlyOwner {
-        // Withdraw protocol fees in all tracked tokens that are whitelisted
-        for (uint256 i = 0; i < protocolFeeTokens.length; i++) {
-            address token = protocolFeeTokens[i];
+        // Withdraw protocol fees in all whitelisted tokens
+        for (uint256 i = 0; i < whitelistedProtocolFeeTokens.length; i++) {
+            address token = whitelistedProtocolFeeTokens[i];
             uint256 feeAmount = protocolFeesByToken[token];
 
-            if (feeAmount > 0 && protocolFeeTokenWhitelist[token]) {
+            if (feeAmount > 0) {
                 protocolFeesByToken[token] = 0; // Reset fee amount before transfer
                 ERC20(token).safeTransfer(owner(), feeAmount);
                 emit FeeWithdrawn(owner(), token, feeAmount);
@@ -454,7 +449,10 @@ contract BullaFrendLendV2 is BullaClaimControllerBase, ERC165, Ownable, IBullaFr
      * @param token The token address to whitelist for withdrawals
      */
     function addToFeeTokenWhitelist(address token) external onlyOwner {
-        protocolFeeTokenWhitelist[token] = true;
+        if (!protocolFeeTokenWhitelist[token]) {
+            protocolFeeTokenWhitelist[token] = true;
+            whitelistedProtocolFeeTokens.push(token);
+        }
 
         emit TokenAddedToFeesWhitelist(token);
     }
@@ -464,7 +462,19 @@ contract BullaFrendLendV2 is BullaClaimControllerBase, ERC165, Ownable, IBullaFr
      * @param token The token address to remove from withdrawal whitelist
      */
     function removeFromFeeTokenWhitelist(address token) external onlyOwner {
-        protocolFeeTokenWhitelist[token] = false;
+        if (protocolFeeTokenWhitelist[token]) {
+            protocolFeeTokenWhitelist[token] = false;
+
+            // Remove from whitelistedProtocolFeeTokens array
+            for (uint256 i = 0; i < whitelistedProtocolFeeTokens.length; i++) {
+                if (whitelistedProtocolFeeTokens[i] == token) {
+                    whitelistedProtocolFeeTokens[i] =
+                        whitelistedProtocolFeeTokens[whitelistedProtocolFeeTokens.length - 1];
+                    whitelistedProtocolFeeTokens.pop();
+                    break;
+                }
+            }
+        }
 
         emit TokenRemovedFromFeesWhitelist(token);
     }
diff --git a/src/BullaInvoice.sol b/src/BullaInvoice.sol
@@ -41,9 +41,8 @@ contract BullaInvoice is BullaClaimControllerBase, ERC165, Ownable, IBullaInvoic
 
     ClaimMetadata public EMPTY_METADATA = ClaimMetadata({attachmentURI: "", tokenURI: ""});
 
-    address[] public protocolFeeTokens;
+    address[] public whitelistedProtocolFeeTokens;
     mapping(address => uint256) public protocolFeesByToken;
-    mapping(address => bool) private _tokenExists;
 
     // Whitelist for protocol fee token withdrawals
     mapping(address => bool) public protocolFeeTokenWhitelist;
@@ -329,10 +328,6 @@ contract BullaInvoice is BullaClaimControllerBase, ERC165, Ownable, IBullaInvoic
                 // Track protocol fee for this token if any interest was paid
                 // No need to track gas fee as it is the balance of the contract
                 if (protocolFee > 0) {
-                    if (!_tokenExists[claim.token]) {
-                        protocolFeeTokens.push(claim.token);
-                        _tokenExists[claim.token] = true;
-                    }
                     protocolFeesByToken[claim.token] += protocolFee;
                 }
                 // Handle ERC20 payments
@@ -496,12 +491,12 @@ contract BullaInvoice is BullaClaimControllerBase, ERC165, Ownable, IBullaInvoic
             emit FeeWithdrawn(owner(), address(0), ethBalance);
         }
 
-        // Withdraw protocol fees in all tracked tokens that are whitelisted
-        for (uint256 i = 0; i < protocolFeeTokens.length; i++) {
-            address token = protocolFeeTokens[i];
+        // Withdraw protocol fees in all whitelisted tokens
+        for (uint256 i = 0; i < whitelistedProtocolFeeTokens.length; i++) {
+            address token = whitelistedProtocolFeeTokens[i];
             uint256 feeAmount = protocolFeesByToken[token];
 
-            if (feeAmount > 0 && protocolFeeTokenWhitelist[token]) {
+            if (feeAmount > 0) {
                 protocolFeesByToken[token] = 0; // Reset fee amount before transfer
                 ERC20(token).safeTransfer(owner(), feeAmount);
                 emit FeeWithdrawn(owner(), token, feeAmount);
@@ -527,7 +522,10 @@ contract BullaInvoice is BullaClaimControllerBase, ERC165, Ownable, IBullaInvoic
      * @param token The token address to whitelist for withdrawals
      */
     function addToFeeTokenWhitelist(address token) external onlyOwner {
-        protocolFeeTokenWhitelist[token] = true;
+        if (!protocolFeeTokenWhitelist[token]) {
+            protocolFeeTokenWhitelist[token] = true;
+            whitelistedProtocolFeeTokens.push(token);
+        }
 
         emit TokenAddedToFeesWhitelist(token);
     }
@@ -537,7 +535,19 @@ contract BullaInvoice is BullaClaimControllerBase, ERC165, Ownable, IBullaInvoic
      * @param token The token address to remove from withdrawal whitelist
      */
     function removeFromFeeTokenWhitelist(address token) external onlyOwner {
-        protocolFeeTokenWhitelist[token] = false;
+        if (protocolFeeTokenWhitelist[token]) {
+            protocolFeeTokenWhitelist[token] = false;
+
+            // Remove from whitelistedProtocolFeeTokens array
+            for (uint256 i = 0; i < whitelistedProtocolFeeTokens.length; i++) {
+                if (whitelistedProtocolFeeTokens[i] == token) {
+                    whitelistedProtocolFeeTokens[i] =
+                        whitelistedProtocolFeeTokens[whitelistedProtocolFeeTokens.length - 1];
+                    whitelistedProtocolFeeTokens.pop();
+                    break;
+                }
+            }
+        }
 
         emit TokenRemovedFromFeesWhitelist(token);
     }
diff --git a/test/foundry/BullaFrendLend/BullaFrendLend.t.sol b/test/foundry/BullaFrendLend/BullaFrendLend.t.sol
@@ -952,21 +952,6 @@ contract TestBullaFrendLend is Test {
         assertEq(bullaFrendLend.protocolFeeBPS(), newProtocolFeeBPS, "Protocol fee not updated correctly");
     }
 
-    // helper function to check if token is in protocol fee tokens
-    function isTokenInProtocolFeeTokens(address token) internal view returns (bool) {
-        // we've defined at most 3 tokens in our test
-        for (uint256 i = 0; i < 3; i++) {
-            try bullaFrendLend.protocolFeeTokens(i) returns (address tokenAddr) {
-                if (tokenAddr == token) {
-                    return true;
-                }
-            } catch {
-                break;
-            }
-        }
-        return false;
-    }
-
     function testProtocolFeeWithMultipleTokens() public {
         vm.startPrank(creditor);
         weth.approve(address(bullaFrendLend), 10 ether);
@@ -1052,10 +1037,6 @@ contract TestBullaFrendLend is Test {
         assertEq(bullaFrendLend.protocolFeesByToken(address(weth)), wethProtocolFee, "WETH fee tracking incorrect");
         assertEq(bullaFrendLend.protocolFeesByToken(address(usdc)), usdcProtocolFee, "USDC fee tracking incorrect");
         assertEq(bullaFrendLend.protocolFeesByToken(address(dai)), daiProtocolFee, "DAI fee tracking incorrect");
-
-        assertTrue(isTokenInProtocolFeeTokens(address(weth)), "WETH not found in protocol fee tokens array");
-        assertTrue(isTokenInProtocolFeeTokens(address(usdc)), "USDC not found in protocol fee tokens array");
-        assertTrue(isTokenInProtocolFeeTokens(address(dai)), "DAI not found in protocol fee tokens array");
     }
 
     function testWithdrawAllFees() public {
@@ -1357,88 +1338,6 @@ contract TestBullaFrendLend is Test {
         }
     }
 
-    function testTokenTrackingUniqueness() public {
-        vm.startPrank(creditor);
-        weth.approve(address(bullaFrendLend), 10 ether);
-        vm.stopPrank();
-
-        vm.startPrank(debtor);
-        weth.approve(address(bullaFrendLend), 10 ether);
-        vm.stopPrank();
-
-        bullaClaim.approvalRegistry().permitCreateClaim({
-            user: debtor,
-            controller: address(bullaFrendLend),
-            approvalType: CreateClaimApprovalType.Approved,
-            approvalCount: 2,
-            isBindingAllowed: true,
-            signature: sigHelper.signCreateClaimPermit({
-                pk: debtorPK,
-                user: debtor,
-                controller: address(bullaFrendLend),
-                approvalType: CreateClaimApprovalType.Approved,
-                approvalCount: 2,
-                isBindingAllowed: true
-            })
-        });
-
-        // Create first WETH loan
-        LoanRequestParams memory wethOffer1 = new LoanRequestParamsBuilder().withCreditor(creditor).withDebtor(debtor)
-            .withDescription("WETH Loan 1").withToken(address(weth)).withLoanAmount(1 ether).withNumberOfPeriodsPerYear(365)
-            .build();
-
-        vm.prank(creditor);
-        uint256 wethLoanId1 = bullaFrendLend.offerLoan(wethOffer1);
-
-        vm.prank(debtor);
-        uint256 wethClaimId1 = bullaFrendLend.acceptLoan{value: FEE}(wethLoanId1);
-
-        // Create second WETH loan
-        LoanRequestParams memory wethOffer2 = new LoanRequestParamsBuilder().withCreditor(creditor).withDebtor(debtor)
-            .withDescription("WETH Loan 2").withToken(address(weth)).withLoanAmount(0.5 ether).withNumberOfPeriodsPerYear(
-            365
-        ).build();
-
-        vm.prank(creditor);
-        uint256 wethLoanId2 = bullaFrendLend.offerLoan(wethOffer2);
-
-        vm.prank(debtor);
-        uint256 wethClaimId2 = bullaFrendLend.acceptLoan{value: FEE}(wethLoanId2);
-
-        vm.warp(block.timestamp + 15 days);
-
-        // Make payments on both loans
-        vm.startPrank(debtor);
-
-        uint256 initialContractBalance = weth.balanceOf(address(bullaFrendLend));
-        (uint256 principal1, uint256 interest1) = bullaFrendLend.getTotalAmountDue(wethClaimId1);
-        bullaFrendLend.payLoan(wethClaimId1, principal1 + interest1);
-        uint256 fee1 = weth.balanceOf(address(bullaFrendLend)) - initialContractBalance;
-
-        uint256 contractBalanceAfterFirst = weth.balanceOf(address(bullaFrendLend));
-        (uint256 principal2, uint256 interest2) = bullaFrendLend.getTotalAmountDue(wethClaimId2);
-        bullaFrendLend.payLoan(wethClaimId2, principal2 + interest2);
-        uint256 fee2 = weth.balanceOf(address(bullaFrendLend)) - contractBalanceAfterFirst;
-        vm.stopPrank();
-
-        // Count WETH tokens in the array
-        uint256 wethTokenCount = 0;
-        for (uint256 i = 0; i < 3; i++) {
-            try bullaFrendLend.protocolFeeTokens(i) returns (address token) {
-                if (token == address(weth)) {
-                    wethTokenCount++;
-                }
-            } catch {
-                break;
-            }
-        }
-
-        assertEq(wethTokenCount, 1, "WETH should only appear once in protocol fee tokens array");
-        assertEq(
-            bullaFrendLend.protocolFeesByToken(address(weth)), fee1 + fee2, "WETH fees should accumulate correctly"
-        );
-    }
-
     function testProtocolFeeVariations() public {
         vm.startPrank(creditor);
         weth.approve(address(bullaFrendLend), 10 ether);
diff --git a/test/foundry/BullaFrendLend/BullaFrendLendTokenWhitelist.t.sol b/test/foundry/BullaFrendLend/BullaFrendLendTokenWhitelist.t.sol
@@ -139,9 +139,6 @@ contract TestBullaFrendLendTokenWhitelist is Test {
         vm.prank(debtor);
         bullaFrendLend.payLoan(claimId, PAYMENT_AMOUNT);
 
-        // Verify token was added to protocol fee tracking
-        assertEq(bullaFrendLend.protocolFeeTokens(0), address(token1));
-
         // Verify protocol fees accumulated for this token
         assertGt(bullaFrendLend.protocolFeesByToken(address(token1)), 0);
     }
diff --git a/test/foundry/BullaInvoice/BullaInvoiceProtocolFee.t.sol b/test/foundry/BullaInvoice/BullaInvoiceProtocolFee.t.sol
@@ -277,19 +277,14 @@ contract TestBullaInvoiceProtocolFee is Test {
 
         Invoice memory invoice = bullaInvoice.getInvoice(invoiceId);
         uint256 accruedInterest = invoice.interestComputationState.accruedInterest;
-
-        vm.expectRevert();
-        bullaInvoice.protocolFeeTokens(0);
-
         // Make payment with interest
         vm.prank(debtor);
         token1.approve(address(bullaInvoice), accruedInterest);
 
         vm.prank(debtor);
         bullaInvoice.payInvoice(invoiceId, accruedInterest);
 
-        // Verify token added to array
-        assertEq(bullaInvoice.protocolFeeTokens(0), address(token1), "Token should be added to protocolFeeTokens array");
+        // Verify protocol fee was tracked for whitelisted token
         assertGt(bullaInvoice.protocolFeesByToken(address(token1)), 0, "Protocol fee should be tracked");
     }
 
@@ -324,11 +319,6 @@ contract TestBullaInvoiceProtocolFee is Test {
             "Total protocol fee should be double since it is the same amount twice"
         );
         assertGt(bullaInvoice.protocolFeesByToken(address(token1)), 0, "Protocol fee is not 0");
-
-        // Verify only one entry in array
-        assertEq(bullaInvoice.protocolFeeTokens(0), address(token1), "Token should still be at index 0");
-        vm.expectRevert();
-        bullaInvoice.protocolFeeTokens(1); // Should revert - no second token
     }
 
     // ==================== 5. ETH PAYMENT TESTS ====================
@@ -724,10 +714,6 @@ contract TestBullaInvoiceProtocolFee is Test {
         assertEq(address(bullaInvoice).balance, expectedEthFee, "ETH fees in contract balance");
         assertEq(bullaInvoice.protocolFeesByToken(address(token1)), expectedToken1Fee, "Token1 fees tracked");
         assertEq(bullaInvoice.protocolFeesByToken(address(token2)), expectedToken2Fee, "Token2 fees tracked");
-
-        // Verify token array contains both tokens
-        assertEq(bullaInvoice.protocolFeeTokens(0), address(token1), "First token in array");
-        assertEq(bullaInvoice.protocolFeeTokens(1), address(token2), "Second token in array");
     }
 
     // ==================== 9. EDGE CASES & ERROR HANDLING ====================
diff --git a/test/foundry/BullaInvoice/BullaInvoiceTokenWhitelist.t.sol b/test/foundry/BullaInvoice/BullaInvoiceTokenWhitelist.t.sol
@@ -128,9 +128,6 @@ contract TestBullaInvoiceTokenWhitelist is Test {
         vm.prank(debtor);
         bullaInvoice.payInvoice(invoiceId, PAYMENT_AMOUNT);
 
-        // Verify token was added to protocol fee tracking
-        assertEq(bullaInvoice.protocolFeeTokens(0), address(token1));
-
         // Verify protocol fees accumulated for this token
         assertGt(bullaInvoice.protocolFeesByToken(address(token1)), 0);
     }

Original file line number	Diff line number	Diff line change
`@@ -139,9 +139,6 @@ contract TestBullaFrendLendTokenWhitelist is Test {`
`139`	`139`	`vm.prank(debtor);`
`140`	`140`	`bullaFrendLend.payLoan(claimId, PAYMENT_AMOUNT);`
`141`	`141`
`142`		`- // Verify token was added to protocol fee tracking`
`143`		`- assertEq(bullaFrendLend.protocolFeeTokens(0), address(token1));`
`144`		`-`
`145`	`142`	`// Verify protocol fees accumulated for this token`
`146`	`143`	`assertGt(bullaFrendLend.protocolFeesByToken(address(token1)), 0);`
`147`	`144`	`}`
Original file line number	Diff line number	Diff line change
`@@ -128,9 +128,6 @@ contract TestBullaInvoiceTokenWhitelist is Test {`
`128`	`128`	`vm.prank(debtor);`
`129`	`129`	`bullaInvoice.payInvoice(invoiceId, PAYMENT_AMOUNT);`
`130`	`130`
`131`		`- // Verify token was added to protocol fee tracking`
`132`		`- assertEq(bullaInvoice.protocolFeeTokens(0), address(token1));`
`133`		`-`
`134`	`131`	`// Verify protocol fees accumulated for this token`
`135`	`132`	`assertGt(bullaInvoice.protocolFeesByToken(address(token1)), 0);`
`136`	`133`	`}`