From 0cbaa249884f045831fc833ab6f787b350ecfbce Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Mon, 18 Nov 2024 10:48:30 +0300
Subject: [PATCH] nft: withdrawAndCall

---
 examples/nft/contracts/Connected.sol | 15 +++++++++++----
 examples/nft/contracts/Universal.sol | 11 ++++++-----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/examples/nft/contracts/Connected.sol b/examples/nft/contracts/Connected.sol
index 11e9138..f75a6f2 100644
--- a/examples/nft/contracts/Connected.sol
+++ b/examples/nft/contracts/Connected.sol
@@ -25,6 +25,7 @@ contract Connected is
     error InvalidAddress();
     error Unauthorized();
     error InvalidGasLimit();
+    error GasTokenTransferFailed();
 
     function setCounterparty(address contractAddress) external onlyOwner {
         if (contractAddress == address(0)) revert InvalidAddress();
@@ -104,13 +105,19 @@ contract Connected is
     ) external payable onlyGateway returns (bytes4) {
         if (context.sender != counterparty) revert Unauthorized();
 
-        (address receiver, uint256 tokenId, string memory uri) = abi.decode(
-            message,
-            (address, uint256, string)
-        );
+        (
+            address receiver,
+            uint256 tokenId,
+            string memory uri,
+            uint256 amount
+        ) = abi.decode(message, (address, uint256, string, uint256));
 
         _safeMint(receiver, tokenId);
         _setTokenURI(tokenId, uri);
+        if (amount > 0) {
+            (bool success, ) = receiver.call{value: amount}("");
+            if (!success) revert GasTokenTransferFailed();
+        }
         emit TokenTransferReceived(receiver, tokenId, uri);
         return "";
     }
diff --git a/examples/nft/contracts/Universal.sol b/examples/nft/contracts/Universal.sol
index f270f33..0df47c3 100644
--- a/examples/nft/contracts/Universal.sol
+++ b/examples/nft/contracts/Universal.sol
@@ -83,7 +83,7 @@ contract Universal is
             !IZRC20(destination).transferFrom(msg.sender, address(this), gasFee)
         ) revert TransferFailed();
         IZRC20(destination).approve(address(gateway), gasFee);
-        bytes memory message = abi.encode(receiver, tokenId, uri);
+        bytes memory message = abi.encode(receiver, tokenId, uri, 0);
 
         CallOptions memory callOptions = CallOptions(gasLimit, false);
 
@@ -143,7 +143,7 @@ contract Universal is
                 gasLimit
             );
 
-            SwapHelperLib.swapExactTokensForTokens(
+            uint256 out = SwapHelperLib.swapExactTokensForTokens(
                 uniswapRouter,
                 zrc20,
                 amount,
@@ -151,11 +151,12 @@ contract Universal is
                 0
             );
 
-            IZRC20(destination).approve(address(gateway), gasFee);
-            gateway.call(
+            IZRC20(destination).approve(address(gateway), out);
+            gateway.withdrawAndCall(
                 abi.encodePacked(counterparty[destination]),
+                out - gasFee,
                 destination,
-                abi.encode(receiver, tokenId, uri),
+                abi.encode(receiver, tokenId, uri, out - gasFee),
                 CallOptions(gasLimit, false),
                 RevertOptions(address(0), false, address(0), "", 0)
             );