fea: buyback support pancake and other routers

contracts/buyback/Buyback.sol

@@ -35,8 +35,8 @@ contract Buyback is
   address public constant SWAP_NATIVE_TOKEN_ADDRESS = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
 
   /* ============ State Variables ============ */
-  // 1Inch router whitelist
-  mapping(address => bool) public oneInchRouterWhitelist;
+  // swap router whitelist
+  mapping(address => bool) public routerWhitelist;
   // swap input token whitelist
   mapping(address => bool) public tokenInWhitelist;
   // swap output token
@@ -48,8 +48,15 @@ contract Buyback is
 
   /* ============ Events ============ */
   event BoughtBack(address indexed tokenIn, address indexed tokenOut, uint256 amountIn, uint256 amountOut);
+  event BoughtBack(
+    address indexed pair,
+    address indexed tokenIn,
+    address indexed tokenOut,
+    uint256 amountIn,
+    uint256 amountOut
+  );
   event ReceiverChanged(address indexed receiver);
-  event OneInchRouterChanged(address indexed oneInchRouter, bool added);
+  event RouterChanged(address indexed router, bool added);
   event TokenInChanged(address indexed token, bool added);
   event EmergencyWithdraw(address token, uint256 amount);
 
@@ -99,7 +106,7 @@ contract Buyback is
     _grantRole(PAUSER, _pauser);
     _grantRole(BOT, _bot);
 
-    oneInchRouterWhitelist[_1InchRouter] = true;
+    routerWhitelist[_1InchRouter] = true;
     tokenOut = _tokenOut;
     receiver = _receiver;
   }
@@ -115,7 +122,7 @@ contract Buyback is
     address _1inchRouter,
     bytes calldata _data
   ) external override onlyRole(BOT) nonReentrant whenNotPaused {
-    require(oneInchRouterWhitelist[_1inchRouter], "Invalid 1Inch router");
+    require(routerWhitelist[_1inchRouter], "router not whitelisted");
     require(bytes4(_data[0:4]) == SWAP_FUNCTION_SELECTOR, "Invalid 1Inch function selector");
 
     (, SwapDescription memory swapDesc, ) = abi.decode(_data[4:], (address, SwapDescription, bytes));
@@ -147,6 +154,50 @@ contract Buyback is
     emit BoughtBack(address(swapDesc.srcToken), address(swapDesc.dstToken), swapDesc.amount, amountOut);
   }
 
