From f02f205994e37a50511fb788c615fa206eff17b1 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:06:19 -0400
Subject: [PATCH 01/10] chore: remove erc777

---
 src/plugins/TokenReceiverPlugin.sol   | 34 +++++----------
 test/mocks/MockERC777.sol             | 51 ----------------------
 test/plugin/TokenReceiverPlugin.t.sol | 63 +++++++--------------------
 3 files changed, 26 insertions(+), 122 deletions(-)
 delete mode 100644 test/mocks/MockERC777.sol

diff --git a/src/plugins/TokenReceiverPlugin.sol b/src/plugins/TokenReceiverPlugin.sol
index f22ae3ba..980d7397 100644
--- a/src/plugins/TokenReceiverPlugin.sol
+++ b/src/plugins/TokenReceiverPlugin.sol
@@ -2,7 +2,6 @@
 pragma solidity ^0.8.25;
 
 import {IERC721Receiver} from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
-import {IERC777Recipient} from "@openzeppelin/contracts/interfaces/IERC777Recipient.sol";
 import {IERC1155Receiver} from "@openzeppelin/contracts/interfaces/IERC1155Receiver.sol";
 
 import {
@@ -18,7 +17,7 @@ import {BasePlugin} from "./BasePlugin.sol";
 /// @author ERC-6900 Authors
 /// @notice This plugin allows modular accounts to receive various types of tokens by implementing
 /// required token receiver interfaces.
-contract TokenReceiverPlugin is BasePlugin, IERC721Receiver, IERC777Recipient, IERC1155Receiver {
+contract TokenReceiverPlugin is BasePlugin, IERC721Receiver, IERC1155Receiver {
     string public constant NAME = "Token Receiver Plugin";
     string public constant VERSION = "1.0.0";
     string public constant AUTHOR = "ERC-6900 Authors";
@@ -27,13 +26,6 @@ contract TokenReceiverPlugin is BasePlugin, IERC721Receiver, IERC777Recipient, I
     // ┃    Execution functions    ┃
     // ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
 
-    function tokensReceived(address, address, address, uint256, bytes calldata, bytes calldata)
-        external
-        pure
-        override
-    // solhint-disable-next-line no-empty-blocks
-    {}
-
     function onERC721Received(address, address, uint256, bytes calldata) external pure override returns (bytes4) {
         return IERC721Receiver.onERC721Received.selector;
     }
@@ -72,11 +64,10 @@ contract TokenReceiverPlugin is BasePlugin, IERC721Receiver, IERC777Recipient, I
     function pluginManifest() external pure override returns (PluginManifest memory) {
         PluginManifest memory manifest;
 
-        manifest.executionFunctions = new bytes4[](4);
-        manifest.executionFunctions[0] = this.tokensReceived.selector;
-        manifest.executionFunctions[1] = this.onERC721Received.selector;
-        manifest.executionFunctions[2] = this.onERC1155Received.selector;
-        manifest.executionFunctions[3] = this.onERC1155BatchReceived.selector;
+        manifest.executionFunctions = new bytes4[](3);
+        manifest.executionFunctions[0] = this.onERC721Received.selector;
+        manifest.executionFunctions[1] = this.onERC1155Received.selector;
+        manifest.executionFunctions[2] = this.onERC1155BatchReceived.selector;
 
         // Only runtime validationFunction is needed since callbacks come from token contracts only
         ManifestFunction memory alwaysAllowRuntime = ManifestFunction({
@@ -84,28 +75,23 @@ contract TokenReceiverPlugin is BasePlugin, IERC721Receiver, IERC777Recipient, I
             functionId: 0, // Unused.
             dependencyIndex: 0 // Unused.
         });
-        manifest.validationFunctions = new ManifestAssociatedFunction[](4);
+        manifest.validationFunctions = new ManifestAssociatedFunction[](3);
         manifest.validationFunctions[0] = ManifestAssociatedFunction({
-            executionSelector: this.tokensReceived.selector,
-            associatedFunction: alwaysAllowRuntime
-        });
-        manifest.validationFunctions[1] = ManifestAssociatedFunction({
             executionSelector: this.onERC721Received.selector,
             associatedFunction: alwaysAllowRuntime
         });
-        manifest.validationFunctions[2] = ManifestAssociatedFunction({
+        manifest.validationFunctions[1] = ManifestAssociatedFunction({
             executionSelector: this.onERC1155Received.selector,
             associatedFunction: alwaysAllowRuntime
         });
-        manifest.validationFunctions[3] = ManifestAssociatedFunction({
+        manifest.validationFunctions[2] = ManifestAssociatedFunction({
             executionSelector: this.onERC1155BatchReceived.selector,
             associatedFunction: alwaysAllowRuntime
         });
 
-        manifest.interfaceIds = new bytes4[](3);
+        manifest.interfaceIds = new bytes4[](2);
         manifest.interfaceIds[0] = type(IERC721Receiver).interfaceId;
-        manifest.interfaceIds[1] = type(IERC777Recipient).interfaceId;
-        manifest.interfaceIds[2] = type(IERC1155Receiver).interfaceId;
+        manifest.interfaceIds[1] = type(IERC1155Receiver).interfaceId;
 
         return manifest;
     }
diff --git a/test/mocks/MockERC777.sol b/test/mocks/MockERC777.sol
deleted file mode 100644
index 18840db2..00000000
--- a/test/mocks/MockERC777.sol
+++ /dev/null
@@ -1,51 +0,0 @@
-// SPDX-License-Identifier: UNLICENSED
-pragma solidity ^0.8.19;
-
-import {IERC777} from "@openzeppelin/contracts/token/ERC777/IERC777.sol";
-import {IERC777Recipient} from "@openzeppelin/contracts/token/ERC777/IERC777Recipient.sol";
-
-contract MockERC777 is IERC777 {
-    string public override name;
-    string public override symbol;
-    uint256 public override granularity;
-    uint256 public override totalSupply;
-    mapping(address => uint256) public override balanceOf;
-
-    function mint(address to, uint256 amount) public {
-        balanceOf[to] += amount;
-    }
-
-    function transfer(address to, uint256 amount) public returns (bool) {
-        return transferFrom(msg.sender, to, amount);
-    }
-
-    function transferFrom(address from, address to, uint256 amount) public returns (bool) {
-        IERC777Recipient(to).tokensReceived(msg.sender, from, to, amount, bytes(""), bytes(""));
-        balanceOf[from] -= amount;
-        balanceOf[to] += amount;
-        return true;
-    }
-
-    function send(address to, uint256 amount, bytes calldata) public override {
-        transferFrom(msg.sender, to, amount);
-    }
-
-    function burn(uint256 amount, bytes calldata) external {
-        transferFrom(msg.sender, address(0), amount);
-    }
-
-    function isOperatorFor(address, address) external pure returns (bool) {
-        return false;
-    }
-
-    // solhint-disable-next-line no-empty-blocks
-    function authorizeOperator(address) external {}
-    // solhint-disable-next-line no-empty-blocks
-    function revokeOperator(address) external {}
-    // solhint-disable-next-line no-empty-blocks
-    function defaultOperators() external view returns (address[] memory a) {}
-    // solhint-disable-next-line no-empty-blocks
-    function operatorSend(address, address, uint256, bytes calldata, bytes calldata) external {}
-    // solhint-disable-next-line no-empty-blocks
-    function operatorBurn(address, uint256, bytes calldata, bytes calldata) external {}
-}
diff --git a/test/plugin/TokenReceiverPlugin.t.sol b/test/plugin/TokenReceiverPlugin.t.sol
index 49692e34..7cb90918 100644
--- a/test/plugin/TokenReceiverPlugin.t.sol
+++ b/test/plugin/TokenReceiverPlugin.t.sol
@@ -5,7 +5,6 @@ import {IEntryPoint} from "@eth-infinitism/account-abstraction/interfaces/IEntry
 import {IERC721Receiver} from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
 import {ERC721PresetMinterPauserAutoId} from
     "@openzeppelin/contracts/token/ERC721/presets/ERC721PresetMinterPauserAutoId.sol";
-import {IERC777Recipient} from "@openzeppelin/contracts/token/ERC777/IERC777Recipient.sol";
 import {IERC1155Receiver} from "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol";
 
 import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
@@ -13,7 +12,6 @@ import {FunctionReference} from "../../src/helpers/FunctionReferenceLib.sol";
 import {TokenReceiverPlugin} from "../../src/plugins/TokenReceiverPlugin.sol";
 
 import {MSCAFactoryFixture} from "../mocks/MSCAFactoryFixture.sol";
-import {MockERC777} from "../mocks/MockERC777.sol";
 import {MockERC1155} from "../mocks/MockERC1155.sol";
 import {OptimizedTest} from "../utils/OptimizedTest.sol";
 
@@ -22,8 +20,7 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
     TokenReceiverPlugin public plugin;
 
     ERC721PresetMinterPauserAutoId public t0;
-    MockERC777 public t1;
-    MockERC1155 public t2;
+    MockERC1155 public t1;
 
     // init dynamic length arrays for use in args
     address[] public defaultOperators;
@@ -44,13 +41,10 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
         t0 = new ERC721PresetMinterPauserAutoId("t0", "t0", "");
         t0.mint(address(this));
 
-        t1 = new MockERC777();
-        t1.mint(address(this), _TOKEN_AMOUNT);
-
-        t2 = new MockERC1155();
-        t2.mint(address(this), _TOKEN_ID, _TOKEN_AMOUNT);
+        t1 = new MockERC1155();
+        t1.mint(address(this), _TOKEN_ID, _TOKEN_AMOUNT);
         for (uint256 i = 1; i < _BATCH_TOKEN_IDS; i++) {
-            t2.mint(address(this), i, _TOKEN_AMOUNT);
+            t1.mint(address(this), i, _TOKEN_AMOUNT);
             tokenIds.push(i);
             tokenAmts.push(_TOKEN_AMOUNT);
             zeroTokenAmts.push(0);
@@ -81,54 +75,33 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
         assertEq(t0.ownerOf(_TOKEN_ID), address(acct));
     }
 
-    function test_failERC777Transfer() public {
-        vm.expectRevert(
-            abi.encodePacked(
-                UpgradeableModularAccount.UnrecognizedFunction.selector,
-                IERC777Recipient.tokensReceived.selector,
-                bytes28(0)
-            )
-        );
-        t1.transfer(address(acct), _TOKEN_AMOUNT);
-    }
-
-    function test_passERC777Transfer() public {
-        _initPlugin();
-
-        assertEq(t1.balanceOf(address(this)), _TOKEN_AMOUNT);
-        assertEq(t1.balanceOf(address(acct)), 0);
-        t1.transfer(address(acct), _TOKEN_AMOUNT);
-        assertEq(t1.balanceOf(address(this)), 0);
-        assertEq(t1.balanceOf(address(acct)), _TOKEN_AMOUNT);
-    }
-
     function test_failERC1155Transfer() public {
         // for 1155, reverts are caught in a try catch and bubbled up with a diff reason
         vm.expectRevert("ERC1155: transfer to non-ERC1155Receiver implementer");
-        t2.safeTransferFrom(address(this), address(acct), _TOKEN_ID, _TOKEN_AMOUNT, "");
+        t1.safeTransferFrom(address(this), address(acct), _TOKEN_ID, _TOKEN_AMOUNT, "");
 
         // for 1155, reverts are caught in a try catch and bubbled up with a diff reason
         vm.expectRevert("ERC1155: transfer to non-ERC1155Receiver implementer");
-        t2.safeBatchTransferFrom(address(this), address(acct), tokenIds, tokenAmts, "");
+        t1.safeBatchTransferFrom(address(this), address(acct), tokenIds, tokenAmts, "");
     }
 
     function test_passERC1155Transfer() public {
         _initPlugin();
 
-        assertEq(t2.balanceOf(address(this), _TOKEN_ID), _TOKEN_AMOUNT);
-        assertEq(t2.balanceOf(address(acct), _TOKEN_ID), 0);
-        t2.safeTransferFrom(address(this), address(acct), _TOKEN_ID, _TOKEN_AMOUNT, "");
-        assertEq(t2.balanceOf(address(this), _TOKEN_ID), 0);
-        assertEq(t2.balanceOf(address(acct), _TOKEN_ID), _TOKEN_AMOUNT);
+        assertEq(t1.balanceOf(address(this), _TOKEN_ID), _TOKEN_AMOUNT);
+        assertEq(t1.balanceOf(address(acct), _TOKEN_ID), 0);
+        t1.safeTransferFrom(address(this), address(acct), _TOKEN_ID, _TOKEN_AMOUNT, "");
+        assertEq(t1.balanceOf(address(this), _TOKEN_ID), 0);
+        assertEq(t1.balanceOf(address(acct), _TOKEN_ID), _TOKEN_AMOUNT);
 
         for (uint256 i = 1; i < _BATCH_TOKEN_IDS; i++) {
-            assertEq(t2.balanceOf(address(this), i), _TOKEN_AMOUNT);
-            assertEq(t2.balanceOf(address(acct), i), 0);
+            assertEq(t1.balanceOf(address(this), i), _TOKEN_AMOUNT);
+            assertEq(t1.balanceOf(address(acct), i), 0);
         }
-        t2.safeBatchTransferFrom(address(this), address(acct), tokenIds, tokenAmts, "");
+        t1.safeBatchTransferFrom(address(this), address(acct), tokenIds, tokenAmts, "");
         for (uint256 i = 1; i < _BATCH_TOKEN_IDS; i++) {
-            assertEq(t2.balanceOf(address(this), i), 0);
-            assertEq(t2.balanceOf(address(acct), i), _TOKEN_AMOUNT);
+            assertEq(t1.balanceOf(address(this), i), 0);
+            assertEq(t1.balanceOf(address(acct), i), _TOKEN_AMOUNT);
         }
     }
 
@@ -137,8 +110,6 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
 
         isSupported = acct.supportsInterface(type(IERC721Receiver).interfaceId);
         assertEq(isSupported, false);
-        isSupported = acct.supportsInterface(type(IERC777Recipient).interfaceId);
-        assertEq(isSupported, false);
         isSupported = acct.supportsInterface(type(IERC1155Receiver).interfaceId);
         assertEq(isSupported, false);
     }
@@ -150,8 +121,6 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
 
         isSupported = acct.supportsInterface(type(IERC721Receiver).interfaceId);
         assertEq(isSupported, true);
-        isSupported = acct.supportsInterface(type(IERC777Recipient).interfaceId);
-        assertEq(isSupported, true);
         isSupported = acct.supportsInterface(type(IERC1155Receiver).interfaceId);
         assertEq(isSupported, true);
     }

From 82490c5a73893bb3a9633d177558ec3ced7e63ce Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:14:36 -0400
Subject: [PATCH 02/10] chore: replace OZ ERC721

---
 test/mocks/MockERC721.sol             | 12 ++++++++++++
 test/plugin/TokenReceiverPlugin.t.sol |  9 ++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)
 create mode 100644 test/mocks/MockERC721.sol

diff --git a/test/mocks/MockERC721.sol b/test/mocks/MockERC721.sol
new file mode 100644
index 00000000..1139adf2
--- /dev/null
+++ b/test/mocks/MockERC721.sol
@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.19;
+
+import {ERC721} from "@openzeppelin/contracts/token/ERC721/ERC721.sol";
+
+contract MockERC721 is ERC721 {
+    constructor(string memory name, string memory symbol) ERC721(name, symbol) {}
+
+    function mint(address account, uint256 tokenId) external {
+        _mint(account, tokenId);
+    }
+}
diff --git a/test/plugin/TokenReceiverPlugin.t.sol b/test/plugin/TokenReceiverPlugin.t.sol
index 7cb90918..add51384 100644
--- a/test/plugin/TokenReceiverPlugin.t.sol
+++ b/test/plugin/TokenReceiverPlugin.t.sol
@@ -3,8 +3,6 @@ pragma solidity ^0.8.19;
 
 import {IEntryPoint} from "@eth-infinitism/account-abstraction/interfaces/IEntryPoint.sol";
 import {IERC721Receiver} from "@openzeppelin/contracts/token/ERC721/IERC721Receiver.sol";
-import {ERC721PresetMinterPauserAutoId} from
-    "@openzeppelin/contracts/token/ERC721/presets/ERC721PresetMinterPauserAutoId.sol";
 import {IERC1155Receiver} from "@openzeppelin/contracts/token/ERC1155/IERC1155Receiver.sol";
 
 import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
@@ -12,6 +10,7 @@ import {FunctionReference} from "../../src/helpers/FunctionReferenceLib.sol";
 import {TokenReceiverPlugin} from "../../src/plugins/TokenReceiverPlugin.sol";
 
 import {MSCAFactoryFixture} from "../mocks/MSCAFactoryFixture.sol";
+import {MockERC721} from "../mocks/MockERC721.sol";
 import {MockERC1155} from "../mocks/MockERC1155.sol";
 import {OptimizedTest} from "../utils/OptimizedTest.sol";
 
@@ -19,7 +18,7 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
     UpgradeableModularAccount public acct;
     TokenReceiverPlugin public plugin;
 
-    ERC721PresetMinterPauserAutoId public t0;
+    MockERC721 public t0;
     MockERC1155 public t1;
 
     // init dynamic length arrays for use in args
@@ -38,8 +37,8 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
         acct = factory.createAccount(address(this), 0);
         plugin = _deployTokenReceiverPlugin();
 
-        t0 = new ERC721PresetMinterPauserAutoId("t0", "t0", "");
-        t0.mint(address(this));
+        t0 = new MockERC721("t0", "t0");
+        t0.mint(address(this), _TOKEN_ID);
 
         t1 = new MockERC1155();
         t1.mint(address(this), _TOKEN_ID, _TOKEN_AMOUNT);

From a3bf1afbecf0207803183f201d0059af413b0327 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:39:34 -0400
Subject: [PATCH 03/10] chore: update usage of OZ initialiable, OZ
 UUPSUpgradeable

---
 README.md                                   |  1 +
 src/account/AccountLoupe.sol                |  1 -
 src/account/AccountStorageInitializable.sol | 45 ---------------------
 src/account/UpgradeableModularAccount.sol   | 11 ++---
 test/libraries/AccountStorage.t.sol         | 43 --------------------
 test/mocks/MockDiamondStorageContract.sol   | 13 ------
 6 files changed, 4 insertions(+), 110 deletions(-)
 delete mode 100644 src/account/AccountStorageInitializable.sol
 delete mode 100644 test/libraries/AccountStorage.t.sol
 delete mode 100644 test/mocks/MockDiamondStorageContract.sol

diff --git a/README.md b/README.md
index fb3e9dbc..95eb4e05 100644
--- a/README.md
+++ b/README.md
@@ -7,6 +7,7 @@ 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
diff --git a/src/account/AccountLoupe.sol b/src/account/AccountLoupe.sol
index 8bff326c..a75def37 100644
--- a/src/account/AccountLoupe.sol
+++ b/src/account/AccountLoupe.sol
@@ -24,7 +24,6 @@ abstract contract AccountLoupe is IAccountLoupe {
 
         if (
             selector == IStandardExecutor.execute.selector || selector == IStandardExecutor.executeBatch.selector
-                || selector == UUPSUpgradeable.upgradeTo.selector
                 || selector == UUPSUpgradeable.upgradeToAndCall.selector
                 || selector == IPluginManager.installPlugin.selector
                 || selector == IPluginManager.uninstallPlugin.selector
diff --git a/src/account/AccountStorageInitializable.sol b/src/account/AccountStorageInitializable.sol
deleted file mode 100644
index 5131978b..00000000
--- a/src/account/AccountStorageInitializable.sol
+++ /dev/null
@@ -1,45 +0,0 @@
-// SPDX-License-Identifier: GPL-3.0
-pragma solidity ^0.8.25;
-
-import {Address} from "@openzeppelin/contracts/utils/Address.sol";
-
-import {AccountStorage, getAccountStorage} from "./AccountStorage.sol";
-
-abstract contract AccountStorageInitializable {
-    error AlreadyInitialized();
-    error AlreadyInitializing();
-
-    /// @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 _storage = getAccountStorage();
-        bool isTopLevelCall = !_storage.initializing;
-        if (
-            isTopLevelCall && _storage.initialized < 1
-                || !Address.isContract(address(this)) && _storage.initialized == 1
-        ) {
-            _storage.initialized = 1;
-            if (isTopLevelCall) {
-                _storage.initializing = true;
-            }
-            _;
-            if (isTopLevelCall) {
-                _storage.initializing = false;
-            }
-        } else {
-            revert AlreadyInitialized();
-        }
-    }
-
-    /// @notice Internal function to disable calls to initialization functions
-    /// @dev Reverts if the contract has already been initialized
-    function _disableInitializers() internal virtual {
-        AccountStorage storage _storage = getAccountStorage();
-        if (_storage.initializing) {
-            revert AlreadyInitializing();
-        }
-        if (_storage.initialized != type(uint8).max) {
-            _storage.initialized = type(uint8).max;
-        }
-    }
-}
diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
index 81496a2a..3448ace7 100644
--- a/src/account/UpgradeableModularAccount.sol
+++ b/src/account/UpgradeableModularAccount.sol
@@ -8,6 +8,7 @@ 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";
@@ -18,13 +19,12 @@ 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,
-    AccountStorageInitializable,
+    Initializable,
     BaseAccount,
     IERC165,
     IPluginExecutor,
@@ -292,11 +292,6 @@ contract UpgradeableModularAccount is
         return getAccountStorage().supportedIfaces[interfaceId] > 0;
     }
 
-    /// @inheritdoc UUPSUpgradeable
-    function upgradeTo(address newImplementation) public override onlyProxy wrapNativeFunction {
-        _upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
-    }
-
     /// @inheritdoc UUPSUpgradeable
     function upgradeToAndCall(address newImplementation, bytes memory data)
         public
@@ -305,7 +300,7 @@ contract UpgradeableModularAccount is
         onlyProxy
         wrapNativeFunction
     {
-        _upgradeToAndCallUUPS(newImplementation, data, true);
+        UUPSUpgradeable.upgradeToAndCall(newImplementation, data);
     }
 
     /// @notice Gets the entry point for this account
diff --git a/test/libraries/AccountStorage.t.sol b/test/libraries/AccountStorage.t.sol
deleted file mode 100644
index 7cab9991..00000000
--- a/test/libraries/AccountStorage.t.sol
+++ /dev/null
@@ -1,43 +0,0 @@
-// 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.AlreadyInitialized.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));
-    }
-}
diff --git a/test/mocks/MockDiamondStorageContract.sol b/test/mocks/MockDiamondStorageContract.sol
deleted file mode 100644
index 81d40dc5..00000000
--- a/test/mocks/MockDiamondStorageContract.sol
+++ /dev/null
@@ -1,13 +0,0 @@
-// 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 {}
-}

From 9502d2c80fd9a0c4bce769c050538923ec62ad05 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:44:54 -0400
Subject: [PATCH 04/10] chore: update usage of OZ ECDSA

---
 src/plugins/owner/SingleOwnerPlugin.sol | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/src/plugins/owner/SingleOwnerPlugin.sol b/src/plugins/owner/SingleOwnerPlugin.sol
index bfff2aa4..b1d5d5e1 100644
--- a/src/plugins/owner/SingleOwnerPlugin.sol
+++ b/src/plugins/owner/SingleOwnerPlugin.sol
@@ -5,6 +5,7 @@ import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import {IERC1271} from "@openzeppelin/contracts/interfaces/IERC1271.sol";
 import {SignatureChecker} from "@openzeppelin/contracts/utils/cryptography/SignatureChecker.sol";
 import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {IPluginManager} from "../../interfaces/IPluginManager.sol";
@@ -38,6 +39,7 @@ import {ISingleOwnerPlugin} from "./ISingleOwnerPlugin.sol";
 /// send user operations through a bundler.
 contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
     using ECDSA for bytes32;
+    using MessageHashUtils for bytes32;
 
     string public constant NAME = "Single Owner Plugin";
     string public constant VERSION = "1.0.0";
@@ -99,7 +101,7 @@ contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
     {
         if (functionId == uint8(FunctionId.VALIDATION_OWNER_OR_SELF)) {
             // Validate the user op signature against the owner.
-            (address signer,) = (userOpHash.toEthSignedMessageHash()).tryRecover(userOp.signature);
+            (address signer,,) = (userOpHash.toEthSignedMessageHash()).tryRecover(userOp.signature);
             if (signer == address(0) || signer != _owners[msg.sender]) {
                 return _SIG_VALIDATION_FAILED;
             }
@@ -154,7 +156,7 @@ contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
             functionId: uint8(FunctionId.VALIDATION_OWNER_OR_SELF),
             dependencyIndex: 0 // Unused.
         });
-        manifest.validationFunctions = new ManifestAssociatedFunction[](8);
+        manifest.validationFunctions = new ManifestAssociatedFunction[](7);
         manifest.validationFunctions[0] = ManifestAssociatedFunction({
             executionSelector: this.transferOwnership.selector,
             associatedFunction: ownerValidationFunction
@@ -176,10 +178,6 @@ contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
             associatedFunction: ownerValidationFunction
         });
         manifest.validationFunctions[5] = ManifestAssociatedFunction({
-            executionSelector: UUPSUpgradeable.upgradeTo.selector,
-            associatedFunction: ownerValidationFunction
-        });
-        manifest.validationFunctions[6] = ManifestAssociatedFunction({
             executionSelector: UUPSUpgradeable.upgradeToAndCall.selector,
             associatedFunction: ownerValidationFunction
         });
@@ -189,7 +187,7 @@ contract SingleOwnerPlugin is BasePlugin, ISingleOwnerPlugin, IERC1271 {
             functionId: 0, // Unused.
             dependencyIndex: 0 // Unused.
         });
