-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Do Tu
authored and
Do Tu
committed
Aug 12, 2024
1 parent
ab9eed2
commit 38f2226
Showing
14 changed files
with
304 additions
and
12 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,184 @@ | ||
| // SPDX-License-Identifier: MIT OR Apache-2.0 | ||
| pragma solidity >=0.6.2 <0.9.0; | ||
| pragma experimental ABIEncoderV2; | ||
|
|
||
| import { EnumerableSet } from "../../dependencies/@openzeppelin-4.9.3/contracts/utils/structs/EnumerableSet.sol"; | ||
| import { JSONParserLib } from "../../dependencies/@solady-0.0.228/src/utils/JSONParserLib.sol"; | ||
| import { Vm } from "../../dependencies/@forge-std-1.9.1/src/Vm.sol"; | ||
| import { StdStorage, stdStorage } from "../../dependencies/@forge-std-1.9.1/src/StdStorage.sol"; | ||
| import { console, vm, vme } from "../utils/Helpers.sol"; | ||
| import { TNetwork } from "../types/TNetwork.sol"; | ||
| import { TContract } from "../types/TContract.sol"; | ||
|
|
||
| interface InitializableOZV4 { | ||
| event Initialized(uint8); | ||
| } | ||
|
|
||
| interface InitializableOZV5 { | ||
| event Initialized(uint64); | ||
| } | ||
|
|
||
| interface IERC1967 { | ||
| event Upgraded(address indexed); | ||
| } | ||
|
|
||
| library LibInitializeGuard { | ||
| using JSONParserLib for *; | ||
| using stdStorage for StdStorage; | ||
| using EnumerableSet for EnumerableSet.AddressSet; | ||
|
|
||
| struct InitializedSlot { | ||
| bool found; | ||
| bytes32 slot; | ||
| uint256 offset; | ||
| } | ||
|
|
||
| struct Cache { | ||
| EnumerableSet.AddressSet _proxies; | ||
| mapping(address addr => uint256) _initVersion; | ||
| mapping(address addr => Vm.ChainInfo) _chainInfo; | ||
| mapping(address proxy => InitializedSlot) _initSlot; | ||
| mapping(address proxy => address logic) _proxy2Logic; | ||
| } | ||
|
|
||
| uint256 internal constant MAX_INIT_VERSION_OZV4 = type(uint8).max; | ||
| uint256 internal constant MAX_INIT_VERSION_OZV5 = type(uint64).max; | ||
|
|
||
| // keccak256(abi.encode(uint256(keccak256("openzeppelin.storage.Initializable")) - 1)) & ~bytes32(uint256(0xff)) | ||
| bytes32 private constant INITIALIZABLE_STORAGE_OZV5 = | ||
| 0xf0c57e16840df040f15088dc2f81fe391c3923bec73e23a9662efc9c229c6a00; | ||
|
|
||
| function start() internal { | ||
| bool started; | ||
| try vme.getUserDefinedConfig("LibInitializeGuard") returns (bytes memory raw) { | ||
| if (raw.length != 0) { | ||
| started = abi.decode(raw, (bool)); | ||
| } | ||
| } catch { } | ||
|
|
||
| require(!started, "LibInitializeGuard: already initialized"); | ||
| vme.setUserDefinedConfig("LibInitializeGuard", abi.encode(true)); | ||
|
|
||
| vm.recordLogs(); | ||
| vm.startStateDiffRecording(); | ||
| } | ||
|
|
||
| function end() internal { | ||
| bool started; | ||
| try vme.getUserDefinedConfig("LibInitializeGuard") returns (bytes memory raw) { | ||
| if (raw.length != 0) { | ||
| started = abi.decode(raw, (bool)); | ||
| } | ||
| } catch { } | ||
|
|
||
| require(started, "LibInitializeGuard: not initialized"); | ||
| vme.setUserDefinedConfig("LibInitializeGuard", abi.encode(false)); | ||
|
|
||
| Vm.Log[] memory logs = vm.getRecordedLogs(); | ||
| Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); | ||
| Cache storage $ = _getCacheStorage(); | ||
|
|
||
| for (uint256 i; i < logs.length; ++i) { | ||
| address emitter = logs[i].emitter; | ||
| bytes32 eventSig = logs[i].topics[0]; | ||
|
|
||
| if (eventSig == InitializableOZV4.Initialized.selector) { | ||
| $._initVersion[emitter] = abi.decode(logs[i].data, (uint8)); | ||
| } | ||
| if (eventSig == InitializableOZV5.Initialized.selector) { | ||
| $._initVersion[emitter] = abi.decode(logs[i].data, (uint64)); | ||
| } | ||
| if (eventSig == IERC1967.Upgraded.selector) { | ||
| $._proxies.add(emitter); | ||
| $._proxy2Logic[emitter] = address(uint160(uint256(logs[i].topics[1]))); | ||
| } | ||
| } | ||
|
|
||
| for (uint256 i; i < stateDiffs.length; ++i) { | ||
| if ($._proxies.contains(stateDiffs[i].account) && !$._initSlot[stateDiffs[i].account].found) { | ||
| address proxy = stateDiffs[i].account; | ||
|
|
||
| $._chainInfo[proxy] = stateDiffs[i].chainInfo; | ||
| $._initSlot[proxy] = _getInitializedSlot($, proxy); | ||
| } | ||
| } | ||
|
|
||
| address[] memory proxies = $._proxies.values(); | ||
| for (uint256 i; i < proxies.length; ++i) { | ||
| uint256 proxyVersion = $._initVersion[proxies[i]]; | ||
| uint256 logicVersion = $._initVersion[$._proxy2Logic[proxies[i]]]; | ||
|
|
||
| require( | ||
| $._initSlot[proxies[i]].found, | ||
| string.concat("LibInitializeGuard: Proxy ", vm.getLabel(proxies[i]), " does not have `_initialized` slot!") | ||
| ); | ||
|
|
||
| require( | ||
| logicVersion == MAX_INIT_VERSION_OZV4 || logicVersion == MAX_INIT_VERSION_OZV5, | ||
| string.concat( | ||
| "LibInitializeGuard: Logic of proxy ", vm.getLabel(proxies[i]), " does not disable initialized version!" | ||
| ) | ||
| ); | ||
|
|
||
| require( | ||
| proxyVersion != 0, string.concat("LibInitializeGuard: Proxy ", vm.getLabel(proxies[i]), " does not initialize!") | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| function _getInitializedSlot(Cache storage $cache, address proxy) internal returns (InitializedSlot memory initSlot) { | ||
| Vm.ChainInfo memory chainInfo = $cache._chainInfo[proxy]; | ||
| TNetwork networkType = vme.getNetworkTypeByForkId(chainInfo.forkId); | ||
| TContract contractType = vme.getContractTypeByRawData(networkType, proxy); | ||
| string memory contractName = vme.getContractName(contractType); | ||
|
|
||
| string[] memory inputs = new string[](4); | ||
| inputs[0] = "forge"; | ||
| inputs[1] = "inspect"; | ||
| inputs[2] = contractName; | ||
| inputs[3] = "storage"; | ||
|
|
||
| string memory ret = string(vm.ffi(inputs)); | ||
| JSONParserLib.Item memory layout = ret.parse().at('"storage"'); | ||
| uint256 layoutSize = layout.size(); | ||
|
|
||
| for (uint256 i; i < layoutSize; ++i) { | ||
| JSONParserLib.Item memory storageSlot = layout.at(i); | ||
|
|
||
| if (keccak256(bytes(storageSlot.at('"label"').value().decodeString())) == keccak256("_initialized")) { | ||
| initSlot.found = true; | ||
| initSlot.slot = vm.parseBytes32(storageSlot.at('"slot"').value().decodeString()); | ||
| initSlot.offset = storageSlot.at('"offset"').value().parseUint(); | ||
|
|
||
| return initSlot; | ||
| } | ||
| } | ||
|
|
||
| uint256 initVersion = $cache._initVersion[proxy]; | ||
| if (!initSlot.found && initVersion != 0) { | ||
| // assume proxy use `Initializable` from OpenZeppelin v5 | ||
| bytes32 slotValue = vm.load(proxy, INITIALIZABLE_STORAGE_OZV5); | ||
| if (slotValue != 0) { | ||
| initSlot.found = true; | ||
| initSlot.slot = INITIALIZABLE_STORAGE_OZV5; | ||
| initSlot.offset = 0; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| function _getCacheStorage() internal pure returns (Cache storage $) { | ||
| bytes32 slot = keccak256("LibInitializeGuard.Cache.storage.slot"); | ||
|
|
||
| assembly ("memory-safe") { | ||
| $.slot := slot | ||
| } | ||
| } | ||
|
|
||
| function _getStdStorage() internal pure returns (StdStorage storage $) { | ||
| bytes32 slot = keccak256("StdStorage.storage.slot"); | ||
|
|
||
| assembly ("memory-safe") { | ||
| $.slot := slot | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,22 @@ | ||
| // SPDX-License-Identifier: MIT | ||
| pragma solidity ^0.8.19; | ||
|
|
||
| import { Initializable } from "@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; | ||
|
|
||
| contract MockInitialize is Initializable { | ||
| address private _addr; | ||
|
|
||
| constructor() { | ||
| _disableInitializers(); | ||
| } | ||
|
|
||
| function initialize() public initializer { } | ||
|
|
||
| function initializeV2() external reinitializer(2) { } | ||
|
|
||
| function initializeV3(address a) external reinitializer(3) { | ||
| _addr = a; | ||
| } | ||
|
|
||
| function initializeV4() external reinitializer(4) { } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.