add not registered check to session key rotation

alchemyplatform · Dec 13, 2023 · 2ba8823 · 2ba8823
1 parent 84f4b3d
commit 2ba8823
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 0 deletions.
diff --git a/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol b/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol
@@ -55,8 +55,15 @@ contract SessionKeyPermissionsPlugin is ISessionKeyPermissionsPlugin, SessionKey
 
     /// @inheritdoc ISessionKeyPermissionsPlugin
     function rotateKey(address oldSessionKey, address newSessionKey) external override {
+        // Assert that the new key is not already registered
+        SessionKeyId newKeyId = _sessionKeyIdOf(msg.sender, newSessionKey);
+        if (SessionKeyId.unwrap(newKeyId) != 0) {
+            revert KeyAlreadyRegistered(newSessionKey);
+        }
+        // Assert that the old key is registered
         SessionKeyId keyId = _sessionKeyIdOf(msg.sender, oldSessionKey);
         _assertRegistered(keyId, oldSessionKey);
+        // Update the key ID of the old key to zero, and the new key to the old key's ID
         _updateSessionKeyId(msg.sender, oldSessionKey, SessionKeyId.wrap(bytes32(0)));
         _updateSessionKeyId(msg.sender, newSessionKey, keyId);
 

diff --git a/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol b/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol
@@ -448,6 +448,35 @@ contract SessionKeyPermissionsPluginTest is Test {
         assertEq(returnedEndTime, endTime);
     }
 
+    function test_rotateKey_invalid_alreadyRegistered() public {
+        // Attempt to rotate the key to itself
+        vm.expectRevert(
+            abi.encodeWithSelector(ISessionKeyPermissionsPlugin.KeyAlreadyRegistered.selector, sessionKey1)
+        );
+        vm.prank(owner1);
+        SessionKeyPermissionsPlugin(address(account1)).rotateKey(sessionKey1, sessionKey1);
+
+        // Add a second session key
+        address sessionKey2 = makeAddr("sessionKey2");
+        address[] memory keysToAdd = new address[](1);
+        keysToAdd[0] = sessionKey2;
+        vm.prank(owner1);
+        SessionKeyPlugin(address(account1)).updateSessionKeys(
+            keysToAdd, new SessionKeyPlugin.SessionKeyToRemove[](0)
+        );
+
+        // Register the second session key
+        vm.prank(owner1);
+        SessionKeyPermissionsPlugin(address(account1)).registerKey(sessionKey2, 0);
+
+        // Attempt to rotate the key to the second session key
+        vm.expectRevert(
+            abi.encodeWithSelector(ISessionKeyPermissionsPlugin.KeyAlreadyRegistered.selector, sessionKey2)
+        );
+        vm.prank(owner1);
+        SessionKeyPermissionsPlugin(address(account1)).rotateKey(sessionKey1, sessionKey2);
+    }
+
     function testFuzz_sessionKeyPermissions_setRequiredPaymaster(address requiredPaymaster) public {
         bytes[] memory updates = new bytes[](1);
         updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setRequiredPaymaster, (requiredPaymaster));