feat: added ability to swap token on and from ZetaChain

omnichain/swap/contracts/Swap.sol

@@ -42,29 +42,38 @@
             params.to = recipient;
         }
 
+        swapAndWithdraw(zrc20, amount, params.target, params.to);
+    }
+
+    function swapAndWithdraw(
+        address inputToken,
+        uint256 amount,
+        address targetToken,
+        bytes memory recipient
+    ) internal {
         uint256 inputForGas;
         address gasZRC20;
         uint256 gasFee;
 
-        (gasZRC20, gasFee) = IZRC20(params.target).withdrawGasFee();
+        (gasZRC20, gasFee) = IZRC20(targetToken).withdrawGasFee();
 
         inputForGas = SwapHelperLib.swapTokensForExactTokens(
             systemContract,
-            zrc20,
+            inputToken,
             gasFee,
             gasZRC20,
             amount
         );
 
         uint256 outputAmount = SwapHelperLib.swapExactTokensForTokens(
             systemContract,
-            zrc20,
+            inputToken,
             amount - inputForGas,
-            params.target,
+            targetToken,
             0
         );
 
-        IZRC20(gasZRC20).approve(params.target, gasFee);
-        IZRC20(params.target).withdraw(params.to, outputAmount);
+        IZRC20(gasZRC20).approve(targetToken, gasFee);
+        IZRC20(targetToken).withdraw(recipient, outputAmount);
     }
 }

omnichain/swap/contracts/SwapToAnyToken.sol

@@ -9,6 +9,7 @@
 import "@zetachain/toolkit/contracts/OnlySystem.sol";
 
 contract SwapToAnyToken is zContract, OnlySystem {
+    error TransferFailed();
     SystemContract public systemContract;
 
     uint256 constant BITCOIN = 18332;
@@ -23,6 +24,8 @@
         bool withdraw;
     }
 
+    receive() external payable {}
+
     function onCrossChainCall(
         zContext calldata context,
         address zrc20,
@@ -37,49 +40,96 @@
 
         if (context.chainID == BITCOIN) {
             params.target = BytesHelperLib.bytesToAddress(message, 0);
-            params.to = abi.encodePacked(BytesHelperLib.bytesToAddress(message, 20));
+            params.to = abi.encodePacked(
+                BytesHelperLib.bytesToAddress(message, 20)
+            );
             if (message.length >= 41) {
                 params.withdraw = BytesHelperLib.bytesToBool(message, 40);
             }
         } else {
-            (address targetToken, bytes memory recipient, bool withdrawFlag) = abi.decode(
-                message,
-                (address, bytes, bool)
-            );
+            (
+                address targetToken,
+                bytes memory recipient,
+                bool withdrawFlag
+            ) = abi.decode(message, (address, bytes, bool));
             params.target = targetToken;
             params.to = recipient;
             params.withdraw = withdrawFlag;
         }
 
+        swapAndWithdraw(
+            zrc20,
+            amount,
+            params.target,
+            params.to,
+            params.withdraw
+        );
+    }
+
+    function swapAndWithdraw(
+        address inputToken,
+        uint256 amount,
+        address targetToken,
+        bytes memory recipient,
+        bool withdraw
+    ) internal {
         uint256 inputForGas;
         address gasZRC20;
         uint256 gasFee;
 
-        if (params.withdraw) {
-            (gasZRC20, gasFee) = IZRC20(params.target).withdrawGasFee();
+        if (withdraw) {
+            (gasZRC20, gasFee) = IZRC20(targetToken).withdrawGasFee();
 
             inputForGas = SwapHelperLib.swapTokensForExactTokens(
                 systemContract,
-                zrc20,
+                inputToken,
                 gasFee,
                 gasZRC20,
                 amount
             );
         }
 
         uint256 outputAmount = SwapHelperLib.swapExactTokensForTokens(
             systemContract,
-            zrc20,
-            params.withdraw ? amount - inputForGas : amount,
-            params.target,
+            inputToken,
+            withdraw ? amount - inputForGas : amount,
+            targetToken,
             0
         );
 
-        if (params.withdraw) {
-            IZRC20(gasZRC20).approve(params.target, gasFee);
-            IZRC20(params.target).withdraw(params.to, outputAmount);
+        if (withdraw) {
+            IZRC20(gasZRC20).approve(targetToken, gasFee);
+            IZRC20(targetToken).withdraw(recipient, outputAmount);
         } else {
-            IWETH9(params.target).transfer(address(uint160(bytes20(params.to))), outputAmount);
+            address wzeta = systemContract.wZetaContractAddress();
+            if (targetToken == wzeta) {
+                IWETH9(wzeta).withdraw(outputAmount);
+                address payable to = payable(
+                    address(uint160(bytes20(recipient)))
+                );
+                (bool success, ) = to.call{value: outputAmount}("");
+                if (!success) revert TransferFailed();
+            } else {
+                address to = address(uint160(bytes20(recipient)));
+                bool success = IWETH9(targetToken).transfer(to, outputAmount);
+                if (!success) revert TransferFailed();
+            }
         }
     }
+
+    function swap(
+        address inputToken,
+        uint256 amount,
+        address targetToken,
+        bytes memory recipient,
+        bool withdraw
+    ) public {
+        bool success = IZRC20(inputToken).transferFrom(
+            msg.sender,
+            address(this),
+            amount
+        );
+        if (!success) revert TransferFailed();
+        swapAndWithdraw(inputToken, amount, targetToken, recipient, withdraw);
+    }
 }

omnichain/swap/tasks/swap.ts

@@ -0,0 +1,47 @@
+import { task } from "hardhat/config";
+import { HardhatRuntimeEnvironment } from "hardhat/types";
+import { parseEther } from "@ethersproject/units";
+import { ethers } from "ethers";
+
+const main = async (args: any, hre: HardhatRuntimeEnvironment) => {
+  const [signer] = await hre.ethers.getSigners();
+
+  if (!/zeta_(testnet|mainnet)/.test(hre.network.name)) {
+    throw new Error('🚨 Please use either "zeta_testnet" or "zeta_mainnet".');
+  }
+
+  const factory = await hre.ethers.getContractFactory("SwapToAnyToken");
+  const contract = factory.attach(args.contract);
+
+  const amount = parseEther(args.amount);
+  const inputToken = args.inputToken;
+  const targetToken = args.targetToken;
+  const recipient = ethers.utils.arrayify(args.recipient);
+  const withdraw = JSON.parse(args.withdraw);
+
+  const erc20Factory = await hre.ethers.getContractFactory("ERC20");
+  const inputTokenContract = erc20Factory.attach(args.inputToken);
+
+  const approval = await inputTokenContract.approve(args.contract, amount);
+  await approval.wait();
+
+  const tx = await contract.swap(
+    inputToken,
+    amount,
+    targetToken,
+    recipient,
+    withdraw
+  );
+
+  await tx.wait();
+  console.log(`Transaction hash: ${tx.hash}`);
+};
+
+task("swap", "Interact with the Swap contract from ZetaChain", main)
+  .addFlag("json", "Output JSON")
+  .addParam("contract", "Contract address")
+  .addParam("amount", "Token amount to send")
+  .addParam("inputToken", "Input token address")
+  .addParam("targetToken", "Target token address")
+  .addParam("recipient", "Recipient address")
+  .addParam("withdraw", "Withdraw flag (true/false)");