simplify permitted call mapping

erc6900 · May 29, 2024 · 30ce3cf · 30ce3cf
1 parent 46a0f31
commit 30ce3cf
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 20 deletions.
diff --git a/src/account/AccountStorage.sol b/src/account/AccountStorage.sol
@@ -57,8 +57,7 @@ struct AccountStorage {
     mapping(address => PluginData) pluginData;
     // Execution functions and their associated functions
     mapping(bytes4 => SelectorData) selectorData;
-    // bytes24 key = address(calling plugin) || bytes4(selector of execution function)
-    mapping(bytes24 => bool) callPermitted;
+    mapping(address caller => mapping(bytes4 selector => bool)) callPermitted;
     // key = address(calling plugin) || target address
     mapping(IPlugin => mapping(address => PermittedExternalCallData)) permittedExternalCalls;
     // For ERC165 introspection
@@ -77,10 +76,6 @@ function getAccountStorage() pure returns (AccountStorage storage _storage) {
     }
 }
 
-function getPermittedCallKey(address addr, bytes4 selector) pure returns (bytes24) {
-    return bytes24(bytes20(addr)) | (bytes24(selector) >> 160);
-}
-
 using EnumerableSet for EnumerableSet.Bytes32Set;
 
 function toSetValue(FunctionReference functionReference) pure returns (bytes32) {

diff --git a/src/account/PluginManagerInternals.sol b/src/account/PluginManagerInternals.sol
@@ -21,7 +21,6 @@ import {
     getAccountStorage,
     SelectorData,
     toSetValue,
-    getPermittedCallKey,
     PermittedExternalCallData
 } from "./AccountStorage.sol";
 
@@ -220,7 +219,7 @@ abstract contract PluginManagerInternals is IPluginManager {
         for (uint256 i = 0; i < length; ++i) {
             // If there are duplicates, this will just enable the flag again. This is not a problem, since the
             // boolean will be set to false twice during uninstall, which is fine.
-            _storage.callPermitted[getPermittedCallKey(plugin, manifest.permittedExecutionSelectors[i])] = true;
+            _storage.callPermitted[plugin][manifest.permittedExecutionSelectors[i]] = true;
         }
 
         // Add the permitted external calls to the account.
@@ -417,7 +416,7 @@ abstract contract PluginManagerInternals is IPluginManager {
 
         length = manifest.permittedExecutionSelectors.length;
         for (uint256 i = 0; i < length; ++i) {
-            _storage.callPermitted[getPermittedCallKey(plugin, manifest.permittedExecutionSelectors[i])] = false;
+            _storage.callPermitted[plugin][manifest.permittedExecutionSelectors[i]] = false;
         }
 
         length = manifest.executionFunctions.length;

diff --git a/src/account/UpgradeableModularAccount.sol b/src/account/UpgradeableModularAccount.sol
@@ -23,7 +23,6 @@ import {AccountLoupe} from "./AccountLoupe.sol";
 import {
     AccountStorage,
     getAccountStorage,
-    getPermittedCallKey,
     SelectorData,
     toSetValue,
     toFunctionReference,
@@ -187,11 +186,9 @@ contract UpgradeableModularAccount is
         bytes4 selector = bytes4(data[:4]);
         address callingPlugin = msg.sender;
 
-        bytes24 execFromPluginKey = getPermittedCallKey(callingPlugin, selector);
-
         AccountStorage storage _storage = getAccountStorage();
 
-        if (!_storage.callPermitted[execFromPluginKey]) {
+        if (!_storage.callPermitted[callingPlugin][selector]) {
             revert ExecFromPluginNotPermitted(callingPlugin, selector);
         }
 

diff --git a/test/libraries/AccountStorage.t.sol b/test/libraries/AccountStorage.t.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.19;
 
 import {Test} from "forge-std/Test.sol";
-import {_ACCOUNT_STORAGE_SLOT, getPermittedCallKey} from "../../src/account/AccountStorage.sol";
+import {_ACCOUNT_STORAGE_SLOT} 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";
@@ -34,10 +34,4 @@ contract AccountStorageTest is Test {
         // 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));
-    }
 }