From 2ba882398509ca2ec2fae3a30a144f031713ef65 Mon Sep 17 00:00:00 2001 From: adam Date: Wed, 13 Dec 2023 12:19:23 -0500 Subject: [PATCH] add not registered check to session key rotation --- .../SessionKeyPermissionsPlugin.sol | 7 +++++ .../SessionKeyPermissionsPlugin.t.sol | 29 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol b/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol index 0bceee98..af118d52 100644 --- 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 index 5833bab2..33ce259d 100644 --- 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));