Merge pull request #80 from propeller-heads/router/hr/ENG-4258-verify-valid-target

foundry/src/executors/UniswapV2Executor.sol

@@ -6,10 +6,17 @@ import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import "@uniswap-v2/contracts/interfaces/IUniswapV2Pair.sol";
 
 error UniswapV2Executor__InvalidDataLength();
+error UniswapV2Executor__InvalidTarget();
 
 contract UniswapV2Executor is IExecutor {
     using SafeERC20 for IERC20;
 
+    bytes32 internal constant POOL_INIT_CODE_HASH =
+        0x96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f;
+
+    address private constant FACTORY =
+        0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f;
+
     // slither-disable-next-line locked-ether
     function swap(uint256 givenAmount, bytes calldata data)
         external
@@ -22,6 +29,9 @@ contract UniswapV2Executor is IExecutor {
         IERC20 tokenIn;
 
         (tokenIn, target, receiver, zeroForOne) = _decodeData(data);
+
+        _verifyPairAddress(target);
+
         calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne);
         tokenIn.safeTransfer(target, givenAmount);
 
@@ -74,4 +84,24 @@ contract UniswapV2Executor is IExecutor {
         uint256 denominator = (uint256(reserveIn) * 1000) + amountInWithFee;
         amount = numerator / denominator;
     }
+
+    function _verifyPairAddress(address target) internal view {
+        address token0 = IUniswapV2Pair(target).token0();
+        address token1 = IUniswapV2Pair(target).token1();
+        bytes32 salt = keccak256(abi.encodePacked(token0, token1));
+        address pair = address(
+            uint160(
+                uint256(
+                    keccak256(
+                        abi.encodePacked(
+                            hex"ff", FACTORY, salt, POOL_INIT_CODE_HASH
+                        )
+                    )
+                )
+            )
+        );
+        if (pair != target) {
+            revert UniswapV2Executor__InvalidTarget();
+        }
+    }
 }

foundry/src/executors/UniswapV3Executor.sol

@@ -9,10 +9,13 @@ import "@interfaces/ICallback.sol";
 
 error UniswapV3Executor__InvalidDataLength();
 error UniswapV3Executor__InvalidFactory();
+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;
@@ -42,6 +45,9 @@ contract UniswapV3Executor is IExecutor, ICallback {
             address target,
             bool zeroForOne
         ) = _decodeData(data);
+
+        _verifyPairAddress(tokenIn, tokenOut, fee, target);
+
         int256 amount0;
         int256 amount1;
         IUniswapV3Pool pool = IUniswapV3Pool(target);
@@ -146,4 +152,31 @@ contract UniswapV3Executor is IExecutor, ICallback {
             tokenIn, tokenOut, fee, self, ICallback.handleCallback.selector
         );
     }
+
+    function _verifyPairAddress(
+        address tokenA,
+        address tokenB,
+        uint24 fee,
+        address target
+    ) internal view {
+        (address token0, address token1) =
+            tokenA < tokenB ? (tokenA, tokenB) : (tokenB, tokenA);
+        address pool = address(
+            uint160(
+                uint256(
+                    keccak256(
+                        abi.encodePacked(
+                            hex"ff",
+                            factory,
+                            keccak256(abi.encode(token0, token1, fee)),
+                            POOL_INIT_CODE_HASH
+                        )
+                    )
+                )
+            )
+        );
+        if (pool != target) {
+            revert UniswapV3Executor__InvalidTarget();
+        }
+    }
 }

foundry/test/LibPrefixLengthEncodedByteArray.t.sol

@@ -51,14 +51,14 @@ contract LibPrefixLengthEncodedByteArrayTest is Test {
         assertEq(this.size(multiple), 3);
     }
 
-    function test_RevertIf_InvalidLength() public {
+    function testInvalidLength() public {
         // Length prefix larger than remaining data
         vm.expectRevert();
         bytes memory invalid = hex"0004414243";
         this.next(invalid);
     }
 
-    function test_RevertIf_IncompletePrefix() public {
+    function testIncompletePrefix() public {
         // Only 1 byte instead of 2 bytes prefix
         vm.expectRevert();
         bytes memory invalid = hex"01";

foundry/test/executors/UniswapV2Executor.t.sol

@@ -26,6 +26,20 @@ contract UniswapV2ExecutorExposed is UniswapV2Executor {
     {
         return _getAmountOut(target, amountIn, zeroForOne);
     }
+
+    function verifyPairAddress(address target) external view {
+        _verifyPairAddress(target);
+    }
+}
+
+contract FakeUniswapV2Pool {
+    address public token0;
+    address public token1;
+
+    constructor(address _tokenA, address _tokenB) {
+        token0 = _tokenA < _tokenB ? _tokenA : _tokenB;
+        token1 = _tokenA < _tokenB ? _tokenB : _tokenA;
+    }
 }
 
 contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
@@ -62,6 +76,16 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         uniswapV2Exposed.decodeParams(invalidParams);
     }
 
+    function testVerifyPairAddress() public view {
+        uniswapV2Exposed.verifyPairAddress(WETH_DAI_POOL);
+    }
+
+    function testInvalidTarget() public {
+        address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR));
+        vm.expectRevert(UniswapV2Executor__InvalidTarget.selector);
+        uniswapV2Exposed.verifyPairAddress(fakePool);
+    }
+
     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,16 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants {
         uint256 finalBalance = DAI.balanceOf(BOB);
         assertGe(finalBalance, amountOut);
     }
+
+    function testSwapFailureInvalidTarget() public {
+        uint256 amountIn = 10 ** 18;
+        bool zeroForOne = false;
+        address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR));
+        bytes memory protocolData =
+            abi.encodePacked(WETH_ADDR, fakePool, 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,15 @@ contract UniswapV3ExecutorExposed is UniswapV3Executor {
     {
         return _decodeData(data);
     }
+
+    function verifyPairAddress(
+        address tokenA,
+        address tokenB,
+        uint24 fee,
+        address target
+    ) external view {
+        _verifyPairAddress(tokenA, tokenB, fee, target);
+    }
 }
 
 contract UniswapV3ExecutorTest is Test, Constants {
@@ -69,6 +78,12 @@ contract UniswapV3ExecutorTest is Test, Constants {
         uniswapV3Exposed.decodeData(invalidParams);
     }
 
+    function testVerifyPairAddress() public view {
+        uniswapV3Exposed.verifyPairAddress(
+            WETH_ADDR, DAI_ADDR, 3000, DAI_WETH_USV3
+        );
+    }
+
     function testUSV3Callback() public {
         uint24 poolFee = 3000;
         uint256 amountOwed = 1000000000000000000;
@@ -113,6 +128,25 @@ contract UniswapV3ExecutorTest is Test, Constants {
         assertGe(IERC20(DAI_ADDR).balanceOf(address(this)), expAmountOut);
     }
 
+    function testSwapFailureInvalidTarget() public {
+        uint256 amountIn = 10 ** 18;
+        deal(WETH_ADDR, address(uniswapV3Exposed), amountIn);
+        bool zeroForOne = false;
+        address fakePool = DUMMY; // Contract with minimal code
+
+        bytes memory protocolData = abi.encodePacked(
+            WETH_ADDR,
+            DAI_ADDR,
+            uint24(3000),
+            address(this),
+            fakePool,
+            zeroForOne
+        );
+
+        vm.expectRevert(UniswapV3Executor__InvalidTarget.selector);
+        uniswapV3Exposed.swap(amountIn, protocolData);
+    }
+
     function encodeUniswapV3Swap(
         address tokenIn,
         address tokenOut,