From a11df6a1c7926d539d81b6d711bff0a67801e033 Mon Sep 17 00:00:00 2001
From: Stanislav Bezkorovainyi <stanislavbezkor@gmail.com>
Date: Thu, 28 Nov 2024 18:56:25 +0100
Subject: [PATCH] Fix H03 (#17)

---
 .../contracts/bridgehub/Bridgehub.sol         | 27 ++++++++++---------
 .../contracts/bridgehub/IBridgehub.sol        |  6 ++---
 .../l1/integration/DeploymentTest.t.sol       |  9 +++++--
 3 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/l1-contracts/contracts/bridgehub/Bridgehub.sol b/l1-contracts/contracts/bridgehub/Bridgehub.sol
index c31bcf285..d62e7bffb 100644
--- a/l1-contracts/contracts/bridgehub/Bridgehub.sol
+++ b/l1-contracts/contracts/bridgehub/Bridgehub.sol
@@ -226,20 +226,9 @@ contract Bridgehub is IBridgehub, ReentrancyGuard, Ownable2StepUpgradeable, Paus
         messageRoot = _messageRoot;
     }
 
-    /// @notice Used for the upgrade to set the baseTokenAssetId previously stored as baseToken.
-    /// @param _chainId the chainId of the chain.
-    function setLegacyBaseTokenAssetId(uint256 _chainId) external override {
-        if (baseTokenAssetId[_chainId] != bytes32(0)) {
-            return;
-        }
-        address token = __DEPRECATED_baseToken[_chainId];
-        require(token != address(0), "BH: token not set");
-        baseTokenAssetId[_chainId] = DataEncoding.encodeNTVAssetId(block.chainid, token);
-    }
-
-    /// @notice Used to set the legacy chain address for the upgrade.
+    /// @notice Used to set the legacy chain data for the upgrade.
     /// @param _chainId The chainId of the legacy chain we are migrating.
-    function setLegacyChainAddress(uint256 _chainId) external override {
+    function registerLegacyChain(uint256 _chainId) external override {
         address ctm = chainTypeManager[_chainId];
         if (ctm == address(0)) {
             revert ChainNotLegacy();
@@ -247,11 +236,23 @@ contract Bridgehub is IBridgehub, ReentrancyGuard, Ownable2StepUpgradeable, Paus
         if (zkChainMap.contains(_chainId)) {
             revert ChainAlreadyLive();
         }
+
+        // From now on, since `zkChainMap` did not contain the chain, we assume
+        // that the chain is a legacy chain in the process of migration, i.e.
+        // its stored `baseTokenAssetId`, etc.
+
+        address token = __DEPRECATED_baseToken[_chainId];
+        require(token != address(0), "BH: token not set");
+
+        baseTokenAssetId[_chainId] = DataEncoding.encodeNTVAssetId(block.chainid, token);
+
         address chainAddress = IChainTypeManager(ctm).getZKChainLegacy(_chainId);
         if (chainAddress == address(0)) {
             revert ChainNotLegacy();
         }
         _registerNewZKChain(_chainId, chainAddress);
+        messageRoot.addNewChain(_chainId);
+        settlementLayer[_chainId] = block.chainid;
     }
 
     //// Registry
diff --git a/l1-contracts/contracts/bridgehub/IBridgehub.sol b/l1-contracts/contracts/bridgehub/IBridgehub.sol
index 3f05bba35..90ef736d6 100644
--- a/l1-contracts/contracts/bridgehub/IBridgehub.sol
+++ b/l1-contracts/contracts/bridgehub/IBridgehub.sol
@@ -227,13 +227,11 @@ interface IBridgehub is IAssetHandler, IL1AssetHandler {
 
     function L1_CHAIN_ID() external view returns (uint256);
 
-    function setLegacyBaseTokenAssetId(uint256 _chainId) external;
-
     function registerAlreadyDeployedZKChain(uint256 _chainId, address _hyperchain) external;
 
-    function setLegacyChainAddress(uint256 _chainId) external;
-
     /// @notice return the ZK chain contract for a chainId
     /// @dev It is a legacy method. Do not use!
     function getHyperchain(uint256 _chainId) external view returns (address);
+
+    function registerLegacyChain(uint256 _chainId) external;
 }
diff --git a/l1-contracts/test/foundry/l1/integration/DeploymentTest.t.sol b/l1-contracts/test/foundry/l1/integration/DeploymentTest.t.sol
index b4cca3bb1..0955cfa90 100644
--- a/l1-contracts/test/foundry/l1/integration/DeploymentTest.t.sol
+++ b/l1-contracts/test/foundry/l1/integration/DeploymentTest.t.sol
@@ -26,6 +26,7 @@ import {IZKChain} from "contracts/state-transition/chain-interfaces/IZKChain.sol
 import {IChainTypeManager} from "contracts/state-transition/IChainTypeManager.sol";
 import {DataEncoding} from "contracts/common/libraries/DataEncoding.sol";
 import {IncorrectBridgeHubAddress} from "contracts/common/L1ContractErrors.sol";
+import {MessageRoot} from "contracts/bridgehub/MessageRoot.sol";
 
 contract DeploymentTests is L1ContractDeployer, ZKChainDeployer, TokenDeployer, L2TxMocker {
     uint256 constant TEST_USERS_COUNT = 10;
@@ -106,8 +107,12 @@ contract DeploymentTests is L1ContractDeployer, ZKChainDeployer, TokenDeployer,
             keccak256(abi.encode(randomChainId, 204)),
             bytes32(uint256(uint160(address(chainTypeManager))))
         );
-        bridgehub.setLegacyBaseTokenAssetId(randomChainId);
-        bridgehub.setLegacyChainAddress(randomChainId);
+        bridgehub.registerLegacyChain(randomChainId);
+
+        assertEq(bridgehub.settlementLayer(randomChainId), block.chainid);
+
+        address messageRoot = address(bridgehub.messageRoot());
+        assertTrue(MessageRoot(messageRoot).chainIndex(randomChainId) != 0);
     }
 
     function test_registerAlreadyDeployedZKChain() public {