refactor withdraw

immutable · Nov 22, 2023 · e14bced · e14bced
1 parent 3097957
commit e14bced
Showing 1 changed file with 84 additions and 59 deletions.
diff --git a/src/child/ChildERC20Bridge.sol b/src/child/ChildERC20Bridge.sol
@@ -298,81 +298,106 @@ contract ChildERC20Bridge is BridgeRoles, IChildERC20BridgeErrors, IChildERC20Br
         }
 
         address rootToken;
-        uint256 feeAmount = msg.value;
         if (childTokenAddr == NATIVE_IMX) {
-            // Native IMX.
-            if (msg.value < amount) {
-                revert InsufficientValue();
-            }
-
-            feeAmount = msg.value - amount;
-            rootToken = rootIMXToken;
+            rootToken = _withdrawIMX(amount);
         } else if (childTokenAddr == wIMXToken) {
-            // Wrapped IMX.
-            // Transfer and unwrap IMX.
-            uint256 expectedBalance = address(this).balance + amount;
+            rootToken = _withdrawWIMX(amount);
+        } else if (childTokenAddr == childETHToken) {
+            rootToken = _withdrawETH(amount);
+        } else {
+            rootToken = _withdrawERC20(rootToken, childTokenAddr, amount);
+        }
 
-            IWIMX wIMX = IWIMX(wIMXToken);
-            if (!wIMX.transferFrom(msg.sender, address(this), amount)) {
-                revert TransferWIMXFailed();
-            }
-            wIMX.withdraw(amount);
+        // Encode the message payload
+        bytes memory payload = abi.encode(WITHDRAW_SIG, rootToken, msg.sender, receiver, amount);
 
-            if (address(this).balance != expectedBalance) {
-                revert BalanceInvariantCheckFailed(address(this).balance, expectedBalance);
-            }
+        // Account for fee amount
+        uint256 feeAmount = (childTokenAddr == NATIVE_IMX) ? msg.value - amount : msg.value;
 
-            rootToken = rootIMXToken;
-        } else if (childTokenAddr == childETHToken) {
-            // Wrapped ETH.
-            IChildERC20 childToken = IChildERC20(childTokenAddr);
-            rootToken = NATIVE_ETH;
+        // Send the message to the bridge adaptor and up to root chain
+        bridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
 
-            if (!childToken.burn(msg.sender, amount)) {
-                revert BurnFailed();
-            }
-        } else {
-            // Other ERC20 Tokens
-            IChildERC20 childToken = IChildERC20(childTokenAddr);
+        _emitWithdrawEvent(rootToken, childTokenAddr, msg.sender, receiver, amount);
+    }
 
-            if (address(childToken).code.length == 0) {
-                revert EmptyTokenContract();
-            }
-            rootToken = childToken.rootToken();
+    function _withdrawIMX(uint256 amount) private returns (address) {
+        if (msg.value < amount) {
+            revert InsufficientValue();
+        }
+        return rootIMXToken;
+    }
 
-            if (rootTokenToChildToken[rootToken] != address(childToken)) {
-                revert NotMapped();
-            }
+    function _withdrawETH(uint256 amount) private returns (address) {
+        if (!IChildERC20(childETHToken).burn(msg.sender, amount)) {
+            revert BurnFailed();
+        }
+        return NATIVE_ETH;
+    }
 
-            // A mapped token should never have root token unset
-            if (rootToken == address(0)) {
-                revert ZeroAddressRootToken();
-            }
+    function _withdrawWIMX(uint256 amount) private returns (address) {
+        // Calculate expected balance
+        uint256 expectedBalance = address(this).balance + amount;
 
-            // A mapped token should never have the bridge unset
-            if (childToken.bridge() != address(this)) {
-                revert IncorrectBridgeAddress();
-            }
+        IWIMX wIMX = IWIMX(wIMXToken);
 
-            if (!childToken.burn(msg.sender, amount)) {
-                revert BurnFailed();
-            }
+        // Transfer to contract
+        if (!wIMX.transferFrom(msg.sender, address(this), amount)) {
+            revert TransferWIMXFailed();
         }
 
-        // Encode the message payload
-        bytes memory payload = abi.encode(WITHDRAW_SIG, rootToken, msg.sender, receiver, amount);
+        // Withdraw
+        wIMX.withdraw(amount);
+
+        // Assert balance
+        if (address(this).balance != expectedBalance) {
+            revert BalanceInvariantCheckFailed(address(this).balance, expectedBalance);
+        }
 
-        // Send the message to the bridge adaptor and up to root chain
-        bridgeAdaptor.sendMessage{value: feeAmount}(payload, msg.sender);
+        return rootIMXToken;
+    }
 
-        if (childTokenAddr == NATIVE_IMX) {
-            emit ChildChainNativeIMXWithdraw(rootToken, msg.sender, receiver, amount);
-        } else if (childTokenAddr == wIMXToken) {
-            emit ChildChainWrappedIMXWithdraw(rootToken, msg.sender, receiver, amount);
-        } else if (childTokenAddr == childETHToken) {
-            emit ChildChainEthWithdraw(msg.sender, receiver, amount);
+    function _withdrawERC20(address rootToken, address childToken, uint256 amount) private returns (address) {
+        // Validate code existance
+        if (address(childToken).code.length == 0) {
+            revert EmptyTokenContract();
+        }
+
+        rootToken = IChildERC20(childToken).rootToken();
+
+        // Assert mapping
+        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 (IChildERC20(childToken).bridge() != address(this)) {
+            revert IncorrectBridgeAddress();
+        }
+
+        // Burn tokens
+        if (!IChildERC20(childToken).burn(msg.sender, amount)) {
+            revert BurnFailed();
+        }
+
+        return rootToken;
+    }
+
+    function _emitWithdrawEvent(address rootToken, address childToken, address sender, address receiver, uint256 amount)
+        private
+    {
+        if (childToken == NATIVE_IMX) {
+            emit ChildChainNativeIMXWithdraw(rootToken, sender, receiver, amount);
+        } else if (childToken == wIMXToken) {
+            emit ChildChainWrappedIMXWithdraw(rootToken, sender, receiver, amount);
+        } else if (childToken == childETHToken) {
+            emit ChildChainEthWithdraw(sender, receiver, amount);
         } else {
-            emit ChildChainERC20Withdraw(rootToken, childTokenAddr, msg.sender, receiver, amount);
+            emit ChildChainERC20Withdraw(rootToken, childToken, sender, receiver, amount);
         }
     }