+  /// @dev buy back tokens using router
+  /// @param _router The address of the router.
+  /// @param _tokenIn The address of the input token.
+  /// @param _tokenOut The address of the output token.
+  /// @param _amountIn The amount to sell.
+  /// @param _amountOutMin The minimum amount to receive.
+  /// @param _swapData The swap data.
+  function buyback(
+    address _router,
+    address _tokenIn,
+    address _tokenOut,
+    uint256 _amountIn,
+    uint256 _amountOutMin,
+    bytes calldata _swapData
+  ) external onlyRole(BOT) nonReentrant whenNotPaused {
+    require(tokenInWhitelist[_tokenIn], "token not whitelisted");
+    require(tokenOut == _tokenOut, "token not whitelisted");
+    require(routerWhitelist[_router], "router not whitelisted");
+
+    uint256 beforeTokenIn = _getTokenBalance(_tokenIn, address(this));
+    uint256 beforeTokenOut = _getTokenBalance(_tokenOut, address(this));
+
+    bool isNativeTokenIn = (_tokenIn == SWAP_NATIVE_TOKEN_ADDRESS);
+    if (!isNativeTokenIn) {
+      IERC20(_tokenIn).safeApprove(_router, _amountIn);
+    }
+    (bool success, ) = _router.call{ value: isNativeTokenIn ? _amountIn : 0 }(_swapData);
+    require(success, "swap failed");
+
+    if (!isNativeTokenIn) {
+      IERC20(_tokenIn).safeApprove(_router, 0);
+    }
+
+    uint256 actualAmountIn = beforeTokenIn - _getTokenBalance(_tokenIn, address(this));
+    uint256 actualAmountOut = _getTokenBalance(_tokenOut, address(this)) - beforeTokenOut;
+
+    require(actualAmountIn <= _amountIn, "exceed amount in");
+    require(actualAmountOut >= _amountOutMin, "not enough profit");
+
+    IERC20(_tokenOut).safeTransfer(receiver, actualAmountOut);
+
+    emit BoughtBack(_router, _tokenIn, _tokenOut, actualAmountIn, actualAmountOut);
+  }
+
   /**
    * @dev change receiver
    * @param _receiver - Address of the receiver
@@ -159,27 +210,14 @@ contract Buyback is
     emit ReceiverChanged(_receiver);
   }
 
-  /**
-   * @dev add 1Inch router to whitelist
-   * @param _1InchRouter - Address of the 1Inch router
-   */
-  function add1InchRouterWhitelist(address _1InchRouter) external onlyRole(MANAGER) {
-    require(_1InchRouter != address(0), "Invalid 1Inch router");
-    require(!oneInchRouterWhitelist[_1InchRouter], "Already whitelisted");
-
-    oneInchRouterWhitelist[_1InchRouter] = true;
-    emit OneInchRouterChanged(_1InchRouter, true);
-  }
-
-  /**
-   * @dev remove 1Inch router from whitelist
-   * @param _1InchRouter - Address of the 1Inch router
-   */
-  function remove1InchRouterWhitelist(address _1InchRouter) external onlyRole(MANAGER) {
-    require(oneInchRouterWhitelist[_1InchRouter], "1Inch router is not in whitelist");
-
-    delete oneInchRouterWhitelist[_1InchRouter];
-    emit OneInchRouterChanged(_1InchRouter, false);
+  /// @dev sets the router whitelist.
+  /// @param _router The address of the router.
+  /// @param status The status of the router.
+  function setRouterWhitelist(address _router, bool status) external onlyRole(MANAGER) {
+    require(_router != address(0), "Invalid router address");
+    require(routerWhitelist[_router] != status, "whitelist same status");
+    routerWhitelist[_router] = status;
+    emit RouterChanged(_router, status);
   }
 
   /**
@@ -239,4 +277,12 @@ contract Buyback is
   // /* ============ Internal Functions ============ */
 
   function _authorizeUpgrade(address newImplementation) internal override onlyRole(DEFAULT_ADMIN_ROLE) {}
+
+  function _getTokenBalance(address _token, address account) internal view returns (uint256) {
+    if (_token == SWAP_NATIVE_TOKEN_ADDRESS) {
+      return account.balance;
+    } else {
+      return IERC20(_token).balanceOf(account);
+    }
+  }
 }

contracts/buyback/interfaces/IBuyback.sol

@@ -16,4 +16,13 @@ interface IBuyback {
   }
 
   function buyback(address _1inchRouter, bytes calldata _data) external;
+
+  function buyback(
+    address router,
+    address tokenIn,
+    address tokenOut,
+    uint256 amountIn,
+    uint256 amountOutMin,
+    bytes calldata swapData
+  ) external;
 }

contracts/dao/ListaAutoBuyback.sol

@@ -15,18 +15,37 @@ import "../buyback/library/RevertReasonParser.sol";
   * @dev result of swap will be sent to receiver address to distribute to users
   */
 contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
-
     using SafeERC20 for IERC20;
 
+    struct SwapDescription {
+        address srcToken;
+        address dstToken;
+        address payable srcReceiver;
+        address payable dstReceiver;
+        uint256 amount;
+        uint256 minReturnAmount;
+        uint256 flags;
+    }
+
     event BoughtBack(address indexed tokenIn, uint256 amountIn, uint256 amountOut);
+    event BoughtBack(
+        address indexed pair,
+        address indexed tokenIn,
+        address indexed tokenOut,
+        uint256 amountIn,
+        uint256 amountOut
+    );
 
     event ReceiverChanged(address indexed receiver);
-
     event RouterChanged(address indexed router, bool added);
+    event TokenWhitelistChanged(address indexed token, bool added);
+    event AdminTransfer(address token, uint256 amount);
+
 
     bytes32 public constant BOT = keccak256("BOT");
 
     bytes4 public constant SWAP_FUNCTION_SELECTOR = bytes4(keccak256("swap(address,(address,address,address,address,uint256,uint256,uint256),bytes)"));
