feat(LibInitializeGuard): add guard for safely upgrade contract with reinitializer

Author: Do Tu

foundry.toml

@@ -27,4 +27,5 @@ localhost = "http://localhost:8545"
 [dependencies]
 "@openzeppelin" = { version = "4.9.3", url = "https://github.com/OpenZeppelin/openzeppelin-contracts/archive/refs/tags/v4.9.3.zip" }
 "@forge-std" = { version = "1.9.1", url = "https://github.com/foundry-rs/forge-std/archive/refs/tags/v1.9.1.zip" }
-"@solady" = { version = "0.0.228", url = "https://github.com/Vectorized/solady/archive/refs/tags/v0.0.228.zip" }
+"@solady" = { version = "0.0.228", url = "https://github.com/Vectorized/solady/archive/refs/tags/v0.0.228.zip" }
+"@openzeppelin-v5" = { version = "5.0.2", url = "https://github.com/OpenZeppelin/openzeppelin-contracts/archive/refs/tags/v5.0.2.zip" }

remappings.txt

@@ -1,4 +1,5 @@
 @fdk/=script/
 @forge-std-1.9.1=dependencies/@forge-std-1.9.1/src/
 @openzeppelin-4.9.3=dependencies/@openzeppelin-4.9.3
-@solady-0.0.228=dependencies/@solady-0.0.228
+@solady-0.0.228=dependencies/@solady-0.0.228
+@openzeppelin-v5-5.0.2=dependencies/@openzeppelin-v5-5.0.2

script/BaseMigration.s.sol

@@ -7,10 +7,12 @@ import { TransparentProxyOZv4_9_5 } from "../src/TransparentProxyOZv4_9_5.sol";
 import { LibString } from "../dependencies/@solady-0.0.228/src/utils/LibString.sol";
 import { console } from "../dependencies/@forge-std-1.9.1/src/console.sol";
 import { StdStyle } from "../dependencies/@forge-std-1.9.1/src/StdStyle.sol";
+import { Vm } from "../dependencies/@forge-std-1.9.1/src/Vm.sol";
 import { ScriptExtended, IScriptExtended } from "./extensions/ScriptExtended.s.sol";
 import { OnchainExecutor } from "./OnchainExecutor.s.sol"; // cheat to load artifact to parent `out` directory
 import { IMigrationScript } from "./interfaces/IMigrationScript.sol";
 import { LibProxy } from "./libraries/LibProxy.sol";
+import { LibInitializeGuard } from "./libraries/LibInitializeGuard.sol";
 import { DefaultContract } from "./utils/DefaultContract.sol";
 import { ProxyInterface, LibDeploy, DeployInfo, UpgradeInfo } from "./libraries/LibDeploy.sol";
 import { cheatBroadcast } from "./utils/Helpers.sol";
@@ -38,6 +40,18 @@ abstract contract BaseMigration is ScriptExtended {
     ProxyInterface /* proxyInterface */
   ) internal virtual { }
 
+  function _beforeRunningScript() internal virtual override {
+    vm.recordLogs();
+    vm.startStateDiffRecording();
+  }
+
+  function _afterRunningScript() internal virtual override {
+    Vm.Log[] memory recordedLogs = vm.getRecordedLogs();
+    Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff();
+
+    LibInitializeGuard.validate({ logs: recordedLogs, stateDiffs: stateDiffs });
+  }
+
   function _sharedArguments() internal virtual returns (bytes memory rawSharedArgs);
 
   function _injectDependencies() internal virtual { }

script/configs/NetworkConfig.sol

@@ -74,12 +74,18 @@ abstract contract NetworkConfig is INetworkConfig {
 
   function setForkMode(bool shouldEnable) public virtual {
     _isForkModeEnabled = shouldEnable;
+    emit ForkModeUpdated(shouldEnable);
   }
 
   function getNetworkData(TNetwork network) public view virtual returns (NetworkData memory) {
     return _networkDataMap[network];
   }
 
+  function getNetworkTypeByForkId(uint256 forkId) public view virtual returns (TNetwork network) {
+    network = _forkId2Network[forkId];
+    if (network == TNetwork.wrap(0x0)) network = DefaultNetwork.LocalHost.key();
+  }
+
   function getDeploymentDirectory(TNetwork network) public view virtual returns (string memory dirPath) {
     string memory dirName = network.dir();
     require(bytes(dirName).length != 0, "NetworkConfig: Deployment directory not found");

script/extensions/ScriptExtended.s.sol

@@ -83,9 +83,13 @@ abstract contract ScriptExtended is BaseScriptExtended, Script, StdAssertions, I
       console.log("ScriptExtended:".blue(), "Prechecking completed in", vm.toString(end - start), "milliseconds.\n");
     }
 
+    _beforeRunningScript();
+
     (bool success, bytes memory data) = address(this).delegatecall(callData);
     success.handleRevert(msg.sig, data);
 
+    _afterRunningScript();
+
     if (vme.getRuntimeConfig().disablePostcheck) {
       console.log("\nPostchecking is disabled.".yellow());
     } else {
@@ -99,6 +103,10 @@ abstract contract ScriptExtended is BaseScriptExtended, Script, StdAssertions, I
     }
   }
 
+  function _beforeRunningScript() internal virtual { }
+
+  function _afterRunningScript() internal virtual { }
+
   function _requireOn(TNetwork networkType) private view {
     require(network() == networkType, string.concat("ScriptExtended: Only allowed on ", vme.getAlias(networkType)));
   }
@@ -137,7 +145,7 @@ abstract contract ScriptExtended is BaseScriptExtended, Script, StdAssertions, I
   }
 
   function prankOrBroadcast(address by) internal virtual {
-    if (vme.isPostChecking()) {
+    if (vme.isPostChecking() || vme.isPreChecking()) {
       vm.prank(by);
     } else {
       vm.broadcast(by);

script/interfaces/configs/INetworkConfig.sol

@@ -13,6 +13,8 @@ interface INetworkConfig {
     string explorer;
   }
 
+  event ForkModeUpdated(bool enabled);
+
   function setNetworkInfo(NetworkData memory networkData) external;
 
   function setForkMode(bool shouldEnable) external;
@@ -25,6 +27,8 @@ interface INetworkConfig {
 
   function getNetworkData(TNetwork network) external view returns (NetworkData memory);
 
+  function getNetworkTypeByForkId(uint256 forkId) external view returns (TNetwork network);
+
   function getForkId(TNetwork network) external view returns (uint256 forkId);
 
   function getForkId(TNetwork, uint256 forkBlockNumber) external view returns (uint256 forkId);

script/libraries/LibErrorHandler.sol

@@ -1,9 +1,9 @@
-// SPDX-License-Identifier: MIT
-pragma solidity ^0.8.0;
+// SPDX-License-Identifier: MIT OR Apache-2.0
+pragma solidity >=0.6.2 <0.9.0;
+pragma experimental ABIEncoderV2;
 
 library LibErrorHandler {
   /// @dev Reserves error definition to upload to signature database.
-  error ExternalCallFailed(bytes4 msgSig, bytes4 callSig);
 
   /// @notice handle low level call revert if call failed,
   /// If external call return empty bytes, reverts with custom error.