sendTokenWithPayloadToEvmWithRefund (#9)

* sendTokenWithPayloadToEvmWithRefund

* naming change

* Tests + bugfix

Author: derpy-duck

src/WormholeRelayerSDK.sol

@@ -146,6 +146,26 @@ abstract contract TokenSender is TokenBase {
         );
     }
 
+    function sendTokenWithPayloadToEvm(
+        uint16 targetChain,
+        address targetAddress,
+        bytes memory payload,
+        uint256 receiverValue,
+        uint256 gasLimit,
+        address token,
+        uint256 amount,
+        uint16 refundChain,
+        address refundAddress
+    ) internal returns (uint64) {
+        VaaKey[] memory vaaKeys = new VaaKey[](1);
+        vaaKeys[0] = transferTokens(token, amount, targetChain, targetAddress);
+
+        (uint256 cost,) = wormholeRelayer.quoteEVMDeliveryPrice(targetChain, receiverValue, gasLimit);
+        return wormholeRelayer.sendVaasToEvm{value: cost}(
+            targetChain, targetAddress, payload, receiverValue, gasLimit, vaaKeys, refundChain, refundAddress
+        );
+    }
+
     function forwardTokenWithPayloadToEvm(
         uint16 targetChain,
         address targetAddress,

src/testing/WormholeRelayerTest.sol

@@ -131,6 +131,8 @@ abstract contract WormholeRelayerTest is Test {
 
         vm.selectFork(firstFork.fork);
         setUpGeneral();
+
+        vm.selectFork(firstFork.fork);
     }
 
     function performDelivery() public {
@@ -161,6 +163,9 @@ abstract contract WormholeRelayerTest is Test {
         bytes memory attestation = home.guardian.fetchSignedMessageFromLogs(log, home.chainId);
 
         for (uint256 i = 0; i < activeForksList.length; ++i) {
+            if(activeForksList[i] == home.chainId) {
+                continue;
+            }
             ActiveFork memory fork = activeForks[activeForksList[i]];
             vm.selectFork(fork.fork);
             fork.tokenBridge.createWrapped(attestation);

test/Fork.t.sol

@@ -35,21 +35,118 @@ contract Toy is Base {
     }
 }
 
+contract TokenToy is TokenSender, TokenReceiver {
+    constructor (address _wormholeRelayer, address _bridge, address _wormhole) TokenBase(_wormholeRelayer, _bridge, _wormhole) {}
+    
+    uint256 constant GAS_LIMIT = 250_000;
+    
+    function quoteCrossChainDeposit(uint16 targetChain) public view returns (uint256 cost) {
+        // Cost of delivering token and payload to targetChain
+        uint256 deliveryCost;
+        (deliveryCost,) = wormholeRelayer.quoteEVMDeliveryPrice(targetChain, 0, GAS_LIMIT);
+
+        // Total cost: delivery cost + cost of publishing the 'sending token' wormhole message
+        cost = deliveryCost + wormhole.messageFee();
+    }
+    
+    function sendCrossChainDeposit(
+        uint16 targetChain,
+        address recipient,
+        uint256 amount,
+        address token
+    ) public payable {
+        uint256 cost = quoteCrossChainDeposit(targetChain);
+        require(msg.value == cost, "msg.value must be quoteCrossChainDeposit(targetChain)");
+
+        IERC20(token).transferFrom(msg.sender, address(this), amount);
+
+        bytes memory payload = abi.encode(recipient);
+        sendTokenWithPayloadToEvm(
+            targetChain, 
+            fromWormholeFormat(registeredSenders[targetChain]), // address (on targetChain) to send token and payload to
+            payload, 
+            0, // receiver value
+            GAS_LIMIT, 
+            token, // address of IERC20 token contract
+            amount
+        );
+    }
+
+    function sendCrossChainDeposit(
+        uint16 targetChain,
+        address recipient,
+        uint256 amount,
+        address token,
+        uint16 refundChain,
+        address refundAddress
+    ) public payable {
+        uint256 cost = quoteCrossChainDeposit(targetChain);
+        require(msg.value == cost, "msg.value must be quoteCrossChainDeposit(targetChain)");
+
+        IERC20(token).transferFrom(msg.sender, address(this), amount);
+
+        bytes memory payload = abi.encode(recipient);
+        sendTokenWithPayloadToEvm(
+            targetChain, 
+            fromWormholeFormat(registeredSenders[targetChain]), // address (on targetChain) to send token and payload to
+            payload, 
+            0, // receiver value
+            GAS_LIMIT, 
+            token, // address of IERC20 token contract
+            amount,
+            refundChain,
+            refundAddress
+        );
+    }
+
+    function receivePayloadAndTokens(
+        bytes memory payload,
+        TokenReceived[] memory receivedTokens,
+        bytes32 sourceAddress,
+        uint16 sourceChain,
+        bytes32 // deliveryHash
+    ) internal override onlyWormholeRelayer isRegisteredSender(sourceChain, sourceAddress) {
+        require(receivedTokens.length == 1, "Expected 1 token transfers");
+        address recipient = abi.decode(payload, (address));
+        IERC20(receivedTokens[0].tokenAddress).transfer(recipient, receivedTokens[0].amount);
+    }
+}
+
 contract WormholeSDKTest is WormholeRelayerBasicTest {
     Toy toySource;
     Toy toyTarget;
+    TokenToy tokenToySource;
+    TokenToy tokenToyTarget;
+    ERC20Mock public token;
 
     function setUpSource() public override {
         toySource = new Toy(address(relayerSource), address(wormholeSource));
         toySource.setRegisteredSender(targetChain, toWormholeFormat(address(this)));
+
+        tokenToySource = new TokenToy(address(relayerSource), address(tokenBridgeSource), address(wormholeSource));
+
+        token = createAndAttestToken(sourceChain);
     }
 
     function setUpTarget() public override {
         toyTarget = new Toy(address(relayerTarget), address(wormholeTarget));
         toyTarget.setRegisteredSender(sourceChain, toWormholeFormat(address(this)));
+
+        tokenToyTarget = new TokenToy(address(relayerTarget), address(tokenBridgeTarget), address(wormholeTarget));
+    }
+
+    function setUpGeneral() public override {
+        vm.selectFork(sourceFork);
+        tokenToySource.setRegisteredSender(targetChain, toWormholeFormat(address(tokenToyTarget)));
+        
+
+        vm.selectFork(targetFork);
+        tokenToyTarget.setRegisteredSender(sourceChain, toWormholeFormat(address(tokenToySource)));
+
     }
 
     function testSendMessage() public {
+
         vm.recordLogs();
         (uint256 cost,) = relayerSource.quoteEVMDeliveryPrice(targetChain, 1e17, 100_000);
         relayerSource.sendPayloadToEvm{value: cost}(targetChain, address(toyTarget), abi.encode(55), 1e17, 100_000);
@@ -70,4 +167,53 @@ contract WormholeSDKTest is WormholeRelayerBasicTest {
         vm.selectFork(sourceFork);
         require(56 == toySource.payloadReceived());
     }
+
+    function testSendToken() public {
+        
+        vm.selectFork(sourceFork);
+
+        uint256 amount = 19e17;
+        token.approve(address(tokenToySource), amount);
+
+        vm.selectFork(targetFork);
+        address recipient = 0x1234567890123456789012345678901234567890;
+
+        vm.selectFork(sourceFork);
+        uint256 cost = tokenToySource.quoteCrossChainDeposit(targetChain);
+
+        vm.recordLogs();
+        tokenToySource.sendCrossChainDeposit{value: cost}(
+            targetChain, recipient, amount, address(token)
+        );
+        performDelivery();
+
+        vm.selectFork(targetFork);
+        address wormholeWrappedToken = tokenBridgeTarget.wrappedAsset(sourceChain, toWormholeFormat(address(token)));
+        assertEq(IERC20(wormholeWrappedToken).balanceOf(recipient), amount);
+    }
+
+    function testSendTokenWithRefund() public {
+        
+        vm.selectFork(sourceFork);
+
+        uint256 amount = 19e17;
+        token.approve(address(tokenToySource), amount);
+
+        vm.selectFork(targetFork);
+        address recipient = 0x1234567890123456789012345678901234567890;
+        address refundAddress = 0x2234567890123456789012345678901234567890;
+        vm.selectFork(sourceFork);
+        uint256 cost = tokenToySource.quoteCrossChainDeposit(targetChain);
+
+        vm.recordLogs();
+        tokenToySource.sendCrossChainDeposit{value: cost}(
+            targetChain, recipient, amount, address(token), targetChain, refundAddress
+        );
+        performDelivery();
+
+        vm.selectFork(targetFork);
+        address wormholeWrappedToken = tokenBridgeTarget.wrappedAsset(sourceChain, toWormholeFormat(address(token)));
+        assertEq(IERC20(wormholeWrappedToken).balanceOf(recipient), amount);
+        assertTrue(refundAddress.balance > 0);
+    }
 }