Skip to content

Commit

Permalink
feat(LibInitializeGuard): add guard for safely upgrade contract with …
Browse files Browse the repository at this point in the history
…reinitializer
  • Loading branch information
Do Tu authored and Do Tu committed Aug 13, 2024
1 parent ab9eed2 commit bb5463c
Show file tree
Hide file tree
Showing 12 changed files with 386 additions and 12 deletions.
3 changes: 2 additions & 1 deletion foundry.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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" }
3 changes: 2 additions & 1 deletion remappings.txt
Original file line number Diff line number Diff line change
@@ -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
14 changes: 14 additions & 0 deletions script/BaseMigration.s.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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 { }
Expand Down
6 changes: 6 additions & 0 deletions script/configs/NetworkConfig.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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");
Expand Down
10 changes: 9 additions & 1 deletion script/extensions/ScriptExtended.s.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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)));
}
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 4 additions & 0 deletions script/interfaces/configs/INetworkConfig.sol
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,8 @@ interface INetworkConfig {
string explorer;
}

event ForkModeUpdated(bool enabled);

function setNetworkInfo(NetworkData memory networkData) external;

function setForkMode(bool shouldEnable) external;
Expand All @@ -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);
Expand Down
6 changes: 3 additions & 3 deletions script/libraries/LibErrorHandler.sol
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
Loading

0 comments on commit bb5463c

Please sign in to comment.