🍬 WETH unwrap

NaniDAO · Feb 14, 2024 · 81e8c3b · 81e8c3b
1 parent fe6cab8
commit 81e8c3b
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 31 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -1,17 +1,18 @@
 IETest:testBalanceInERC20() (gas: 49207)
-IETest:testBalanceInETH() (gas: 40333)
+IETest:testBalanceInETH() (gas: 40311)
 IETest:testCommandSendERC0() (gas: 102992)
 IETest:testCommandSendETH() (gas: 69474)
 IETest:testCommandSendUSDC() (gas: 135722)
-IETest:testCommandSwapDAI() (gas: 101991)
-IETest:testCommandSwapETH() (gas: 113467)
-IETest:testCommandSwapUSDC() (gas: 157045)
-IETest:testDeploy() (gas: 2599275)
-IETest:testENSNameOwnership() (gas: 84088)
-IETest:testIENameSetting() (gas: 8142)
+IETest:testCommandSwapDAI() (gas: 101807)
+IETest:testCommandSwapETH() (gas: 112631)
+IETest:testCommandSwapForETH() (gas: 123827)
+IETest:testCommandSwapUSDC() (gas: 156200)
+IETest:testDeploy() (gas: 2653217)
+IETest:testENSNameOwnership() (gas: 84066)
+IETest:testIENameSetting() (gas: 8209)
 IETest:testPreviewCommandSendDecimals() (gas: 91984)
-IETest:testPreviewCommandSendUSDC() (gas: 66113)
+IETest:testPreviewCommandSendUSDC() (gas: 66091)
 IETest:testPreviewSend() (gas: 42673)
-IETest:testPreviewSendCommand() (gas: 54504)
+IETest:testPreviewSendCommand() (gas: 54549)
 IETest:testSendETH() (gas: 59695)
 IETest:testTotalSupply() (gas: 14844)
diff --git a/src/IE.sol b/src/IE.sol
@@ -299,45 +299,57 @@ contract IE {
         if (_tokenIn == address(0)) _tokenIn = tokens[tokenIn];
         address _tokenOut = _returnTokenConstant(bytes32(bytes(tokenOut)));
         if (_tokenOut == address(0)) _tokenOut = tokens[tokenOut];
-        bool isETH = _tokenIn == ETH;
-        if (isETH) _tokenIn = WETH;
-        uint256 _amountIn = _stringToUint(amountIn, isETH ? 18 : _tokenIn.readDecimals());
+        bool ETHIn = _tokenIn == ETH;
+        bool ETHOut = _tokenOut == ETH;
+        if (ETHIn) _tokenIn = WETH;
+        if (ETHOut) _tokenOut = WETH;
+        uint256 _amountIn = _stringToUint(amountIn, ETHIn ? 18 : _tokenIn.readDecimals());
         if (_amountIn >= 1 << 255) revert Overflow();
         (address pool, bool zeroForOne) = _computePoolAddress(_tokenIn, _tokenOut);
         ISwapRouter(pool).swap(
-            msg.sender,
+            !ETHOut ? msg.sender : address(this),
             zeroForOne,
             int256(_amountIn),
             zeroForOne ? MIN_SQRT_RATIO_PLUS_ONE : MAX_SQRT_RATIO_MINUS_ONE,
-            abi.encodePacked(isETH, msg.sender, _tokenIn, _tokenOut)
+            abi.encodePacked(ETHIn, ETHOut, msg.sender, _tokenIn, _tokenOut)
         );
     }
 
     /// @dev Fallback `uniswapV3SwapCallback`.
     /// If ETH is swapped, WETH is forwarded.
