From 6adc2fa81e0e363a3608bc0dd3b5feff9fb7886c Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Thu, 27 Jun 2024 12:04:43 +0300
Subject: [PATCH] feat: added ability to swap token on and from ZetaChain

---
 omnichain/swap/contracts/Swap.sol           | 21 ++++--
 omnichain/swap/contracts/SwapToAnyToken.sol | 76 ++++++++++++++++-----
 2 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/omnichain/swap/contracts/Swap.sol b/omnichain/swap/contracts/Swap.sol
index fa36a853..e6400ee8 100644
--- a/omnichain/swap/contracts/Swap.sol
+++ b/omnichain/swap/contracts/Swap.sol
@@ -42,15 +42,24 @@ contract Swap is zContract, OnlySystem {
             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
@@ -58,13 +67,13 @@ contract Swap is zContract, OnlySystem {
 
         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);
     }
 }
diff --git a/omnichain/swap/contracts/SwapToAnyToken.sol b/omnichain/swap/contracts/SwapToAnyToken.sol
index 4a22dc6d..d1196b36 100644
--- a/omnichain/swap/contracts/SwapToAnyToken.sol
+++ b/omnichain/swap/contracts/SwapToAnyToken.sol
@@ -23,6 +23,8 @@ contract SwapToAnyToken is zContract, OnlySystem {
         bool withdraw;
     }
 
+    receive() external payable {}
+
     function onCrossChainCall(
         zContext calldata context,
         address zrc20,
@@ -37,49 +39,91 @@ contract SwapToAnyToken is zContract, OnlySystem {
 
         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 outputAmount;
         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(
+        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 recipientAddress = payable(
+                    address(uint160(bytes20(recipient)))
+                );
+                recipientAddress.transfer(outputAmount);
+            } else {
+                address recipientAddress = address(uint160(bytes20(recipient)));
+                IWETH9(targetToken).transfer(recipientAddress, outputAmount);
+            }
         }
     }
+
+    function swap(
+        address inputToken,
+        uint256 amount,
+        address targetToken,
+        bytes memory recipient,
+        bool withdraw
+    ) public {
+        IZRC20(inputToken).transferFrom(msg.sender, address(this), amount);
+
+        swapAndWithdraw(inputToken, amount, targetToken, recipient, withdraw);
+    }
 }