fix: Pay in LZToken (SC-1152)

src/forwarders/LZForwarder.sol

@@ -1,6 +1,9 @@
 // SPDX-License-Identifier: AGPL-3.0-or-later
 pragma solidity ^0.8.0;
 
+import { IERC20 }    from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { SafeERC20 } from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
+
 struct MessagingParams {
     uint32  dstEid;
     bytes32 receiver;
@@ -21,10 +24,12 @@ struct MessagingFee {
 }
 
 interface ILayerZeroEndpointV2 {
+    function lzToken() external view returns (address);
     function send(
         MessagingParams calldata _params,
         address                  _refundAddress
     ) external payable returns (MessagingReceipt memory);
+    function setLzToken(address _lzToken) external;
     function quote(
         MessagingParams calldata _params,
         address                  _sender
@@ -33,6 +38,8 @@ interface ILayerZeroEndpointV2 {
 
 library LZForwarder {
 
+    error LzTokenUnavailable();
+
     uint32 public constant ENDPOINT_ID_BASE     = 30184;
     uint32 public constant ENDPOINT_ID_BNB      = 30102;
     uint32 public constant ENDPOINT_ID_ETHEREUM = 30101;
@@ -51,19 +58,30 @@ library LZForwarder {
         ILayerZeroEndpointV2 endpoint,
         bytes         memory _message,
         bytes         memory _options,
-        address              _refundAddress
+        address              _refundAddress,
+        bool                 _payInLzToken
     ) internal {
         MessagingParams memory params = MessagingParams({
             dstEid:       _dstEid,
             receiver:     _receiver,
             message:      _message,
             options:      _options,
-            payInLzToken: false
+            payInLzToken: _payInLzToken
         });
 
         MessagingFee memory fee = endpoint.quote(params, address(this));
+        if (fee.lzTokenFee > 0) _payLzToken(endpoint, fee.lzTokenFee);
 
         endpoint.send{ value: fee.nativeFee }(params, _refundAddress);
     }
 
+    function _payLzToken(ILayerZeroEndpointV2 endpoint, uint256 _lzTokenFee) internal {
+        // @dev Cannot cache the token because it is not immutable in the endpoint.
+        address lzToken = endpoint.lzToken();
+        if (lzToken == address(0)) revert LzTokenUnavailable();
+
+        // Pay LZ token fee by sending tokens to the endpoint.
+        SafeERC20.safeTransfer(IERC20(lzToken), address(endpoint), _lzTokenFee);
+    }
+
 }

src/receivers/LZReceiver.sol

@@ -4,13 +4,13 @@ pragma solidity ^0.8.0;
 import { Address } from "openzeppelin-contracts/contracts/utils/Address.sol";
 import { Ownable } from "openzeppelin-contracts/contracts/access/Ownable.sol";
 
-import { OApp, Origin } from "layerzerolabs/oapp-evm/contracts/oapp/OApp.sol";
+import { OAppReceiver, Origin, OAppCore } from "layerzerolabs/oapp-evm/contracts/oapp/OApp.sol";
 
 /**
  * @title  LZReceiver
  * @notice Receive messages from LayerZero-style bridge.
  */
-contract LZReceiver is OApp {
+contract LZReceiver is OAppReceiver {
 
     using Address for address;
 
@@ -27,7 +27,7 @@ contract LZReceiver is OApp {
         address _target,
         address _delegate,
         address _owner
-    ) OApp(_destinationEndpoint, _delegate) Ownable(_owner) {
+    ) OAppCore(_destinationEndpoint, _delegate) Ownable(_owner) {
         target          = _target;
         sourceAuthority = _sourceAuthority;
         srcEid          = _srcEid;

test/IntegrationBase.t.sol

@@ -47,7 +47,7 @@ abstract contract IntegrationBaseTest is Test {
 
     Bridge bridge;
 
-    function setUp() public {
+    function setUp() public virtual {
         source = getChain("mainnet").createFork();
     }
 

test/LZIntegration.t.sol

@@ -199,7 +199,8 @@ contract LZIntegrationTest is IntegrationBaseTest {
             ILayerZeroEndpointV2(bridge.sourceCrossChainMessenger),
             message,
             options,
-            sourceAuthority
+            sourceAuthority,
+            false
         );
     }
 
@@ -214,7 +215,8 @@ contract LZIntegrationTest is IntegrationBaseTest {
             ILayerZeroEndpointV2(bridge.destinationCrossChainMessenger),
             message,
             options,
-            destinationAuthority
+            destinationAuthority,
+            false
         );
     }
 

test/LZIntegrationWithLZToken.t.sol

@@ -0,0 +1,140 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+pragma solidity >=0.8.0;
+
+import { IERC20 } from "forge-std/interfaces/IERC20.sol";
+
+import { OptionsBuilder } from "layerzerolabs/oapp-evm/contracts/oapp/libs/OptionsBuilder.sol";
+
+import { LZBridgeTesting }                   from "src/testing/bridges/LZBridgeTesting.sol";
+import { LZForwarder, ILayerZeroEndpointV2 } from "src/forwarders/LZForwarder.sol";
+import { LZReceiver, Origin }                from "src/receivers/LZReceiver.sol";
+import { RecordedLogs }                      from "src/testing/utils/RecordedLogs.sol";
+
+import "./IntegrationBase.t.sol";
+
+interface ITreasury {
+    function setLzTokenEnabled(bool _lzTokenEnabled) external;
+    function setLzTokenFee(uint256 _lzTokenFee) external;
+}
+
+contract LZIntegrationTestWithLZToken is IntegrationBaseTest {
+
+    using DomainHelpers   for *;
+    using LZBridgeTesting for *;
+    using OptionsBuilder  for bytes;
+
+    uint32 sourceEndpointId = LZForwarder.ENDPOINT_ID_ETHEREUM;
+    uint32 destinationEndpointId;
+
+    address sourceEndpoint = LZForwarder.ENDPOINT_ETHEREUM;
+    address destinationEndpoint;
+
+    address lzToken  = 0x6985884C4392D348587B19cb9eAAf157F13271cd;
+    address lzOwner  = 0xBe010A7e3686FdF65E93344ab664D065A0B02478;
+    address treasury = 0x5ebB3f2feaA15271101a927869B3A56837e73056;
+
+    Domain destination2;
+    Bridge bridge2;
+
+    function setUp() public override {
+        super.setUp();
+
+        source.selectFork();
+
+        vm.startPrank(lzOwner);
+        ILayerZeroEndpointV2(sourceEndpoint).setLzToken(lzToken);
+        ITreasury(treasury).setLzTokenEnabled(true);
+        ITreasury(treasury).setLzTokenFee(1e18);
+        vm.stopPrank();
+    }
+
+    function test_base() public {
+        destinationEndpointId = LZForwarder.ENDPOINT_ID_BASE;
+        destinationEndpoint   = LZForwarder.ENDPOINT_BASE;
+
+        runCrossChainTests(getChain("base").createFork());
+    }
+
+    function test_binance() public {
+        destinationEndpointId = LZForwarder.ENDPOINT_ID_BNB;
+        destinationEndpoint   = LZForwarder.ENDPOINT_BNB;
+
+        runCrossChainTests(getChain("bnb_smart_chain").createFork());
+    }
+
+    function initSourceReceiver() internal override returns (address) {
+        return address(new LZReceiver(
+            sourceEndpoint,
+            destinationEndpointId,
+            bytes32(uint256(uint160(destinationAuthority))),
+            address(moSource),
+            makeAddr("delegate"),
+            makeAddr("owner")
+        ));
+    }
+
+    function initDestinationReceiver() internal override returns (address) {
+        return address(new LZReceiver(
+            destinationEndpoint,
+            sourceEndpointId,
+            bytes32(uint256(uint160(sourceAuthority))),
+            address(moDestination),
+            makeAddr("delegate"),
+            makeAddr("owner")
+        ));
+    }
+
+    function initBridgeTesting() internal override returns (Bridge memory) {
+        return LZBridgeTesting.createLZBridge(source, destination);
+    }
+
+    function queueSourceToDestination(bytes memory message) internal override {
+        // Gas to queue message
+        vm.deal(sourceAuthority, 1 ether);
+        deal(lzToken, sourceAuthority, 1 ether);
+
+        bytes memory options = OptionsBuilder.newOptions().addExecutorLzReceiveOption(200_000, 0);
+
+        assertEq(IERC20(lzToken).balanceOf(address(sourceAuthority)), 1 ether);
+        assertEq(address(sourceAuthority).balance,                    1 ether);
+
+        LZForwarder.sendMessage(
+            destinationEndpointId,
+            bytes32(uint256(uint160(destinationReceiver))),
+            ILayerZeroEndpointV2(bridge.sourceCrossChainMessenger),
+            message,
+            options,
+            sourceAuthority,
+            true
+        );
+
+        // LZ token and ETH spent
+        assertLt(IERC20(lzToken).balanceOf(address(sourceAuthority)), 1 ether);
+        assertLt(address(sourceAuthority).balance,                    1 ether);
+    }
+
+    function queueDestinationToSource(bytes memory message) internal override {
+        vm.deal(destinationAuthority, 1 ether); // Gas to queue message
+
+        bytes memory options = OptionsBuilder.newOptions().addExecutorLzReceiveOption(200_000, 0);
+
+        LZForwarder.sendMessage(
+            sourceEndpointId,
+            bytes32(uint256(uint160(sourceReceiver))),
+            ILayerZeroEndpointV2(bridge.destinationCrossChainMessenger),
+            message,
+            options,
+            destinationAuthority,
+            false
+        );
+    }
+
+    function relaySourceToDestination() internal override {
+        bridge.relayMessagesToDestination(true, sourceAuthority, destinationReceiver);
+    }
+
+    function relayDestinationToSource() internal override {
+        bridge.relayMessagesToSource(true, destinationAuthority, sourceReceiver);
+    }
+
+}