From 4b77128df23647bbd247be9e6536b68ea9454056 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Sat, 22 Feb 2025 00:26:15 +0530 Subject: [PATCH] chore: rename _computePairAddress to _verifyPairAddress, add fake v2 pool in v2 test file --- foundry/src/executors/UniswapV2Executor.sol | 24 +++++------ foundry/src/executors/UniswapV3Executor.sol | 24 ++++++----- .../test/executors/UniswapV2Executor.t.sol | 43 +++++++++---------- .../test/executors/UniswapV3Executor.t.sol | 27 ++++++------ foundry/test/mock/MockUniswapV2Pool.sol | 13 ------ 5 files changed, 60 insertions(+), 71 deletions(-) delete mode 100644 foundry/test/mock/MockUniswapV2Pool.sol diff --git a/foundry/src/executors/UniswapV2Executor.sol b/foundry/src/executors/UniswapV2Executor.sol index 750bb18..cddad3f 100644 --- a/foundry/src/executors/UniswapV2Executor.sol +++ b/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)) { - revert UniswapV2Executor__InvalidTarget(); - } + _verifyPairAddress(target); + calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne); tokenIn.safeTransfer(target, givenAmount); @@ -83,27 +85,23 @@ contract UniswapV2Executor is IExecutor { amount = numerator / denominator; } - function _computePairAddress(address target) - internal - view - returns (address pair) - { + function _verifyPairAddress(address target) internal view { 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", - FACTORY, - salt, - hex"96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f" + hex"ff", FACTORY, salt, POOL_INIT_CODE_HASH ) ) ) ) ); + if (pair != target) { + revert UniswapV2Executor__InvalidTarget(); + } } } diff --git a/foundry/src/executors/UniswapV3Executor.sol b/foundry/src/executors/UniswapV3Executor.sol index 9caff33..431e2d7 100644 --- a/foundry/src/executors/UniswapV3Executor.sol +++ b/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)) { - revert UniswapV3Executor__InvalidTarget(); - } + _verifyPairAddress(tokenIn, tokenOut, fee, target); int256 amount0; int256 amount1; @@ -153,14 +153,15 @@ contract UniswapV3Executor is IExecutor, ICallback { ); } - function _computePairAddress(address tokenA, address tokenB, uint24 fee) - internal - view - returns (address pool) - { + function _verifyPairAddress( + address tokenA, + address tokenB, + 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(); + } } } diff --git a/foundry/test/executors/UniswapV2Executor.t.sol b/foundry/test/executors/UniswapV2Executor.t.sol index 88973a8..bb8c553 100644 --- a/foundry/test/executors/UniswapV2Executor.t.sol +++ b/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) - external - view - returns (address pair) - { - return _computePairAddress(target); + 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; } } @@ -71,19 +76,14 @@ 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 testVerifyPairAddress() public view { + uniswapV2Exposed.verifyPairAddress(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 test_RevertIf_InvalidTarget() public { + address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR)); + vm.expectRevert(UniswapV2Executor__InvalidTarget.selector); + uniswapV2Exposed.verifyPairAddress(fakePool); } 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(new MockUniswapV2Pool(WETH_ADDR, DAI_ADDR)); + address fakePool = address(new FakeUniswapV2Pool(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); diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index 14a6982..1691c96 100644 --- a/foundry/test/executors/UniswapV3Executor.t.sol +++ b/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) - external - view - returns (address) - { - return _computePairAddress(tokenA, tokenB, fee); + function verifyPairAddress( + address tokenA, + address tokenB, + uint24 fee, + address target + ) external view { + _verifyPairAddress(tokenA, tokenB, fee, target); } } @@ -77,10 +78,10 @@ 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 testVerifyPairAddress() public view { + uniswapV3Exposed.verifyPairAddress( + WETH_ADDR, DAI_ADDR, 3000, 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 ); diff --git a/foundry/test/mock/MockUniswapV2Pool.sol b/foundry/test/mock/MockUniswapV2Pool.sol deleted file mode 100644 index 80a9b92..0000000 --- a/foundry/test/mock/MockUniswapV2Pool.sol +++ /dev/null @@ -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; - } -}