feat: Add new pool initialization action and validation at PM call (#26)

* feat: update universal router with new action and position call protection

* test: fix test case

* feat: reduce optimizer runs

* cleanup: remove unused code

* doc: update command comments

.forge-snapshots/BinNativePancakeSwapV4Test#test_v4BinSwap_ExactInSingle_NativeIn.snap

@@ -1 +1 @@
-145983
+146503

.forge-snapshots/BinNativePancakeSwapV4Test#test_v4BinSwap_ExactInSingle_NativeOut_RouterRecipient.snap

@@ -1 +1 @@
-123041
+123561

.forge-snapshots/BinNativePancakeSwapV4Test#test_v4BinSwap_v4InitializeBinPool.snap

@@ -0,0 +1 @@
+138906

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_ExactInSingle.snap

@@ -1 +1 @@
-146945
+147468

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_ExactIn_MultiHop.snap

@@ -1 +1 @@
-177946
+178629

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_ExactIn_SingleHop.snap

@@ -1 +1 @@
-148741
+149270

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_ExactOut_MultiHop.snap

@@ -1 +1 @@
-181787
+182474

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_ExactOut_SingleHop.snap

@@ -1 +1 @@
-153114
+153645

.forge-snapshots/BinPancakeSwapV4Test#test_v4BinSwap_InitializeBinPool.snap

@@ -0,0 +1 @@
+159096

.forge-snapshots/BinPancakeSwapV4Test#test_v4ClSwap_ExactOutSingle.snap

@@ -1 +1 @@
-151324
+151845

.forge-snapshots/CLNativePancakeSwapV4Test#test_v4ClSwap_ExactInSingle_NativeIn.snap

@@ -1 +1 @@
-171713
+172182

.forge-snapshots/CLNativePancakeSwapV4Test#test_v4ClSwap_ExactInSingle_NativeOut.snap

@@ -1 +1 @@
-173985
+174508

.forge-snapshots/CLNativePancakeSwapV4Test#test_v4ClSwap_v4InitializeClPool.snap

@@ -0,0 +1 @@
+140212

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactInSingle.snap

@@ -1 +1 @@
-181084
+181622

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactIn_MultiHop.snap

@@ -1 +1 @@
-245737
+246450

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactIn_SingleHop.snap

@@ -1 +1 @@
-182624
+183168

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactOutSingle.snap

@@ -1 +1 @@
-185423
+185957

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactOut_MultiHop.snap

@@ -1 +1 @@
-249496
+250209

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_ExactOut_SingleHop.snap

@@ -1 +1 @@
-186956
+187500

.forge-snapshots/CLPancakeSwapV4Test#test_v4ClSwap_v4InitializeClPool.snap

@@ -0,0 +1 @@
+160352

.forge-snapshots/PancakeSwapV2Test#test_v2Swap_exactInput0For1.snap

@@ -1 +1 @@
-100125
+100195

.forge-snapshots/PancakeSwapV2Test#test_v2Swap_exactOutput0For1.snap

@@ -1 +1 @@
-100754
+100807

.forge-snapshots/PancakeSwapV3Test#test_v3Swap_ExactInput0For1.snap

@@ -1 +1 @@
-151862
+151937

.forge-snapshots/PancakeSwapV3Test#test_v3Swap_ExactInput0For1_ContractBalance.snap

@@ -1 +1 @@
-154767
+154851

.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactInput_MultiHop.snap

@@ -1 +1 @@
-243258
+243405

.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput0For1.snap

@@ -1 +1 @@
-150557
+150632

.forge-snapshots/PancakeSwapV3Test#test_v3Swap_exactOutput_MultiHop.snap

@@ -1 +1 @@
-254163
+254310

.forge-snapshots/StableSwapTest#test_stableSwap_ExactInput0For1.snap

@@ -1 +1 @@
-193973
+194005

.forge-snapshots/StableSwapTest#test_stableSwap_ExactInput1For0.snap

@@ -1 +1 @@
-194041
+194073

.forge-snapshots/UniversalRouterBytecodeSize.snap

@@ -1 +1 @@
-24566
+24411

.forge-snapshots/UniversalRouterTest#test_sweep_token.snap

@@ -1 +1 @@
-55441
+55456

.forge-snapshots/V3ToV4MigrationNativeTest#test_v4CLPositionmanager_Mint_Native.snap

@@ -1 +1 @@
-558784
+561402

.forge-snapshots/V3ToV4MigrationTest#test_v3PositionManager_burn.snap

@@ -1 +1 @@
-291590
+291611

.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity.snap

@@ -1 +1 @@
-594723
+597077

.forge-snapshots/V3ToV4MigrationTest#test_v4BinPositionmanager_BinAddLiquidity_Native.snap

@@ -1 +1 @@
-570540
+572762

.forge-snapshots/V3ToV4MigrationTest#test_v4CLPositionmanager_Mint.snap

@@ -1 +1 @@
-582977
+585727

foundry.toml

@@ -3,7 +3,7 @@ src = "src"
 out = 'foundry-out'
 libs = ["lib"]
 via_ir = true
-optimizer_runs = 30_000
+optimizer_runs = 10_000
 ffi = true
 fs_permissions = [
     { access = "read-write", path = ".forge-snapshots/" },

src/base/Dispatcher.sol

@@ -17,6 +17,9 @@ import {IERC721Permit} from "pancake-v4-periphery/src/pool-cl/interfaces/IERC721
 import {ActionConstants} from "pancake-v4-periphery/src/libraries/ActionConstants.sol";
 import {BaseActionsRouter} from "pancake-v4-periphery/src/base/BaseActionsRouter.sol";
 import {CalldataDecoder} from "pancake-v4-periphery/src/libraries/CalldataDecoder.sol";
+import {PoolKey} from "pancake-v4-core/src/types/PoolKey.sol";
+import {ICLPoolManager} from "pancake-v4-core/src/pool-cl/interfaces/ICLPoolManager.sol";
+import {IBinPoolManager} from "pancake-v4-core/src/pool-bin/interfaces/IBinPoolManager.sol";
 
 /// @title Decodes and Executes Commands
 /// @notice Called by the UniversalRouter contract to efficiently decode and execute a singular command
@@ -34,8 +37,6 @@ abstract contract Dispatcher is
 
     error InvalidCommandType(uint256 commandType);
     error BalanceTooLow();
-    error InvalidAction(bytes4 action);
-    error NotAuthorizedForToken(uint256 tokenId);
 
     /// @notice Executes encoded commands along with provided inputs.
     /// @param commands A set of concatenated commands, each 1 byte in length
@@ -305,14 +306,34 @@ abstract contract Dispatcher is
                     /// @dev ensure there's follow-up action if v3 position's removed token are sent to router contract
                     (success, output) = address(V3_POSITION_MANAGER).call(inputs);
                     return (success, output);
+                } else if (command == Commands.V4_CL_INITIALIZE_POOL) {
+                    PoolKey calldata poolKey;
+                    uint160 sqrtPriceX96;
+                    assembly {
+                        poolKey := inputs.offset
+                        sqrtPriceX96 := calldataload(add(inputs.offset, 0xc0)) // poolKey has 6 variable, so it takes 192 space = 0xc0
+                    }
+                    // <wip> remove "" hookData once we updated universal-router dependencies
+                    (success, output) = address(clPoolManager).call(
+                        abi.encodeCall(ICLPoolManager.initialize, (poolKey, sqrtPriceX96, ""))
+                    );
+                } else if (command == Commands.V4_BIN_INITIALIZE_POOL) {
+                    PoolKey calldata poolKey;
+                    uint24 activeId;
+                    assembly {
+                        poolKey := inputs.offset
+                        activeId := calldataload(add(inputs.offset, 0xc0)) // poolKey has 6 variable, so it takes 192 space = 0xc0
+                    }
+                    // <wip> remove "" hookData once we updated universal-router dependencies
+                    (success, output) = address(binPoolManager).call(
+                        abi.encodeCall(IBinPoolManager.initialize, (poolKey, activeId, ""))
+                    );
                 } else if (command == Commands.V4_CL_POSITION_CALL) {
-                    // should only call modifyLiquidities() with Actions.CL_MINT_POSITION
-                    // do not permit or approve this contract over a v4 position or someone could use this command to decrease, burn, or transfer your position
+                    _checkV4ClPositionManagerCall(inputs);
                     (success, output) = address(V4_CL_POSITION_MANAGER).call{value: address(this).balance}(inputs);
                     return (success, output);
                 } else if (command == Commands.V4_BIN_POSITION_CALL) {
-                    // should only call modifyLiquidities() with Actions.BIN_ADD_LIQUIDITY
-                    // do not permit or approve this contract over a v4 position or someone could use this command to decrease, burn, or transfer your position
+                    _checkV4BinPositionManagerCall(inputs);
                     (success, output) = address(V4_BIN_POSITION_MANAGER).call{value: address(this).balance}(inputs);
                     return (success, output);
                 } else {

src/libraries/Commands.sol

@@ -35,9 +35,11 @@ library Commands {
     uint256 constant V4_SWAP = 0x10;
     uint256 constant V3_POSITION_MANAGER_PERMIT = 0x11;
     uint256 constant V3_POSITION_MANAGER_CALL = 0x12;
-    uint256 constant V4_CL_POSITION_CALL = 0x13;
-    uint256 constant V4_BIN_POSITION_CALL = 0x14;
-    // COMMAND_PLACEHOLDER = 0x15 -> 0x20
+    uint256 constant V4_CL_INITIALIZE_POOL = 0x13;
+    uint256 constant V4_BIN_INITIALIZE_POOL = 0x14;
+    uint256 constant V4_CL_POSITION_CALL = 0x15;
+    uint256 constant V4_BIN_POSITION_CALL = 0x16;
+    // COMMAND_PLACEHOLDER = 0x17 -> 0x20
 
     // Command Types where 0x21<=value<=0x3f
     uint256 constant EXECUTE_SUB_PLAN = 0x21;

src/modules/V3ToV4Migrator.sol

@@ -4,10 +4,21 @@ pragma solidity ^0.8.0;
 import {RouterImmutables} from "../base/RouterImmutables.sol";
 import {IV3NonfungiblePositionManager} from
     "pancake-v4-periphery/src/interfaces/external/IV3NonfungiblePositionManager.sol";
+import {Actions} from "pancake-v4-periphery/src/libraries/Actions.sol";
+import {CalldataDecoder} from "pancake-v4-periphery/src/libraries/CalldataDecoder.sol";
+import {IPositionManager} from "pancake-v4-periphery/src/interfaces/IPositionManager.sol";
+import {console2} from "forge-std/console2.sol";
 
 /// @title V3 to V4 Migrator
 /// @notice A contract that migrates liquidity from PancakeSwap V3 to V4
 abstract contract V3ToV4Migrator is RouterImmutables {
+    using CalldataDecoder for bytes;
+
+    error NotAuthorizedForToken(uint256 tokenId);
+    error InvalidAction(bytes4 action);
+    error OnlyMintAllowed();
+    error OnlyAddLiqudityAllowed();
+
     /// @dev validate if an action is decreaseLiquidity, collect, or burn
     function isValidAction(bytes4 selector) internal pure returns (bool) {
         return selector == IV3NonfungiblePositionManager.decreaseLiquidity.selector
@@ -21,4 +32,68 @@ abstract contract V3ToV4Migrator is RouterImmutables {
         return caller == owner || V3_POSITION_MANAGER.getApproved(tokenId) == caller
             || V3_POSITION_MANAGER.isApprovedForAll(owner, caller);
     }
+
+    /// @dev check that the v4 position manager call is a safe call
+    /// of the position-altering Actions, we only allow Actions.MINT
+    /// this is because, if a user could be tricked into approving the UniversalRouter for
+    /// their position, an attacker could take their fees, or drain their entire position
+    function _checkV4ClPositionManagerCall(bytes calldata inputs) internal view {
+        bytes4 selector;
+        assembly {
+            selector := calldataload(inputs.offset)
+        }
+        if (selector != V4_CL_POSITION_MANAGER.modifyLiquidities.selector) {
+            revert InvalidAction(selector);
+        }
+
+        // slice is `abi.encode(bytes unlockData, uint256 deadline)`
+        bytes calldata slice = inputs[4:];
+        // the first bytes(0) extracts the unlockData parameter from modifyLiquidities
+        // unlockData = `abi.encode(bytes actions, bytes[] params)`
+        // the second bytes(0) extracts the actions parameter from unlockData
+        bytes calldata actions = slice.toBytes(0).toBytes(0);
+
+        uint256 numActions = actions.length;
+
+        for (uint256 actionIndex = 0; actionIndex < numActions; actionIndex++) {
+            uint256 action = uint8(actions[actionIndex]);
+
+            if (
+                action == Actions.CL_INCREASE_LIQUIDITY || action == Actions.CL_DECREASE_LIQUIDITY
+                    || action == Actions.CL_BURN_POSITION
+            ) {
+                revert OnlyMintAllowed();
+            }
+        }
+    }
+
+    /// @dev check that the v4 position manager call is a safe call
+    /// of the position-altering Actions, we only allow Actions.BIN_ADD_LIQUIDITY
+    /// this is because, if a user could be tricked into approving the UniversalRouter for
+    /// their position, an attacker could drain their entire position
+    function _checkV4BinPositionManagerCall(bytes calldata inputs) internal view {
+        bytes4 selector;
+        assembly {
+            selector := calldataload(inputs.offset)
+        }
+        if (selector != V4_BIN_POSITION_MANAGER.modifyLiquidities.selector) {
+            revert InvalidAction(selector);
+        }
+
+        // slice is `abi.encode(bytes unlockData, uint256 deadline)`
+        bytes calldata slice = inputs[4:];
+        // the first bytes(0) extracts the unlockData parameter from modifyLiquidities
+        // unlockData = `abi.encode(bytes actions, bytes[] params)`
+        // the second bytes(0) extracts the actions parameter from unlockData
+        bytes calldata actions = slice.toBytes(0).toBytes(0);
+
+        uint256 numActions = actions.length;
+
+        for (uint256 actionIndex = 0; actionIndex < numActions; actionIndex++) {
+            uint256 action = uint8(actions[actionIndex]);
+            if (action == Actions.BIN_REMOVE_LIQUIDITY) {
+                revert OnlyAddLiqudityAllowed();
+            }
+        }
+    }
 }

test/UniversalRouter.t.sol

@@ -631,7 +631,7 @@ contract UniversalRouterTest is Test, GasSnapshot, Permit2SignatureHelpers, Depl
         // if valid commands, return
         if (command >= 0x00 && command <= 0x06) return;
         if (command >= 0x08 && command <= 0x0e) return;
-        if (command >= 0x10 && command <= 0x14) return;
+        if (command >= 0x10 && command <= 0x16) return;
         if (command >= 0x21 && command <= 0x23) return;
 
         bytes memory commands = abi.encodePacked(bytes1(uint8(command)));