Whitelist polygon

contracts/aggregator/polygon/flashloan/helpers.sol

@@ -225,7 +225,7 @@ contract Helper is Variables {
      * @notice Returns fee for the passed route in BPS. 1 BPS == 0.01%.
      * @param _route route number for flashloan.
      */
-    function calculateFeeBPS(uint256 _route)
+    function calculateFeeBPS(uint256 _route, address account_)
         public
         view
         returns (uint256 BPS_)
@@ -244,7 +244,7 @@ contract Helper is Variables {
             revert("Invalid source");
         }
 
-        if (BPS_ < InstaFeeBPS) {
+        if (!isWhitelisted[account_] && BPS_ < InstaFeeBPS) {
             BPS_ = InstaFeeBPS;
         }
     }

contracts/aggregator/polygon/flashloan/main.sol

@@ -4,7 +4,50 @@ pragma solidity ^0.8.0;
 import "./helpers.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
 
-contract FlashAggregatorPolygon is Helper {
+contract AdminModule is Helper {
+    event updateOwnerLog(address indexed oldOwner, address indexed newOwner);
+
+    event updateWhitelistLog(
+        address indexed account,
+        bool indexed isWhitelisted_
+    );
+
+    /**
+     * @dev owner gaurd.
+     * @notice owner gaurd.
+     */
+    modifier onlyOwner() {
+        require(msg.sender == owner, "not-owner");
+        _;
+    }
+
+    /**
+     * @dev Update owner.
+     * @notice Update owner.
+     * @param newOwner_ address of new owner.
+     */
+    function updateOwner(address newOwner_) external onlyOwner {
+        address oldOwner_ = owner;
+        owner = newOwner_;
+        emit updateOwnerLog(oldOwner_, newOwner_);
+    }
+
+    /**
+     * @dev Update whitelist.
+     * @notice Update whitelist.
+     * @param account_ address to update the whitelist for.
+     * @param whitelist_ want to whitelist -> true, else false.
+     */
+    function updateWhitelist(address account_, bool whitelist_)
+        external
+        onlyOwner
+    {
+        isWhitelisted[account_] = whitelist_;
+        emit updateWhitelistLog(account_, whitelist_);
+    }
+}
+
+contract FlashAggregatorPolygon is AdminModule {
     using SafeERC20 for IERC20;
 
     event LogFlashloan(
@@ -44,7 +87,7 @@ contract FlashAggregatorPolygon is Helper {
         instaLoanVariables_._amounts = _amounts;
         instaLoanVariables_._instaFees = calculateFees(
             _amounts,
-            calculateFeeBPS(1)
+            calculateFeeBPS(1, sender_)
         );
         instaLoanVariables_._iniBals = calculateBalances(
             _assets,
@@ -112,7 +155,7 @@ contract FlashAggregatorPolygon is Helper {
         );
         instaLoanVariables_._instaFees = calculateFees(
             amounts_,
-            calculateFeeBPS(route_)
+            calculateFeeBPS(route_, sender_)
         );
 
         if (route_ == 5) {

contracts/aggregator/polygon/flashloan/variables.sol

@@ -50,4 +50,8 @@ contract Variables is ConstantVariables {
         address token1;
         uint24 fee;
     }
+
+    address public owner = address(0); // TODO: update this before deploying
+
+    mapping(address => bool) public isWhitelisted;
 }

contracts/resolver/polygon/interfaces.sol

@@ -4,7 +4,10 @@ pragma solidity ^0.8.0;
 interface InstaFlashloanAggregatorInterface {
     function getRoutes() external pure returns (uint16[] memory);
 
-    function calculateFeeBPS(uint256 _route) external view returns (uint256);
+    function calculateFeeBPS(uint256 _route, address account_)
+        external
+        view
+        returns (uint256 BPS_);
 
     function tokenToCToken(address) external view returns (address);
 

contracts/resolver/polygon/main.sol

@@ -55,7 +55,8 @@ contract FlashResolverPolygon is Helper {
                 }
             } else if (routesWithAvailability_[i] != 0) {
                 uint256 routeFeeBPS_ = flashloanAggregator.calculateFeeBPS(
-                    routesWithAvailability_[i]
+                    routesWithAvailability_[i],
+                    randomAddr_
                 );
                 if (feeBPS_ > routeFeeBPS_) {
                     feeBPS_ = routeFeeBPS_;

contracts/resolver/polygon/variables.sol

@@ -14,14 +14,17 @@ contract Variables {
     address public constant balancerLendingAddr =
         0xBA12222222228d8Ba445958a75a0704d566BF2C8;
 
-    address private flashloanAggregatorAddr =
+    address public constant flashloanAggregatorAddr =
         0xB2A7F20D10A006B0bEA86Ce42F2524Fde5D6a0F4;
-    InstaFlashloanAggregatorInterface internal flashloanAggregator =
+    InstaFlashloanAggregatorInterface internal constant flashloanAggregator =
         InstaFlashloanAggregatorInterface(flashloanAggregatorAddr);
 
-    address public constant wethAddr = 0x8f3Cf7ad23Cd3CaDbD9735AFf958023239c6A063;
-    address public constant usdcAddr = 0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174;
-    address public constant wmaticAddr = 0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0;
+    address public constant wethAddr =
+        0x8f3Cf7ad23Cd3CaDbD9735AFf958023239c6A063;
+    address public constant usdcAddr =
+        0x2791Bca1f2de4661ED88A30C99A7a9449Aa84174;
+    address public constant wmaticAddr =
+        0x7D1AfA7B718fb893dB30A3aBc0Cfc608AaCfeBB0;
     address public constant uniswapFactoryAddr =
         0x1F98431c8aD98523631AE4a59f267346ea31F984;
     bytes32 internal constant POOL_INIT_CODE_HASH =
@@ -32,4 +35,6 @@ contract Variables {
         address token1;
         uint24 fee;
     }
+
+    address internal constant randomAddr_ = address(0);
 }