From 5338e6428e591167671418859b1bc20094a0bef2 Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Wed, 27 Nov 2024 11:43:33 +0300
Subject: [PATCH] gasZRC20 check

---
 contracts/nft/contracts/zetachain/UniversalNFT.sol  | 12 ++++++------
 .../token/contracts/zetachain/UniversalToken.sol    | 13 +++++++------
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/contracts/nft/contracts/zetachain/UniversalNFT.sol b/contracts/nft/contracts/zetachain/UniversalNFT.sol
index d1086cf..181f031 100644
--- a/contracts/nft/contracts/zetachain/UniversalNFT.sol
+++ b/contracts/nft/contracts/zetachain/UniversalNFT.sol
@@ -71,9 +71,9 @@ abstract contract UniversalNFT is
         string memory uri = tokenURI(tokenId);
         _burn(tokenId);
 
-        (, uint256 gasFee) = IZRC20(destination).withdrawGasFeeWithGasLimit(
-            gasLimitAmount
-        );
+        (address gasZRC20, uint256 gasFee) = IZRC20(destination)
+            .withdrawGasFeeWithGasLimit(gasLimitAmount);
+        if (destination != gasZRC20) revert InvalidAddress();
         if (
             !IZRC20(destination).transferFrom(msg.sender, address(this), gasFee)
         ) revert TransferFailed();
@@ -143,9 +143,9 @@ abstract contract UniversalNFT is
             _setTokenURI(tokenId, uri);
             emit TokenTransferReceived(receiver, tokenId, uri);
         } else {
-            (, uint256 gasFee) = IZRC20(destination).withdrawGasFeeWithGasLimit(
-                gasLimitAmount
-            );
+            (address gasZRC20, uint256 gasFee) = IZRC20(destination)
+                .withdrawGasFeeWithGasLimit(gasLimitAmount);
+            if (destination != gasZRC20) revert InvalidAddress();
 
             uint256 out = SwapHelperLib.swapExactTokensForTokens(
                 uniswapRouter,
diff --git a/contracts/token/contracts/zetachain/UniversalToken.sol b/contracts/token/contracts/zetachain/UniversalToken.sol
index f67e45d..c358bc7 100644
--- a/contracts/token/contracts/zetachain/UniversalToken.sol
+++ b/contracts/token/contracts/zetachain/UniversalToken.sol
@@ -65,9 +65,10 @@ abstract contract UniversalToken is
         if (receiver == address(0)) revert InvalidAddress();
         _burn(msg.sender, amount);
 
-        (, uint256 gasFee) = IZRC20(destination).withdrawGasFeeWithGasLimit(
-            gasLimitAmount
-        );
+        (address gasZRC20, uint256 gasFee) = IZRC20(destination)
+            .withdrawGasFeeWithGasLimit(gasLimitAmount);
+        if (destination != gasZRC20) revert InvalidAddress();
+
         if (
             !IZRC20(destination).transferFrom(msg.sender, address(this), gasFee)
         ) revert TransferFailed();
@@ -116,9 +117,9 @@ abstract contract UniversalToken is
         if (destination == address(0)) {
             _mint(receiver, tokenAmount);
         } else {
-            (, uint256 gasFee) = IZRC20(destination).withdrawGasFeeWithGasLimit(
-                gasLimitAmount
-            );
+            (address gasZRC20, uint256 gasFee) = IZRC20(destination)
+                .withdrawGasFeeWithGasLimit(gasLimitAmount);
+            if (destination != gasZRC20) revert InvalidAddress();
             uint256 out = SwapHelperLib.swapExactTokensForTokens(
                 uniswapRouter,
                 zrc20,