From f9ca7f8959fed2234412d2aa4a47906b8d644be0 Mon Sep 17 00:00:00 2001
From: lsqrl <lsqrl91@gmail.com>
Date: Mon, 18 Dec 2023 14:45:13 +0100
Subject: [PATCH] added weth reward internal function used twice

---
 src/services/debit/GmxService.sol | 58 ++++++++++++++++++++++---------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/src/services/debit/GmxService.sol b/src/services/debit/GmxService.sol
index f2fab36..90409cd 100644
--- a/src/services/debit/GmxService.sol
+++ b/src/services/debit/GmxService.sol
@@ -88,9 +88,7 @@ contract GmxService is Whitelisted, AuctionRateModel, DebitService {
     function _close(uint256 tokenID, Agreement memory agreement, bytes memory data) internal override {
         uint256 minAmountOut = abi.decode(data, (uint256));
         uint256 userVirtualDeposit = virtualDeposit[tokenID];
-        // delete virtual deposits
         delete virtualDeposit[tokenID];
-        totalVirtualDeposits -= userVirtualDeposit;
 
         routerV2.unstakeAndRedeemGlp(
             agreement.loans[0].token,
@@ -104,14 +102,18 @@ contract GmxService is Whitelisted, AuctionRateModel, DebitService {
         // register rewards
         uint256 finalBalance = weth.balanceOf(address(this));
         uint256 newRewards = totalRewards + (finalBalance - initialBalance);
-        // calculate share of rewards to give to the user
-        uint256 totalWithdraw = ((newRewards + totalVirtualDeposits) * agreement.collaterals[0].amount) /
-            totalCollateral;
-        // Subtracting the virtual deposit we get the weth part: this is the weth the user is entitled to
-        // Due to integer arithmetic, we may get underflow if we do not make checks
-        uint256 toTransfer = totalWithdraw >= userVirtualDeposit
-            ? totalWithdraw - userVirtualDeposit <= finalBalance ? totalWithdraw - userVirtualDeposit : finalBalance
-            : 0;
+
+        uint256 toTransfer = _wethReward(
+            agreement.collaterals[0].amount,
+            userVirtualDeposit,
+            newRewards,
+            totalVirtualDeposits,
+            totalCollateral,
+            finalBalance
+        );
+
+        // delete virtual deposits here since _wethReward already subtracts user deposit from total
+        totalVirtualDeposits -= userVirtualDeposit;
         // update totalRewards and totalCollateral with an extra check to avoid integer arithmetic underflows
         totalRewards = toTransfer < newRewards ? newRewards - toTransfer : 0;
         totalCollateral -= agreement.collaterals[0].amount;
@@ -120,12 +122,17 @@ contract GmxService is Whitelisted, AuctionRateModel, DebitService {
     }
 
     function wethReward(uint256 tokenID) public view returns (uint256) {
-        uint256 collateral = agreements[tokenID].collaterals[0].amount;
-        uint256 newRewards = totalRewards + rewardTracker.claimableReward(address(this));
-        // calculate share of rewards to give to the user
-        uint256 totalWithdraw = ((newRewards + totalVirtualDeposits) * collateral) / totalCollateral;
-        // Subtracting the virtual deposit we get the weth part: this is the weth the user is entitled to
-        return totalWithdraw - virtualDeposit[tokenID];
+        uint256 claimableReward = rewardTracker.claimableReward(address(this));
+        uint256 finalBalance = weth.balanceOf(address(this)) + claimableReward;
+        return
+            _wethReward(
+                agreements[tokenID].collaterals[0].amount,
+                virtualDeposit[tokenID],
+                totalRewards + claimableReward,
+                totalVirtualDeposits,
+                totalCollateral,
+                finalBalance
+            );
     }
 
     function quote(Agreement memory agreement) public view override returns (uint256[] memory) {
@@ -149,4 +156,23 @@ contract GmxService is Whitelisted, AuctionRateModel, DebitService {
 
         return results;
     }
+
+    function _wethReward(
+        uint256 collateral,
+        uint256 userVirtualDeposit,
+        uint256 totalRewards,
+        uint256 totalVirtualDeposits,
+        uint256 totalCollateral,
+        uint256 finalBalance
+    ) internal pure returns (uint256) {
+        // calculate share of rewards to give to the user
+        uint256 totalWithdraw = ((totalRewards + totalVirtualDeposits - userVirtualDeposit) * collateral) /
+            totalCollateral;
+        return
+            // Subtracting the virtual deposit we get the weth part: this is the weth the user is entitled to
+            // Due to integer arithmetic, we may get underflow if we do not make checks
+            totalWithdraw >= userVirtualDeposit
+                ? totalWithdraw - userVirtualDeposit <= finalBalance ? totalWithdraw - userVirtualDeposit : finalBalance
+                : 0;
+    }
 }