diff --git a/script/sample/contracts/SampleProxyDeploy.s.sol b/script/sample/contracts/SampleProxyDeploy.s.sol index 9c22d71..75ba96c 100644 --- a/script/sample/contracts/SampleProxyDeploy.s.sol +++ b/script/sample/contracts/SampleProxyDeploy.s.sol @@ -16,7 +16,7 @@ contract SampleProxyDeploy is SampleMigration { instance = SampleProxy(_deployProxy(Contract.SampleProxy.key())); assertEq(instance.getMessage(), ISharedArgument(address(vme)).sharedArguments().proxyMessage); - vm.prank(sender()); - instance.initializeV4(); + // vm.prank(sender()); + // instance.initializeV4(); } } diff --git a/script/sample/utils/Contract.sol b/script/sample/utils/Contract.sol index 5514669..dc58b20 100644 --- a/script/sample/utils/Contract.sol +++ b/script/sample/utils/Contract.sol @@ -13,7 +13,8 @@ enum Contract { tWETH, Sample, SampleClone, - SampleProxy + SampleProxy, + SampleProxyForTestingPurpose } using { key, name } for Contract global; @@ -31,5 +32,6 @@ function name(Contract contractEnum) pure returns (string memory) { if (contractEnum == Contract.tWRON) return "tWRON"; if (contractEnum == Contract.SampleClone) return "SampleClone"; if (contractEnum == Contract.SampleProxy) return "SampleProxy"; + if (contractEnum == Contract.SampleProxyForTestingPurpose) return "SampleProxyForTestingPurpose"; revert("Contract: Unknown contract"); } diff --git a/src/mocks/ForTesting/InitializableTesting.sol b/src/mocks/ForTesting/InitializableTesting.sol new file mode 100644 index 0000000..eeee1f0 --- /dev/null +++ b/src/mocks/ForTesting/InitializableTesting.sol @@ -0,0 +1,115 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; +import "@openzeppelin-4.9.3/contracts/utils/Address.sol"; + +abstract contract InitializableTesting { + /** + * @dev Indicates that the contract has been initialized. + * @custom:oz-retyped-from bool + */ + uint8 internal _initialized; + + /** + * @dev Indicates that the contract is in the process of being initialized. + */ + bool internal _initializing; + + /** + * @dev Triggered when the contract has been initialized or reinitialized. + */ + event Initialized(uint8 version); + + /** + * @dev A modifier that defines a protected initializer function that can be invoked at most once. In its scope, + * `onlyInitializing` functions can be used to initialize parent contracts. + * + * Similar to `reinitializer(1)`, except that functions marked with `initializer` can be nested in the context of a + * constructor. + * + * Emits an {Initialized} event. + */ + modifier initializer() { + bool isTopLevelCall = !_initializing; + require( + (isTopLevelCall && _initialized < 1) || (!Address.isContract(address(this)) && _initialized == 1), + "Initializable: contract is already initialized" + ); + _initialized = 1; + if (isTopLevelCall) { + _initializing = true; + } + _; + if (isTopLevelCall) { + _initializing = false; + emit Initialized(1); + } + } + + /** + * @dev A modifier that defines a protected reinitializer function that can be invoked at most once, and only if the + * contract hasn't been initialized to a greater version before. In its scope, `onlyInitializing` functions can be + * used to initialize parent contracts. + * + * A reinitializer may be used after the original initialization step. This is essential to configure modules that + * are added through upgrades and that require initialization. + * + * When `version` is 1, this modifier is similar to `initializer`, except that functions marked with `reinitializer` + * cannot be nested. If one is invoked in the context of another, execution will revert. + * + * Note that versions can jump in increments greater than 1; this implies that if multiple reinitializers coexist in + * a contract, executing them in the right order is up to the developer or operator. + * + * WARNING: setting the version to 255 will prevent any future reinitialization. + * + * Emits an {Initialized} event. + */ + modifier reinitializer(uint8 version) { + require(!_initializing && _initialized < version, "Initializable: contract is already initialized"); + _initialized = version; + _initializing = true; + _; + _initializing = false; + emit Initialized(version); + } + + /** + * @dev Modifier to protect an initialization function so that it can only be invoked by functions with the + * {initializer} and {reinitializer} modifiers, directly or indirectly. + */ + modifier onlyInitializing() { + require(_initializing, "Initializable: contract is not initializing"); + _; + } + + /** + * @dev Locks the contract, preventing any future reinitialization. This cannot be part of an initializer call. + * Calling this in the constructor of a contract will prevent that contract from being initialized or reinitialized + * to any version. It is recommended to use this to lock implementation contracts that are designed to be called + * through proxies. + * + * Emits an {Initialized} event the first time it is successfully executed. + */ + function _disableInitializers() internal virtual { + require(!_initializing, "Initializable: contract is initializing"); + if (_initialized != type(uint8).max) { + _initialized = type(uint8).max; + emit Initialized(type(uint8).max); + } + } + + /** + * @dev Returns the highest version that has been initialized. See {reinitializer}. + */ + function _getInitializedVersion() internal view returns (uint8) { + return _initialized; + } + + /** + * @dev Returns `true` if the contract is currently initializing. See {onlyInitializing}. + */ + function _isInitializing() internal view returns (bool) { + return _initializing; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose.sol new file mode 100644 index 0000000..6ac0036 --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose.sol @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Initializable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { Initializable as InitializableV5 } from + "../../../dependencies/@openzeppelin-v5-5.0.2/contracts/proxy/utils/Initializable.sol"; +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; + +contract SampleProxyForTestingPurpose is Ownable, InitializableV5 { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + + constructor() { + // _disableInitializers(); + } + + function initialize(string calldata message) external initializer { + _message = message; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose2.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose2.sol new file mode 100644 index 0000000..7bd272f --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose2.sol @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; +import "./InitializableTesting.sol"; + +contract SampleProxyForTestingPurpose2 is Ownable, InitializableTesting { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) public { + _message = message; + } + + function abcXYZ(string calldata message) public { } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose3.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose3.sol new file mode 100644 index 0000000..d797333 --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose3.sol @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; +import "./InitializableTesting.sol"; + +contract SampleProxyForTestingPurpose3 is Ownable, InitializableTesting { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) public initializer { + _message = message; + _initialized = type(uint8).max - 10; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose4.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose4.sol new file mode 100644 index 0000000..29a39f3 --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose4.sol @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Initializable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { Initializable as InitializableV5 } from + "../../../dependencies/@openzeppelin-v5-5.0.2/contracts/proxy/utils/Initializable.sol"; +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; + +contract SampleProxyForTestingPurpose4 is Ownable, Initializable { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) external initializer { + _message = message; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose5.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose5.sol new file mode 100644 index 0000000..a53b1ec --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose5.sol @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Initializable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { Initializable as InitializableV5 } from + "../../../dependencies/@openzeppelin-v5-5.0.2/contracts/proxy/utils/Initializable.sol"; +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; + +contract SampleProxyForTestingPurpose5 is Ownable, Initializable { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + uint256 internal _newVariable; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) external initializer { + _message = message; + } + + function initializeV2(uint256 newValues) external reinitializer(2) { + _newVariable = newValues; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose6.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose6.sol new file mode 100644 index 0000000..5984835 --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose6.sol @@ -0,0 +1,40 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Initializable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { Initializable as InitializableV5 } from + "../../../dependencies/@openzeppelin-v5-5.0.2/contracts/proxy/utils/Initializable.sol"; +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; + +contract SampleProxyForTestingPurpose6 is Ownable, Initializable { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + uint256 internal _newVariable; + uint256 internal _newVariable2; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) external initializer { + _message = message; + } + + function initializeV2(uint256 newValues) external reinitializer(2) { + _newVariable = newValues; + } + + function initializeV3(uint256 newValues) external reinitializer(3) { + _newVariable2 = newValues; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/src/mocks/ForTesting/SampleProxyForTestingPurpose7.sol b/src/mocks/ForTesting/SampleProxyForTestingPurpose7.sol new file mode 100644 index 0000000..b04a3f5 --- /dev/null +++ b/src/mocks/ForTesting/SampleProxyForTestingPurpose7.sol @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.25; + +import { Initializable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { Initializable as InitializableV5 } from + "../../../dependencies/@openzeppelin-v5-5.0.2/contracts/proxy/utils/Initializable.sol"; +import { Ownable } from "../../../dependencies/@openzeppelin-4.9.3/contracts/access/Ownable.sol"; + +contract SampleProxyForTestingPurpose7 is Ownable, Initializable { + uint256[50] private __gap; + + string internal _message; + address internal _addr; + uint256 internal _newVariable; + uint256 internal _newVariable2; + + constructor() { + _disableInitializers(); + } + + function initialize(string calldata message) external initializer { + _message = message; + } + + function initializeV2(uint256 newValues) external reinitializer(2) { + _newVariable = newValues; + } + + function initializeV3(uint256 newValues) external reinitializer(3) { + _newVariable2 = newValues; + } + + function initializeV4(uint256 newValues) external reinitializer(3) { + _newVariable2 = newValues; + } + + function initializeV5(uint256 newValues) external reinitializer(5) { + _newVariable2 = newValues; + } + + function setMessage(string memory message) public { + _message = message; + } + + function getMessage() public view returns (string memory) { + return _message; + } +} diff --git a/test/LibInitializeGuard.t.sol b/test/LibInitializeGuard.t.sol new file mode 100644 index 0000000..af46528 --- /dev/null +++ b/test/LibInitializeGuard.t.sol @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.0; + +import { Test } from "../dependencies/@forge-std-1.9.1/src/Test.sol"; +import { vme, Vm } from "script/utils/Constants.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 { BaseGeneralConfig } from "script/BaseGeneralConfig.sol"; +import { BaseMigration } from "script/BaseMigration.s.sol"; +import { Initializable } from "../dependencies/@openzeppelin-4.9.3/contracts/proxy/utils/Initializable.sol"; +import { LibInitializeGuard } from "@fdk/libraries/LibInitializeGuard.sol"; +import { SampleProxy } from "src/mocks/SampleProxy.sol"; +import { SampleProxyDeploy } from "script/sample/contracts/SampleProxyDeploy.s.sol"; +import { LibProxy } from "script/libraries/LibProxy.sol"; +import { MockConfig } from "./MockConfig.sol"; +import { SampleProxyForTestingPurpose2 } from "src/mocks/ForTesting/SampleProxyForTestingPurpose2.sol"; +import { SampleProxyForTestingPurpose5 } from "src/mocks/ForTesting/SampleProxyForTestingPurpose5.sol"; +import { SampleProxyForTestingPurpose6 } from "src/mocks/ForTesting/SampleProxyForTestingPurpose6.sol"; +import { SampleProxyForTestingPurpose7 } from "src/mocks/ForTesting/SampleProxyForTestingPurpose7.sol"; + +interface ITransparentUpgradeableProxy { + function upgradeTo(address) external; + function upgradeToAndCall(address, bytes memory) external payable; +} + +contract ValidateWrapper { + function runValidate(Vm.Log[] memory logs, Vm.AccountAccess[] memory stateDiffs) public { + LibInitializeGuard.validate(logs, stateDiffs); + } +} + +contract LibInitializeGuardTest is Test { + using LibProxy for address; + using StdStyle for *; + + SampleProxy _sample; + + function setUp() public { + deployCodeTo("MockConfig.sol:MockConfig", abi.encode(""), 0, address(vme)); + } + + function testConcrete_DeployProxy_ForTheFirstTime() public { + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose4"); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + _sample = new SampleProxyDeploy().run(); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + LibInitializeGuard.validate(logs, stateDiffs); + } + + function testConcrete_Upgrade_WithoutInitialization() public { + testConcrete_DeployProxy_ForTheFirstTime(); + _sample = new SampleProxyDeploy().run(); + address admin = address(_sample).getProxyAdmin(); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + address newLogic = address(new SampleProxyForTestingPurpose2()); + vm.prank(admin); + ITransparentUpgradeableProxy(address(_sample)).upgradeTo(newLogic); + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose2"); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + LibInitializeGuard.validate(logs, stateDiffs); + } + + function testConcrete_UpgradeFrom_V1_To_V2_ByCallingInitializeV2() public { + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose4"); + + _sample = new SampleProxyDeploy().run(); + address admin = address(_sample).getProxyAdmin(); + bytes memory callData = abi.encodeCall(SampleProxyForTestingPurpose5.initializeV2, (100)); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + address newLogic = address(new SampleProxyForTestingPurpose5()); + vm.prank(admin); + ITransparentUpgradeableProxy(address(_sample)).upgradeToAndCall(newLogic, callData); + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose5"); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + LibInitializeGuard.validate(logs, stateDiffs); + } + + function testConcrete_UpgradeFrom_V2_To_V3_ByCallingInitializeV3() public { + testConcrete_UpgradeFrom_V1_To_V2_ByCallingInitializeV2(); + address admin = address(_sample).getProxyAdmin(); + bytes memory callData = abi.encodeCall(SampleProxyForTestingPurpose6.initializeV3, (100)); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + address newLogic = address(new SampleProxyForTestingPurpose6()); + vm.prank(admin); + ITransparentUpgradeableProxy(address(_sample)).upgradeToAndCall(newLogic, callData); + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose6"); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + LibInitializeGuard.validate(logs, stateDiffs); + } + + function testRevert_When_UpgradeFrom_V3_To_V5_ByCallingInitilizeV5() public { + testConcrete_UpgradeFrom_V2_To_V3_ByCallingInitializeV3(); + address admin = address(_sample).getProxyAdmin(); + bytes memory callData = abi.encodeCall(SampleProxyForTestingPurpose7.initializeV5, (100)); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + address newLogic = address(new SampleProxyForTestingPurpose7()); + vm.prank(admin); + ITransparentUpgradeableProxy(address(_sample)).upgradeToAndCall(newLogic, callData); + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose7"); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + ValidateWrapper _wrapper = new ValidateWrapper(); + vm.expectRevert("LibInitializeGuard: Version does not correctly increment!"); + _wrapper.runValidate(logs, stateDiffs); + } + + function testRevert_When_FoundFourInitializeFunctions_But_ActualInitVerIsOne() public { + vm.recordLogs(); + vm.startStateDiffRecording(); + + _sample = new SampleProxyDeploy().run(); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + ValidateWrapper _wrapper = new ValidateWrapper(); + vm.expectRevert(bytes(string.concat("LibInitializeGuard: Invalid initialized version! Expected: 4 Got: 1"))); + _wrapper.runValidate(logs, stateDiffs); + } + + function testRevert_When_NotDisableInitializedVersion() public { + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose"); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + _sample = new SampleProxyDeploy().run(); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + ValidateWrapper _wrapper = new ValidateWrapper(); + vm.expectRevert( + bytes( + string.concat( + "LibInitializeGuard: Logic ", + vm.getLabel(address(_sample).getProxyImplementation()), + " did not disable initialized version!" + ) + ) + ); + _wrapper.runValidate(logs, stateDiffs); + } + + function testRevert_When_ProxyDoesNot_Initialize() public { + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose2"); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + _sample = new SampleProxyDeploy().run(); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + ValidateWrapper _wrapper = new ValidateWrapper(); + vm.expectRevert( + bytes(string.concat("LibInitializeGuard: Proxy ", vm.getLabel(address(_sample)), " does not initialize!".red())) + ); + _wrapper.runValidate(logs, stateDiffs); + } + + function testRevert_When_Does_Not_Correctly_Increment() public { + MockConfig(address(vme)).updateSampleProxyLogicForTesting("SampleProxyForTestingPurpose3"); + + vm.recordLogs(); + vm.startStateDiffRecording(); + + _sample = new SampleProxyDeploy().run(); + + Vm.Log[] memory logs = vm.getRecordedLogs(); + Vm.AccountAccess[] memory stateDiffs = vm.stopAndReturnStateDiff(); + + ValidateWrapper _wrapper = new ValidateWrapper(); + vm.expectRevert("LibInitializeGuard: Version does not correctly increment!"); + _wrapper.runValidate(logs, stateDiffs); + } +} diff --git a/test/MockConfig.sol b/test/MockConfig.sol new file mode 100644 index 0000000..56e6bd2 --- /dev/null +++ b/test/MockConfig.sol @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity >=0.6.2 <0.9.0; +pragma experimental ABIEncoderV2; + +import { console } from "../dependencies/@forge-std-1.9.1/src/console.sol"; + +import "script/sample/SampleGeneralConfig.sol"; + +contract MockConfig is SampleGeneralConfig { + function updateSampleProxyLogicForTesting(string memory contractName) public { + console.log("Cheating contract logic for testing..."); + _contractNameMap[Contract.SampleProxy.key()] = contractName; + } +}