add shared runtime validation

erc6900 · Jun 5, 2024 · 73f6768 · 73f6768
1 parent 18c20d7
commit 73f6768
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 28 deletions.
diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
@@ -79,6 +79,7 @@ contract UpgradeableModularAccount is
     error UnexpectedAggregator(address plugin, uint8 functionId, address aggregator);
     error UnrecognizedFunction(bytes4 selector);
     error UserOpValidationFunctionMissing(bytes4 selector);
+    error ValidationDoesNotApply(bytes4 selector, address plugin, uint8 functionId, bool shared);
 
     // Wraps execution of a native function with runtime validation and hooks
     // Used for upgradeTo, upgradeToAndCall, execute, executeBatch, installPlugin, uninstallPlugin
@@ -280,11 +281,12 @@ contract UpgradeableModularAccount is
         if (_storage.selectorData[execSelector].denyExecutionCount > 0) {
             revert AlwaysDenyRule();
         }
-        if (!_storage.selectorData[execSelector].validations.contains(toSetValue(runtimeValidationFunction))) {
-            revert RuntimeValidationFunctionMissing(execSelector);
-        }
 
-        _doRuntimeValidation(runtimeValidationFunction, data, authorization[21:]);
+        // Check if the runtime validation function is allowed to be called
+        bool isSharedValidation = uint8(authorization[21]) == 1;
+        _checkIfValidationApplies(execSelector, runtimeValidationFunction, isSharedValidation);
+
+        _doRuntimeValidation(runtimeValidationFunction, data, authorization[22:]);
 
         // If runtime validation passes, execute the call
 
@@ -434,24 +436,7 @@ contract UpgradeableModularAccount is
         FunctionReference userOpValidationFunction = FunctionReference.wrap(bytes21(userOp.signature[:21]));
         bool isSharedValidation = uint8(userOp.signature[21]) == 1;
 
-        // Check that the provided validation function is applicable to the selector
-        if (isSharedValidation) {
-            if (
-                !_sharedValidationAllowed(selector)
-                    || !_storage.defaultValidations.contains(toSetValue(userOpValidationFunction))
-            ) {
-                revert UserOpValidationFunctionMissing(selector);
-            }
-        } else {
-            // Not shared validation, but per-selector
-            if (
-                !getAccountStorage().selectorData[selector].validations.contains(
-                    toSetValue(userOpValidationFunction)
-                )
-            ) {
-                revert UserOpValidationFunctionMissing(selector);
-            }
-        }
+        _checkIfValidationApplies(selector, userOpValidationFunction, isSharedValidation);
 
         validationData =
             _doUserOpValidation(selector, userOpValidationFunction, userOp, userOp.signature[22:], userOpHash);
@@ -622,6 +607,28 @@ contract UpgradeableModularAccount is
     // solhint-disable-next-line no-empty-blocks
     function _authorizeUpgrade(address newImplementation) internal override {}
 
+    function _checkIfValidationApplies(bytes4 selector, FunctionReference validationFunction, bool shared)
+        internal
+        view
+    {
+        AccountStorage storage _storage = getAccountStorage();
+
+        // Check that the provided validation function is applicable to the selector
+        if (shared) {
+            if (
+                !_sharedValidationAllowed(selector)
+                    || !_storage.defaultValidations.contains(toSetValue(validationFunction))
+            ) {
+                revert UserOpValidationFunctionMissing(selector);
+            }
+        } else {
+            // Not shared validation, but per-selector
+            if (!getAccountStorage().selectorData[selector].validations.contains(toSetValue(validationFunction))) {
+                revert UserOpValidationFunctionMissing(selector);
+            }
+        }
+    }
+
     function _sharedValidationAllowed(bytes4 selector) internal view returns (bool) {
         if (
             selector == this.execute.selector || selector == this.executeBatch.selector

diff --git a/test/account/AccountReturnData.t.sol b/test/account/AccountReturnData.t.sol
@@ -59,7 +59,9 @@ contract AccountReturnDataTest is AccountTestBase {
                 account1.execute,
                 (address(regularResultContract), 0, abi.encodeCall(RegularResultContract.foo, ()))
             ),
-            abi.encodePacked(singleOwnerPlugin, ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER)
+            abi.encodePacked(
+                singleOwnerPlugin, ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER, SELECTOR_ASSOCIATED_VALIDATION
+            )
         );
 
         bytes32 result = abi.decode(abi.decode(returnData, (bytes)), (bytes32));
@@ -83,7 +85,9 @@ contract AccountReturnDataTest is AccountTestBase {
 
         bytes memory retData = account1.executeWithAuthorization(
             abi.encodeCall(account1.executeBatch, (calls)),
-            abi.encodePacked(singleOwnerPlugin, ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER)
+            abi.encodePacked(
+                singleOwnerPlugin, ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER, SELECTOR_ASSOCIATED_VALIDATION
+            )
         );
 
         bytes[] memory returnDatas = abi.decode(retData, (bytes[]));

diff --git a/test/account/MultiValidation.t.sol b/test/account/MultiValidation.t.sol
@@ -67,13 +67,21 @@ contract MultiValidationTest is AccountTestBase {
         );
         account1.executeWithAuthorization(
             abi.encodeCall(IStandardExecutor.execute, (address(0), 0, "")),
-            abi.encodePacked(address(validator2), uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER))
+            abi.encodePacked(
+                address(validator2),
+                uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER),
+                SELECTOR_ASSOCIATED_VALIDATION
+            )
         );
 
         vm.prank(owner2);
         account1.executeWithAuthorization(
             abi.encodeCall(IStandardExecutor.execute, (address(0), 0, "")),
-            abi.encodePacked(address(validator2), uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER))
+            abi.encodePacked(
+                address(validator2),
+                uint8(ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER),
+                SELECTOR_ASSOCIATED_VALIDATION
+            )
         );
     }
 

diff --git a/test/account/SharedValidationTest.t.sol b/test/account/SharedValidationTest.t.sol
@@ -38,7 +38,7 @@ contract SharedValidationTestTest is AccountTestBase {
         );
     }
 
-    function test_sharedValidation_simple() public {
+    function test_sharedValidation_userOp_simple() public {
         PackedUserOperation memory userOp = PackedUserOperation({
             sender: address(account1),
             nonce: 0,
@@ -66,4 +66,17 @@ contract SharedValidationTestTest is AccountTestBase {
 
         assertEq(ethRecipient.balance, 2 wei);
     }
+
+    function test_sharedValidation_runtime_simple() public {
+        // Deploy the account first
+        sharedValidationFactoryFixture.createAccount(owner1, 0);
+
+        vm.prank(owner1);
+        account1.executeWithAuthorization(
+            abi.encodeCall(UpgradeableModularAccount.execute, (ethRecipient, 1 wei, "")),
+            abi.encodePacked(ownerValidation, SHARED_VALIDATION)
+        );
+
+        assertEq(ethRecipient.balance, 2 wei);
+    }
 }
diff --git a/test/utils/AccountTestBase.sol b/test/utils/AccountTestBase.sol
@@ -50,7 +50,11 @@ abstract contract AccountTestBase is OptimizedTest {
                     abi.encodeCall(SingleOwnerPlugin.transferOwnership, (address(this)))
                 )
             ),
-            abi.encodePacked(address(singleOwnerPlugin), ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER)
+            abi.encodePacked(
+                address(singleOwnerPlugin),
+                ISingleOwnerPlugin.FunctionId.VALIDATION_OWNER,
+                SELECTOR_ASSOCIATED_VALIDATION
+            )
         );
     }