-    fallback() external {
+    fallback() external payable {
         uint256 amount0Delta;
         uint256 amount1Delta;
-        bool isETH;
+        bool ETHIn;
+        bool ETHOut;
         address payer;
         address tokenIn;
         address tokenOut;
         assembly ("memory-safe") {
             amount0Delta := calldataload(0x4)
             amount1Delta := calldataload(0x24)
-            isETH := byte(0, calldataload(0x84))
-            payer := shr(96, calldataload(add(0x84, 1)))
-            tokenIn := shr(96, calldataload(add(0x84, 21)))
-            tokenOut := shr(96, calldataload(add(0x84, 41)))
+            ETHIn := byte(0, calldataload(0x84))
+            ETHOut := byte(0, calldataload(add(0x84, 1)))
+            payer := shr(96, calldataload(add(0x84, 2)))
+            tokenIn := shr(96, calldataload(add(0x84, 22)))
+            tokenOut := shr(96, calldataload(add(0x84, 42)))
         }
         (address pool, bool zeroForOne) = _computePoolAddress(tokenIn, tokenOut);
-        if (msg.sender != pool) revert Unauthorized();
-        if (isETH) WETH.safeTransferETH(zeroForOne ? amount0Delta : amount1Delta);
-        isETH
+        if (msg.sender != pool) revert Unauthorized(); // Only pair pool can call.
+        if (ETHIn) WETH.safeTransferETH(zeroForOne ? amount0Delta : amount1Delta);
+        ETHIn
             ? WETH.safeTransfer(msg.sender, zeroForOne ? amount0Delta : amount1Delta)
             : tokenIn.safeTransferFrom(payer, msg.sender, zeroForOne ? amount0Delta : amount1Delta);
+        if (ETHOut) {
+            uint256 amount = WETH.balanceOf(address(this));
+            IWETH(WETH).withdraw(amount);
+            payer.safeTransferETH(amount);
+        }
     }
 
+    /// @dev ETH receiver fallback.
+    receive() external payable {}
+
     /// @dev Computes the create2 address for given token pair.
     function _computePoolAddress(address tokenA, address tokenB)
         internal
@@ -633,24 +645,29 @@ contract IE {
     }
 }
 
+/// @dev ENS name resolution helper contracts interface.
+interface IENSHelper {
+    function addr(bytes32) external view returns (address);
+    function owner(bytes32) external view returns (address);
+    function ownerOf(uint256) external view returns (address);
+    function resolver(bytes32) external view returns (address);
+}
+
 /// @dev Simple token transfer interface.
 interface IToken {
     function transfer(address, uint256) external returns (bool);
 }
 
+/// @dev Simple wrapped ether (WETH) token interface.
+interface IWETH {
+    function withdraw(uint256) external;
+}
+
 /// @notice Simple calldata executor interface.
 interface IExecutor {
     function execute(address, uint256, bytes calldata) external payable returns (bytes memory);
 }
 
-/// @dev ENS name resolution helper contracts interface.
-interface IENSHelper {
-    function addr(bytes32) external view returns (address);
-    function owner(bytes32) external view returns (address);
-    function ownerOf(uint256) external view returns (address);
-    function resolver(bytes32) external view returns (address);
-}
-
 /// @dev Simple Uniswap V3 swapping interface.
 interface ISwapRouter {
     function swap(address, bool, int256, uint160, bytes calldata)

diff --git a/test/IE.t.sol b/test/IE.t.sol
@@ -177,6 +177,17 @@ contract IETest is Test {
         ie.command{value: 10 ether}("swap 10 eth for dai");
     }
 
+    function testCommandSwapForETH() public payable {
+        uint256 startBalETH = DAI_WHALE.balance;
+        uint256 startBalDAI = IERC20(DAI).balanceOf(DAI_WHALE);
+        vm.prank(DAI_WHALE);
+        IERC20(DAI).approve(address(ie), 100 ether);
+        vm.prank(DAI_WHALE);
+        ie.command("swap 100 dai for eth");
+        assert(startBalETH < DAI_WHALE.balance);
+        assertEq(startBalDAI - 100 ether, IERC20(DAI).balanceOf(DAI_WHALE));
+    }
+
     function testCommandSwapDAI() public payable {
         vm.prank(DAI_WHALE);
         IERC20(DAI).approve(address(ie), 100 ether);