test: add target verification tests for usv2, usv3

foundry/src/executors/UniswapV2Executor.sol

@@ -15,10 +15,11 @@ contract UniswapV2Executor is IExecutor {
         0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f;
 
     // slither-disable-next-line locked-ether
-    function swap(
+    function swap(uint256 givenAmount, bytes calldata data)
-        uint256 givenAmount,
+        external
-        bytes calldata data
+        payable
-    ) external payable returns (uint256 calculatedAmount) {
+        returns (uint256 calculatedAmount)
+    {
         address target;
         address receiver;
         bool zeroForOne;
@@ -40,9 +41,7 @@ contract UniswapV2Executor is IExecutor {
         }
     }
 
-    function _decodeData(
+    function _decodeData(bytes calldata data)
-        bytes calldata data
-    )
         internal
         pure
         returns (
@@ -61,11 +60,11 @@ contract UniswapV2Executor is IExecutor {
         zeroForOne = uint8(data[60]) > 0;
     }
 
-    function _getAmountOut(
+    function _getAmountOut(address target, uint256 amountIn, bool zeroForOne)
-        address target,
+        internal
-        uint256 amountIn,
+        view
-        bool zeroForOne
+        returns (uint256 amount)
-    ) internal view returns (uint256 amount) {
+    {
         IUniswapV2Pair pair = IUniswapV2Pair(target);
         uint112 reserveIn;
         uint112 reserveOut;
@@ -84,9 +83,11 @@ contract UniswapV2Executor is IExecutor {
         amount = numerator / denominator;
     }
 
-    function _computePairAddress(
+    function _computePairAddress(address target)
-        address target
+        internal
-    ) internal view returns (address pair) {
+        view
+        returns (address pair)
+    {
         address token0 = IUniswapV2Pair(target).token0();
         address token1 = IUniswapV2Pair(target).token1();
         bytes32 salt = keccak256(abi.encodePacked(token0, token1));

foundry/src/executors/UniswapV3Executor.sol

@@ -30,10 +30,11 @@ contract UniswapV3Executor is IExecutor, ICallback {
     }
 
     // slither-disable-next-line locked-ether
-    function swap(
+    function swap(uint256 amountIn, bytes calldata data)
-        uint256 amountIn,
+        external
-        bytes calldata data
+        payable
-    ) external payable returns (uint256 amountOut) {
+        returns (uint256 amountOut)
+    {
         (
             address tokenIn,
             address tokenOut,
@@ -71,9 +72,10 @@ contract UniswapV3Executor is IExecutor, ICallback {
         }
     }
 
-    function handleCallback(
+    function handleCallback(bytes calldata msgData)
-        bytes calldata msgData
+        public
-    ) public returns (bytes memory result) {
+        returns (bytes memory result)
+    {
         // The data has the following layout:
         // - amount0Delta (32 bytes)
         // - amount1Delta (32 bytes)
@@ -81,18 +83,15 @@ contract UniswapV3Executor is IExecutor, ICallback {
         // - dataLength (32 bytes)
         // - protocolData (variable length)
 
-        (int256 amount0Delta, int256 amount1Delta) = abi.decode(
+        (int256 amount0Delta, int256 amount1Delta) =
-            msgData[:64],
+            abi.decode(msgData[:64], (int256, int256));
-            (int256, int256)
-        );
 
         address tokenIn = address(bytes20(msgData[128:148]));
 
         verifyCallback(msgData[128:]);
 
-        uint256 amountOwed = amount0Delta > 0
+        uint256 amountOwed =
-            ? uint256(amount0Delta)
+            amount0Delta > 0 ? uint256(amount0Delta) : uint256(amount1Delta);
-            : uint256(amount1Delta);
 
         IERC20(tokenIn).safeTransfer(msg.sender, amountOwed);
         return abi.encode(amountOwed, tokenIn);
@@ -104,32 +103,24 @@ contract UniswapV3Executor is IExecutor, ICallback {
         uint24 poolFee = uint24(bytes3(data[40:43]));
 
         // slither-disable-next-line unused-return
-        CallbackValidationV2.verifyCallback(
+        CallbackValidationV2.verifyCallback(factory, tokenIn, tokenOut, poolFee);
-            factory,
-            tokenIn,
-            tokenOut,
-            poolFee
-        );
     }
 
     function uniswapV3SwapCallback(
-        int256 /* amount0Delta */,
+        int256, /* amount0Delta */
-        int256 /* amount1Delta */,
+        int256, /* amount1Delta */
         bytes calldata /* data */
     ) external {
         uint256 dataOffset = 4 + 32 + 32 + 32; // Skip selector + 2 ints + data_offset
-        uint256 dataLength = uint256(
+        uint256 dataLength =
-            bytes32(msg.data[dataOffset:dataOffset + 32])
+            uint256(bytes32(msg.data[dataOffset:dataOffset + 32]));
-        );
 
         bytes calldata fullData = msg.data[4:dataOffset + 32 + dataLength];
 
         handleCallback(fullData);
     }
 
-    function _decodeData(
+    function _decodeData(bytes calldata data)
-        bytes calldata data
-    )
         internal
         pure
         returns (
@@ -152,29 +143,23 @@ contract UniswapV3Executor is IExecutor, ICallback {
         zeroForOne = uint8(data[83]) > 0;
     }
 
-    function _makeV3CallbackData(
+    function _makeV3CallbackData(address tokenIn, address tokenOut, uint24 fee)
-        address tokenIn,
+        internal
-        address tokenOut,
+        view
-        uint24 fee
+        returns (bytes memory)
-    ) internal view returns (bytes memory) {
+    {
-        return
+        return abi.encodePacked(
-            abi.encodePacked(
+            tokenIn, tokenOut, fee, self, ICallback.handleCallback.selector
-                tokenIn,
-                tokenOut,
-                fee,
-                self,
-                ICallback.handleCallback.selector
         );
     }
 
-    function _computePairAddress(
+    function _computePairAddress(address tokenA, address tokenB, uint24 fee)
-        address tokenA,
+        internal
-        address tokenB,
+        view
-        uint24 fee
+        returns (address pool)
-    ) internal view returns (address pool) {
+    {
-        (address token0, address token1) = tokenA < tokenB
+        (address token0, address token1) =
-            ? (tokenA, tokenB)
+            tokenA < tokenB ? (tokenA, tokenB) : (tokenB, tokenA);
-            : (tokenB, tokenA);
         pool = address(
             uint160(
                 uint256(

foundry/test/executors/UniswapV2Executor.t.sol

@@ -4,6 +4,7 @@ pragma solidity ^0.8.26;
 import "@src/executors/UniswapV2Executor.sol";
 import {Test} from "../../lib/forge-std/src/Test.sol";
 import {Constants} from "../Constants.sol";
+import {MockUniswapV2Pool} from "../mock/MockUniswapV2Pool.sol";
 
 contract UniswapV2ExecutorExposed is UniswapV2Executor {
     function decodeParams(bytes calldata data)
@@ -26,6 +27,14 @@ contract UniswapV2ExecutorExposed is UniswapV2Executor {
     {
         return _getAmountOut(target, amountIn, zeroForOne);
     }
+
+    function computePairAddress(address target)
+        external
+        view
+        returns (address pair)
+    {
+        return _computePairAddress(target);
+    }
 }
 
 contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
@@ -62,6 +71,21 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         uniswapV2Exposed.decodeParams(invalidParams);
     }
 
+    function testComputePairAddress() public view {
+        address computedPair =
+            uniswapV2Exposed.computePairAddress(WETH_DAI_POOL);
+        assertEq(computedPair, WETH_DAI_POOL);
+    }
+
+    function testComputePairAddressInvalid() public {
+        address tokenA = WETH_ADDR;
+        address tokenB = DAI_ADDR;
+        address maliciousPool = address(new MockUniswapV2Pool(tokenA, tokenB));
+        address computedPair =
+            uniswapV2Exposed.computePairAddress(maliciousPool);
+        assertNotEq(computedPair, maliciousPool);
+    }
+
     function testAmountOut() public view {
         uint256 amountOut =
             uniswapV2Exposed.getAmountOut(WETH_DAI_POOL, 10 ** 18, false);
@@ -80,7 +104,7 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         assertGe(amountOut, 0);
     }
 
-    function testSwapUniswapV2() public {
+    function testSwap() public {
         uint256 amountIn = 10 ** 18;
         uint256 amountOut = 1847751195973566072891;
         bool zeroForOne = false;
@@ -120,4 +144,17 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         uint256 finalBalance = DAI.balanceOf(BOB);
         assertGe(finalBalance, amountOut);
     }
+
+    function test_RevertIf_InvalidTarget() public {
+        uint256 amountIn = 10 ** 18;
+        bool zeroForOne = false;
+        address maliciousPool =
+            address(new MockUniswapV2Pool(WETH_ADDR, DAI_ADDR));
+        bytes memory protocolData =
+            abi.encodePacked(WETH_ADDR, maliciousPool, BOB, zeroForOne);
+
+        deal(WETH_ADDR, address(uniswapV2Exposed), amountIn);
+        vm.expectRevert(UniswapV2Executor__InvalidTarget.selector);
+        uniswapV2Exposed.swap(amountIn, protocolData);
+    }
 }

foundry/test/executors/UniswapV3Executor.t.sol

@@ -22,6 +22,14 @@ contract UniswapV3ExecutorExposed is UniswapV3Executor {
     {
         return _decodeData(data);
     }
+
+    function computePairAddress(address tokenA, address tokenB, uint24 fee)
+        external
+        view
+        returns (address)
+    {
+        return _computePairAddress(tokenA, tokenB, fee);
+    }
 }
 
 contract UniswapV3ExecutorTest is Test, Constants {
@@ -69,6 +77,20 @@ contract UniswapV3ExecutorTest is Test, Constants {
         uniswapV3Exposed.decodeData(invalidParams);
     }
 
+    function testComputePairAddress() public view {
+        address computedPair =
+            uniswapV3Exposed.computePairAddress(WETH_ADDR, DAI_ADDR, 3000);
+        assertEq(computedPair, DAI_WETH_USV3);
+    }
+
+    function testComputePairAddressInvalid() public view {
+        address maliciousPool = DUMMY; // Contract with malicious behavior
+
+        address computedPair =
+            uniswapV3Exposed.computePairAddress(WETH_ADDR, DAI_ADDR, 3000);
+        assertNotEq(computedPair, maliciousPool);
+    }
+
     function testUSV3Callback() public {
         uint24 poolFee = 3000;
         uint256 amountOwed = 1000000000000000000;
@@ -113,6 +135,25 @@ contract UniswapV3ExecutorTest is Test, Constants {
         assertGe(IERC20(DAI_ADDR).balanceOf(address(this)), expAmountOut);
     }
 
+    function test_RevertIf_InvalidTargetV3() public {
+        uint256 amountIn = 10 ** 18;
+        deal(WETH_ADDR, address(uniswapV3Exposed), amountIn);
+        bool zeroForOne = false;
+        address maliciousPool = DUMMY;
+
+        bytes memory protocolData = abi.encodePacked(
+            WETH_ADDR,
+            DAI_ADDR,
+            uint24(3000),
+            address(this),
+            maliciousPool,
+            zeroForOne
+        );
+
+        vm.expectRevert(UniswapV3Executor__InvalidTarget.selector);
+        uniswapV3Exposed.swap(amountIn, protocolData);
+    }
+
     function encodeUniswapV3Swap(
         address tokenIn,
         address tokenOut,

foundry/test/mock/MockUniswapV2Pool.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: Unlicense
+pragma solidity ^0.8.26;
+
+// Mock for the UniswapV2Pool contract, it is expected to have malicious behavior
+contract MockUniswapV2Pool {
+    address public token0;
+    address public token1;
+
+    constructor(address _tokenA, address _tokenB) {
+        token0 = _tokenA < _tokenB ? _tokenA : _tokenB;
+        token1 = _tokenA < _tokenB ? _tokenB : _tokenA;
+    }
+}