From 7936ba1c943a616a143dc6f8ecb1da61073b05a8 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Fri, 21 Feb 2025 20:00:38 +0530 Subject: [PATCH] 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" + ) + ) + ) + ) ); } }