use estimated gas as part of compose message

Author: PeterMPhillips

script/EnsoReceiverDeployer.s.sol

@@ -12,10 +12,7 @@ contract EnsoReceiverDeployer is Script {
     address OWNER = 0x826e0BB2276271eFdF2a500597f37b94f6c153bA;
     address BACKEND_SIGNER = 0xFE503EE14863F6aCEE10BCdc66aC5e2301b3A946;
 
-    function run()
-        public
-        returns (EnsoReceiver implementation, ERC4337CloneFactory factory)
-    {
+    function run() public returns (EnsoReceiver implementation, ERC4337CloneFactory factory) {
         uint256 deployerPrivateKey = vm.envUint("PRIVATE_KEY");
 
         vm.startBroadcast(deployerPrivateKey);

script/LayerZeroDeployer.s.sol

@@ -45,7 +45,7 @@ contract LayerZeroDeployer is Script {
             router = 0xF75584eF6673aD213a685a1B58Cc0330B8eA22Cf;
         }
 
-        lzReceiver = address(new LayerZeroReceiver{ salt: "LayerZeroReceiver" }(endpoint, router, deployer, 100_000));
+        lzReceiver = address(new LayerZeroReceiver{ salt: "LayerZeroReceiver" }(endpoint, router, deployer));
 
         vm.stopBroadcast();
     }

src/bridge/LayerZeroReceiver.sol

@@ -13,13 +13,10 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
     using SafeERC20 for IERC20;
 
     address private constant _NATIVE_ASSET = address(0);
-    uint256 private constant _TRY_CATCH_GAS = 2000;
 
     address public immutable endpoint;
     IEnsoRouter public immutable router;
 
-    uint256 public reserveGas;
-
     mapping(address => bool) public validOFT;
     mapping(address => bool) public validRegistrar;
     mapping(bytes32 => bool) public messageExecuted;
@@ -33,7 +30,7 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
     event ReserveGasUpdated(uint256 amount);
     event FundsCollected(address token, uint256 amount);
 
-    error InsufficientGas(bytes32 guid);
+    error InsufficientGas(bytes32 guid, uint256 estimatedGas, uint256 availableGas);
     error NotEndpoint(address sender);
     error NotRegistrar(address sender);
     error NotSelf();
@@ -45,15 +42,13 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
     error InvalidMsgValue(uint256 actual, uint256 expected);
     error MessageExecuted(bytes32 key);
 
-    constructor(address _endpoint, address _router, address _owner, uint256 _reserveGas) Ownable(_owner) {
+    constructor(address _endpoint, address _router, address _owner) Ownable(_owner) {
         if (_endpoint == address(0)) revert EndpointNotSet();
         if (_router == address(0)) revert RouterNotSet();
         endpoint = _endpoint;
         router = IEnsoRouter(_router);
         validRegistrar[_owner] = true;
         emit RegistrarAdded(_owner);
-        reserveGas = _reserveGas;
-        emit ReserveGasUpdated(_reserveGas);
     }
 
     // layer zero callback
@@ -76,20 +71,16 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
 
         uint256 amount = _message.amountLD();
         bytes memory composeMsg = _message.composeMsg();
-        (address receiver, uint256 nativeDrop, bytes memory shortcutData) = abi.decode(composeMsg, (address, uint256, bytes));
+        (address receiver, uint256 nativeDrop, uint256 estimatedGas, bytes memory shortcutData) =
+            abi.decode(composeMsg, (address, uint256, uint256, bytes));
         if (msg.value != nativeDrop) revert InvalidMsgValue(msg.value, nativeDrop);
-
         uint256 availableGas = gasleft();
-        if (availableGas < reserveGas) revert InsufficientGas(_guid);
+        if (availableGas < estimatedGas) revert InsufficientGas(_guid, estimatedGas, availableGas);
+
         // try to execute shortcut
-        try this.execute{ gas: availableGas - reserveGas }(token, amount, shortcutData, msg.value) {
+        try this.execute(token, amount, shortcutData, msg.value) {
             emit ShortcutExecutionSuccessful(_guid);
         } catch (bytes memory err) {
-            if (err.length == 0 && gasleft() < reserveGas - _TRY_CATCH_GAS) {
-                // assume that the shortcut failed due to an out of gas error,
-                // to discourage griefing we will revert instead of transferring funds.
-                revert InsufficientGas(_guid);
-            }
             // if shortcut fails send funds to receiver
             emit ShortcutExecutionFailed(_guid, err);
             _transfer(token, receiver, amount);
@@ -153,17 +144,8 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
         emit RegistrarRemoved(account);
     }
 
-    function setReserveGas(uint256 amount) external onlyOwner {
-        reserveGas = amount;
-        emit ReserveGasUpdated(amount);
-    }
-
     // sweep funds to the contract owner in order to refund user
-    function sweep(
-        bytes32 messageKey,
-        address token,
-        uint256 amount
-    ) external onlyOwner {
+    function sweep(bytes32 messageKey, address token, uint256 amount) external onlyOwner {
         // message key is passed to block subsequent calls to lzCompose in case a failing message becomes executable.
         // internal message data is not validated in case the message itself is malformed or incorrect
         // (e.g. oft returns incorrect token)
@@ -174,11 +156,7 @@ contract LayerZeroReceiver is Ownable, ILayerZeroComposer {
         emit FundsCollected(token, amount);
     }
 
-    function getMessageKey(
-        address from,
-        bytes32 guid,
-        bytes calldata message
-    ) public pure returns (bytes32) {
+    function getMessageKey(address from, bytes32 guid, bytes calldata message) public pure returns (bytes32) {
         return keccak256(abi.encode(from, guid, message));
     }
 

test/Bridge.t.sol

@@ -46,7 +46,7 @@ contract BridgeTest is Test {
         vm.selectFork(_ethereumFork);
         router = new EnsoRouter();
         shortcuts = EnsoShortcuts(payable(router.shortcuts()));
-        lzReceiver = new LayerZeroReceiver(address(this), address(router), address(this), 100_000);
+        lzReceiver = new LayerZeroReceiver(address(this), address(router), address(this));
 
         address[] memory ofts = new address[](2);
         ofts[0] = ethPool;
@@ -60,7 +60,7 @@ contract BridgeTest is Test {
         uint256 balanceBefore = weth.balanceOf(address(this));
 
         (bytes32[] memory commands, bytes[] memory state) = _buildWethDeposit(ETH_AMOUNT);
-        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, commands, state);
+        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, 0, commands, state);
 
         // transfer funds
         (bool success,) = address(lzReceiver).call{ value: ETH_AMOUNT }("");
@@ -78,7 +78,7 @@ contract BridgeTest is Test {
 
         // TOO MUCH VALUE ATTEMPTED TO TRANSFER
         (bytes32[] memory commands, bytes[] memory state) = _buildWethDeposit(ETH_AMOUNT * 100);
-        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, commands, state);
+        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, 0, commands, state);
 
         // transfer funds
         (bool success,) = address(lzReceiver).call{ value: ETH_AMOUNT }("");
@@ -91,32 +91,21 @@ contract BridgeTest is Test {
         assertEq(balanceBefore, address(this).balance);
     }
 
-    function testEthBridgeWithShortcutOutOfGas() public {
+    function testEthBridgeWithLessThanEstimateGas() public {
         vm.selectFork(_ethereumFork);
 
         (bytes32[] memory commands, bytes[] memory state) = _buildWethDeposit(ETH_AMOUNT);
-        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, commands, state);
+        // exact gas amount needed for execution
+        uint256 estimatedGas = 94_910;
+        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, estimatedGas, commands, state); 
 
         // transfer funds
         (bool success,) = address(lzReceiver).call{ value: ETH_AMOUNT }("");
         if (!success) revert TransferFailed();
         // trigger compose with insufficient gas
-        vm.expectRevert();
-        lzReceiver.lzCompose{ gas: 108_000 }(ethPool, bytes32(0), message, address(0), "");
-    }
-
-    function testEthBridgeWithLessThanReserveGas() public {
-        vm.selectFork(_ethereumFork);
-
-        (bytes32[] memory commands, bytes[] memory state) = _buildWethDeposit(ETH_AMOUNT);
-        bytes memory message = _buildLzComposeMessage(ETH_AMOUNT, 0, commands, state);
-
-        // transfer funds
-        (bool success,) = address(lzReceiver).call{ value: ETH_AMOUNT }("");
-        if (!success) revert TransferFailed();
-        // trigger compose with insufficient gas
-        vm.expectRevert();
-        lzReceiver.lzCompose{ gas: 99_000 }(ethPool, bytes32(0), message, address(0), "");
+        vm.expectRevert(abi.encodeWithSelector(LayerZeroReceiver.InsufficientGas.selector, bytes32(0), estimatedGas, estimatedGas - 1));
+        // exactly 1 less gas than needed for lz compose
+        lzReceiver.lzCompose{ gas: 105_317 }(ethPool, bytes32(0), message, address(0), ""); 
     }
 
     function testUsdcBridge() public {
@@ -125,7 +114,7 @@ contract BridgeTest is Test {
         uint256 balanceBefore = IERC20(usdc).balanceOf(vitalik);
 
         (bytes32[] memory commands, bytes[] memory state) = _buildTransfer(usdc, vitalik, USDC_AMOUNT);
-        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, 0, commands, state);
+        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, 0, 0, commands, state);
 
         // transfer funds
         vm.startPrank(usdcPool);
@@ -145,7 +134,7 @@ contract BridgeTest is Test {
 
         (bytes32[] memory commands, bytes[] memory state) =
             _buildTokenAndValueTransfer(usdc, vitalik, USDC_AMOUNT, ETH_AMOUNT);
-        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, ETH_AMOUNT, commands, state);
+        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, ETH_AMOUNT, 0, commands, state);
 
         // transfer funds
         vm.startPrank(usdcPool);
@@ -168,7 +157,7 @@ contract BridgeTest is Test {
 
         // TOO MUCH VALUE ATTEMPTED TO TRANSFER
         (bytes32[] memory commands, bytes[] memory state) = _buildTransfer(usdc, vitalik, USDC_AMOUNT * 100);
-        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, 0, commands, state);
+        bytes memory message = _buildLzComposeMessage(USDC_AMOUNT, 0, 0, commands, state);
 
         // transfer funds
         vm.startPrank(usdcPool);
@@ -215,6 +204,7 @@ contract BridgeTest is Test {
     function _buildLzComposeMessage(
         uint256 amount,
         uint256 nativeDrop,
+        uint256 estimatedGas,
         bytes32[] memory commands,
         bytes[] memory state
     )
@@ -226,7 +216,7 @@ contract BridgeTest is Test {
         bytes memory shortcutData =
             abi.encodeWithSelector(shortcuts.executeShortcut.selector, bytes32(0), bytes32(0), commands, state);
         // encode callback data
-        bytes memory callbackData = abi.encode(address(this), nativeDrop, shortcutData);
+        bytes memory callbackData = abi.encode(address(this), nativeDrop, estimatedGas, shortcutData);
         // encode message
         message = OFTComposeMsgCodec.encode(
             uint64(0),