From 13f3299b9c90a9e39e8c2d7b080c7d25181f8ae7 Mon Sep 17 00:00:00 2001
From: Denis Fadeev <denis@fadeev.org>
Date: Mon, 18 Nov 2024 11:12:46 +0300
Subject: [PATCH] nft: return extra tokens to sender, handle a case where
 receiver is not the same as sender

---
 examples/nft/contracts/Connected.sol | 21 ++++++++++++++-------
 examples/nft/contracts/Universal.sol |  8 +++++---
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/examples/nft/contracts/Connected.sol b/examples/nft/contracts/Connected.sol
index f75a6f2..482b815 100644
--- a/examples/nft/contracts/Connected.sol
+++ b/examples/nft/contracts/Connected.sol
@@ -75,7 +75,13 @@ contract Connected is
 
         string memory uri = tokenURI(tokenId);
         _burn(tokenId);
-        bytes memory message = abi.encode(destination, receiver, tokenId, uri);
+        bytes memory message = abi.encode(
+            destination,
+            receiver,
+            tokenId,
+            uri,
+            msg.sender
+        );
         if (destination == address(0)) {
             gateway.call(
                 counterparty,
@@ -90,7 +96,7 @@ contract Connected is
                     address(this),
                     true,
                     address(0),
-                    abi.encode(receiver, tokenId, uri),
+                    abi.encode(receiver, tokenId, uri, msg.sender),
                     gasLimit
                 )
             );
@@ -109,13 +115,14 @@ contract Connected is
             address receiver,
             uint256 tokenId,
             string memory uri,
-            uint256 amount
-        ) = abi.decode(message, (address, uint256, string, uint256));
+            uint256 amount,
+            address sender
+        ) = abi.decode(message, (address, uint256, string, uint256, address));
 
         _safeMint(receiver, tokenId);
         _setTokenURI(tokenId, uri);
         if (amount > 0) {
-            (bool success, ) = receiver.call{value: amount}("");
+            (bool success, ) = sender.call{value: amount}("");
             if (!success) revert GasTokenTransferFailed();
         }
         emit TokenTransferReceived(receiver, tokenId, uri);
@@ -123,9 +130,9 @@ contract Connected is
     }
 
     function onRevert(RevertContext calldata context) external onlyGateway {
-        (address sender, uint256 tokenId, string memory uri) = abi.decode(
+        (, uint256 tokenId, string memory uri, address sender) = abi.decode(
             context.revertMessage,
-            (address, uint256, string)
+            (address, uint256, string, address)
         );
 
         _safeMint(sender, tokenId);
diff --git a/examples/nft/contracts/Universal.sol b/examples/nft/contracts/Universal.sol
index 0df47c3..07faddb 100644
--- a/examples/nft/contracts/Universal.sol
+++ b/examples/nft/contracts/Universal.sol
@@ -125,14 +125,16 @@ contract Universal is
         uint256 amount,
         bytes calldata message
     ) external override onlyGateway {
+        revert("Method not implemented");
         if (context.sender != counterparty[zrc20]) revert Unauthorized();
 
         (
             address destination,
             address receiver,
             uint256 tokenId,
-            string memory uri
-        ) = abi.decode(message, (address, address, uint256, string));
+            string memory uri,
+            address sender
+        ) = abi.decode(message, (address, address, uint256, string, address));
 
         if (destination == address(0)) {
             _safeMint(receiver, tokenId);
@@ -156,7 +158,7 @@ contract Universal is
                 abi.encodePacked(counterparty[destination]),
                 out - gasFee,
                 destination,
-                abi.encode(receiver, tokenId, uri, out - gasFee),
+                abi.encode(receiver, tokenId, uri, out - gasFee, sender),
                 CallOptions(gasLimit, false),
                 RevertOptions(address(0), false, address(0), "", 0)
             );