Improve fee withdraw logic

xtokenmarket · Jan 11, 2023 · be4965d · be4965d
1 parent eb1a833
commit be4965d
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 7 deletions.
diff --git a/contracts/FungibleOriginationPool.sol b/contracts/FungibleOriginationPool.sol
@@ -751,18 +751,31 @@ contract FungibleOriginationPool is
             require(success);
             // send fees to core
             originationCore.receiveFees{value: originationCoreFees}();
+            // reset accrued origination fees amount
+            originationCoreFees = 0;
         } else {
             transferAmount =
                 purchaseToken.balanceOf(address(this)) -
                 originationCoreFees;
             purchaseToken.safeTransfer(owner(), transferAmount);
-            purchaseToken.safeTransfer(
-                address(originationCore),
-                originationCoreFees
-            );
+            // handle timelocked tokens
+            try
+                purchaseToken.transfer(
+                    address(originationCore),
+                    originationCoreFees
+                )
+            {
+                // reset accrued origination fees amount
+                originationCoreFees = 0;
+            } catch {
+                // if transfer to core fails don't reset the fee amount
+                // approve tokens to origination core to be transferred
+                purchaseToken.safeIncreaseAllowance(
+                    address(originationCore),
+                    originationCoreFees
+                );
+            }
         }
-        // reset accrued origination fees amount
-        originationCoreFees = 0;
 
         emit PurchaseTokenClaim(owner(), transferAmount);
     }

diff --git a/contracts/OriginationCore.sol b/contracts/OriginationCore.sol
@@ -215,6 +215,25 @@ contract OriginationCore is
         }
     }
 
+    /**
+     * @dev Withdraw unclaimed fees from origination pool
+     * @dev Callable only if there are any unclaimed fees
+     * @dev Callable only after owner claims sale purchase tokens
+     *
+     * @param _pool the pool address
+     * @param _feeToken The token address of the fee to claim
+     */
+    function withdrawFees(address _pool, address _feeToken) external {
+        require(
+            xTokenManager.isRevenueController(msg.sender),
+            "Only callable by revenue controller."
+        );
+
+        uint256 feeAmount = IFungibleOriginationPool(_pool).originationCoreFees();
+
+        IERC20(_feeToken).transferFrom(_pool, msg.sender, feeAmount);
+    }
+
     /**
      * @dev Function used by origination pools to send the origination fees to this contract
      * @dev Only used after a token sale has finished successfully on claiming the tokens

diff --git a/package.json b/package.json
@@ -11,7 +11,7 @@
     "chai": "^4.3.4",
     "ethereum-waffle": "^3.4.0",
     "ethers": "^5.5.2",
-    "hardhat": "^2.9.2",
+    "hardhat": "^2.12.6",
     "hardhat-deploy": "^0.9.23",
     "hardhat-deploy-ethers": "^0.3.0-beta.12",
     "prettier": "^2.5.1",