Skip to content

Commit

Permalink
tmp
Browse files Browse the repository at this point in the history
  • Loading branch information
@TuDo1403
TuDo1403 committed Aug 12, 2024
1 parent 38f2226 commit fe18da6
Show file tree
Hide file tree
Showing 2 changed files with 93 additions and 34 deletions.
109 changes: 93 additions & 16 deletions script/libraries/LibInitializeGuard.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,8 @@ 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 { LibString } from "../../dependencies/@solady-0.0.228/src/utils/LibString.sol";
import { Vm, VmSafe } 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";
Expand All @@ -24,25 +25,31 @@ interface IERC1967 {

library LibInitializeGuard {
using JSONParserLib for *;
using LibString for *;
using stdStorage for StdStorage;
using EnumerableSet for EnumerableSet.AddressSet;

struct InitializedSlot {
bool found;
bytes32 slot;
uint256 offset;
uint256 bitOffset;
uint256 nBits;
}

struct Cache {
EnumerableSet.AddressSet _proxies;
mapping(address addr => uint256) _initVersion;
EnumerableSet.AddressSet _logics;
mapping(address addr => uint256[]) _initVersions;
mapping(address addr => Vm.ChainInfo) _chainInfo;
mapping(address proxy => InitializedSlot) _initSlot;
mapping(address proxy => address logic) _proxy2Logic;
mapping(address logic => address proxy) _logic2Proxy;
}

uint256 internal constant MAX_INIT_VERSION_OZV4 = type(uint8).max;
uint256 internal constant MAX_INIT_VERSION_OZV5 = type(uint64).max;
uint256 internal constant NUM_BIT_OZV4 = 8;
uint256 internal constant NUM_BIT_OZV5 = 64;

// keccak256(abi.encode(uint256(keccak256("openzeppelin.storage.Initializable")) - 1)) & ~bytes32(uint256(0xff))
bytes32 private constant INITIALIZABLE_STORAGE_OZV5 =
Expand Down Expand Up @@ -83,30 +90,64 @@ library LibInitializeGuard {
bytes32 eventSig = logs[i].topics[0];

if (eventSig == InitializableOZV4.Initialized.selector) {
$._initVersion[emitter] = abi.decode(logs[i].data, (uint8));
$._initVersions[emitter].push(abi.decode(logs[i].data, (uint8)));
}
if (eventSig == InitializableOZV5.Initialized.selector) {
$._initVersion[emitter] = abi.decode(logs[i].data, (uint64));
$._initVersions[emitter].push(abi.decode(logs[i].data, (uint64)));
}
if (eventSig == IERC1967.Upgraded.selector) {
$._proxies.add(emitter);
$._proxy2Logic[emitter] = address(uint160(uint256(logs[i].topics[1])));
address logic = address(uint160(uint256(logs[i].topics[1])));
$._proxy2Logic[emitter] = logic;
$._logics.add(logic);
$._logic2Proxy[logic] = emitter;
}
}

for (uint256 i; i < stateDiffs.length; ++i) {
if ($._proxies.contains(stateDiffs[i].account) && !$._initSlot[stateDiffs[i].account].found) {
if ($._proxies.contains(stateDiffs[i].account)) {
address proxy = stateDiffs[i].account;

$._chainInfo[proxy] = stateDiffs[i].chainInfo;
$._initSlot[proxy] = _getInitializedSlot($, proxy);
if (!$._initSlot[proxy].found) {
$._chainInfo[proxy] = stateDiffs[i].chainInfo;
$._initSlot[proxy] = _getInitializedSlot($, proxy);
}
}
}

for (uint256 i; i < stateDiffs.length; ++i) {
if ($._logics.contains(stateDiffs[i].account)) {
if (stateDiffs[i].kind != VmSafe.AccountAccessKind.DelegateCall) continue;

Vm.StorageAccess[] memory sAccesses = stateDiffs[i].storageAccesses;
address proxy = $._logic2Proxy[stateDiffs[i].account];

for (uint256 j; j < sAccesses.length; ++j) {
if (!(sAccesses[j].isWrite && sAccesses[j].account == proxy && sAccesses[j].slot == $._initSlot[proxy].slot))
{
continue;
}

uint256 mask = (1 << $._initSlot[proxy].nBits) - 1;
uint256 bitOffset = $._initSlot[proxy].bitOffset;

uint256 prevInitVersion = (uint256(sAccesses[j].previousValue) >> bitOffset) & mask;
uint256 newInitVersion = (uint256(sAccesses[j].newValue) >> bitOffset) & mask;

if (prevInitVersion == newInitVersion) continue;

require(newInitVersion == prevInitVersion + 1, "LibInitializeGuard: Invalid initialized version increment!");
console.log(
"[INIT] %s: v%s -> v%s", vm.getLabel(proxy), vm.toString(prevInitVersion), vm.toString(newInitVersion)
);
}
}
}

address[] memory proxies = $._proxies.values();
for (uint256 i; i < proxies.length; ++i) {
uint256 proxyVersion = $._initVersion[proxies[i]];
uint256 logicVersion = $._initVersion[$._proxy2Logic[proxies[i]]];
uint256 proxyVersion = $._initVersions[proxies[i]][0];
uint256 logicVersion = $._initVersions[$._proxy2Logic[proxies[i]]][0];

require(
$._initSlot[proxies[i]].found,
Expand All @@ -123,6 +164,40 @@ library LibInitializeGuard {
require(
proxyVersion != 0, string.concat("LibInitializeGuard: Proxy ", vm.getLabel(proxies[i]), " does not initialize!")
);

uint256 initCount = _getInitializeFnCount($, proxies[i]);
require(
proxyVersion == initCount,
string.concat(
"LibInitializeGuard: Invalid initialized version!",
"Expected: ",
vm.toString(initCount),
" Got: ",
vm.toString(proxyVersion)
)
);
}
}

function _getInitializeFnCount(Cache storage $cache, address proxy) internal returns (uint256 count) {
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] = "methodIdentifiers";

string memory ret = vm.toLowercase(string(vm.ffi(inputs)));
string[] memory allKeys = vm.parseJsonKeys(ret, ".");

for (uint256 i; i < allKeys.length; ++i) {
if (allKeys[i].contains("initialize")) {
count++;
}
}
}

Expand All @@ -147,21 +222,23 @@ library LibInitializeGuard {

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();
initSlot.slot = bytes32(vm.parseUint(storageSlot.at('"slot"').value().decodeString()));
initSlot.bitOffset = storageSlot.at('"offset"').value().parseUint() * 8;
initSlot.nBits = NUM_BIT_OZV4;

return initSlot;
}
}

uint256 initVersion = $cache._initVersion[proxy];
if (!initSlot.found && initVersion != 0) {
uint256 vTail = $cache._initVersions[proxy].length - 1;
if ($cache._initVersions[proxy][vTail] != 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;
initSlot.bitOffset = 0;
initSlot.nBits = NUM_BIT_OZV5;
}
}
}
Expand Down
18 changes: 0 additions & 18 deletions soldeer.lock
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,24 +17,6 @@ version = "0.0.228"
source = "https://github.com/Vectorized/solady/archive/refs/tags/v0.0.228.zip"
checksum = "90c6c562403b67338e3fa5fbad36a80b129df333d89db6759c5fd3ed15d4f0da"

[[dependencies]]
name = "@openzeppelin-5.0.2"
version = "5.0.2"
source = "https://github.com/OpenZeppelin/openzeppelin-contracts/archive/refs/tags/v5.0.2.zip"
checksum = "ca49e0776066328da0087977864dfaef0d5e54a0ea9859236a7cd9ad02abb9e5"

[[dependencies]]
name = "@openzeppelin"
version = "5.0.2"
source = "https://github.com/OpenZeppelin/openzeppelin-contracts/archive/refs/tags/v5.0.2.zip"
checksum = "ca49e0776066328da0087977864dfaef0d5e54a0ea9859236a7cd9ad02abb9e5"

[[dependencies]]
name = "@openzeppelinV5"
version = "5.0.2"
source = "https://github.com/OpenZeppelin/openzeppelin-contracts/archive/refs/tags/v5.0.2.zip"
checksum = "ca49e0776066328da0087977864dfaef0d5e54a0ea9859236a7cd9ad02abb9e5"

[[dependencies]]
name = "@openzeppelin-v5"
version = "5.0.2"
Expand Down

0 comments on commit fe18da6

Please sign in to comment.