-        manifest.validationFunctions[7] = ManifestAssociatedFunction({
+        manifest.validationFunctions[6] = ManifestAssociatedFunction({
             executionSelector: this.isValidSignature.selector,
             associatedFunction: alwaysAllowRuntime
         });

From 36134b411521d278cd0525090d5ceb2a482e609c Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Sat, 20 Apr 2024 16:50:07 -0400
Subject: [PATCH 05/10] chore: cleanup, fix errs

---
 src/account/AccountStorage.sol               | 3 ---
 test/account/UpgradeableModularAccount.t.sol | 2 ++
 test/comparison/CompareSimpleAccount.t.sol   | 2 ++
 test/plugin/SingleOwnerPlugin.t.sol          | 2 ++
 4 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/account/AccountStorage.sol b/src/account/AccountStorage.sol
index 35b3eabb..e198042f 100644
--- a/src/account/AccountStorage.sol
+++ b/src/account/AccountStorage.sol
@@ -48,9 +48,6 @@ struct SelectorData {
 }
 
 struct AccountStorage {
-    // AccountStorageInitializable variables
-    uint8 initialized;
-    bool initializing;
     // Plugin metadata storage
     EnumerableSet.AddressSet plugins;
     mapping(address => PluginData) pluginData;
diff --git a/test/account/UpgradeableModularAccount.t.sol b/test/account/UpgradeableModularAccount.t.sol
index ae491c13..426a54e9 100644
--- a/test/account/UpgradeableModularAccount.t.sol
+++ b/test/account/UpgradeableModularAccount.t.sol
@@ -5,6 +5,7 @@ import {console} from "forge-std/Test.sol";
 
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
 import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
+import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 
 import {PluginManagerInternals} from "../../src/account/PluginManagerInternals.sol";
 import {UpgradeableModularAccount} from "../../src/account/UpgradeableModularAccount.sol";
@@ -23,6 +24,7 @@ import {AccountTestBase} from "../utils/AccountTestBase.sol";
 
 contract UpgradeableModularAccountTest is AccountTestBase {
     using ECDSA for bytes32;
+    using MessageHashUtils for bytes32;
 
     TokenReceiverPlugin public tokenReceiverPlugin;
 
diff --git a/test/comparison/CompareSimpleAccount.t.sol b/test/comparison/CompareSimpleAccount.t.sol
index 7c3182b2..1179dac3 100644
--- a/test/comparison/CompareSimpleAccount.t.sol
+++ b/test/comparison/CompareSimpleAccount.t.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.19;
 import {Test} from "forge-std/Test.sol";
 
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
+import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 import {EntryPoint} from "@eth-infinitism/account-abstraction/core/EntryPoint.sol";
 import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
 
@@ -14,6 +15,7 @@ import {Counter} from "../mocks/Counter.sol";
 
 contract CompareSimpleAccountTest is Test {
     using ECDSA for bytes32;
+    using MessageHashUtils for bytes32;
 
     EntryPoint public entryPoint;
     address payable public beneficiary;
diff --git a/test/plugin/SingleOwnerPlugin.t.sol b/test/plugin/SingleOwnerPlugin.t.sol
index 73e45648..dea5b2f9 100644
--- a/test/plugin/SingleOwnerPlugin.t.sol
+++ b/test/plugin/SingleOwnerPlugin.t.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.19;
 import {EntryPoint} from "@eth-infinitism/account-abstraction/core/EntryPoint.sol";
 import {PackedUserOperation} from "@eth-infinitism/account-abstraction/interfaces/PackedUserOperation.sol";
 import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
+import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 
 import {SingleOwnerPlugin} from "../../src/plugins/owner/SingleOwnerPlugin.sol";
 import {ISingleOwnerPlugin} from "../../src/plugins/owner/ISingleOwnerPlugin.sol";
@@ -12,6 +13,7 @@ import {OptimizedTest} from "../utils/OptimizedTest.sol";
 
 contract SingleOwnerPluginTest is OptimizedTest {
     using ECDSA for bytes32;
+    using MessageHashUtils for bytes32;
 
     SingleOwnerPlugin public plugin;
     EntryPoint public entryPoint;

From 2472fd91d9a3272bbb19d8c3db47edfc34e23097 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Tue, 23 Apr 2024 13:52:30 -0400
Subject: [PATCH 06/10] fix: cleanup

---
 src/account/UpgradeableModularAccount.sol | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
index 3448ace7..9c843b8d 100644
--- a/src/account/UpgradeableModularAccount.sol
+++ b/src/account/UpgradeableModularAccount.sol
@@ -300,7 +300,7 @@ contract UpgradeableModularAccount is
         onlyProxy
         wrapNativeFunction
     {
-        UUPSUpgradeable.upgradeToAndCall(newImplementation, data);
+        upgradeToAndCall(newImplementation, data);
     }
 
     /// @notice Gets the entry point for this account

From 01a70198ecb59acf22f7065f5efc5b08b19e6234 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Mon, 29 Apr 2024 11:43:10 -0400
Subject: [PATCH 07/10] fix: revert to custom initializable, make changes from
 oz 5.0

---
 README.md                                   |  1 -
 src/account/AccountStorage.sol              |  3 +
 src/account/AccountStorageInitializable.sol | 65 +++++++++++++++++++++
 src/account/UpgradeableModularAccount.sol   |  4 +-
 test/libraries/AccountStorage.t.sol         | 43 ++++++++++++++
 test/mocks/MockDiamondStorageContract.sol   | 13 +++++
 6 files changed, 126 insertions(+), 3 deletions(-)
 create mode 100644 src/account/AccountStorageInitializable.sol
 create mode 100644 test/libraries/AccountStorage.t.sol
 create mode 100644 test/mocks/MockDiamondStorageContract.sol

diff --git a/README.md b/README.md
index 95eb4e05..fb3e9dbc 100644
--- a/README.md
+++ b/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
diff --git a/src/account/AccountStorage.sol b/src/account/AccountStorage.sol
index e198042f..35b3eabb 100644
--- a/src/account/AccountStorage.sol
+++ b/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;
diff --git a/src/account/AccountStorageInitializable.sol b/src/account/AccountStorageInitializable.sol
new file mode 100644
index 00000000..efd7f582
--- /dev/null
+++ b/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);
+        }
+    }
+}
diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
index 9c843b8d..a4364d64 100644
--- a/src/account/UpgradeableModularAccount.sol
+++ b/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,
diff --git a/test/libraries/AccountStorage.t.sol b/test/libraries/AccountStorage.t.sol
new file mode 100644
index 00000000..f59049a3
--- /dev/null
+++ b/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));
+    }
+}
diff --git a/test/mocks/MockDiamondStorageContract.sol b/test/mocks/MockDiamondStorageContract.sol
new file mode 100644
index 00000000..81d40dc5
--- /dev/null
+++ b/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 {}
+}

