From 08867af74dc0b3532903578444ac9c6cbe42a632 Mon Sep 17 00:00:00 2001 From: ChefMist <133624774+ChefMist@users.noreply.github.com> Date: Fri, 13 Sep 2024 09:44:16 +0800 Subject: [PATCH] feat: prevent permit2.permit frontrun --- ...apV3Test#test_v3Swap_exactOutput0For1.snap | 2 +- .../UniversalRouterBytecodeSize.snap | 2 +- .../UniversalRouterTest#test_sweep_token.snap | 2 +- ...#test_v4CLPositionmanager_Mint_Native.snap | 2 +- ..._v4BinPositionmanager_BinAddLiquidity.snap | 2 +- ...ositionmanager_BinAddLiquidity_Native.snap | 2 +- ...ionTest#test_v4CLPositionmanager_Mint.snap | 2 +- src/base/Dispatcher.sol | 12 +- test/UniversalRouter.t.sol | 207 +++++++++++++++++- 9 files changed, 222 insertions(+), 11 deletions(-) diff --git a/.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput0For1.snap b/.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput0For1.snap index 56302cc..73d7c6c 100644 --- a/.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput0For1.snap +++ b/.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput0For1.snap @@ -1 +1 @@ -152658 \ No newline at end of file +152659 \ No newline at end of file diff --git a/.forge-snapshots/UniversalRouterBytecodeSize.snap b/.forge-snapshots/UniversalRouterBytecodeSize.snap index 4cdd517..484ebbf 100644 --- a/.forge-snapshots/UniversalRouterBytecodeSize.snap +++ b/.forge-snapshots/UniversalRouterBytecodeSize.snap @@ -1 +1 @@ -24108 \ No newline at end of file +24160 \ No newline at end of file diff --git a/.forge-snapshots/UniversalRouterTest#test_sweep_token.snap b/.forge-snapshots/UniversalRouterTest#test_sweep_token.snap index 0efb7fb..63f357d 100644 --- a/.forge-snapshots/UniversalRouterTest#test_sweep_token.snap +++ b/.forge-snapshots/UniversalRouterTest#test_sweep_token.snap @@ -1 +1 @@ -55428 \ No newline at end of file +55441 \ No newline at end of file diff --git a/.forge-snapshots/V3ToV4MigrationNativeTest#test_v4CLPositionmanager_Mint_Native.snap b/.forge-snapshots/V3ToV4MigrationNativeTest#test_v4CLPositionmanager_Mint_Native.snap index 43f8a2a..8a3f431 100644 --- a/.forge-snapshots/V3ToV4MigrationNativeTest#test_v4CLPositionmanager_Mint_Native.snap +++ b/.forge-snapshots/V3ToV4MigrationNativeTest#test_v4CLPositionmanager_Mint_Native.snap @@ -1 +1 @@ -559336 \ No newline at end of file +559373 \ No newline at end of file diff --git a/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity.snap b/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity.snap index 1c644c8..ad6e1a7 100644 --- a/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity.snap +++ b/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity.snap @@ -1 +1 @@ -594165 \ No newline at end of file +594191 \ No newline at end of file diff --git a/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity_Native.snap b/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity_Native.snap index 6c91c9a..23aaec4 100644 --- a/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity_Native.snap +++ b/.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity_Native.snap @@ -1 +1 @@ -569991 \ No newline at end of file +570028 \ No newline at end of file diff --git a/.forge-snapshots/V3ToV4MigrationTest#test_v4CLPositionmanager_Mint.snap b/.forge-snapshots/V3ToV4MigrationTest#test_v4CLPositionmanager_Mint.snap index fff8307..4613764 100644 --- a/.forge-snapshots/V3ToV4MigrationTest#test_v4CLPositionmanager_Mint.snap +++ b/.forge-snapshots/V3ToV4MigrationTest#test_v4CLPositionmanager_Mint.snap @@ -1 +1 @@ -583518 \ No newline at end of file +583544 \ No newline at end of file diff --git a/src/base/Dispatcher.sol b/src/base/Dispatcher.sol index 58e98b6..5f7a035 100755 --- a/src/base/Dispatcher.sol +++ b/src/base/Dispatcher.sol @@ -119,7 +119,11 @@ abstract contract Dispatcher is permitBatch := add(inputs.offset, calldataload(inputs.offset)) } bytes calldata data = inputs.toBytes(1); - PERMIT2.permit(msgSender(), permitBatch, data); + try PERMIT2.permit(msgSender(), permitBatch, data) {} + catch (bytes memory reason) { + output = reason; + success = false; + } return (success, output); } else if (command == Commands.SWEEP) { // equivalent: abi.decode(inputs, (address, address, uint256)) @@ -204,7 +208,11 @@ abstract contract Dispatcher is permitSingle := inputs.offset } bytes calldata data = inputs.toBytes(6); // PermitSingle takes first 6 slots (0..5) - PERMIT2.permit(msgSender(), permitSingle, data); + try PERMIT2.permit(msgSender(), permitSingle, data) {} + catch (bytes memory reason) { + output = reason; + success = false; + } return (success, output); } else if (command == Commands.WRAP_ETH) { // equivalent: abi.decode(inputs, (address, uint256)) diff --git a/test/UniversalRouter.t.sol b/test/UniversalRouter.t.sol index fb86d11..300de47 100644 --- a/test/UniversalRouter.t.sol +++ b/test/UniversalRouter.t.sol @@ -6,7 +6,9 @@ import {GasSnapshot} from "forge-gas-snapshot/GasSnapshot.sol"; import {ERC20} from "solmate/src/tokens/ERC20.sol"; import {IERC165} from "@openzeppelin/contracts/interfaces/IERC165.sol"; import {ActionConstants} from "pancake-v4-periphery/src/libraries/ActionConstants.sol"; +import {Permit2SignatureHelpers} from "pancake-v4-periphery/test/shared/Permit2SignatureHelpers.sol"; import {IAllowanceTransfer} from "permit2/src/interfaces/IAllowanceTransfer.sol"; +import {DeployPermit2} from "permit2/test/utils/DeployPermit2.sol"; import {WETH} from "solmate/src/tokens/WETH.sol"; import {UniversalRouter} from "../src/UniversalRouter.sol"; @@ -19,21 +21,26 @@ import {MockERC721} from "./mock/MockERC721.sol"; import {MockERC1155} from "./mock/MockERC1155.sol"; import {RouterParameters} from "../src/base/RouterImmutables.sol"; -contract UniversalRouterTest is Test, GasSnapshot { +contract UniversalRouterTest is Test, GasSnapshot, Permit2SignatureHelpers, DeployPermit2 { error ContractSizeTooLarge(uint256 diff); + error InvalidNonce(); address RECIPIENT = makeAddr("alice"); uint256 constant AMOUNT = 10 ** 18; UniversalRouter router; MockERC20 erc20; + MockERC20 erc20_2; MockERC721 erc721; MockERC1155 erc1155; WETH weth9 = new WETH(); + IAllowanceTransfer permit2; function setUp() public { + permit2 = IAllowanceTransfer(deployPermit2()); + RouterParameters memory params = RouterParameters({ - permit2: address(0), + permit2: address(permit2), weth9: address(weth9), v2Factory: address(0), v3Factory: address(0), @@ -53,6 +60,7 @@ contract UniversalRouterTest is Test, GasSnapshot { router = new UniversalRouter(params); erc20 = new MockERC20(); + erc20_2 = new MockERC20(); erc721 = new MockERC721(); erc1155 = new MockERC1155(); } @@ -197,4 +205,199 @@ contract UniversalRouterTest is Test, GasSnapshot { vm.expectRevert(Payments.InsufficientETH.selector); router.execute{value: 1 ether}(commands, inputs); } + + function test_permit2Single() public { + // pre-req: + (address charlie, uint256 charliePK) = makeAddrAndKey("charlie"); + uint160 permitAmount = type(uint160).max; + uint48 permitExpiration = uint48(block.timestamp + 10e18); + uint48 permitNonce = 0; + + IAllowanceTransfer.PermitSingle memory permit = + defaultERC20PermitAllowance(address(erc20), permitAmount, permitExpiration, permitNonce); + permit.spender = address(router); + bytes memory sig = getPermitSignature(permit, charliePK, permit2.DOMAIN_SEPARATOR()); + + // before verify + (uint160 _amount, uint48 _expiration, uint48 _nonce) = + permit2.allowance(charlie, address(erc20), address(router)); + assertEq(_amount, 0); + assertEq(_expiration, 0); + assertEq(_nonce, 0); + + // execute + vm.startPrank(charlie); + bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.PERMIT2_PERMIT))); + bytes[] memory inputs = new bytes[](1); + inputs[0] = abi.encode(permit, sig); + router.execute(commands, inputs); + + // after verify + (_amount, _expiration, _nonce) = permit2.allowance(charlie, address(erc20), address(router)); + assertEq(_amount, permitAmount); + assertEq(_expiration, permitExpiration); + assertEq(_nonce, permitNonce + 1); + } + + function test_permit2Batch() public { + // pre-req: + (address charlie, uint256 charliePK) = makeAddrAndKey("charlie"); + uint160 permitAmount = type(uint160).max; + uint48 permitExpiration = uint48(block.timestamp + 10e18); + uint48 permitNonce = 0; + + address[] memory tokens = new address[](2); + tokens[0] = address(erc20); + tokens[1] = address(erc20_2); + + IAllowanceTransfer.PermitBatch memory permit = + defaultERC20PermitBatchAllowance(tokens, permitAmount, permitExpiration, permitNonce); + permit.spender = address(router); + bytes memory sig = getPermitBatchSignature(permit, charliePK, permit2.DOMAIN_SEPARATOR()); + + // before verify + for (uint256 i; i < tokens.length; i++) { + (uint160 _amount, uint48 _expiration, uint48 _nonce) = + permit2.allowance(charlie, tokens[i], address(router)); + assertEq(_amount, 0); + assertEq(_expiration, 0); + assertEq(_nonce, 0); + } + + // execute + vm.startPrank(charlie); + bytes memory commands = abi.encodePacked(bytes1(uint8(Commands.PERMIT2_PERMIT_BATCH))); + bytes[] memory inputs = new bytes[](1); + inputs[0] = abi.encode(permit, sig); + router.execute(commands, inputs); + + // after verify + for (uint256 i; i < tokens.length; i++) { + (uint160 _amount, uint48 _expiration, uint48 _nonce) = + permit2.allowance(charlie, address(tokens[i]), address(router)); + assertEq(_amount, permitAmount); + assertEq(_expiration, permitExpiration); + assertEq(_nonce, permitNonce + 1); + } + } + + /// @dev test showing that if permit command have ALLOW_REVERT flag and was front-run, the next command can still execute + function test_permit2Single_frontrun() public { + // pre-req + address bob = makeAddr("bob"); + (address charlie, uint256 charliePK) = makeAddrAndKey("charlie"); + uint160 permitAmount = type(uint160).max; + uint48 permitExpiration = uint48(block.timestamp + 10e18); + uint48 permitNonce = 0; + + IAllowanceTransfer.PermitSingle memory permit = + defaultERC20PermitAllowance(address(erc20), permitAmount, permitExpiration, permitNonce); + permit.spender = address(router); + bytes memory sig = getPermitSignature(permit, charliePK, permit2.DOMAIN_SEPARATOR()); + + // bob front-runs the permits + vm.prank(bob); + permit2.permit(charlie, permit, sig); + + // bob's front-run was successful + (uint160 _amount, uint48 _expiration, uint48 _nonce) = + permit2.allowance(charlie, address(erc20), address(router)); + assertEq(_amount, permitAmount); + assertEq(_expiration, permitExpiration); + assertEq(_nonce, permitNonce + 1); + + // before + assertEq(weth9.balanceOf(address(router)), 0); + + // charlie tries to call universal router permit2_permit and wrap_eth command + vm.deal(charlie, 1 ether); + vm.startPrank(charlie); + + bytes[] memory inputs = new bytes[](2); + inputs[0] = abi.encode(permit, sig); + inputs[1] = abi.encode(ActionConstants.ADDRESS_THIS, ActionConstants.CONTRACT_BALANCE); + + bytes memory commands; + + // attempt 1: execute and expect revert + commands = abi.encodePacked(bytes1(uint8(Commands.PERMIT2_PERMIT)), bytes1(uint8(Commands.WRAP_ETH))); + vm.expectRevert( + abi.encodeWithSelector( + IUniversalRouter.ExecutionFailed.selector, 0, abi.encodePacked(InvalidNonce.selector) + ) + ); + router.execute{value: 1 ether}(commands, inputs); + + // attempt 2: execute with allow revert flag and no revert expected + commands = abi.encodePacked( + bytes1(uint8(Commands.PERMIT2_PERMIT)) | Commands.FLAG_ALLOW_REVERT, bytes1(uint8(Commands.WRAP_ETH)) + ); + router.execute{value: 1 ether}(commands, inputs); + + // after + assertEq(weth9.balanceOf(address(router)), 1 ether); + } + + /// @dev test showing that if permit command have ALLOW_REVERT flag and was front-run, the next command can still execute + function test_permit2Batch_frontrun() public { + // pre-req + address bob = makeAddr("bob"); + (address charlie, uint256 charliePK) = makeAddrAndKey("charlie"); + uint160 permitAmount = type(uint160).max; + uint48 permitExpiration = uint48(block.timestamp + 10e18); + uint48 permitNonce = 0; + + address[] memory tokens = new address[](2); + tokens[0] = address(erc20); + tokens[1] = address(erc20_2); + + IAllowanceTransfer.PermitBatch memory permit = + defaultERC20PermitBatchAllowance(tokens, permitAmount, permitExpiration, permitNonce); + permit.spender = address(router); + bytes memory sig = getPermitBatchSignature(permit, charliePK, permit2.DOMAIN_SEPARATOR()); + + // bob front-runs the permits + vm.prank(bob); + permit2.permit(charlie, permit, sig); + + // bob's front-run was successful + for (uint256 i; i < tokens.length; i++) { + (uint160 _amount, uint48 _expiration, uint48 _nonce) = + permit2.allowance(charlie, address(tokens[i]), address(router)); + assertEq(_amount, permitAmount); + assertEq(_expiration, permitExpiration); + assertEq(_nonce, permitNonce + 1); + } + + // before + assertEq(weth9.balanceOf(address(router)), 0); + + // charlie tries to call universal router permit2_permit and wrap_eth command + vm.deal(charlie, 1 ether); + vm.startPrank(charlie); + + bytes[] memory inputs = new bytes[](2); + inputs[0] = abi.encode(permit, sig); + inputs[1] = abi.encode(ActionConstants.ADDRESS_THIS, ActionConstants.CONTRACT_BALANCE); + + bytes memory commands; + + // attempt 1: execute and expect revert + commands = abi.encodePacked(bytes1(uint8(Commands.PERMIT2_PERMIT_BATCH)), bytes1(uint8(Commands.WRAP_ETH))); + vm.expectRevert( + abi.encodeWithSelector( + IUniversalRouter.ExecutionFailed.selector, 0, abi.encodePacked(InvalidNonce.selector) + ) + ); + router.execute{value: 1 ether}(commands, inputs); + + // attempt 2: execute with allow revert flag and no revert expected + commands = abi.encodePacked( + bytes1(uint8(Commands.PERMIT2_PERMIT_BATCH)) | Commands.FLAG_ALLOW_REVERT, bytes1(uint8(Commands.WRAP_ETH)) + ); + router.execute{value: 1 ether}(commands, inputs); + + // after + assertEq(weth9.balanceOf(address(router)), 1 ether); + } }