fix: add equality check, amountInOrOut check, update _decodeData

foundry/src/executors/UniswapV4Executor.sol

@@ -2,14 +2,9 @@
 pragma solidity ^0.8.26;
 
 import "@interfaces/IExecutor.sol";
-import {
+import {IERC20, SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
-    IERC20,
-    SafeERC20
-} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {IPoolManager} from "@uniswap/v4-core/src/interfaces/IPoolManager.sol";
-import {
+import {Currency, CurrencyLibrary} from "@uniswap/v4-core/src/types/Currency.sol";
-    Currency, CurrencyLibrary
-} from "@uniswap/v4-core/src/types/Currency.sol";
 import {PoolKey} from "@uniswap/v4-core/src/types/PoolKey.sol";
 import {IHooks} from "@uniswap/v4-core/src/interfaces/IHooks.sol";
 import {V4Router} from "@uniswap/v4-periphery/src/V4Router.sol";
@@ -26,50 +21,60 @@ contract UniswapV4Executor is IExecutor, V4Router {
 
     constructor(IPoolManager _poolManager) V4Router(_poolManager) {}
 
-    function swap(uint256, bytes calldata data)
+    function swap(
-        external
+        uint256,
-        payable
+        bytes calldata data
-        returns (uint256 amountOut)
+    ) external payable returns (uint256 amountInOrOut) {
-    {
+        (
-        (address tokenOut, address receiver, bool isExactInput, uint256 amount)
+            address tokenIn,
-        = _decodeData(data);
+            address tokenOut,
+            address receiver,
+            bool isExactInput,
+            uint256 amount
+        ) = _decodeData(data);
 
-        uint256 balanceBefore = IERC20(tokenOut).balanceOf(receiver);
+        uint256 tokenOutBefore = IERC20(tokenOut).balanceOf(receiver);
+        uint256 tokenInBefore = IERC20(tokenIn).balanceOf(address(this));
 
         this.executeActions(data);
 
-        uint256 balanceAfter = IERC20(tokenOut).balanceOf(receiver);
+        uint256 tokenOutAfter = IERC20(tokenOut).balanceOf(receiver);
+        uint256 tokenInAfter = IERC20(tokenIn).balanceOf(address(this));
 
         if (isExactInput) {
-            amountOut = balanceAfter - balanceBefore;
+            amountInOrOut = tokenOutAfter - tokenOutBefore;
         } else {
-            amountOut = amount;
+            amountInOrOut = tokenInBefore - tokenInAfter;
         }
 
-        // Checks if the amountOut is not 0.
+        // slither-disable-next-line incorrect-equality
-        // Slither does not allow strict equality checks.
+        if (amountInOrOut == 0) {
-        if (amountOut < 1) {
             revert UniswapV4Executor__SwapFailed();
         }
-        return amountOut;
+        return amountInOrOut;
     }
 
     function executeActions(bytes calldata actions) public {
         _executeActions(actions);
     }
 
-    function _decodeData(bytes calldata data)
+    function _decodeData(
+        bytes calldata data
+    )
         internal
         pure
         returns (
+            address tokenIn,
             address tokenOut,
             address receiver,
             bool isExactInput,
             uint256 amount
         )
     {
-        (bytes memory actions, bytes[] memory params) =
+        (bytes memory actions, bytes[] memory params) = abi.decode(
-            abi.decode(data, (bytes, bytes[]));
+            data,
+            (bytes, bytes[])
+        );
 
         // First byte of actions determines the swap type
         uint8 action = uint8(bytes1(actions[0]));
@@ -78,46 +83,65 @@ contract UniswapV4Executor is IExecutor, V4Router {
         (, receiver, ) = abi.decode(params[2], (Currency, address, uint256));
 
         if (action == uint8(Actions.SWAP_EXACT_IN_SINGLE)) {
-            IV4Router.ExactInputSingleParams memory swapParams =
+            IV4Router.ExactInputSingleParams memory swapParams = abi.decode(
-                abi.decode(params[0], (IV4Router.ExactInputSingleParams));
+                params[0],
+                (IV4Router.ExactInputSingleParams)
+            );
 
+            tokenIn = swapParams.zeroForOne
+                ? address(uint160(swapParams.poolKey.currency0.toId()))
+                : address(uint160(swapParams.poolKey.currency1.toId()));
             tokenOut = swapParams.zeroForOne
                 ? address(uint160(swapParams.poolKey.currency1.toId()))
                 : address(uint160(swapParams.poolKey.currency0.toId()));
             isExactInput = true;
             amount = swapParams.amountIn;
         } else if (action == uint8(Actions.SWAP_EXACT_OUT_SINGLE)) {
-            IV4Router.ExactOutputSingleParams memory swapParams =
+            IV4Router.ExactOutputSingleParams memory swapParams = abi.decode(
-                abi.decode(params[0], (IV4Router.ExactOutputSingleParams));
+                params[0],
+                (IV4Router.ExactOutputSingleParams)
+            );
 
+            tokenIn = swapParams.zeroForOne
+                ? address(uint160(swapParams.poolKey.currency0.toId()))
+                : address(uint160(swapParams.poolKey.currency1.toId()));
             tokenOut = swapParams.zeroForOne
                 ? address(uint160(swapParams.poolKey.currency1.toId()))
                 : address(uint160(swapParams.poolKey.currency0.toId()));
             isExactInput = false;
             amount = swapParams.amountOut;
         } else if (action == uint8(Actions.SWAP_EXACT_IN)) {
-            IV4Router.ExactInputParams memory swapParams =
+            IV4Router.ExactInputParams memory swapParams = abi.decode(
-                abi.decode(params[0], (IV4Router.ExactInputParams));
+                params[0],
+                (IV4Router.ExactInputParams)
+            );
 
-            PathKey memory lastPath =
+            tokenIn = address(uint160(swapParams.currencyIn.toId()));
-                swapParams.path[swapParams.path.length - 1];
+            PathKey memory lastPath = swapParams.path[
+                swapParams.path.length - 1
+            ];
             tokenOut = address(uint160(lastPath.intermediateCurrency.toId()));
             isExactInput = true;
             amount = swapParams.amountIn;
         } else if (action == uint8(Actions.SWAP_EXACT_OUT)) {
-            IV4Router.ExactOutputParams memory swapParams =
+            IV4Router.ExactOutputParams memory swapParams = abi.decode(
-                abi.decode(params[0], (IV4Router.ExactOutputParams));
+                params[0],
+                (IV4Router.ExactOutputParams)
+            );
 
+            PathKey memory firstPath = swapParams.path[0];
+            tokenIn = address(uint160(firstPath.intermediateCurrency.toId()));
             tokenOut = address(uint160(swapParams.currencyOut.toId()));
             isExactInput = false;
             amount = swapParams.amountOut;
         }
     }
 
-    function _pay(Currency token, address payer, uint256 amount)
+    function _pay(
-        internal
+        Currency token,
-        override
+        address payer,
-    {
+        uint256 amount
+    ) internal override {
         token.transfer(payer, amount);
     }
 

foundry/test/executors/UniswapV4Executor.t.sol

@@ -9,10 +9,13 @@ import {console} from "forge-std/console.sol";
 contract UniswapV4ExecutorExposed is UniswapV4Executor {
     constructor(IPoolManager _poolManager) UniswapV4Executor(_poolManager) {}
 
-    function decodeData(bytes calldata data)
+    function decodeData(
+        bytes calldata data
+    )
         external
         pure
         returns (
+            address tokenIn,
             address tokenOut,
             address receiver,
             bool isExactInput,
@@ -34,8 +37,9 @@ contract UniswapV4ExecutorTest is Test, Constants {
     function setUp() public {
         uint256 forkBlock = 21817316;
         vm.createSelectFork(vm.rpcUrl("mainnet"), forkBlock);
-        uniswapV4Exposed =
+        uniswapV4Exposed = new UniswapV4ExecutorExposed(
-            new UniswapV4ExecutorExposed(IPoolManager(poolManager));
+            IPoolManager(poolManager)
+        );
     }
 
     function testDecodeParams() public view {
@@ -53,9 +57,15 @@ contract UniswapV4ExecutorTest is Test, Constants {
             expectedAmount
         );
 
-        (address tokenOut, address receiver, bool isExactInput, uint256 amount)
+        (
-        = uniswapV4Exposed.decodeData(data);
+            address tokenIn,
+            address tokenOut,
+            address receiver,
+            bool isExactInput,
+            uint256 amount
+        ) = uniswapV4Exposed.decodeData(data);
 
+        assertEq(tokenIn, USDE_ADDR);
         assertEq(tokenOut, USDT_ADDR);
         assertEq(receiver, expectedReceiver);
         assertTrue(isExactInput);
@@ -63,15 +73,21 @@ contract UniswapV4ExecutorTest is Test, Constants {
     }
 
     function testSwap() public {
-        vm.startPrank(BOB);
         uint256 amountIn = 100 ether;
         deal(USDE_ADDR, address(uniswapV4Exposed), amountIn);
         uint256 usdeBalanceBeforePool = USDE.balanceOf(poolManager);
-        uint256 usdeBalanceBeforeSwapExecutor =
+        uint256 usdeBalanceBeforeSwapExecutor = USDE.balanceOf(
-            USDE.balanceOf(address(uniswapV4Exposed));
+            address(uniswapV4Exposed)
+        );
 
         bytes memory data = _encodeExactInputSingle(
-            USDE_ADDR, USDT_ADDR, 100, BOB, true, 1, uint128(amountIn)
+            USDE_ADDR,
+            USDT_ADDR,
+            100,
+            BOB,
+            true,
+            1,
+            uint128(amountIn)
         );
 
         uint256 amountOut = uniswapV4Exposed.swap(amountIn, data);