Improve fee withdraw logic

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);
     }

contracts/OriginationCore.sol

@@ -5,6 +5,7 @@ import "@openzeppelin/contracts/proxy/transparent/TransparentUpgradeableProxy.so
 import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
 import "./interface/IPoolDeployer.sol";
 import "./interface/INFTDeployer.sol";
@@ -21,6 +22,7 @@ contract OriginationCore is
     Initializable,
     OwnableUpgradeable
 {
+    using SafeERC20 for IERC20;
     //--------------------------------------------------------------------------
     // State variables
     //--------------------------------------------------------------------------
@@ -61,6 +63,8 @@ contract OriginationCore is
     event SetListingFee(uint256 fee);
     event CustomListingFeeEnabled(address indexed deployer, uint256 customFee);
     event CustomListingFeeDisabled(address indexed deployer);
+    event TokenFeeWithdraw(address indexed token, uint256 amount);
+    event EthFeeWithdraw(uint256 amount);
 
     //--------------------------------------------------------------------------
     // Constructor / Initializer
@@ -109,10 +113,11 @@ contract OriginationCore is
      * @dev Must pay the listing fee
      *
      * @param saleParams The token sale params
+     * @return originationPool address of the deployed pool
      */
     function createFungibleListing(
         IFungibleOriginationPool.SaleParams calldata saleParams
-    ) external payable {
+    ) external payable returns (address originationPool) {
         uint256 feeOwed = customListingFeeEnabled[msg.sender]
             ? customListingFee[msg.sender]
             : listingFee;
@@ -127,7 +132,7 @@ contract OriginationCore is
         );
 
         // Deploy the pool
-        address originationPool = poolDeployer.deployFungibleOriginationPool(
+        originationPool = poolDeployer.deployFungibleOriginationPool(
             address(proxyAdmin)
         );
         // Deploy the vesting entry nft if there is a vesting period
@@ -206,15 +211,37 @@ contract OriginationCore is
                 ""
             );
             require(success);
+            emit EthFeeWithdraw(address(this).balance);
         } else {
-            bool success = IERC20(_feeToken).transfer(
+            uint256 fees = IERC20(_feeToken).balanceOf(address(this));
+            IERC20(_feeToken).safeTransfer(
                 msg.sender,
-                IERC20(_feeToken).balanceOf(address(this))
+                fees
             );
-            require(success);
+            emit TokenFeeWithdraw(_feeToken, fees);
         }
     }
 
+    /**
+     * @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);
+        emit TokenFeeWithdraw(_feeToken, 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

contracts/interface/IFungibleOriginationPool.sol

@@ -32,4 +32,6 @@ interface IFungibleOriginationPool {
         address vestingEntryNFT,
         SaleParams calldata saleParams
     ) external;
+
+    function originationCoreFees() external view returns (uint256);
 }

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",