IMX Withdraw L2

src/child/ChildERC20Bridge.sol

@@ -40,6 +40,7 @@ contract ChildERC20Bridge is
     bytes32 public constant DEPOSIT_SIG = keccak256("DEPOSIT");
     bytes32 public constant WITHDRAW_SIG = keccak256("WITHDRAW");
     address public constant NATIVE_ETH = address(0xeee);
+    address public constant NATIVE_IMX = address(0xfff);
 
     IChildERC20BridgeAdaptor public bridgeAdaptor;
 
@@ -134,29 +135,65 @@ contract ChildERC20Bridge is
         _withdraw(childToken, receiver, amount);
     }
 
-    function _withdraw(IChildERC20 childToken, address receiver, uint256 amount) private {
-        if (address(childToken).code.length == 0) {
-            revert EmptyTokenContract();
-        }
+    function withdrawIMX(uint256 amount) external payable {
+        _withdrawIMX(msg.sender, amount);
+    }
 
-        address rootToken = childToken.rootToken();
+    function withdrawToIMX(address receiver, uint256 amount) external payable {
+        _withdrawIMX(receiver, amount);
+    }
 
-        if (rootTokenToChildToken[rootToken] != address(childToken)) {
-            revert NotMapped();
+    function _withdrawIMX(address receiver, uint256 amount) private {
+        if (msg.value < amount) {
+            revert InsufficientValue();
         }
 
-        // A mapped token should never have root token unset
-        if (rootToken == address(0)) {
-            revert ZeroAddressRootToken();
+        uint256 expectedBalance = address(this).balance - (msg.value - amount);
+
+        _withdraw(IChildERC20(NATIVE_IMX), receiver, amount);
+
+        if (address(this).balance != expectedBalance) {
+            revert BalanceInvariantCheckFailed(address(this).balance, expectedBalance);
         }
+    }
 
-        // A mapped token should never have the bridge unset
-        if (childToken.bridge() != address(this)) {
-            revert BridgeNotSet();
+    function _withdraw(IChildERC20 childToken, address receiver, uint256 amount) private {
+        if (address(childToken) == address(0)) {
+            revert ZeroAddress();
+        }
+        if (amount == 0) {
+            revert ZeroAmount();
         }
 
-        if (!childToken.burn(msg.sender, amount)) {
-            revert BurnFailed();
+        address rootToken;
+        uint256 feeAmount = msg.value;
+
+        if (address(childToken) == NATIVE_IMX) {
+            feeAmount = msg.value - amount;
+            rootToken = rootIMXToken;
+        } else {
+            if (address(childToken).code.length == 0) {
+                revert EmptyTokenContract();
+            }
+            rootToken = childToken.rootToken();
+
+            if (rootTokenToChildToken[rootToken] != address(childToken)) {
+                revert NotMapped();
+            }
+
+            // A mapped token should never have root token unset
+            if (rootToken == address(0)) {
+                revert ZeroAddressRootToken();
+            }
+
+            // A mapped token should never have the bridge unset
+            if (childToken.bridge() != address(this)) {
+                revert BridgeNotSet();
+            }
+
+            if (!childToken.burn(msg.sender, amount)) {
+                revert BurnFailed();
+            }
         }
 
         // TODO Should we enforce receiver != 0? old poly contracts don't
@@ -165,9 +202,14 @@ contract ChildERC20Bridge is
         bytes memory payload = abi.encode(WITHDRAW_SIG, rootToken, msg.sender, receiver, amount);
 
         // Send the message to the bridge adaptor and up to root chain
-        bridgeAdaptor.sendMessage{value: msg.value}(payload, msg.sender);
 
-        emit ChildChainERC20Withdraw(rootToken, address(childToken), msg.sender, receiver, amount);
+        bridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
+
+        if (address(childToken) == NATIVE_IMX) {
+            emit ChildChainNativeIMXWithdraw(rootToken, msg.sender, receiver, amount);
+        } else {
+            emit ChildChainERC20Withdraw(rootToken, address(childToken), msg.sender, receiver, amount);
+        }
     }
 
     function _mapToken(bytes calldata data) private {

src/interfaces/child/IChildERC20Bridge.sol

@@ -32,6 +32,9 @@ interface IChildERC20BridgeEvents {
         address indexed receiver,
         uint256 amount
     );
+    event ChildChainNativeIMXWithdraw(
+        address indexed rootToken, address depositor, address indexed receiver, uint256 amount
+    );
 
     event ChildChainERC20Deposit(
         address indexed rootToken,
@@ -52,6 +55,10 @@ interface IChildERC20BridgeEvents {
 
 // TODO add parameters to errors if it makes sense
 interface IChildERC20BridgeErrors {
+    /// @notice Error when the amount requested is less than the value sent.
+    error InsufficientValue();
+    /// @notice Error when there is no gas payment received.
+    error ZeroAmount();
     /// @notice Error when the contract to mint had no bytecode.
     error EmptyTokenContract();
     /// @notice Error when the mint operation failed.
@@ -84,4 +91,6 @@ interface IChildERC20BridgeErrors {
     error BridgeNotSet();
     /// @notice Error when a call to the given child token's `burn` function fails.
     error BurnFailed();
+    /// @notice Error when token balance invariant check fails.
+    error BalanceInvariantCheckFailed(uint256 actualBalance, uint256 expectedBalance);
 }

test/integration/child/withdrawals/ChildAxelarBridgeWithdrawIMX.t.sol

@@ -0,0 +1,133 @@
+// SPDX-License-Identifier: Apache 2.0
+pragma solidity ^0.8.21;
+
+import {Test, console2} from "forge-std/Test.sol";
+import {ERC20PresetMinterPauser} from "@openzeppelin/contracts/token/ERC20/presets/ERC20PresetMinterPauser.sol";
+import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
+import {Strings} from "@openzeppelin/contracts/utils/Strings.sol";
+import {MockAxelarGateway} from "../../../../src/test/root/MockAxelarGateway.sol";
+import {MockAxelarGasService} from "../../../../src/test/root/MockAxelarGasService.sol";
+import {ChildERC20Bridge, IChildERC20BridgeEvents} from "../../../../src/child/ChildERC20Bridge.sol";
+import {
+    ChildAxelarBridgeAdaptor,
+    IChildAxelarBridgeAdaptorEvents,
+    IChildAxelarBridgeAdaptorErrors
+} from "../../../../src/child/ChildAxelarBridgeAdaptor.sol";
+import {Utils} from "../../../utils.t.sol";
+import {WETH} from "../../../../src/test/root/WETH.sol";
+import {ChildERC20} from "../../../../src/child/ChildERC20.sol";
+
+contract ChildERC20BridgeWithdrawIMXIntegrationTest is
+    Test,
+    IChildERC20BridgeEvents,
+    IChildAxelarBridgeAdaptorEvents,
+    IChildAxelarBridgeAdaptorErrors,
+    Utils
+{
+    address constant CHILD_BRIDGE = address(3);
+    address constant CHILD_BRIDGE_ADAPTOR = address(4);
+    string constant CHILD_CHAIN_NAME = "test";
+    address constant ROOT_IMX_TOKEN = address(555555);
+    address constant NATIVE_ETH = address(0xeee);
+    address constant WRAPPED_ETH = address(0xddd);
+
+    ChildERC20Bridge public childBridge;
+    ChildAxelarBridgeAdaptor public axelarAdaptor;
+    address public rootToken;
+    address public rootImxToken;
+    ChildERC20 public childTokenTemplate;
+    MockAxelarGasService public axelarGasService;
+    MockAxelarGateway public mockAxelarGateway;
+
+    function setUp() public {
+        (childBridge, axelarAdaptor, rootToken, rootImxToken, childTokenTemplate, axelarGasService, mockAxelarGateway) =
+            childIntegrationSetup();
+    }
+
+    function test_WithdrawIMX_CallsBridgeAdaptor() public {
+        uint256 withdrawFee = 300;
+        uint256 withdrawAmount = 7 ether;
+
+        bytes memory predictedPayload =
+            abi.encode(WITHDRAW_SIG, ROOT_IMX_TOKEN, address(this), address(this), withdrawAmount);
+        vm.expectCall(
+            address(axelarAdaptor),
+            withdrawFee,
+            abi.encodeWithSelector(axelarAdaptor.sendMessage.selector, predictedPayload, address(this))
+        );
+
+        childBridge.withdrawIMX{value: withdrawFee + withdrawAmount}(withdrawAmount);
+    }
+
+    function test_WithdrawIMX_CallsAxelarGateway() public {
+        uint256 withdrawFee = 300;
+        uint256 withdrawAmount = 7 ether;
+
+        bytes memory predictedPayload =
+            abi.encode(WITHDRAW_SIG, ROOT_IMX_TOKEN, address(this), address(this), withdrawAmount);
+        vm.expectCall(
+            address(mockAxelarGateway),
+            0,
+            abi.encodeWithSelector(
+                mockAxelarGateway.callContract.selector,
+                childBridge.rootChain(),
+                childBridge.rootERC20BridgeAdaptor(),
+                predictedPayload
+            )
+        );
+
+        childBridge.withdrawIMX{value: withdrawFee + withdrawAmount}(withdrawAmount);
+    }
+
+    function test_WithdrawIMX_CallsGasService() public {
+        uint256 withdrawFee = 300;
+        uint256 withdrawAmount = 7 ether;
+
+        bytes memory predictedPayload =
+            abi.encode(WITHDRAW_SIG, ROOT_IMX_TOKEN, address(this), address(this), withdrawAmount);
+
+        vm.expectCall(
+            address(axelarGasService),
+            withdrawFee,
+            abi.encodeWithSelector(
+                axelarGasService.payNativeGasForContractCall.selector,
+                address(axelarAdaptor),
+                childBridge.rootChain(),
+                childBridge.rootERC20BridgeAdaptor(),
+                predictedPayload,
+                address(this)
+            )
+        );
+
+        childBridge.withdrawIMX{value: withdrawFee + withdrawAmount}(withdrawAmount);
+    }
+
+    function test_WithdrawIMXEmitsAxelarMessageEvent() public {
+        uint256 withdrawFee = 300;
+        uint256 withdrawAmount = 7 ether;
+
+        bytes memory predictedPayload =
+            abi.encode(WITHDRAW_SIG, ROOT_IMX_TOKEN, address(this), address(this), withdrawAmount);
+
+        vm.expectEmit(address(axelarAdaptor));
+        emit AxelarMessage(childBridge.rootChain(), childBridge.rootERC20BridgeAdaptor(), predictedPayload);
+
+        childBridge.withdrawIMX{value: withdrawFee + withdrawAmount}(withdrawAmount);
+    }
+
+    function test_WithdrawIMX_ReducesBalance() public {
+        uint256 withdrawFee = 300;
+        uint256 withdrawAmount = 7 ether;
+
+        uint256 preBal = address(this).balance;
+        uint256 preGasBal = address(axelarGasService).balance;
+
+        childBridge.withdrawIMX{value: withdrawFee + withdrawAmount}(withdrawAmount);
+
+        uint256 postBal = address(this).balance;
+        uint256 postGasBal = address(axelarGasService).balance;
+
+        assertEq(postBal, preBal - withdrawFee - withdrawAmount, "Balance not reduced");
+        assertEq(postGasBal, preGasBal + withdrawFee, "Gas service not getting paid");
+    }
+}