From 7e58e9993255ebf3e238b9330502d77da96249e3 Mon Sep 17 00:00:00 2001
From: adam-alchemy <127769144+adam-alchemy@users.noreply.github.com>
Date: Tue, 12 Dec 2023 12:45:45 -0800
Subject: [PATCH] fix: [spearbit-76] Session key data storage collision (#14)

---
 .../permissions/SessionKeyPermissionsBase.sol | 10 +---
 .../SessionKeyPermissionsPlugin.t.sol         | 53 +++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/src/plugins/session/permissions/SessionKeyPermissionsBase.sol b/src/plugins/session/permissions/SessionKeyPermissionsBase.sol
index d41a15f8..f9e95cbd 100644
--- a/src/plugins/session/permissions/SessionKeyPermissionsBase.sol
+++ b/src/plugins/session/permissions/SessionKeyPermissionsBase.sol
@@ -54,12 +54,6 @@ abstract contract SessionKeyPermissionsBase is ISessionKeyPermissionsPlugin {
         uint256 limitUsed;
     }
 
-    // PluginStorageLib KEY DEFINITIONS
-    // When adding a new permission type, you must:
-    // 1. Add new prefixes here for all stored structs
-    // 2. Define the key derivation and checking functions in a new file
-    // 3. Use the checking function in PermissionsCheckerPlugin
-
     // Prefixes:
     bytes4 internal constant SESSION_KEY_ID_PREFIX = bytes4(0x1a01dae4); // bytes4(keccak256("SessionKeyId"))
     bytes4 internal constant SESSION_KEY_DATA_PREFIX = bytes4(0x16bff296); // bytes4(keccak256("SessionKeyData"))
@@ -84,7 +78,7 @@ abstract contract SessionKeyPermissionsBase is ISessionKeyPermissionsPlugin {
 
     // ContractData (128 bytes)
     // 12 padding zeros || associated address || CONTRACT_DATA_PREFIX || batch index || sessionKeyId
-    // || 12 padding zero bytes || contractAddress
+    // || contractAddress  || 12 padding zero bytes
 
     // FunctionData (128 bytes)
     // 12 padding zeros || associated address || FUNCTION_DATA_PREFIX || batch index || sessionKeyId || selector
@@ -130,7 +124,7 @@ abstract contract SessionKeyPermissionsBase is ISessionKeyPermissionsPlugin {
         bytes memory associatedStorageKey =
             PluginStorageLib.allocateAssociatedStorageKey(associated, prefixAndBatchIndex, 1);
 
-        bytes32 sessionKeyDataKey = bytes32(abi.encodePacked(SESSION_KEY_DATA_PREFIX, SessionKeyId.unwrap(id)));
+        bytes32 sessionKeyDataKey = SessionKeyId.unwrap(id);
         return _toSessionKeyData(PluginStorageLib.associatedStorageLookup(associatedStorageKey, sessionKeyDataKey));
     }
 
diff --git a/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol b/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol
index 6cc525ab..5833bab2 100644
--- a/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol
+++ b/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol
@@ -591,6 +591,59 @@ contract SessionKeyPermissionsPluginTest is Test {
         SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates);
     }
 
+    function test_sessionKeyPerms_independentKeyStorage() public {
+        address sessionKey2 = makeAddr("sessionKey2");
+
+        address[] memory sessionKeysToAdd = new address[](1);
+        sessionKeysToAdd[0] = sessionKey2;
+
+        vm.prank(owner1);
+        SessionKeyPlugin(address(account1)).updateSessionKeys(
+            sessionKeysToAdd, new SessionKeyPlugin.SessionKeyToRemove[](0)
+        );
+        vm.prank(owner1);
+        SessionKeyPermissionsPlugin(address(account1)).registerKey(sessionKey2, 0);
+
+        ISessionKeyPermissionsPlugin.ContractAccessControlType accessControlType1;
+        ISessionKeyPermissionsPlugin.ContractAccessControlType accessControlType2;
+
+        accessControlType1 = sessionKeyPermissionsPlugin.getAccessControlType(address(account1), sessionKey1);
+        accessControlType2 = sessionKeyPermissionsPlugin.getAccessControlType(address(account1), sessionKey2);
+
+        assertEq(
+            uint8(accessControlType1),
+            uint8(ISessionKeyPermissionsPlugin.ContractAccessControlType.ALLOWLIST),
+            "sessionKey1 should start with an allowlist"
+        );
+        assertEq(
+            uint8(accessControlType2),
+            uint8(ISessionKeyPermissionsPlugin.ContractAccessControlType.ALLOWLIST),
+            "sessionKey2 should start with an allowlist"
+        );
+
+        bytes[] memory updates = new bytes[](1);
+        updates[0] = abi.encodeCall(
+            ISessionKeyPermissionsUpdates.setAccessListType,
+            (ISessionKeyPermissionsPlugin.ContractAccessControlType.NONE)
+        );
+        vm.prank(owner1);
+        SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates);
+
+        accessControlType1 = sessionKeyPermissionsPlugin.getAccessControlType(address(account1), sessionKey1);
+        accessControlType2 = sessionKeyPermissionsPlugin.getAccessControlType(address(account1), sessionKey2);
+
+        assertEq(
+            uint8(accessControlType1),
+            uint8(ISessionKeyPermissionsPlugin.ContractAccessControlType.NONE),
+            "sessionKey1 should now have no allowlist"
+        );
+        assertEq(
+            uint8(accessControlType2),
+            uint8(ISessionKeyPermissionsPlugin.ContractAccessControlType.ALLOWLIST),
+            "sessionKey2 should still have an allowlist"
+        );
+    }
+
     function _runSessionKeyExecUserOp(
         address target,
         address sessionKey,