✔️ Pool check in callback

NaniDAO · Feb 13, 2024 · 4f92f67 · 4f92f67
1 parent bb09ff0
commit 4f92f67
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 34 deletions.
diff --git a/.gas-snapshot b/.gas-snapshot
@@ -1,11 +1,11 @@
-IETest:testBalanceInERC20() (gas: 62602)
+IETest:testBalanceInERC20() (gas: 62613)
 IETest:testBalanceInETH() (gas: 53716)
 IETest:testCommandSendERC0() (gas: 114314)
 IETest:testCommandSendETH() (gas: 80786)
 IETest:testCommandSendUSDC() (gas: 147022)
-IETest:testCommandSwapDAI() (gas: 101295)
-IETest:testCommandSwapUSDC() (gas: 156369)
-IETest:testDeploy() (gas: 2105248)
+IETest:testCommandSwapDAI() (gas: 101091)
+IETest:testCommandSwapUSDC() (gas: 156474)
+IETest:testDeploy() (gas: 2121089)
 IETest:testENSNameFromENSHelper() (gas: 29103)
 IETest:testENSNameOwnership() (gas: 109286)
 IETest:testIENameSetting() (gas: 8187)

diff --git a/src/IE.sol b/src/IE.sol
@@ -290,7 +290,7 @@ contract IE {
             zeroForOne,
             int256(_amountIn),
             zeroForOne ? MIN_SQRT_RATIO_PLUS_ONE : MAX_SQRT_RATIO_MINUS_ONE,
-            abi.encodePacked(isETH, zeroForOne, _tokenIn, msg.sender)
+            abi.encodePacked(isETH, zeroForOne, _tokenIn, msg.sender, pool)
         );
     }
 
@@ -319,25 +319,26 @@ contract IE {
     /// @dev Fallback `uniswapV3SwapCallback`.
     /// If ETH is swapped, WETH is forwarded.
     fallback() external {
-        int256 amount0Delta;
-        int256 amount1Delta;
+        uint256 amount0Delta;
+        uint256 amount1Delta;
         bool isETH;
         bool zeroForOne;
         address tokenIn;
         address payer;
+        address pool;
         assembly ("memory-safe") {
             amount0Delta := calldataload(0x4)
             amount1Delta := calldataload(0x24)
             isETH := byte(0, calldataload(0x84))
             zeroForOne := byte(0, calldataload(add(0x84, 1)))
             tokenIn := shr(96, calldataload(add(0x84, 2)))
             payer := shr(96, calldataload(add(0x84, 22)))
+            pool := shr(96, calldataload(add(0x84, 42)))
+            if iszero(eq(caller(), pool)) { revert(codesize(), 0x00) }
         }
         isETH
-            ? WETH.safeTransfer(msg.sender, uint256(zeroForOne ? amount0Delta : amount1Delta))
-            : tokenIn.safeTransferFrom(
-                payer, msg.sender, uint256(zeroForOne ? amount0Delta : amount1Delta)
-            );
+            ? WETH.safeTransfer(msg.sender, zeroForOne ? amount0Delta : amount1Delta)
+            : tokenIn.safeTransferFrom(payer, msg.sender, zeroForOne ? amount0Delta : amount1Delta);
     }
 
     /// ================== BALANCE & SUPPLY HELPERS ================== ///
@@ -354,7 +355,7 @@ contract IE {
         address _token = _returnConstant(bytes32(bytes(normalized)));
         if (_token == address(0)) _token = tokens[token];
         bool isETH = _token == ETH;
-        balance = isETH ? _name.balance : _balanceOf(_token, _name);
+        balance = isETH ? _name.balance : _token.balanceOf(_name);
         balanceAdjusted = balance / 10 ** (isETH ? 18 : _token.readDecimals());
     }
 
@@ -372,21 +373,6 @@ contract IE {
         supplyAdjusted = supply / 10 ** _token.readDecimals();
     }
 
-    /// @dev Returns the amount of ERC20/721 `token` owned by `account`.
-    function _balanceOf(address token, address account)
-        internal
-        view
-        virtual
-        returns (uint256 amount)
-    {
-        assembly ("memory-safe") {
-            mstore(0x00, 0x70a08231000000000000000000000000) // `balanceOf(address)`.
-            mstore(0x14, account) // Store the `account` argument.
-            if iszero(staticcall(gas(), token, 0x10, 0x24, 0x20, 0x20)) { revert(codesize(), 0x00) }
-            amount := mload(0x20)
-        }
-    }
-
     /// @dev Returns the total supply of ERC20/721 `token`.
     function _totalSupply(address token) internal view virtual returns (uint256 supply) {
         assembly ("memory-safe") {
@@ -548,11 +534,7 @@ interface IENSHelper {
 
 /// @dev Simple Uniswap V3 swapping interface.
 interface ISwapRouter {
-    function swap(
-        address recipient,
-        bool zeroForOne,
-        int256 amountSpecified,
-        uint160 sqrtPriceLimitX96,
-        bytes calldata data
-    ) external returns (int256 amount0, int256 amount1);
+    function swap(address, bool, int256, uint160, bytes calldata)
+        external
+        returns (int256, int256);
 }