Merge branch 'main' into feature/registry-deploy

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@rhinestone/modulekit",
-  "version": "0.4.10",
+  "version": "0.4.11",
   "description": "A development kit for building and testing smart account modules.",
   "license": "GPL-3.0",
   "author": {

src/ModuleKit.sol

@@ -2,7 +2,12 @@
 pragma solidity ^0.8.23;
 
 /* solhint-disable no-unused-import */
-import { UserOpData, AccountInstance, RhinestoneModuleKit } from "./test/RhinestoneModuleKit.sol";
+import {
+    UserOpData,
+    AccountInstance,
+    RhinestoneModuleKit,
+    AccountType
+} from "./test/RhinestoneModuleKit.sol";
 import { ModuleKitHelpers } from "./test/ModuleKitHelpers.sol";
 import { ModuleKitUserOp } from "./test/ModuleKitUserOp.sol";
 import { PackedUserOperation } from "./external/ERC4337.sol";

src/test/ModuleKitHelpers.sol

@@ -23,8 +23,13 @@ import {
     writeAccountEnv,
     getFactory,
     getHelper as getHelperFromStorage,
-    getAccountEnv as getAccountEnvFromStorage
+    getAccountEnv as getAccountEnvFromStorage,
+    getInstalledModules as getInstalledModulesFromStorage,
+    writeInstalledModule as writeInstalledModuleToStorage,
+    removeInstalledModule as removeInstalledModuleFromStorage,
+    InstalledModule
 } from "./utils/Storage.sol";
+import { recordLogs, VmSafe, getRecordedLogs } from "./utils/Vm.sol";
 
 library ModuleKitHelpers {
     /*//////////////////////////////////////////////////////////////////////////
@@ -165,7 +170,6 @@ library ModuleKitHelpers {
         address module
     )
         internal
-        view
         returns (bool)
     {
         return HelperBase(instance.accountHelper).isModuleInstalled(instance, moduleTypeId, module);
@@ -178,7 +182,6 @@ library ModuleKitHelpers {
         bytes memory data
     )
         internal
-        view
         returns (bool)
     {
         return HelperBase(instance.accountHelper).isModuleInstalled(
@@ -262,12 +265,60 @@ library ModuleKitHelpers {
         userOpData.entrypoint = instance.aux.entrypoint;
     }
 
+    function getInstalledModules(
+        AccountInstance memory instance
+    )
+        internal
+        view
+        returns (InstalledModule[] memory)
+    {
+        return getInstalledModulesFromStorage(instance.account);
+    }
+
+    function writeInstalledModule(
+        AccountInstance memory instance,
+        InstalledModule memory module
+    )
+        internal
+    {
+        writeInstalledModuleToStorage(module, instance.account);
+    }
+
+    function removeInstalledModule(
+        AccountInstance memory instance,
+        uint256 moduleType,
+        address moduleAddress
+    )
+        internal
+    {
+        // Get installed modules for account
+        InstalledModule[] memory installedModules = getInstalledModules(instance);
+        // Find module to remove (not super scalable at high module counts)
+        for (uint256 i; i < installedModules.length; i++) {
+            if (
+                installedModules[i].moduleType == moduleType
+                    && installedModules[i].moduleAddress == moduleAddress
+            ) {
+                // Remove module from storage
+                removeInstalledModuleFromStorage(i, instance.account);
+                return;
+            }
+        }
+    }
     /*//////////////////////////////////////////////////////////////////////////
                                 CONTROL FLOW
     //////////////////////////////////////////////////////////////////////////*/
 
     function expect4337Revert(AccountInstance memory) internal {
-        writeExpectRevert(1);
+        writeExpectRevert("");
+    }
+
+    function expect4337Revert(AccountInstance memory, bytes4 selector) internal {
+        writeExpectRevert(abi.encodePacked(selector));
+    }
+
+    function expect4337Revert(AccountInstance memory, bytes memory message) internal {
+        writeExpectRevert(message);
     }
 
     /**
@@ -332,7 +383,22 @@ library ModuleKitHelpers {
     }
 
     function deployAccount(AccountInstance memory instance) internal {
+        // Record logs to track installed modules
+        recordLogs();
+        // Deploy account
         HelperBase(instance.accountHelper).deployAccount(instance);
+        // Parse logs and determine if a module was installed
+        VmSafe.Log[] memory logs = getRecordedLogs();
+        for (uint256 i; i < logs.length; i++) {
+            // ModuleInstalled(uint256, address)
+            if (
+                logs[i].topics[0]
+                    == 0xd21d0b289f126c4b473ea641963e766833c2f13866e4ff480abd787c100ef123
+            ) {
+                (uint256 moduleType, address module) = abi.decode(logs[i].data, (uint256, address));
+                writeInstalledModuleToStorage(InstalledModule(moduleType, module), logs[i].emitter);
+            }
+        }
     }
 
     function setAccountType(AccountInstance memory, AccountType env) internal {

src/test/helpers/HelperBase.sol

@@ -270,8 +270,8 @@ abstract contract HelperBase {
         address module
     )
         public
-        view
         virtual
+        deployAccountForAction(instance)
         returns (bool)
     {
         return isModuleInstalled(instance, moduleTypeId, module, "");
@@ -284,8 +284,8 @@ abstract contract HelperBase {
         bytes memory additionalContext
     )
         public
-        view
         virtual
+        deployAccountForAction(instance)
         returns (bool)
     {
         return IERC7579Account(instance.account).isModuleInstalled(

src/test/helpers/KernelHelpers.sol

@@ -394,9 +394,9 @@ contract KernelHelpers is HelperBase {
         bytes memory data
     )
         public
-        view
         virtual
         override
+        deployAccountForAction(instance)
         returns (bool)
     {
         if (moduleTypeId == MODULE_TYPE_HOOK) {

src/test/helpers/SafeHelpers.sol

@@ -245,9 +245,9 @@ contract SafeHelpers is HelperBase {
         bytes memory data
     )
         public
-        view
         virtual
         override
+        deployAccountForAction(instance)
         returns (bool)
     {
         if (moduleTypeId == MODULE_TYPE_HOOK) {

src/test/utils/ERC4337Helpers.sol

@@ -16,17 +16,24 @@ import { GasParser } from "./gas/GasParser.sol";
 import {
     getSimulateUserOp,
     getExpectRevert,
-    writeExpectRevert,
+    getExpectRevertMessage,
+    clearExpectRevert,
     getGasIdentifier,
-    writeGasIdentifier
+    writeGasIdentifier,
+    writeInstalledModule,
+    getInstalledModules,
+    removeInstalledModule,
+    InstalledModule
 } from "./Storage.sol";
 
 library ERC4337Helpers {
     using Simulator for PackedUserOperation;
 
     error UserOperationReverted(
-        bytes32 userOpHash, address sender, uint256 nonce, bytes revertReason
+        bytes32 userOpHash, address sender, string senderLabel, uint256 nonce, bytes revertReason
     );
+    error InvalidRevertMessage(bytes4 expected, bytes4 reason);
+    error InvalidRevertMessageBytes(bytes expected, bytes reason);
 
     function exec4337(PackedUserOperation[] memory userOps, IEntryPoint onEntryPoint) internal {
         uint256 isExpectRevert = getExpectRevert();
@@ -45,10 +52,12 @@ library ERC4337Helpers {
         // Execute userOps
         address payable beneficiary = payable(address(0x69));
         bytes memory userOpCalldata = abi.encodeCall(IEntryPoint.handleOps, (userOps, beneficiary));
-        (bool success,) = address(onEntryPoint).call(userOpCalldata);
+        (bool success, bytes memory returnData) = address(onEntryPoint).call(userOpCalldata);
 
         if (isExpectRevert == 0) {
             require(success, "UserOperation execution failed");
+        } else if (isExpectRevert == 2 && !success) {
+            checkRevertMessage(returnData);
         }
 
         // Parse logs and determine if a revert happened
@@ -64,14 +73,45 @@ library ERC4337Helpers {
                     abi.decode(logs[i].data, (uint256, bool, uint256, uint256));
                 totalUserOpGas = actualGasUsed;
                 if (!userOpSuccess) {
+                    bytes32 userOpHash = logs[i].topics[1];
                     if (isExpectRevert == 0) {
-                        bytes32 userOpHash = logs[i].topics[1];
                         bytes memory revertReason = getUserOpRevertReason(logs, userOpHash);
+                        address account = address(bytes20(logs[i].topics[2]));
                         revert UserOperationReverted(
-                            userOpHash, address(bytes20(logs[i].topics[2])), nonce, revertReason
+                            userOpHash, account, getLabel(account), nonce, revertReason
                         );
                     } else {
-                        writeExpectRevert(0);
+                        if (isExpectRevert == 2) {
+                            checkRevertMessage(getUserOpRevertReason(logs, userOpHash));
+                        }
+                        clearExpectRevert();
+                    }
+                }
+            }
+            // ModuleInstalled(uint256, address)
+            else if (
+                logs[i].topics[0]
+                    == 0xd21d0b289f126c4b473ea641963e766833c2f13866e4ff480abd787c100ef123
+            ) {
+                (uint256 moduleType, address module) = abi.decode(logs[i].data, (uint256, address));
+                writeInstalledModule(InstalledModule(moduleType, module), logs[i].emitter);
+            }
+            // ModuleUninstalled(uint256, address)
+            else if (
+                logs[i].topics[0]
+                    == 0x341347516a9de374859dfda710fa4828b2d48cb57d4fbe4c1149612b8e02276e
+            ) {
+                (uint256 moduleType, address module) = abi.decode(logs[i].data, (uint256, address));
+                // Get all installed modules
+                InstalledModule[] memory installedModules = getInstalledModules(logs[i].emitter);
+                // Remove the uninstalled module from the list of installed modules
+                for (uint256 j; j < installedModules.length; j++) {
+                    if (
+                        installedModules[j].moduleAddress == module
+                            && installedModules[j].moduleType == moduleType
+                    ) {
+                        removeInstalledModule(j, logs[i].emitter);
+                        break;
                     }
                 }
             }
@@ -84,7 +124,7 @@ library ERC4337Helpers {
                 require(!success, "UserOperation execution did not fail as expected");
             }
         }
-        writeExpectRevert(0);
+        clearExpectRevert();
 
         // Calculate gas for userOp
         string memory gasIdentifier = getGasIdentifier();
@@ -108,7 +148,7 @@ library ERC4337Helpers {
 
     function getUserOpRevertReason(
         VmSafe.Log[] memory logs,
-        bytes32 /* userOpHash */
+        bytes32 userOpHash
     )
         internal
         pure
@@ -119,12 +159,29 @@ library ERC4337Helpers {
             if (
                 logs[i].topics[0]
                     == 0x1c4fada7374c0a9ee8841fc38afe82932dc0f8e69012e927f061a8bae611a201
+                    && logs[i].topics[1] == userOpHash
             ) {
                 (, revertReason) = abi.decode(logs[i].data, (uint256, bytes));
             }
         }
     }
 
+    function checkRevertMessage(bytes memory actualReason) internal view {
+        bytes memory revertMessage = getExpectRevertMessage();
+
+        if (revertMessage.length == 4) {
+            bytes4 expected = bytes4(revertMessage);
+            bytes4 actual = bytes4(actualReason);
+            if (expected != actual) {
+                revert InvalidRevertMessage(expected, actual);
+            }
+        } else {
+            if (revertMessage.length != actualReason.length) {
+                revert InvalidRevertMessageBytes(revertMessage, actualReason);
+            }
+        }
+    }
+
     function calculateGas(
         PackedUserOperation[] memory userOps,
         IEntryPoint onEntryPoint,