add withdrawer role

revert-finance · Jul 11, 2023 · 7faa9b9 · 7faa9b9
1 parent 2d08c9c
commit 7faa9b9
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 11 deletions.
diff --git a/src/automators/AutoExit.sol b/src/automators/AutoExit.sol
@@ -32,8 +32,8 @@ contract AutoExit is Automator {
         uint64 token1SlippageX64
     );
 
-    constructor(INonfungiblePositionManager _npm, address _operator, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) 
-        Automator(_npm, _operator, _TWAPSeconds, _maxTWAPTickDifference,  _protocolRewardX64, _swapRouterOptions) {
+    constructor(INonfungiblePositionManager _npm, address _operator, address _withdrawer, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) 
+        Automator(_npm, _operator, _withdrawer, _TWAPSeconds, _maxTWAPTickDifference,  _protocolRewardX64, _swapRouterOptions) {
     }
 
     // define how stoploss / limit should be handled

diff --git a/src/automators/AutoRange.sol b/src/automators/AutoRange.sol
@@ -27,8 +27,8 @@ contract AutoRange is Automator {
         uint64 token1SlippageX64
     );
 
-    constructor(INonfungiblePositionManager _npm, address _operator, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) 
-        Automator(_npm, _operator, _TWAPSeconds, _maxTWAPTickDifference, _protocolRewardX64, _swapRouterOptions) {
+    constructor(INonfungiblePositionManager _npm, address _operator, address _withdrawer, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) 
+        Automator(_npm, _operator, _withdrawer, _TWAPSeconds, _maxTWAPTickDifference, _protocolRewardX64, _swapRouterOptions) {
     }
 
     // defines when and how a position can be changed by operator

diff --git a/src/automators/Automator.sol b/src/automators/Automator.sol
@@ -44,16 +44,18 @@ abstract contract Automator is Ownable {
 
     // admin events
     event OperatorChanged(address newOperator, bool active);
+    event WithdrawerChanged(address newWithdrawer);
     event TWAPConfigChanged(uint32 TWAPSeconds, uint16 maxTWAPTickDifference);
     event SwapRouterChanged(uint8 swapRouterIndex);
 
     // configurable by owner
     mapping(address => bool) public operators;
+    address public withdrawer;
     uint32 public TWAPSeconds;
     uint16 public maxTWAPTickDifference;
     uint8 public swapRouterIndex; // default is 0
 
-    constructor(INonfungiblePositionManager npm, address _operator, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) {
+    constructor(INonfungiblePositionManager npm, address _operator, address _withdrawer, uint32 _TWAPSeconds, uint16 _maxTWAPTickDifference, uint64 _protocolRewardX64, address[] memory _swapRouterOptions) {
 
         nonfungiblePositionManager = npm;
         weth = IWETH9(npm.WETH9());
@@ -67,6 +69,7 @@ abstract contract Automator is Ownable {
         emit SwapRouterChanged(0);
 
         setOperator(_operator, true);
+        setWithdrawer(_withdrawer);
 
         setTWAPConfig(_maxTWAPTickDifference, _TWAPSeconds);
 
@@ -88,6 +91,15 @@ abstract contract Automator is Ownable {
         swapRouterIndex = _swapRouterIndex;
     }
 
+    /**
+     * @notice Owner controlled function to set withdrawer address
+     * @param _withdrawer withdrawer
+     */
+    function setWithdrawer(address _withdrawer) public onlyOwner {
+        emit WithdrawerChanged(_withdrawer);
+        withdrawer = _withdrawer;
+    }
+
     /**
      * @notice Owner controlled function to activate/deactivate operator address
      * @param _operator operator
@@ -119,7 +131,12 @@ abstract contract Automator is Ownable {
      * @param tokens Addresses of tokens to withdraw
      * @param to Address to send to
      */
-    function withdrawBalances(address[] calldata tokens, address to) external onlyOwner {
+    function withdrawBalances(address[] calldata tokens, address to) external {
+
+        if (msg.sender != withdrawer) {
+            revert Unauthorized();
+        }
+
         uint i;
         uint count = tokens.length;
         for(;i < count;++i) {
@@ -134,7 +151,12 @@ abstract contract Automator is Ownable {
      * @notice Withdraws ETH balance
      * @param to Address to send to
      */
-    function withdrawETH(address to) external onlyOwner {
+    function withdrawETH(address to) external {
+
+        if (msg.sender != withdrawer) {
+            revert Unauthorized();
+        }
+
         uint256 balance = address(this).balance;
         if (balance > 0) {
             (bool sent,) = to.call{value: balance}("");

diff --git a/test/IntegrationTestBase.sol b/test/IntegrationTestBase.sol
@@ -19,6 +19,7 @@ abstract contract IntegrationTestBase is Test {
 
     address constant WHALE_ACCOUNT = 0xF977814e90dA44bFA03b6295A0616a897441aceC;
     address constant OPERATOR_ACCOUNT = 0xF977814e90dA44bFA03b6295A0616a897441aceC;
+    address constant WITHDRAWER_ACCOUNT = 0xF977814e90dA44bFA03b6295A0616a897441aceC;
 
 
     address FACTORY = 0x1F98431c8aD98523631AE4a59f267346ea31F984;

diff --git a/test/integration/automators/AutoExit.t.sol b/test/integration/automators/AutoExit.t.sol
@@ -11,7 +11,7 @@ contract AutoExitTest is IntegrationTestBase {
 
     function setUp() external {
         _setupBase();
-        autoExit = new AutoExit(NPM, OPERATOR_ACCOUNT, 60, 100, uint64(Q64 / 400), _getSwapRouterOptions());
+        autoExit = new AutoExit(NPM, OPERATOR_ACCOUNT, WITHDRAWER_ACCOUNT, 60, 100, uint64(Q64 / 400), _getSwapRouterOptions());
     }
 
     function _setConfig(
@@ -104,6 +104,7 @@ contract AutoExitTest is IntegrationTestBase {
         address[] memory addresses = new address[](2);
         addresses[0] = address(DAI);
         addresses[1] = address(USDC);
+        vm.prank(WITHDRAWER_ACCOUNT); 
         autoExit.withdrawBalances(addresses, address(this));
         uint balanceAfter = DAI.balanceOf(address(this));
 
@@ -136,6 +137,7 @@ contract AutoExitTest is IntegrationTestBase {
         // protocol fee
         balanceBefore = USDC.balanceOf(address(this));
 
+        vm.prank(WITHDRAWER_ACCOUNT); 
         autoExit.withdrawBalances(addresses, address(this));
 
         balanceAfter = USDC.balanceOf(address(this));
@@ -257,7 +259,7 @@ contract AutoExitTest is IntegrationTestBase {
     function testOracleCheck() external {
 
         // create range adjustor with more strict oracle config    
-        autoExit = new AutoExit(NPM, OPERATOR_ACCOUNT, 60 * 30, 4, uint64(Q64 / 400), _getSwapRouterOptions());
+        autoExit = new AutoExit(NPM, OPERATOR_ACCOUNT, WITHDRAWER_ACCOUNT, 60 * 30, 4, uint64(Q64 / 400), _getSwapRouterOptions());
 
         vm.prank(TEST_NFT_2_ACCOUNT);
         NPM.setApprovalForAll(address(autoExit), true);

diff --git a/test/integration/automators/AutoRange.t.sol b/test/integration/automators/AutoRange.t.sol
@@ -13,7 +13,7 @@ contract AutoRangeTest is IntegrationTestBase {
 
     function setUp() external {
         _setupBase();
-        autoRange = new AutoRange(NPM, OPERATOR_ACCOUNT, 60, 100, uint64(Q64 / 400), _getSwapRouterOptions());
+        autoRange = new AutoRange(NPM, OPERATOR_ACCOUNT, WITHDRAWER_ACCOUNT, 60, 100, uint64(Q64 / 400), _getSwapRouterOptions());
     }
 
     function testSetTWAPSeconds() external {
@@ -330,7 +330,7 @@ contract AutoRangeTest is IntegrationTestBase {
     function testOracleCheck() external {
 
         // create range adjustor with more strict oracle config    
-        autoRange = new AutoRange(NPM, OPERATOR_ACCOUNT, 60 * 30, 4, uint64(Q64 / 400), _getSwapRouterOptions());
+        autoRange = new AutoRange(NPM, OPERATOR_ACCOUNT, WITHDRAWER_ACCOUNT, 60 * 30, 4, uint64(Q64 / 400), _getSwapRouterOptions());
 
         vm.prank(TEST_NFT_2_ACCOUNT);
         NPM.setApprovalForAll(address(autoRange), true);