fix: revert to custom initializable, make changes from oz 5.0

README.md

@@ -7,7 +7,6 @@ The implementation includes an upgradable modular account with two plugins (`Sin
 ## Important Callouts
 
 - **Not audited and should NOT be used in production**.
-- The Reference Implementation Account uses the storage slot provided by OpenZeppelin's Initializable library. This would prevent upgrades to and from other accounts that use the same library.
 - Not optimized in both deployments and execution. We’ve explicitly removed some optimizations for reader comprehension.
 
 ## Development

src/account/AccountStorage.sol

@@ -48,6 +48,9 @@ struct SelectorData {
 }
 
 struct AccountStorage {
+    // AccountStorageInitializable variables
+    uint8 initialized;
+    bool initializing;
     // Plugin metadata storage
     EnumerableSet.AddressSet plugins;
     mapping(address => PluginData) pluginData;

src/account/AccountStorageInitializable.sol

@@ -0,0 +1,65 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.25;
+
+import {AccountStorage, getAccountStorage} from "./AccountStorage.sol";
+
+abstract contract AccountStorageInitializable {
+    /**
+     * @dev The contract is already initialized.
+     */
+    error InvalidInitialization();
+
+    /**
+     * @dev The contract is not initializing.
+     */
+    error NotInitializing();
+
+    /**
+     * @dev Triggered when the contract has been initialized or reinitialized.
+     */
+    event Initialized(uint64 version);
+
+    /// @notice Modifier to put on function intended to be called only once per implementation
+    /// @dev Reverts if the contract has already been initialized
+    modifier initializer() {
+        AccountStorage storage $ = getAccountStorage();
+
+        // Cache values to avoid duplicated sloads
+        bool isTopLevelCall = !$.initializing;
+        uint64 initialized = $.initialized;
+
+        // Allowed calls:
+        // - initialSetup: the contract is not in the initializing state and no previous version was
+        //                 initialized
+        // - construction: the contract is initialized at version 1 (no reininitialization) and the
+        //                 current contract is just being deployed
+        bool initialSetup = initialized == 0 && isTopLevelCall;
+        bool construction = initialized == 1 && address(this).code.length == 0;
+
+        if (!initialSetup && !construction) {
+            revert InvalidInitialization();
+        }
+        $.initialized = 1;
+        if (isTopLevelCall) {
+            $.initializing = true;
+        }
+        _;
+        if (isTopLevelCall) {
+            $.initializing = false;
+            emit Initialized(1);
+        }
+    }
+
+    /// @notice Internal function to disable calls to initialization functions
+    /// @dev Reverts if the contract has already been initialized
+    function _disableInitializers() internal virtual {
+        AccountStorage storage $ = getAccountStorage();
+        if ($.initializing) {
+            revert InvalidInitialization();
+        }
+        if ($.initialized != type(uint8).max) {
+            $.initialized = type(uint8).max;
+            emit Initialized(type(uint8).max);
+        }
+    }
+}

src/account/UpgradeableModularAccount.sol

@@ -8,7 +8,6 @@ import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeab
 import {IERC165} from "@openzeppelin/contracts/utils/introspection/IERC165.sol";
 import {EnumerableMap} from "@openzeppelin/contracts/utils/structs/EnumerableMap.sol";
 import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
 
 import {FunctionReferenceLib} from "../helpers/FunctionReferenceLib.sol";
 import {_coalescePreValidation, _coalesceValidation} from "../helpers/ValidationDataHelpers.sol";
@@ -19,12 +18,13 @@ import {IStandardExecutor, Call} from "../interfaces/IStandardExecutor.sol";
 import {AccountExecutor} from "./AccountExecutor.sol";
 import {AccountLoupe} from "./AccountLoupe.sol";
 import {AccountStorage, getAccountStorage, getPermittedCallKey, SelectorData} from "./AccountStorage.sol";
+import {AccountStorageInitializable} from "./AccountStorageInitializable.sol";
 import {PluginManagerInternals} from "./PluginManagerInternals.sol";
 
 contract UpgradeableModularAccount is
     AccountExecutor,
     AccountLoupe,
-    Initializable,
+    AccountStorageInitializable,
     BaseAccount,
     IERC165,
     IPluginExecutor,

test/libraries/AccountStorage.t.sol

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {Test} from "forge-std/Test.sol";
+import {_ACCOUNT_STORAGE_SLOT, getPermittedCallKey} from "../../src/account/AccountStorage.sol";
+import {AccountStorageInitializable} from "../../src/account/AccountStorageInitializable.sol";
+import {MockDiamondStorageContract} from "../mocks/MockDiamondStorageContract.sol";
+import {ERC1967Proxy} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
+
+// Test implementation of AccountStorageInitializable which is contained in UpgradeableModularAccount
+contract AccountStorageTest is Test {
+    MockDiamondStorageContract public impl;
+    address public proxy;
+
+    function setUp() external {
+        impl = new MockDiamondStorageContract();
+        proxy = address(new ERC1967Proxy(address(impl), ""));
+    }
+
+    function test_storageSlotImpl() external {
+        // disable initializers sets value to uint8(max)
+        assertEq(uint256(vm.load(address(impl), _ACCOUNT_STORAGE_SLOT)), type(uint8).max);
+
+        // should revert if we try to initialize again
+        vm.expectRevert(AccountStorageInitializable.InvalidInitialization.selector);
+        impl.initialize();
+    }
+
+    function test_storageSlotProxy() external {
+        // before init, proxy's slot should be empty
+        assertEq(uint256(vm.load(proxy, _ACCOUNT_STORAGE_SLOT)), uint256(0));
+
+        MockDiamondStorageContract(proxy).initialize();
+        // post init slot should contains: packed(uint8 initialized = 1, bool initializing = 0)
+        assertEq(uint256(vm.load(proxy, _ACCOUNT_STORAGE_SLOT)), uint256(1));
+    }
+
+    function testFuzz_permittedCallKey(address addr, bytes4 selector) public {
+        bytes24 key = getPermittedCallKey(addr, selector);
+        assertEq(bytes20(addr), bytes20(key));
+        assertEq(bytes4(selector), bytes4(key << 160));
+    }
+}

test/mocks/MockDiamondStorageContract.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {AccountStorageInitializable} from "../../src/account/AccountStorageInitializable.sol";
+
+contract MockDiamondStorageContract is AccountStorageInitializable {
+    constructor() {
+        _disableInitializers();
+    }
+
+    // solhint-disable-next-line no-empty-blocks
+    function initialize() external initializer {}
+}