+    address public constant SWAP_NATIVE_TOKEN_ADDRESS = 0xEeeeeEeeeEeEeeEeEeEeeEEEeeeeEeeeeeeeEEeE;
 
     // The offset of the dstReceiver in the call data
     // 4 bytes for the function selector + 32 bytes for executor + 32 bytes for srcToken +
@@ -38,15 +57,19 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
 
     address public defaultReceiver;
 
-    mapping(address => bool) public oneInchRouterWhitelist;
+    mapping(address => bool) public routerWhitelist;
 
     mapping(uint256 => uint256) public dailyBought;
 
+    mapping(address => bool) public tokenWhitelist;
+
     /// @custom:oz-upgrades-unsafe-allow constructor
     constructor() {
         _disableInitializers();
     }
 
+    receive() external payable {}
+
     function initialize(
         address _admin,
         address _bot,
@@ -62,7 +85,7 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
         _setupRole(BOT, _bot);
 
         defaultReceiver = _initReceiver;
-        oneInchRouterWhitelist[_initRouter] = true;
+        routerWhitelist[_initRouter] = true;
     }
 
     /**
@@ -77,8 +100,16 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
         onlyRole(BOT)
     {
         require(_amountIn > 0, "amountIn is zero");
-        require(oneInchRouterWhitelist[_1inchRouter], "router not whitelisted");
+        require(routerWhitelist[_1inchRouter], "router not whitelisted");
         require(_getFunctionSelector(_data) == SWAP_FUNCTION_SELECTOR, "invalid function selector of _data");
+        require(tokenWhitelist[_tokenIn], "token in not whitelisted");
+        (, SwapDescription memory swapDesc, ) = abi.decode(_data[4:], (address, SwapDescription, bytes));
+
+        require(_tokenIn == swapDesc.srcToken, "Invalid swap input token");
+        require(tokenWhitelist[swapDesc.dstToken], "token out not whitelisted");
+        require(address(swapDesc.dstReceiver) == defaultReceiver, "Invalid receiver");
+        require(swapDesc.amount > 0, "Invalid swap input amount");
+
         require(_extractDstReceiver(_data) == defaultReceiver, "invalid dst receiver of _data");
         require(IERC20(_tokenIn).balanceOf(address(this)) >= _amountIn, "insufficient balance");
         // Approves the 1inch router contract to spend the specified amount of _tokenIn
@@ -92,14 +123,64 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
         }
 
         (uint256 amountOut,) = abi.decode(result, (uint256, uint256));
+        require(amountOut >= swapDesc.minReturnAmount, "insufficient output amount");
         uint256 today = block.timestamp / DAY * DAY;
         dailyBought[today] = dailyBought[today] + amountOut;
 
         emit BoughtBack(_tokenIn, _amountIn, amountOut);
     }
 
+    /// @dev buy back tokens using router
+    /// @param _router The address of the router.
+    /// @param _tokenIn The address of the input token.
+    /// @param _tokenOut The address of the output token.
+    /// @param _amountIn The amount to sell.
+    /// @param _amountOutMin The minimum amount to receive.
+    /// @param _swapData The swap data.
+    function buyback(
+        address _router,
+        address _tokenIn,
+        address _tokenOut,
+        uint256 _amountIn,
+        uint256 _amountOutMin,
+        bytes calldata _swapData
+    ) external onlyRole(BOT) {
+        require(tokenWhitelist[_tokenIn], "token not whitelisted");
+        require(tokenWhitelist[_tokenOut], "token not whitelisted");
+        require(routerWhitelist[_router], "router not whitelisted");
+
+        uint256 beforeTokenIn = _getTokenBalance(_tokenIn, address(this));
+        uint256 beforeTokenOut = _getTokenBalance(_tokenOut, address(this));
+
+        bool isNativeTokenIn = (_tokenIn == SWAP_NATIVE_TOKEN_ADDRESS);
+        if (!isNativeTokenIn) {
+            IERC20(_tokenIn).safeApprove(_router, _amountIn);
+        }
+        (bool success, ) = _router.call{value: isNativeTokenIn ? _amountIn : 0}(_swapData);
+        require(success, "swap failed");
+        if (!isNativeTokenIn) {
+            IERC20(_tokenIn).safeApprove(_router, 0);
+        }
+
+        uint256 actualAmountIn = beforeTokenIn - _getTokenBalance(_tokenIn, address(this));
+        uint256 actualAmountOut = _getTokenBalance(_tokenOut, address(this)) - beforeTokenOut;
+
+        require(actualAmountIn <= _amountIn, "exceed amount in");
+        require(actualAmountOut >= _amountOutMin, "not enough profit");
+
+        IERC20(_tokenOut).safeTransfer(defaultReceiver, actualAmountOut);
+
+        emit BoughtBack(_router, _tokenIn, _tokenOut, actualAmountIn, actualAmountOut);
+    }
+
     function adminTransfer(address _token, uint256 _amount) external onlyRole(DEFAULT_ADMIN_ROLE) {
-        IERC20(_token).safeTransfer(msg.sender, _amount);
+        if (_token == SWAP_NATIVE_TOKEN_ADDRESS) {
+            (bool success, ) = payable(msg.sender).call{ value: _amount }("");
+            require(success, "Withdraw failed");
+        } else {
+            IERC20(_token).safeTransfer(msg.sender, _amount);
+        }
+        emit AdminTransfer(_token, _amount);
     }
 
     function changeDefaultReceiver(address _receiver) external onlyRole(DEFAULT_ADMIN_ROLE) {
@@ -110,20 +191,27 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
         emit ReceiverChanged(defaultReceiver);
     }
 
-    function add1InchRouterWhitelist(address _router) external onlyRole(DEFAULT_ADMIN_ROLE) {
-        require(!oneInchRouterWhitelist[_router], "router already whitelisted");
-
-        oneInchRouterWhitelist[_router] = true;
-        emit RouterChanged(_router, true);
+    /// @dev sets the router whitelist.
+    /// @param _router The address of the router.
+    /// @param status The status of the router.
+    function setRouterWhitelist(address _router, bool status) external onlyRole(DEFAULT_ADMIN_ROLE) {
+        require(_router != address(0), "Invalid router address");
+        require(routerWhitelist[_router] != status, "whitelist same status");
+        routerWhitelist[_router] = status;
+        emit RouterChanged(_router, status);
     }
 
-    function remove1InchRouterWhitelist(address _router) external onlyRole(DEFAULT_ADMIN_ROLE) {
-        require(oneInchRouterWhitelist[_router], "router not whitelisted");
-
-        delete oneInchRouterWhitelist[_router];
-        emit RouterChanged(_router, false);
+    /// @dev sets the token whitelist.
+    /// @param token The address of the token.
+    /// @param status The status of the token.
+    function setTokenWhitelist(address token, bool status) external onlyRole(DEFAULT_ADMIN_ROLE) {
+        require(token != address(0), "Invalid token");
+        require(tokenWhitelist[token] != status, "whitelist same status");
+        tokenWhitelist[token] = status;
+        emit TokenWhitelistChanged(token, status);
     }
 
+
     function _getFunctionSelector(bytes calldata _data) private pure returns (bytes4) {
         return bytes4(_data[0:4]);
     }
@@ -135,4 +223,12 @@ contract ListaAutoBuyback is Initializable, AccessControlUpgradeable {
             dstReceiver := calldataload(add(_data.offset, SWAP_DST_RECEIVER_OFFSET))
         }
     }
+
+    function _getTokenBalance(address _token, address account) internal view returns (uint256) {
+        if (_token == SWAP_NATIVE_TOKEN_ADDRESS) {
+            return account.balance;
+        } else {
+            return IERC20(_token).balanceOf(account);
+        }
+    }
 }

scripts/foundry/buyback/deploy_buyback_impl.sol

@@ -0,0 +1,20 @@
+pragma solidity ^0.8.10;
+
+import { Script, console } from "forge-std/Script.sol";
+import { ERC1967Proxy } from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+import { Buyback } from "../../../contracts/buyback/Buyback.sol";
+
+contract BuybackDeploy is Script {
+    function run() public {
+        uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
+        address deployer = vm.addr(deployerPrivateKey);
+        console.log("Deployer: ", deployer);
+        vm.startBroadcast(deployerPrivateKey);
+
+        // Deploy implementation
+        Buyback impl = new Buyback();
+        console.log("Implementation: ", address(impl));
+
+        vm.stopBroadcast();
+    }
+}