test: add target verification tests for usv2, usv3

2025-02-21 22:49:10 +05:30
parent 7936ba1c94
commit 2f1507dd0e
5 changed files with 144 additions and 67 deletions
--- a/foundry/src/executors/UniswapV2Executor.sol
+++ b/foundry/src/executors/UniswapV2Executor.sol
@@ -15,10 +15,11 @@ contract UniswapV2Executor is IExecutor {
        0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f;

    // slither-disable-next-line locked-ether
-    function swap(
-        uint256 givenAmount,
-        bytes calldata data
-    ) external payable returns (uint256 calculatedAmount) {
+    function swap(uint256 givenAmount, bytes calldata data)
+        external
+        payable
+        returns (uint256 calculatedAmount)
+    {
        address target;
        address receiver;
        bool zeroForOne;
@@ -40,9 +41,7 @@ contract UniswapV2Executor is IExecutor {
        }
    }

-    function _decodeData(
-        bytes calldata data
-    )
+    function _decodeData(bytes calldata data)
        internal
        pure
        returns (
@@ -61,20 +60,20 @@ contract UniswapV2Executor is IExecutor {
        zeroForOne = uint8(data[60]) > 0;
    }

-    function _getAmountOut(
-        address target,
-        uint256 amountIn,
-        bool zeroForOne
-    ) internal view returns (uint256 amount) {
+    function _getAmountOut(address target, uint256 amountIn, bool zeroForOne)
+        internal
+        view
+        returns (uint256 amount)
+    {
        IUniswapV2Pair pair = IUniswapV2Pair(target);
        uint112 reserveIn;
        uint112 reserveOut;
        if (zeroForOne) {
            // slither-disable-next-line unused-return
-            (reserveIn, reserveOut, ) = pair.getReserves();
+            (reserveIn, reserveOut,) = pair.getReserves();
        } else {
            // slither-disable-next-line unused-return
-            (reserveOut, reserveIn, ) = pair.getReserves();
+            (reserveOut, reserveIn,) = pair.getReserves();
        }

        require(reserveIn > 0 && reserveOut > 0, "L");
@@ -84,9 +83,11 @@ contract UniswapV2Executor is IExecutor {
        amount = numerator / denominator;
    }

-    function _computePairAddress(
-        address target
-    ) internal view returns (address pair) {
+    function _computePairAddress(address target)
+        internal
+        view
+        returns (address pair)
+    {
        address token0 = IUniswapV2Pair(target).token0();
        address token1 = IUniswapV2Pair(target).token1();
        bytes32 salt = keccak256(abi.encodePacked(token0, token1));
--- a/foundry/src/executors/UniswapV3Executor.sol
+++ b/foundry/src/executors/UniswapV3Executor.sol
@@ -30,10 +30,11 @@ contract UniswapV3Executor is IExecutor, ICallback {
    }

    // slither-disable-next-line locked-ether
-    function swap(
-        uint256 amountIn,
-        bytes calldata data
-    ) external payable returns (uint256 amountOut) {
+    function swap(uint256 amountIn, bytes calldata data)
+        external
+        payable
+        returns (uint256 amountOut)
+    {
        (
            address tokenIn,
            address tokenOut,
@@ -71,9 +72,10 @@ contract UniswapV3Executor is IExecutor, ICallback {
        }
    }

-    function handleCallback(
-        bytes calldata msgData
-    ) public returns (bytes memory result) {
+    function handleCallback(bytes calldata msgData)
+        public
+        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(
-            msgData[:64],
-            (int256, int256)
-        );
+        (int256 amount0Delta, int256 amount1Delta) =
+            abi.decode(msgData[:64], (int256, int256));

        address tokenIn = address(bytes20(msgData[128:148]));

        verifyCallback(msgData[128:]);

-        uint256 amountOwed = amount0Delta > 0
-            ? uint256(amount0Delta)
-            : uint256(amount1Delta);
+        uint256 amountOwed =
+            amount0Delta > 0 ? uint256(amount0Delta) : 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(
-            factory,
-            tokenIn,
-            tokenOut,
-            poolFee
-        );
+        CallbackValidationV2.verifyCallback(factory, tokenIn, tokenOut, poolFee);
    }

    function uniswapV3SwapCallback(
-        int256 /* amount0Delta */,
-        int256 /* amount1Delta */,
+        int256, /* amount0Delta */
+        int256, /* amount1Delta */
        bytes calldata /* data */
    ) external {
        uint256 dataOffset = 4 + 32 + 32 + 32; // Skip selector + 2 ints + data_offset
-        uint256 dataLength = uint256(
-            bytes32(msg.data[dataOffset:dataOffset + 32])
-        );
+        uint256 dataLength =
+            uint256(bytes32(msg.data[dataOffset:dataOffset + 32]));

        bytes calldata fullData = msg.data[4:dataOffset + 32 + dataLength];

        handleCallback(fullData);
    }

-    function _decodeData(
-        bytes calldata data
-    )
+    function _decodeData(bytes calldata data)
        internal
        pure
        returns (
@@ -152,29 +143,23 @@ contract UniswapV3Executor is IExecutor, ICallback {
        zeroForOne = uint8(data[83]) > 0;
    }

-    function _makeV3CallbackData(
-        address tokenIn,
-        address tokenOut,
-        uint24 fee
-    ) internal view returns (bytes memory) {
-        return
-            abi.encodePacked(
-                tokenIn,
-                tokenOut,
-                fee,
-                self,
-                ICallback.handleCallback.selector
-            );
+    function _makeV3CallbackData(address tokenIn, address tokenOut, uint24 fee)
+        internal
+        view
+        returns (bytes memory)
+    {
+        return abi.encodePacked(
+            tokenIn, tokenOut, fee, self, ICallback.handleCallback.selector
+        );
    }

-    function _computePairAddress(
-        address tokenA,
-        address tokenB,
-        uint24 fee
-    ) internal view returns (address pool) {
-        (address token0, address token1) = tokenA < tokenB
-            ? (tokenA, tokenB)
-            : (tokenB, tokenA);
+    function _computePairAddress(address tokenA, address tokenB, uint24 fee)
+        internal
+        view
+        returns (address pool)
+    {
+        (address token0, address token1) =
+            tokenA < tokenB ? (tokenA, tokenB) : (tokenB, tokenA);
        pool = address(
            uint160(
                uint256(
--- a/foundry/test/executors/UniswapV2Executor.t.sol
+++ b/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);
+    }
 }
--- a/foundry/test/executors/UniswapV3Executor.t.sol
+++ b/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,
--- a/foundry/test/mock/MockUniswapV2Pool.sol
+++ b/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;
+    }
+}