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

This commit is contained in:
royvardhan
2025-02-22 00:26:15 +05:30
parent 40bd37a1a4
commit 4b77128df2
5 changed files with 60 additions and 71 deletions

View File

@@ -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();
}
}
}

View File

@@ -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();
}
}
}

View File

@@ -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);

View File

@@ -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
);

View File

@@ -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;
}
}