From 7936ba1c943a616a143dc6f8ecb1da61073b05a8 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Fri, 21 Feb 2025 20:00:38 +0530 Subject: [PATCH 1/5] feat: add target verification for usv2 and usv3 using _computePairAddress --- foundry/src/executors/UniswapV2Executor.sol | 57 +++++++++--- foundry/src/executors/UniswapV3Executor.sol | 96 +++++++++++++++------ 2 files changed, 114 insertions(+), 39 deletions(-) diff --git a/foundry/src/executors/UniswapV2Executor.sol b/foundry/src/executors/UniswapV2Executor.sol index 7239a8a..737f4ae 100644 --- a/foundry/src/executors/UniswapV2Executor.sol +++ b/foundry/src/executors/UniswapV2Executor.sol @@ -6,22 +6,29 @@ 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; + address private constant FACTORY = + 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; IERC20 tokenIn; (tokenIn, target, receiver, zeroForOne) = _decodeData(data); + + if (target != _computePairAddress(target)) { + revert UniswapV2Executor__InvalidTarget(); + } calculatedAmount = _getAmountOut(target, givenAmount, zeroForOne); tokenIn.safeTransfer(target, givenAmount); @@ -33,7 +40,9 @@ contract UniswapV2Executor is IExecutor { } } - function _decodeData(bytes calldata data) + function _decodeData( + bytes calldata data + ) internal pure returns ( @@ -52,20 +61,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"); @@ -74,4 +83,26 @@ contract UniswapV2Executor is IExecutor { uint256 denominator = (uint256(reserveIn) * 1000) + amountInWithFee; amount = numerator / denominator; } + + 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)); + pair = address( + uint160( + uint256( + keccak256( + abi.encodePacked( + hex"ff", + FACTORY, + salt, + hex"96e8ac4277198ff8b6f785478aa9a39f403cb768dd02cbee326c3e7da348845f" + ) + ) + ) + ) + ); + } } diff --git a/foundry/src/executors/UniswapV3Executor.sol b/foundry/src/executors/UniswapV3Executor.sol index d1719fe..74ffaea 100644 --- a/foundry/src/executors/UniswapV3Executor.sol +++ b/foundry/src/executors/UniswapV3Executor.sol @@ -9,6 +9,7 @@ import "@interfaces/ICallback.sol"; error UniswapV3Executor__InvalidDataLength(); error UniswapV3Executor__InvalidFactory(); +error UniswapV3Executor__InvalidTarget(); contract UniswapV3Executor is IExecutor, ICallback { using SafeERC20 for IERC20; @@ -29,11 +30,10 @@ 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, @@ -42,6 +42,11 @@ contract UniswapV3Executor is IExecutor, ICallback { address target, bool zeroForOne ) = _decodeData(data); + + if (target != _computePairAddress(tokenIn, tokenOut, fee)) { + revert UniswapV3Executor__InvalidTarget(); + } + int256 amount0; int256 amount1; IUniswapV3Pool pool = IUniswapV3Pool(target); @@ -66,10 +71,9 @@ 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) @@ -77,15 +81,18 @@ 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); @@ -97,24 +104,32 @@ 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 ( @@ -137,13 +152,42 @@ 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); + pool = address( + uint160( + uint256( + keccak256( + abi.encodePacked( + hex"ff", + factory, + keccak256(abi.encode(token0, token1, fee)), + hex"e34f199b19b2b4f47f68442619d555527d244f78a3297ea89325f843f87b8b54" + ) + ) + ) + ) ); } } From 2f1507dd0ef540de1819b2489cedbdd50643538c Mon Sep 17 00:00:00 2001 From: royvardhan Date: Fri, 21 Feb 2025 22:49:10 +0530 Subject: [PATCH 2/5] test: add target verification tests for usv2, usv3 --- foundry/src/executors/UniswapV2Executor.sol | 35 ++++---- foundry/src/executors/UniswapV3Executor.sol | 83 ++++++++----------- .../test/executors/UniswapV2Executor.t.sol | 39 ++++++++- .../test/executors/UniswapV3Executor.t.sol | 41 +++++++++ foundry/test/mock/MockUniswapV2Pool.sol | 13 +++ 5 files changed, 144 insertions(+), 67 deletions(-) create mode 100644 foundry/test/mock/MockUniswapV2Pool.sol diff --git a/foundry/src/executors/UniswapV2Executor.sol b/foundry/src/executors/UniswapV2Executor.sol index 737f4ae..750bb18 100644 --- 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)); diff --git a/foundry/src/executors/UniswapV3Executor.sol b/foundry/src/executors/UniswapV3Executor.sol index 74ffaea..9caff33 100644 --- 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( diff --git a/foundry/test/executors/UniswapV2Executor.t.sol b/foundry/test/executors/UniswapV2Executor.t.sol index f377f5a..88973a8 100644 --- 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); + } } diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index 86b999b..013fc86 100644 --- 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, diff --git a/foundry/test/mock/MockUniswapV2Pool.sol b/foundry/test/mock/MockUniswapV2Pool.sol new file mode 100644 index 0000000..80a9b92 --- /dev/null +++ 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; + } +} From 40bd37a1a4e2e96fe36d44e025c5f3a046a9ecb5 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Fri, 21 Feb 2025 23:13:06 +0530 Subject: [PATCH 3/5] chore: rm redundant usv3 test --- foundry/test/executors/UniswapV3Executor.t.sol | 8 -------- 1 file changed, 8 deletions(-) diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index 013fc86..14a6982 100644 --- a/foundry/test/executors/UniswapV3Executor.t.sol +++ b/foundry/test/executors/UniswapV3Executor.t.sol @@ -83,14 +83,6 @@ contract UniswapV3ExecutorTest is Test, Constants { 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; From 4b77128df23647bbd247be9e6536b68ea9454056 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Sat, 22 Feb 2025 00:26:15 +0530 Subject: [PATCH 4/5] 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; - } -} From f4f5b841e717bdcd974ca37a9b912f0a56a8f822 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Sat, 22 Feb 2025 01:39:11 +0530 Subject: [PATCH 5/5] chore: keep test naming consistent --- foundry/test/LibPrefixLengthEncodedByteArray.t.sol | 4 ++-- foundry/test/executors/UniswapV2Executor.t.sol | 4 ++-- foundry/test/executors/UniswapV3Executor.t.sol | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) 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 bb8c553..01722eb 100644 --- a/foundry/test/executors/UniswapV2Executor.t.sol +++ b/foundry/test/executors/UniswapV2Executor.t.sol @@ -80,7 +80,7 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants { uniswapV2Exposed.verifyPairAddress(WETH_DAI_POOL); } - function test_RevertIf_InvalidTarget() public { + function testInvalidTarget() public { address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR)); vm.expectRevert(UniswapV2Executor__InvalidTarget.selector); uniswapV2Exposed.verifyPairAddress(fakePool); @@ -145,7 +145,7 @@ contract UniswapV2ExecutorTest is UniswapV2ExecutorExposed, Test, Constants { assertGe(finalBalance, amountOut); } - function test_RevertIf_Swap_InvalidTarget() public { + function testSwapFailureInvalidTarget() public { uint256 amountIn = 10 ** 18; bool zeroForOne = false; address fakePool = address(new FakeUniswapV2Pool(WETH_ADDR, DAI_ADDR)); diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index 1691c96..3e39c7c 100644 --- a/foundry/test/executors/UniswapV3Executor.t.sol +++ b/foundry/test/executors/UniswapV3Executor.t.sol @@ -128,7 +128,7 @@ contract UniswapV3ExecutorTest is Test, Constants { assertGe(IERC20(DAI_ADDR).balanceOf(address(this)), expAmountOut); } - function test_RevertIf_InvalidTarget() public { + function testSwapFailureInvalidTarget() public { uint256 amountIn = 10 ** 18; deal(WETH_ADDR, address(uniswapV3Exposed), amountIn); bool zeroForOne = false;