chore: rename _computePairAddress to _verifyPairAddress, add fake v2 pool in v2 test file

foundry/src/executors/UniswapV2Executor.sol

@@ -11,6 +11,9 @@ error UniswapV2Executor__InvalidTarget();
 contract UniswapV2Executor is IExecutor {
     using SafeERC20 for IERC20;
 
+    bytes32 internal constant POOL_INIT_CODE_HASH =
+        0x96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f;
+
     address private constant FACTORY =
         0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f;
 
@@ -27,9 +30,8 @@ contract UniswapV2Executor is IExecutor {
 
         (tokenIn, target, receiver, zeroForOne) = _decodeData(data);
 
-        if (target != _computePairAddress(target)) {
+        _verifyPairAddress(target);
-            revert UniswapV2Executor__InvalidTarget();
+
-        }
         calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne);
         tokenIn.safeTransfer(target, givenAmount);
 
@@ -83,27 +85,23 @@ contract UniswapV2Executor is IExecutor {
         amount = numerator / denominator;
     }
 
-    function _computePairAddress(address target)
+    function _verifyPairAddress(address target) internal view {
-        internal
-        view
-        returns (address pair)
-    {
         address token0 = IUniswapV2Pair(target).token0();
         address token1 = IUniswapV2Pair(target).token1();
         bytes32 salt = keccak256(abi.encodePacked(token0, token1));
-        pair = address(
+        address pair = address(
             uint160(
                 uint256(
                     keccak256(
                         abi.encodePacked(
-                            hex"ff",
+                            hex"ff", FACTORY, salt, POOL_INIT_CODE_HASH
-                            FACTORY,
-                            salt,
-                            hex"96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f"
                         )
                     )
                 )
             )
         );
+        if (pair != target) {
+            revert UniswapV2Executor__InvalidTarget();
+        }
     }
 }

foundry/src/executors/UniswapV3Executor.sol

@@ -14,6 +14,8 @@ error UniswapV3Executor__InvalidTarget();
 contract UniswapV3Executor is IExecutor, ICallback {
     using SafeERC20 for IERC20;
 
+    bytes32 internal constant POOL_INIT_CODE_HASH =
+        0xe34f199b19b2b4f47f68442619d555527d244f78a3297ea89325f843f87b8b54;
     uint160 private constant MIN_SQRT_RATIO = 4295128739;
     uint160 private constant MAX_SQRT_RATIO =
         1461446703485210103287273052203988822378723970342;
@@ -44,9 +46,7 @@ contract UniswapV3Executor is IExecutor, ICallback {
             bool zeroForOne
         ) = _decodeData(data);
 
-        if (target != _computePairAddress(tokenIn, tokenOut, fee)) {
+        _verifyPairAddress(tokenIn, tokenOut, fee, target);
-            revert UniswapV3Executor__InvalidTarget();
-        }
 
         int256 amount0;
         int256 amount1;
@@ -153,14 +153,15 @@ contract UniswapV3Executor is IExecutor, ICallback {
         );
     }
 
-    function _computePairAddress(address tokenA, address tokenB, uint24 fee)
+    function _verifyPairAddress(
-        internal
+        address tokenA,
-        view
+        address tokenB,
-        returns (address pool)
+        uint24 fee,
-    {
+        address target
+    ) internal view {
         (address token0, address token1) =
             tokenA < tokenB ? (tokenA, tokenB) : (tokenB, tokenA);
-        pool = address(
+        address pool = address(
             uint160(
                 uint256(
                     keccak256(
@@ -168,11 +169,14 @@ contract UniswapV3Executor is IExecutor, ICallback {
                             hex"ff",
                             factory,
                             keccak256(abi.encode(token0, token1, fee)),
-                            hex"e34f199b19b2b4f47f68442619d555527d244f78a3297ea89325f843f87b8b54"
+                            POOL_INIT_CODE_HASH
                         )
                     )
                 )
             )
         );
+        if (pool != target) {
+            revert UniswapV3Executor__InvalidTarget();
+        }
     }
 }

foundry/test/executors/UniswapV2Executor.t.sol

@@ -4,7 +4,6 @@ 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)
@@ -28,12 +27,18 @@ contract UniswapV2ExecutorExposed is UniswapV2Executor {
         return _getAmountOut(target, amountIn, zeroForOne);
     }
 
-    function computePairAddress(address target)
+    function verifyPairAddress(address target) external view {
-        external
+        _verifyPairAddress(target);
-        view
+    }
-        returns (address pair)
+}
-    {
+
-        return _computePairAddress(target);
+contract FakeUniswapV2Pool {
+    address public token0;
+    address public token1;
+
+    constructor(address _tokenA, address _tokenB) {
+        token0 = _tokenA < _tokenB ? _tokenA : _tokenB;
+        token1 = _tokenA < _tokenB ? _tokenB : _tokenA;
     }
 }
 
@@ -71,19 +76,14 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         uniswapV2Exposed.decodeParams(invalidParams);
     }
 
-    function testComputePairAddress() public view {
+    function testVerifyPairAddress() public view {
-        address computedPair =
+        uniswapV2Exposed.verifyPairAddress(WETH_DAI_POOL);
-            uniswapV2Exposed.computePairAddress(WETH_DAI_POOL);
-        assertEq(computedPair, WETH_DAI_POOL);
     }
 
-    function testComputePairAddressInvalid() public {
+    function test_RevertIf_InvalidTarget() public {
-        address tokenA = WETH_ADDR;
+        address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR));
-        address tokenB = DAI_ADDR;
+        vm.expectRevert(UniswapV2Executor__InvalidTarget.selector);
-        address maliciousPool = address(new MockUniswapV2Pool(tokenA, tokenB));
+        uniswapV2Exposed.verifyPairAddress(fakePool);
-        address computedPair =
-            uniswapV2Exposed.computePairAddress(maliciousPool);
-        assertNotEq(computedPair, maliciousPool);
     }
 
     function testAmountOut() public view {
@@ -145,13 +145,12 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         assertGe(finalBalance, amountOut);
     }
 
-    function test_RevertIf_InvalidTarget() public {
+    function test_RevertIf_Swap_InvalidTarget() public {
         uint256 amountIn = 10 ** 18;
         bool zeroForOne = false;
-        address maliciousPool =
+        address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR));
-            address(new MockUniswapV2Pool(WETH_ADDR, DAI_ADDR));
         bytes memory protocolData =
-            abi.encodePacked(WETH_ADDR, maliciousPool, BOB, zeroForOne);
+            abi.encodePacked(WETH_ADDR, fakePool, BOB, zeroForOne);
 
         deal(WETH_ADDR, address(uniswapV2Exposed), amountIn);
         vm.expectRevert(UniswapV2Executor__InvalidTarget.selector);

foundry/test/executors/UniswapV3Executor.t.sol

@@ -23,12 +23,13 @@ contract UniswapV3ExecutorExposed is UniswapV3Executor {
         return _decodeData(data);
     }
 
-    function computePairAddress(address tokenA, address tokenB, uint24 fee)
+    function verifyPairAddress(
-        external
+        address tokenA,
-        view
+        address tokenB,
-        returns (address)
+        uint24 fee,
-    {
+        address target
-        return _computePairAddress(tokenA, tokenB, fee);
+    ) external view {
+        _verifyPairAddress(tokenA, tokenB, fee, target);
     }
 }
 
@@ -77,10 +78,10 @@ contract UniswapV3ExecutorTest is Test, Constants {
         uniswapV3Exposed.decodeData(invalidParams);
     }
 
-    function testComputePairAddress() public view {
+    function testVerifyPairAddress() public view {
-        address computedPair =
+        uniswapV3Exposed.verifyPairAddress(
-            uniswapV3Exposed.computePairAddress(WETH_ADDR, DAI_ADDR, 3000);
+            WETH_ADDR, DAI_ADDR, 3000, DAI_WETH_USV3
-        assertEq(computedPair, DAI_WETH_USV3);
+        );
     }
 
     function testUSV3Callback() public {
@@ -127,18 +128,18 @@ contract UniswapV3ExecutorTest is Test, Constants {
         assertGe(IERC20(DAI_ADDR).balanceOf(address(this)), expAmountOut);
     }
 
-    function test_RevertIf_InvalidTargetV3() public {
+    function test_RevertIf_InvalidTarget() public {
         uint256 amountIn = 10 ** 18;
         deal(WETH_ADDR, address(uniswapV3Exposed), amountIn);
         bool zeroForOne = false;
-        address maliciousPool = DUMMY;
+        address fakePool = DUMMY; // Contract with minimal code
 
         bytes memory protocolData = abi.encodePacked(
             WETH_ADDR,
             DAI_ADDR,
             uint24(3000),
             address(this),
-            maliciousPool,
+            fakePool,
             zeroForOne
         );
 

foundry/test/mock/MockUniswapV2Pool.sol

@@ -1,13 +0,0 @@
-// 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;
-    }
-}