Merge pull request #32 from skip-mev/jw/update-axelar-handler

Update Axelar handler to support new swap router

Author: thal0x

AxelarHandler/src/AxelarHandler.sol

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: UNLICENSED
-pragma solidity 0.8.18;
+pragma solidity ^0.8.0;
 
 import {IWETH} from "./interfaces/IWETH.sol";
 
@@ -13,7 +13,6 @@ import {Ownable2StepUpgradeable} from
 import {UUPSUpgradeable} from "lib/openzeppelin-contracts-upgradeable/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {Initializable} from "lib/openzeppelin-contracts-upgradeable/contracts/proxy/utils/Initializable.sol";
 
-import {ISwapRouter02} from "./interfaces/ISwapRouter02.sol";
 import {SkipSwapRouter} from "./libraries/SkipSwapRouter.sol";
 
 /// @title AxelarHandler
@@ -32,7 +31,7 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
     error NonNativeCannotBeUnwrapped();
     error NativePaymentFailed();
     error WrappingNotEnabled();
-    error SwapFailed();
+    error SwapFailedError();
     error InsufficientSwapOutput();
     error InsufficientNativeToken();
     error ETHSendFailed();
@@ -45,8 +44,7 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
     enum Commands {
         SendToken,
         SendNative,
-        Swap,
-        MultiSwap
+        Swap
     }
 
     bytes32 private _wETHSymbolHash;
@@ -58,7 +56,7 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
 
     bytes32 public constant DISABLED_SYMBOL = keccak256(abi.encodePacked("DISABLED"));
 
-    ISwapRouter02 public swapRouter;
+    address public swapRouter;
 
     bool internal reentrant;
 
@@ -96,7 +94,7 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
     function setSwapRouter(address _swapRouter) external onlyOwner {
         if (_swapRouter == address(0)) revert ZeroAddress();
 
-        swapRouter = ISwapRouter02(_swapRouter);
+        swapRouter = _swapRouter;
     }
 
     /// @notice Sends native currency to other chains through the axelar gateway.
@@ -230,73 +228,13 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
         if (bytes(symbol).length == 0) revert EmptySymbol();
 
         // Get the address of the output token based on the symbol provided
-        IERC20 outputToken = IERC20(_getTokenAddress(symbol));
+        address outputToken = _getTokenAddress(symbol);
 
         uint256 outputAmount;
         if (inputToken == address(0)) {
-            // Native Token
-            if (amount + gasPaymentAmount != msg.value) revert InsufficientNativeToken();
-
-            // Get the contract's balances previous to the swap
-            uint256 preInputBalance = address(this).balance - msg.value;
-            uint256 preOutputBalance = outputToken.balanceOf(address(this));
-
-            // Call the swap router and perform the swap
-            (bool success,) = address(swapRouter).call{value: amount}(swapCalldata);
-            if (!success) revert SwapFailed();
-
-            // Get the contract's balances after the swap
-            uint256 postInputBalance = address(this).balance;
-            uint256 postOutputBalance = outputToken.balanceOf(address(this));
-
-            // Check that the contract's native token balance has increased
-            if (preOutputBalance >= postOutputBalance) revert InsufficientSwapOutput();
-            outputAmount = postOutputBalance - preOutputBalance;
-
-            // Refund the remaining ETH
-            uint256 dust = postInputBalance - preInputBalance - gasPaymentAmount;
-            if (dust != 0) {
-                (bool ethSuccess,) = msg.sender.call{value: dust}("");
-                if (!ethSuccess) revert ETHSendFailed();
-            }
-
-            emit SwapSuccess(inputToken, address(outputToken), amount, outputAmount);
+            outputAmount = _swapAndRefundNative(amount, outputToken, gasPaymentAmount, swapCalldata);
         } else {
-            // ERC20 Token
-            if (gasPaymentAmount != msg.value) revert();
-
-            // Transfer input ERC20 tokens to the contract
-            IERC20 token = IERC20(inputToken);
-            token.safeTransferFrom(msg.sender, address(this), amount);
-
-            // Approve the swap router to spend the input tokens
-            token.safeApprove(address(swapRouter), amount);
-
-            // Get the contract's balances previous to the swap
-            uint256 preInputBalance = token.balanceOf(address(this));
-            uint256 preOutputBalance = outputToken.balanceOf(address(this));
-
-            // Call the swap router and perform the swap
-            (bool success,) = address(swapRouter).call(swapCalldata);
-            if (!success) revert SwapFailed();
-
-            // Get the contract's balances after the swap
-            uint256 dust = token.balanceOf(address(this)) + amount - preInputBalance;
-            uint256 postOutputBalance = outputToken.balanceOf(address(this));
-
-            // Check that the contract's output token balance has increased
-            if (preOutputBalance >= postOutputBalance) revert InsufficientSwapOutput();
-            outputAmount = postOutputBalance - preOutputBalance;
-
-            // Refund the remaining amount
-            if (dust != 0) {
-                token.safeTransfer(msg.sender, dust);
-
-                // Revoke approval
-                token.safeApprove(address(swapRouter), 0);
-            }
-
-            emit SwapSuccess(inputToken, address(outputToken), amount, outputAmount);
+            outputAmount = _swapAndRefund(amount, inputToken, outputToken, gasPaymentAmount, swapCalldata);
         }
 
         // Pay the gas for the GMP transfer
