feat(protocol): introduce ForkManager to improve protocol fork manage…

…ment (#18508)

Co-authored-by: dantaik <[email protected]>
Co-authored-by: David <[email protected]>

packages/protocol/contract_layout_layer1.md

@@ -1855,3 +1855,24 @@
 ╰-------------------------+-------------------------------------------------+------+--------+-------+------------------------------------------------------------╯
 
 
+## ForkManager
+
+╭---------------+-------------+------+--------+-------+---------------------------------------------------╮
+| Name          | Type        | Slot | Offset | Bytes | Contract                                          |
++=========================================================================================================+
+| _initialized  | uint8       | 0    | 0      | 1     | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| _initializing | bool        | 0    | 1      | 1     | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| __gap         | uint256[50] | 1    | 0      | 1600  | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| _owner        | address     | 51   | 0      | 20    | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| __gap         | uint256[49] | 52   | 0      | 1568  | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| _pendingOwner | address     | 101  | 0      | 20    | ForkManager |
+|---------------+-------------+------+--------+-------+---------------------------------------------------|
+| __gap         | uint256[49] | 102  | 0      | 1568  | ForkManager |
+╰---------------+-------------+------+--------+-------+---------------------------------------------------╯
+
+

packages/protocol/contracts/layer1/fork/ForkManager.sol

@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
+import "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";
+
+/// @title ForkManager
+/// @custom:security-contact [email protected]
+/// @notice This contract serves as a base contract for managing up to two forks within the Taiko
+/// protocol. By default, all function calls are routed to the newFork address.
+/// Sub-contracts should override the shouldRouteToOldFork function to route specific function calls
+/// to the old fork address.
+/// These sub-contracts should be placed between a proxy and the actual fork implementations. When
+/// calling upgradeTo, the proxy should always upgrade to a new ForkManager implementation, not an
+/// actual fork implementation.
+/// It is strongly advised to name functions differently for the same functionality across the two
+/// forks, as it is not possible to route the same function to two different forks.
+///
+///                         +--> newFork
+/// PROXY -> FORK_MANAGER --|
+///                         +--> oldFork
+contract ForkManager is UUPSUpgradeable, Ownable2StepUpgradeable {
+    address public immutable oldFork;
+    address public immutable newFork;
+
+    error ForkAddressIsZero();
+    error InvalidParams();
+
+    constructor(address _oldFork, address _currFork) {
+        require(_currFork != address(0) && _currFork != _oldFork, InvalidParams());
+        oldFork = _oldFork;
+        newFork = _currFork;
+    }
+
+    fallback() external payable virtual {
+        _fallback();
+    }
+
+    receive() external payable virtual {
+        _fallback();
+    }
+
+    function isForkManager() public pure returns (bool) {
+        return true;
+    }
+
+    function _fallback() internal virtual {
+        address fork = shouldRouteToOldFork(msg.sig) ? oldFork : newFork;
+        require(fork != address(0), ForkAddressIsZero());
+
+        assembly {
+            calldatacopy(0, 0, calldatasize())
+            let result := delegatecall(gas(), fork, 0, calldatasize(), 0, 0)
+            returndatacopy(0, 0, returndatasize())
+
+            switch result
+            case 0 { revert(0, returndatasize()) }
+            default { return(0, returndatasize()) }
+        }
+    }
+
+    function _authorizeUpgrade(address) internal virtual override onlyOwner { }
+
+    /// @notice Determines if the call should be routed to the old fork.
+    /// @dev This function is intended to be overridden in derived contracts to provide custom
+    /// routing logic.
+    /// @param  _selector The function selector of the call.
+    /// @return A boolean value indicating whether the call should be routed to the old fork.
+    function shouldRouteToOldFork(bytes4 _selector) internal pure virtual returns (bool) { }
+}

packages/protocol/script/gen-layouts.sh

@@ -56,6 +56,7 @@ contracts_layer1=(
 "contracts/layer1/team/tokenunlock/TokenUnlock.sol:TokenUnlock"
 "contracts/layer1/provers/ProverSet.sol:ProverSet"
 "contracts/layer1/provers/GuardianProver.sol:GuardianProver"
+"contracts/layer1/fork/ForkManager.sol:ForkManager"
 )
 
 # Layer 2 contracts

packages/protocol/test/layer1/fork/ForkManager.t.sol

@@ -0,0 +1,69 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import "../TaikoL1Test.sol";
+import "src/layer1/fork/ForkManager.sol";
+
+contract Fork is EssentialContract {
+    bytes32 private immutable __name;
+
+    constructor(bytes32 _name) {
+        __name = _name;
+    }
+
+    function init() external initializer {
+        __Essential_init(msg.sender);
+    }
+
+    function name() public view returns (bytes32) {
+        return __name;
+    }
+}
+
+contract ForkManager_RouteToOldFork is ForkManager {
+    constructor(address _fork1, address _fork2) ForkManager(_fork1, _fork2) { }
+
+    function shouldRouteToOldFork(bytes4 _selector) internal pure override returns (bool) {
+        return _selector == Fork.name.selector;
+    }
+}
+
+contract TestForkManager is TaikoL1Test {
+    address fork1 = address(new Fork("fork1"));
+    address fork2 = address(new Fork("fork2"));
+
+    function test_ForkManager_default_routing() public {
+        address proxy = deployProxy({
+            name: "main_proxy",
+            impl: address(new ForkManager(address(0), fork1)),
+            data: abi.encodeCall(Fork.init, ())
+        });
+
+        assertTrue(ForkManager(payable(proxy)).isForkManager());
+        assertEq(Fork(proxy).name(), "fork1");
+
+        // If we upgrade the proxy's impl to a fork, then alling isForkManager will throw,
+        // so we should never do this in production.
+        Fork(proxy).upgradeTo(fork2);
+        vm.expectRevert();
+        ForkManager(payable(proxy)).isForkManager();
+
+        Fork(proxy).upgradeTo(address(new ForkManager(fork1, fork2)));
+        assertEq(Fork(proxy).name(), "fork2");
+    }
+
+    function test_ForkManager_routing_to_old_fork() public {
+        address proxy = deployProxy({
+            name: "main_proxy",
+            impl: address(new ForkManager_RouteToOldFork(fork1, fork2)),
+            data: abi.encodeCall(Fork.init, ())
+        });
+
+        assertTrue(ForkManager(payable(proxy)).isForkManager());
+        assertEq(Fork(proxy).name(), "fork1");
+
+        Fork(proxy).upgradeTo(address(new ForkManager(fork1, fork2)));
+        assertTrue(ForkManager(payable(proxy)).isForkManager());
+        assertEq(Fork(proxy).name(), "fork2");
+    }
+}