From 19c3e65df6db1c01d1edd4f81ee993280b895b52 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Mon, 29 Apr 2024 11:47:13 -0400
Subject: [PATCH 08/10] chore: add comment for attribution, lint order

---
 src/account/AccountStorageInitializable.sol | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/account/AccountStorageInitializable.sol b/src/account/AccountStorageInitializable.sol
index efd7f582..debcee40 100644
--- a/src/account/AccountStorageInitializable.sol
+++ b/src/account/AccountStorageInitializable.sol
@@ -3,7 +3,14 @@ pragma solidity ^0.8.25;
 
 import {AccountStorage, getAccountStorage} from "./AccountStorage.sol";
 
+/// @title AccountStorageInitializable
+/// @dev Bulk of the impl is lifted from OZ 5.0 Initializible
 abstract contract AccountStorageInitializable {
+    /**
+     * @dev Triggered when the contract has been initialized or reinitialized.
+     */
+    event Initialized(uint64 version);
+
     /**
      * @dev The contract is already initialized.
      */
@@ -14,11 +21,6 @@ abstract contract AccountStorageInitializable {
      */
     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() {

From 66736f2ef14644b493a8db7bd095b489a4c608e1 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Wed, 1 May 2024 14:32:20 -0400
Subject: [PATCH 09/10] fix: make sure to use the right uupsupgradeable

---
 src/account/UpgradeableModularAccount.sol    |  2 +-
 test/account/UpgradeableModularAccount.t.sol | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
index a4364d64..f29d17c7 100644
--- a/src/account/UpgradeableModularAccount.sol
+++ b/src/account/UpgradeableModularAccount.sol
@@ -300,7 +300,7 @@ contract UpgradeableModularAccount is
         onlyProxy
         wrapNativeFunction
     {
-        upgradeToAndCall(newImplementation, data);
+        super.upgradeToAndCall(newImplementation, data);
     }
 
     /// @notice Gets the entry point for this account
diff --git a/test/account/UpgradeableModularAccount.t.sol b/test/account/UpgradeableModularAccount.t.sol
index 426a54e9..3ef6f107 100644
--- a/test/account/UpgradeableModularAccount.t.sol
+++ b/test/account/UpgradeableModularAccount.t.sol
@@ -405,6 +405,20 @@ contract UpgradeableModularAccountTest is AccountTestBase {
         vm.stopPrank();
     }
 
+    function test_upgradeToAndCall() public {
+        vm.startPrank(owner1);
+        UpgradeableModularAccount account3 = new UpgradeableModularAccount(entryPoint);
+        bytes32 slot = account3.proxiableUUID();
+
+        // account has impl from factory
+        assertEq(
+            address(factory.accountImplementation()), address(uint160(uint256(vm.load(address(account1), slot))))
+        );
+        account1.upgradeToAndCall(address(account3), bytes(""));
+        // account has new impl
+        assertEq(address(account3), address(uint160(uint256(vm.load(address(account1), slot)))));
+    }
+
     // Internal Functions
 
     function _printStorageReadsAndWrites(address addr) internal {

From 528f9e760df9fd80acadcb55558758d817667a4e Mon Sep 17 00:00:00 2001
From: howy <132113803+howydev@users.noreply.github.com>
Date: Wed, 1 May 2024 14:39:15 -0400
Subject: [PATCH 10/10] chore: update to 4337 v0.7 (3/3)  (#50)

---
 test/comparison/CompareSimpleAccount.t.sol |  2 +-
 test/plugin/TokenReceiverPlugin.t.sol      | 20 ++++++++++++++++----
 test/utils/AccountTestBase.sol             |  2 +-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/test/comparison/CompareSimpleAccount.t.sol b/test/comparison/CompareSimpleAccount.t.sol
index 1179dac3..93b75ebe 100644
--- a/test/comparison/CompareSimpleAccount.t.sol
+++ b/test/comparison/CompareSimpleAccount.t.sol
@@ -39,7 +39,7 @@ contract CompareSimpleAccountTest is Test {
 
     // helper function to compress 2 gas values into a single bytes32
     function _encodeGas(uint256 g1, uint256 g2) internal pure returns (bytes32) {
-        return bytes32(uint256(g1 << 128 + uint128(g2)));
+        return bytes32(uint256((g1 << 128) + uint128(g2)));
     }
 
     function setUp() public {
diff --git a/test/plugin/TokenReceiverPlugin.t.sol b/test/plugin/TokenReceiverPlugin.t.sol
index add51384..7a7433af 100644
--- a/test/plugin/TokenReceiverPlugin.t.sol
+++ b/test/plugin/TokenReceiverPlugin.t.sol
@@ -75,12 +75,24 @@ contract TokenReceiverPluginTest is OptimizedTest, IERC1155Receiver {
     }
 
     function test_failERC1155Transfer() public {
-        // for 1155, reverts are caught in a try catch and bubbled up with a diff reason
-        vm.expectRevert("ERC1155: transfer to non-ERC1155Receiver implementer");
+        // for 1155, reverts are caught in a try catch and bubbled up with the reason from the account
+        vm.expectRevert(
+            abi.encodePacked(
+                UpgradeableModularAccount.UnrecognizedFunction.selector,
+                IERC1155Receiver.onERC1155Received.selector,
+                bytes28(0)
+            )
+        );
         t1.safeTransferFrom(address(this), address(acct), _TOKEN_ID, _TOKEN_AMOUNT, "");
 
-        // for 1155, reverts are caught in a try catch and bubbled up with a diff reason
-        vm.expectRevert("ERC1155: transfer to non-ERC1155Receiver implementer");
+        // for 1155, reverts are caught in a try catch and bubbled up with the reason from the account
+        vm.expectRevert(
+            abi.encodePacked(
+                UpgradeableModularAccount.UnrecognizedFunction.selector,
+                IERC1155Receiver.onERC1155BatchReceived.selector,
+                bytes28(0)
+            )
+        );
         t1.safeBatchTransferFrom(address(this), address(acct), tokenIds, tokenAmts, "");
     }
 
diff --git a/test/utils/AccountTestBase.sol b/test/utils/AccountTestBase.sol
index e23ed1ee..5d9880f9 100644
--- a/test/utils/AccountTestBase.sol
+++ b/test/utils/AccountTestBase.sol
@@ -42,6 +42,6 @@ abstract contract AccountTestBase is OptimizedTest {
 
     // helper function to compress 2 gas values into a single bytes32
     function _encodeGas(uint256 g1, uint256 g2) internal pure returns (bytes32) {
-        return bytes32(uint256(g1 << 128 + uint128(g2)));
+        return bytes32(uint256((g1 << 128) + uint128(g2)));
     }
 }