diff --git a/foundry/src/executors/UniswapV2Executor.sol b/foundry/src/executors/UniswapV2Executor.sol index 7239a8a..cddad3f 100644 --- a/foundry/src/executors/UniswapV2Executor.sol +++ b/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(); + } + } } diff --git a/foundry/src/executors/UniswapV3Executor.sol b/foundry/src/executors/UniswapV3Executor.sol index d1719fe..431e2d7 100644 --- a/foundry/src/executors/UniswapV3Executor.sol +++ b/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(); + } + } } diff --git a/foundry/test/LibPrefixLengthEncodedByteArray.t.sol b/foundry/test/LibPrefixLengthEncodedByteArray.t.sol index 811dcc1..ad13463 100644 --- a/foundry/test/LibPrefixLengthEncodedByteArray.t.sol +++ b/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"; diff --git a/foundry/test/executors/UniswapV2Executor.t.sol b/foundry/test/executors/UniswapV2Executor.t.sol index f377f5a..01722eb 100644 --- a/foundry/test/executors/UniswapV2Executor.t.sol +++ b/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); + } } diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index 86b999b..3e39c7c 100644 --- a/foundry/test/executors/UniswapV3Executor.t.sol +++ b/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,