From cccb252bf2194af55b983f73815edb0cf1782776 Mon Sep 17 00:00:00 2001 From: royvardhan Date: Fri, 14 Feb 2025 23:27:44 +0530 Subject: [PATCH] feat: update handleCallback in USV3 to do verification --- foundry/src/executors/UniswapV3Executor.sol | 73 +++++++++---------- .../test/executors/UniswapV3Executor.t.sol | 43 ++++------- 2 files changed, 48 insertions(+), 68 deletions(-) diff --git a/foundry/src/executors/UniswapV3Executor.sol b/foundry/src/executors/UniswapV3Executor.sol index 906c771..bfa3eae 100644 --- a/foundry/src/executors/UniswapV3Executor.sol +++ b/foundry/src/executors/UniswapV3Executor.sol @@ -5,7 +5,6 @@ import "@interfaces/IExecutor.sol"; import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol"; import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol"; import "@uniswap/v3-updated/CallbackValidationV2.sol"; -import "forge-std/console.sol"; error UniswapV3Executor__InvalidDataLength(); error UniswapV3Executor__InvalidFactory(); @@ -29,11 +28,10 @@ contract UniswapV3Executor is IExecutor { } // 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, @@ -66,41 +64,36 @@ contract UniswapV3Executor is IExecutor { } } - function handleCallback(bytes calldata msgData) - external - returns (bytes memory result) - { - (int256 amount0Delta, int256 amount1Delta) = - abi.decode(msgData[:64], (int256, int256)); + function handleCallback( + bytes calldata msgData + ) external returns (bytes memory result) { + (int256 amount0Delta, int256 amount1Delta) = abi.decode( + msgData[:64], + (int256, int256) + ); - bytes calldata remainingData = msgData[64:]; + address tokenIn = address(bytes20(msgData[64:84])); + address tokenOut = address(bytes20(msgData[84:104])); + uint24 poolFee = uint24(bytes3(msgData[104:107])); - (uint256 amountOwed, address tokenOwed) = - _verifyUSV3Callback(amount0Delta, amount1Delta, remainingData); + CallbackValidationV2.verifyCallback( + factory, + tokenIn, + tokenOut, + poolFee + ); - IERC20(tokenOwed).safeTransfer(msg.sender, amountOwed); - return abi.encode(amountOwed, tokenOwed); + uint256 amountOwed = amount0Delta > 0 + ? uint256(amount0Delta) + : uint256(amount1Delta); + + IERC20(tokenIn).safeTransfer(msg.sender, amountOwed); + return abi.encode(amountOwed, tokenIn); } - function _verifyUSV3Callback( - int256 amount0Delta, - int256 amount1Delta, + function _decodeData( bytes calldata data - ) internal view returns (uint256 amountIn, address tokenIn) { - tokenIn = address(bytes20(data[0:20])); - address tokenOut = address(bytes20(data[20:40])); - uint24 poolFee = uint24(bytes3(data[40:43])); - - // slither-disable-next-line unused-return - CallbackValidationV2.verifyCallback(factory, tokenIn, tokenOut, poolFee); - - amountIn = - amount0Delta > 0 ? uint256(amount0Delta) : uint256(amount1Delta); - - return (amountIn, tokenIn); - } - - function _decodeData(bytes calldata data) + ) internal pure returns ( @@ -123,11 +116,11 @@ contract UniswapV3Executor is IExecutor { zeroForOne = uint8(data[83]) > 0; } - function _makeV3CallbackData(address tokenIn, address tokenOut, uint24 fee) - internal - view - returns (bytes memory) - { + function _makeV3CallbackData( + address tokenIn, + address tokenOut, + uint24 fee + ) internal view returns (bytes memory) { return abi.encodePacked(tokenIn, tokenOut, fee, self); } } diff --git a/foundry/test/executors/UniswapV3Executor.t.sol b/foundry/test/executors/UniswapV3Executor.t.sol index e0d1638..fe50674 100644 --- a/foundry/test/executors/UniswapV3Executor.t.sol +++ b/foundry/test/executors/UniswapV3Executor.t.sol @@ -8,7 +8,9 @@ import {Constants} from "../Constants.sol"; contract UniswapV3ExecutorExposed is UniswapV3Executor { constructor(address _factory) UniswapV3Executor(_factory) {} - function decodeData(bytes calldata data) + function decodeData( + bytes calldata data + ) external pure returns ( @@ -42,7 +44,12 @@ contract UniswapV3ExecutorTest is Test, Constants { function testDecodeParams() public view { uint24 expectedPoolFee = 500; bytes memory data = abi.encodePacked( - WETH_ADDR, DAI_ADDR, expectedPoolFee, address(2), address(3), false + WETH_ADDR, + DAI_ADDR, + expectedPoolFee, + address(2), + address(3), + false ); ( @@ -63,8 +70,11 @@ contract UniswapV3ExecutorTest is Test, Constants { } function testDecodeParamsInvalidDataLength() public { - bytes memory invalidParams = - abi.encodePacked(WETH_ADDR, address(2), address(3)); + bytes memory invalidParams = abi.encodePacked( + WETH_ADDR, + address(2), + address(3) + ); vm.expectRevert(UniswapV3Executor__InvalidDataLength.selector); uniswapV3Exposed.decodeData(invalidParams); @@ -77,7 +87,7 @@ contract UniswapV3ExecutorTest is Test, Constants { uint256 initialPoolReserve = IERC20(WETH_ADDR).balanceOf(DAI_WETH_USV3); vm.startPrank(DAI_WETH_USV3); - bytes memory callbackData = _encodeUSV3CallbackData( + bytes memory callbackData = abi.encodePacked( int256(amountOwed), // amount0Delta int256(0), // amount1Delta WETH_ADDR, @@ -90,27 +100,4 @@ contract UniswapV3ExecutorTest is Test, Constants { uint256 finalPoolReserve = IERC20(WETH_ADDR).balanceOf(DAI_WETH_USV3); assertEq(finalPoolReserve - initialPoolReserve, amountOwed); } - - function _encodeUSV3CallbackData( - int256 amount0Delta, - int256 amount1Delta, - address tokenIn, - address tokenOut, - uint24 fee - ) internal pure returns (bytes memory) { - // Dummy selector for handleCallback - bytes4 selector = - bytes4(keccak256("handleCallback(int256,int256,bytes)")); - - bytes memory tokenData = abi.encodePacked(tokenIn, tokenOut, fee); - - // [0:4] - function selector - // [4:68] - abi.encode(amount0Delta, amount1Delta) - // [68:end] - abi.encode(tokenData) where tokenData is the packed bytes - return abi.encodePacked( - selector, - abi.encode(amount0Delta, amount1Delta), - abi.encode(tokenData) - ); - } }