Merge pull request #98 from lista-dao/feature/buyback-v3

fea: add withdraw methods for buyback

Author: ricklista

contracts/buyback/Buyback.sol

@@ -258,6 +258,25 @@ contract Buyback is
     emit EmergencyWithdraw(_token, _amount);
   }
 
+  /**
+   * @dev allows bot to withdraw tokens to the receiver address
+   * @param _token token address
+   * @param _amount token amount
+   */
+  function withdraw(address _token, uint256 _amount) external onlyRole(BOT) {
+    require(
+      tokenInWhitelist[_token] || _token == SWAP_NATIVE_TOKEN_ADDRESS || _token == tokenOut,
+      "Token not whitelisted"
+    );
+    require(_amount > 0, "Invalid amount");
+    if (_token == SWAP_NATIVE_TOKEN_ADDRESS) {
+      (bool success, ) = payable(receiver).call{ value: _amount }("");
+      require(success, "Withdraw failed");
+    } else {
+      IERC20(_token).safeTransfer(receiver, _amount);
+    }
+  }
+
   /**
    * @dev pause the contract
    */

contracts/dao/ListaAutoBuyback.sol

@@ -183,6 +183,22 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
         emit AdminTransfer(_token, _amount);
     }
 
+    /**
+     * @dev allow bot to withdraw token to the defaultReceiver address
+     * @param _token token address
+     * @param _amount token amount
+     */
+    function withdraw(address _token, uint256 _amount) external onlyRole(BOT) {
+        require(tokenWhitelist[_token] || _token == SWAP_NATIVE_TOKEN_ADDRESS, "Token not whitelisted");
+        require(_amount > 0, "Invalid amount");
+        if (_token == SWAP_NATIVE_TOKEN_ADDRESS) {
+            (bool success, ) = payable(defaultReceiver).call{ value: _amount }("");
+            require(success, "Withdraw failed");
+        } else {
+            IERC20(_token).safeTransfer(defaultReceiver, _amount);
+        }
+    }
+
     function changeDefaultReceiver(address _receiver) external onlyRole(DEFAULT_ADMIN_ROLE) {
         require(_receiver != address(0), "_receiver is the zero address");
         require(_receiver != defaultReceiver, "_receiver is the same");

test/buyback/Buyback.t.sol

@@ -358,4 +358,20 @@ contract BuybackTest is Test {
     buyback.buyback(pancakeRouter, buyback.SWAP_NATIVE_TOKEN_ADDRESS(), tokenOut, 1 ether, 0, data);
     vm.stopPrank();
   }
+
+  function test_withdraw() public {
+    deal(tokenIn, address(buyback), 1 ether);
+    vm.deal(address(buyback), 1 ether);
+
+    uint256 lisUSDBefore = IERC20(tokenIn).balanceOf(buyback.receiver());
+    uint256 bnbBefore = buyback.receiver().balance;
+
+    vm.startPrank(bot);
+    buyback.withdraw(tokenIn, 1 ether);
+    assertEq(IERC20(tokenIn).balanceOf(buyback.receiver()) - lisUSDBefore, 1 ether, "lisUSD withdraw failed");
+
+    buyback.withdraw(buyback.SWAP_NATIVE_TOKEN_ADDRESS(), 1 ether);
+    assertEq(buyback.receiver().balance - bnbBefore, 1 ether, "BNB withdraw failed");
+    vm.stopPrank();
+  }
 }

test/dao/ListaAutoBuyback.t.sol

@@ -225,4 +225,23 @@ contract ListaAutoBuyBackTest is Test {
         listaAutoBuyBack.buyback(pancakeRouter, listaAutoBuyBack.SWAP_NATIVE_TOKEN_ADDRESS(), USDT, 1 ether, 0, data);
         vm.stopPrank();
     }
+
+    function test_withdraw() public {
+        deal(lisUSD, address(listaAutoBuyBack), 1 ether);
+        vm.deal(address(listaAutoBuyBack), 1 ether);
+        vm.startPrank(admin);
+        listaAutoBuyBack.setTokenWhitelist(lisUSD, true);
+        vm.stopPrank();
+
+        uint256 lisUSDBefore = IERC20(lisUSD).balanceOf(listaAutoBuyBack.defaultReceiver());
+        uint256 bnbBefore = listaAutoBuyBack.defaultReceiver().balance;
+
+        vm.startPrank(bot);
+        listaAutoBuyBack.withdraw(lisUSD, 1 ether);
+        assertEq(IERC20(lisUSD).balanceOf(listaAutoBuyBack.defaultReceiver()) - lisUSDBefore, 1 ether, "lisUSD withdraw failed");
+
+        listaAutoBuyBack.withdraw(listaAutoBuyBack.SWAP_NATIVE_TOKEN_ADDRESS(), 1 ether);
+        assertEq(listaAutoBuyBack.defaultReceiver().balance - bnbBefore, 1 ether, "BNB withdraw failed");
+        vm.stopPrank();
+    }
 }