@@ -390,30 +328,13 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
 
             _sendNative(token, amount, destination);
         } else if (command == Commands.Swap) {
-            (address destination, bool unwrapOut, bytes memory swap) = abi.decode(data, (address, bool, bytes));
-
-            try SkipSwapRouter.swap(swapRouter, destination, tokenIn, amount, swap) returns (
-                IERC20 tokenOut, uint256 amountOut
-            ) {
-                if (unwrapOut) {
-                    _sendNative(address(tokenOut), amountOut, destination);
-                    emit SwapSuccess(address(tokenIn), address(0), amount, amountOut);
-                } else {
-                    _sendToken(address(tokenOut), amountOut, destination);
-                    emit SwapSuccess(address(tokenIn), address(tokenOut), amount, amountOut);
-                }
-            } catch {
-                _sendToken(token, amount, destination);
-                emit SwapFailed();
-            }
-        } else if (command == Commands.MultiSwap) {
-            (address destination, bool unwrapOut, bytes[] memory swaps) = abi.decode(data, (address, bool, bytes[]));
+            (address tokenOut, address destination, bytes memory swap) = abi.decode(data, (address, address, bytes));
 
-            try SkipSwapRouter.multiSwap(swapRouter, destination, tokenIn, amount, swaps) returns (
-                IERC20 tokenOut, uint256 amountOut
+            try SkipSwapRouter.swap(swapRouter, destination, address(tokenIn), tokenOut, amount, swap) returns (
+                uint256 amountOut
             ) {
-                if (unwrapOut) {
-                    _sendNative(address(tokenOut), amountOut, destination);
+                if (tokenOut == address(0)) {
+                    address(destination).call{value: amountOut}("");
                     emit SwapSuccess(address(tokenIn), address(0), amount, amountOut);
                 } else {
                     _sendToken(address(tokenOut), amountOut, destination);
@@ -428,6 +349,85 @@ contract AxelarHandler is AxelarExecutableUpgradeable, Ownable2StepUpgradeable,
         }
     }
 
+    function _swapAndRefund(
+        uint256 amount,
+        address inputToken,
+        address outputToken,
+        uint256 gasPaymentAmount,
+        bytes memory swapCalldata
+    ) internal returns (uint256 outputAmount) {
+        // ERC20 Token
+        if (gasPaymentAmount != msg.value) revert();
+
+        // Transfer input ERC20 tokens to the contract
+        IERC20 token = IERC20(inputToken);
+        token.safeTransferFrom(msg.sender, address(this), amount);
+
+        // Approve the swap router to spend the input tokens
+        token.safeApprove(address(swapRouter), amount);
+
+        // Get the contract's balances previous to the swap
+        uint256 preInputBalance = token.balanceOf(address(this));
+        uint256 preOutputBalance = IERC20(outputToken).balanceOf(address(this));
+
+        // Call the swap router and perform the swap
+        (bool success,) = address(swapRouter).call(swapCalldata);
+        if (!success) revert SwapFailedError();
+
+        // Get the contract's balances after the swap
+        uint256 dust = token.balanceOf(address(this)) + amount - preInputBalance;
+        uint256 postOutputBalance = IERC20(outputToken).balanceOf(address(this));
+
+        // Check that the contract's output token balance has increased
+        if (preOutputBalance >= postOutputBalance) revert InsufficientSwapOutput();
+        outputAmount = postOutputBalance - preOutputBalance;
+
+        // Refund the remaining amount
+        if (dust != 0) {
+            token.safeTransfer(msg.sender, dust);
+
+            // Revoke approval
+            token.safeApprove(address(swapRouter), 0);
+        }
+
+        emit SwapSuccess(inputToken, address(outputToken), amount, outputAmount);
+    }
+
+    function _swapAndRefundNative(
+        uint256 amount,
+        address outputToken,
+        uint256 gasPaymentAmount,
+        bytes memory swapCalldata
+    ) internal returns (uint256 outputAmount) {
+        // Native Token
+        if (amount + gasPaymentAmount != msg.value) revert InsufficientNativeToken();
+
+        // Get the contract's balances previous to the swap
+        uint256 preInputBalance = address(this).balance - msg.value;
+        uint256 preOutputBalance = IERC20(outputToken).balanceOf(address(this));
+
+        // Call the swap router and perform the swap
+        (bool success,) = address(swapRouter).call{value: amount}(swapCalldata);
+        if (!success) revert SwapFailedError();
+
+        // Get the contract's balances after the swap
+        uint256 postInputBalance = address(this).balance;
+        uint256 postOutputBalance = IERC20(outputToken).balanceOf(address(this));
+
+        // Check that the contract's native token balance has increased
+        if (preOutputBalance >= postOutputBalance) revert InsufficientSwapOutput();
+        outputAmount = postOutputBalance - preOutputBalance;
+
+        // Refund the remaining ETH
+        uint256 dust = postInputBalance - preInputBalance - gasPaymentAmount;
+        if (dust != 0) {
+            (bool ethSuccess,) = msg.sender.call{value: dust}("");
+            if (!ethSuccess) revert ETHSendFailed();
+        }
+
+        emit SwapSuccess(address(0), address(outputToken), amount, outputAmount);
+    }
+
     function _sendToken(address token, uint256 amount, address destination) internal {
         IERC20(token).safeTransfer(destination, amount);
     }

AxelarHandler/src/libraries/SkipSwapRouter.sol

@@ -14,105 +14,38 @@ library SkipSwapRouter {
     error InsufficientOutputAmount();
     error NativePaymentFailed();
 
-    enum SwapCommands {
-        ExactInputSingle,
-        ExactInput,
-        ExactTokensForTokens,
-        ExactOutputSingle,
-        ExactOutput,
-        TokensForExactTokens
-    }
-
-    function multiSwap(
-        ISwapRouter02 router,
+    function swap(
+        address router,
         address destination,
-        IERC20 inputToken,
+        address inputToken,
+        address outputToken,
         uint256 amountIn,
-        bytes[] memory swaps
-    ) external returns (IERC20 outputToken, uint256 amountOut) {
-        outputToken = inputToken;
-        amountOut = amountIn;
-
-        uint256 numSwaps = swaps.length;
-        for (uint256 i; i < numSwaps; i++) {
-            // The output token and amount of each iteration is the input token and amount of the next.
-            (outputToken, amountOut) = swap(router, destination, outputToken, amountOut, swaps[i]);
-        }
-    }
+        bytes memory swapData
+    ) public returns (uint256 amountOut) {
+        uint256 preBalIn = IERC20(inputToken).balanceOf(address(this)) - amountIn;
+
+        uint256 preBalOut =
+            outputToken == address(0) ? address(this).balance : IERC20(outputToken).balanceOf(address(this));
+
+        IERC20(inputToken).forceApprove(router, amountIn);
 
-    function swap(ISwapRouter02 router, address destination, IERC20 inputToken, uint256 amountIn, bytes memory payload)
-        public
-        returns (IERC20 outputToken, uint256 outputAmount)
-    {
-        (SwapCommands command, address tokenOut, uint256 amountOut, bytes memory swapData) =
-            abi.decode(payload, (SwapCommands, address, uint256, bytes));
-
-        outputToken = IERC20(tokenOut);
-
-        uint256 preBalIn = inputToken.balanceOf(address(this)) - amountIn;
-        uint256 preBalOut = outputToken.balanceOf(address(this));
-
-        inputToken.forceApprove(address(router), amountIn);
-
-        if (command == SwapCommands.ExactInputSingle) {
-            ISwapRouter02.ExactInputSingleParams memory params;
-            params.tokenIn = address(inputToken);
-            params.tokenOut = tokenOut;
-            params.recipient = address(this);
-            params.amountIn = amountIn;
-            params.amountOutMinimum = amountOut;
-
-            (params.fee, params.sqrtPriceLimitX96) = abi.decode(swapData, (uint24, uint160));
-
-            router.exactInputSingle(params);
-        } else if (command == SwapCommands.ExactInput) {
-            ISwapRouter02.ExactInputParams memory params;
-            params.path = _fixPath(address(inputToken), tokenOut, swapData);
-            params.path = swapData;
-            params.recipient = address(this);
-            params.amountIn = amountIn;
-            params.amountOutMinimum = amountOut;
-
-            router.exactInput(params);
-        } else if (command == SwapCommands.ExactTokensForTokens) {
-            address[] memory path = _fixPath(address(inputToken), tokenOut, abi.decode(swapData, (address[])));
-
-            router.swapExactTokensForTokens(amountIn, amountOut, path, address(this));
-        } else if (command == SwapCommands.ExactOutputSingle) {
-            ISwapRouter02.ExactOutputSingleParams memory params;
-            params.tokenIn = address(inputToken);
-            params.tokenOut = tokenOut;
-            params.recipient = address(this);
-            params.amountInMaximum = amountIn;
-            params.amountOut = amountOut;
-
-            (params.fee, params.sqrtPriceLimitX96) = abi.decode(swapData, (uint24, uint160));
-
-            router.exactOutputSingle(params);
-        } else if (command == SwapCommands.ExactOutput) {
-            ISwapRouter02.ExactOutputParams memory params;
-            params.path = _fixPath(tokenOut, address(inputToken), swapData);
-            params.path = swapData;
-            params.recipient = address(this);
-            params.amountInMaximum = amountIn;
-            params.amountOut = amountOut;
-
-            router.exactOutput(params);
-        } else if (command == SwapCommands.TokensForExactTokens) {
-            address[] memory path = _fixPath(address(inputToken), address(tokenOut), abi.decode(swapData, (address[])));
-
-            router.swapTokensForExactTokens(amountOut, amountIn, path, address(this));
+        (bool success, bytes memory returnData) = router.call(swapData);
+
+        if (!success) {
+            _revertWithData(returnData);
         }
 
-        outputAmount = outputToken.balanceOf(address(this)) - preBalOut;
-        if (outputAmount < amountOut) {
-            revert InsufficientOutputAmount();
+        if (outputToken == address(0)) {
+            amountOut = address(this).balance - preBalOut;
+        } else {
+            amountOut = IERC20(outputToken).balanceOf(address(this)) - preBalOut;
         }
 
-        uint256 dust = inputToken.balanceOf(address(this)) - preBalIn;
+        uint256 dust = IERC20(inputToken).balanceOf(address(this)) - preBalIn;
+
         if (dust != 0) {
-            inputToken.forceApprove(address(router), 0);
-            inputToken.safeTransfer(destination, dust);
+            IERC20(inputToken).forceApprove(router, 0);
+            IERC20(inputToken).safeTransfer(destination, dust);
         }
     }
 
@@ -155,4 +88,10 @@ library SkipSwapRouter {
 
         return path;
     }
+
+    function _revertWithData(bytes memory data) private pure {
+        assembly {
+            revert(add(data, 32), mload(data))
+        }
+    }
 }