Fix aspect wallet aspect operation check issue

artela-network · Sep 14, 2023 · 159e1bc · 159e1bc
1 parent d7c125a
commit 159e1bc
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/contracts/samples/AspectEnabledSimpleAccount.sol b/contracts/samples/AspectEnabledSimpleAccount.sol
@@ -54,7 +54,7 @@ contract AspectEnabledSimpleAccount is SimpleAccount {
      * @dev add a set of Aspect to whitelist
      */
     function approveAspects(address[] calldata aspectIds) external {
-        _requireFromEntryPointOrOwner();
+        _requireFromSelfOrOwner();
         for (uint256 i = 0; i < aspectIds.length; ++i) {
             _aspectWhitelist[aspectIds[i]] = true;
         }
@@ -64,7 +64,7 @@ contract AspectEnabledSimpleAccount is SimpleAccount {
      * @dev remove a set of Aspect from the whitelist
      */
     function disApproveAspects(address[] calldata aspectIds) external {
-        _requireFromEntryPointOrOwner();
+        _requireFromSelfOrOwner();
         for (uint256 i = 0; i < aspectIds.length; ++i) {
             delete _aspectWhitelist[aspectIds[i]];
         }
@@ -94,5 +94,10 @@ contract AspectEnabledSimpleAccount is SimpleAccount {
         }
         return result;
     }
+
+    // Require the function call went through Self or owner
+    function _requireFromSelfOrOwner() internal view {
+        require(msg.sender == address(this) || msg.sender == owner, "account: not Owner or Self");
+    }
 }