From 6998958ef256ee7b3e21dc6a0459ae9d200bb225 Mon Sep 17 00:00:00 2001
From: Stanislav Breadless <stanislavbezkor@gmail.com>
Date: Mon, 8 Jan 2024 12:53:50 +0100
Subject: [PATCH] make upgrade validation more strict

---
 .../contracts/upgrades/BaseZkSyncUpgrade.sol  | 30 +++++++++++++++++--
 .../contracts/upgrades/DefaultUpgrade.sol     | 17 -----------
 2 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/ethereum/contracts/upgrades/BaseZkSyncUpgrade.sol b/ethereum/contracts/upgrades/BaseZkSyncUpgrade.sol
index 71e57badc..8965aaf3c 100644
--- a/ethereum/contracts/upgrades/BaseZkSyncUpgrade.sol
+++ b/ethereum/contracts/upgrades/BaseZkSyncUpgrade.sol
@@ -65,6 +65,22 @@ abstract contract BaseZkSyncUpgrade is Base {
         // on the L2 side would be inaccurate. The effects of this "back-dating" of L2 upgrade batches will be reduced
         // as the permitted delay window is reduced in the future.
         require(block.timestamp >= _proposedUpgrade.upgradeTimestamp, "Upgrade is not ready yet");
+
+        _setNewProtocolVersion(_proposedUpgrade.newProtocolVersion);
+        _upgradeL1Contract(_proposedUpgrade.l1ContractsUpgradeCalldata);
+        _upgradeVerifier(_proposedUpgrade.verifier, _proposedUpgrade.verifierParams);
+        _setBaseSystemContracts(_proposedUpgrade.bootloaderHash, _proposedUpgrade.defaultAccountHash);
+
+        bytes32 txHash;
+        txHash = _setL2SystemContractUpgrade(
+            _proposedUpgrade.l2ProtocolUpgradeTx,
+            _proposedUpgrade.factoryDeps,
+            _proposedUpgrade.newProtocolVersion
+        );
+
+        _postUpgrade(_proposedUpgrade.postUpgradeCalldata);
+
+        emit UpgradeComplete(_proposedUpgrade.newProtocolVersion, txHash, _proposedUpgrade);
     }
 
     /// @notice Change default account bytecode hash, that is used on L2
@@ -104,13 +120,15 @@ abstract contract BaseZkSyncUpgrade is Base {
     /// @notice Change the address of the verifier smart contract
     /// @param _newVerifier Verifier smart contract address
     function _setVerifier(IVerifier _newVerifier) private {
+        if (_newVerifier == IVerifier(address(0))) {
+            return;
+        }
+
         // An upgrade to the verifier must be done carefully to ensure there aren't batches in the committed state
         // during the transition. If verifier is upgraded, it will immediately be used to prove all committed batches.
         // Batches committed expecting the old verifier will fail. Ensure all commited batches are finalized before the
         // verifier is upgraded.
-        if (_newVerifier == IVerifier(address(0))) {
-            return;
-        }
+        require(s.totalBatchesCommitted == s.totalBatchesVerified, "All committed batches must be verified first");
 
         IVerifier oldVerifier = s.verifier;
         s.verifier = _newVerifier;
@@ -128,6 +146,12 @@ abstract contract BaseZkSyncUpgrade is Base {
             return;
         }
 
+        // An upgrade to the verifier params must be done carefully to ensure there aren't batches in the committed state
+        // during the transition. If verifier verifier are upgraded, it will immediately be used to prove all committed batches.
+        // Batches committed expecting the old verifier params will fail. Ensure all commited batches are finalized before the
+        // verifier is upgraded.
+        require(s.totalBatchesCommitted == s.totalBatchesVerified, "All committed batches must be verified first");
+
         VerifierParams memory oldVerifierParams = s.verifierParams;
         s.verifierParams = _newVerifierParams;
         emit NewVerifierParams(oldVerifierParams, _newVerifierParams);
diff --git a/ethereum/contracts/upgrades/DefaultUpgrade.sol b/ethereum/contracts/upgrades/DefaultUpgrade.sol
index cd2bdd29f..790a9aa7d 100644
--- a/ethereum/contracts/upgrades/DefaultUpgrade.sol
+++ b/ethereum/contracts/upgrades/DefaultUpgrade.sol
@@ -24,23 +24,6 @@ contract DefaultUpgrade is BaseZkSyncUpgrade {
     /// @param _proposedUpgrade The upgrade to be executed.
     function upgrade(ProposedUpgrade calldata _proposedUpgrade) public override returns (bytes32) {
         super.upgrade(_proposedUpgrade);
-
-        _setNewProtocolVersion(_proposedUpgrade.newProtocolVersion);
-        _upgradeL1Contract(_proposedUpgrade.l1ContractsUpgradeCalldata);
-        _upgradeVerifier(_proposedUpgrade.verifier, _proposedUpgrade.verifierParams);
-        _setBaseSystemContracts(_proposedUpgrade.bootloaderHash, _proposedUpgrade.defaultAccountHash);
-
-        bytes32 txHash;
-        txHash = _setL2SystemContractUpgrade(
-            _proposedUpgrade.l2ProtocolUpgradeTx,
-            _proposedUpgrade.factoryDeps,
-            _proposedUpgrade.newProtocolVersion
-        );
-
-        _postUpgrade(_proposedUpgrade.postUpgradeCalldata);
-
-        emit UpgradeComplete(_proposedUpgrade.newProtocolVersion, txHash, _proposedUpgrade);
-
         return Diamond.DIAMOND_INIT_SUCCESS_RETURN_VALUE;
     }
 }