From 457385e6929c44be30dab6e3650d3177c9c2512a Mon Sep 17 00:00:00 2001
From: tim <timolsoncrypto@gmail.com>
Date: Tue, 7 Oct 2025 15:46:51 -0400
Subject: [PATCH] flashLoan protocol fee

---
 src/PartyPool.sol | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/PartyPool.sol b/src/PartyPool.sol
index 39ec08a..8ccbc02 100644
--- a/src/PartyPool.sol
+++ b/src/PartyPool.sol
@@ -412,9 +412,25 @@ contract PartyPool is PartyPoolBase, ERC20External, IPartyPool {
         IERC20 token = IERC20(tokenAddr);
         require(amount <= token.balanceOf(address(this)));
         (uint256 fee, ) = _computeFee(amount, FLASH_FEE_PPM);
+
+        // Compute protocol share of flash fee
+        if (PROTOCOL_FEE_PPM > 0 && fee > 0) {
+            uint256 protoShare = (fee * PROTOCOL_FEE_PPM) / 1_000_000; // floor
+            if (protoShare > 0) {
+                uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
+                protocolFeesOwed[tokenIndex] += protoShare;
+            }
+        }
+
         require(token.transfer(address(receiver), amount));
         require(receiver.onFlashLoan(msg.sender, address(token), amount, fee, data) == FLASH_CALLBACK_SUCCESS);
         require(token.transferFrom(address(receiver), address(this), amount + fee));
+
+        // Update cached balance for the borrowed token
+        uint256 tokenIndex = tokenAddressToIndexPlusOne[token] - 1;
+        uint256 balAfter = token.balanceOf(address(this));
+        _recordCachedBalance(tokenIndex, balAfter);
+
         return true;
     }