From 23314725a2669bd282694ab8b956c3ba1fa9ab60 Mon Sep 17 00:00:00 2001 From: adam Date: Sun, 19 Nov 2023 23:19:07 -0500 Subject: [PATCH] Expand session key test coverage --- .../SessionKeyPermissionsPlugin.sol | 8 +- .../SessionKeyERC20SpendLimits.t.sol | 56 ++++++++++- .../permissions/SessionKeyGasLimits.t.sol | 17 ++++ .../SessionKeyNativeTokenSpendLimits.t.sol | 96 +++++++++++++++++++ .../SessionKeyPermissionsPlugin.t.sol | 53 ++++++++++ 5 files changed, 225 insertions(+), 5 deletions(-) diff --git a/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol b/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol index 6db2b1ab..ce9c26cc 100644 --- a/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol +++ b/src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol @@ -561,17 +561,17 @@ contract SessionKeyPermissionsPlugin is ISessionKeyPermissionsPlugin, SessionKey // Must re-check the limits to handle changes due to other user ops. // We manually check for overflows here to give a more informative error message. - uint256 sum; + uint256 newTotalUsage; unchecked { - sum = newUsage + currentUsage; + newTotalUsage = newUsage + currentUsage; } - if (sum < newUsage || sum > spendLimit) { + if (newTotalUsage < newUsage || newTotalUsage > spendLimit) { // If we overflow, or if the limit is exceeded, fail here and revert in the parent context. return false; } // We won't update the refresh interval last used variable now, so just update the spend limit. - limit.limitUsed += newUsage; + limit.limitUsed = newTotalUsage; } else { // We have a interval active that is currently resetting. // Must re-check the amount to handle changes due to other user ops. diff --git a/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol b/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol index bf4b846b..545afa5c 100644 --- a/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol @@ -885,7 +885,7 @@ contract SessionKeyERC20SpendLimitsTest is Test { // Give the account a starting balance token1.mint(address(account1), 100 ether); - // Set the limit to 1 ether + // Set the limit to 1 ether over 1 day bytes[] memory updates = new bytes[](1); updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setERC20SpendLimit, (address(token1), 1 ether, 1 days)); @@ -1036,6 +1036,60 @@ contract SessionKeyERC20SpendLimitsTest is Test { assertEq(spendLimitInfo2.lastUsedTime, time0 + 10 days); } + function test_sessionKeyERC20Limits_refreshInterval_failWithExceedNextLimit() public { + // Set the time to the a unix timestamp + uint256 time0 = 1698708080; + vm.warp(time0); + + // Give the account a starting balance + token1.mint(address(account1), 100 ether); + + // Set the limit to 1 ether over 1 day + bytes[] memory updates = new bytes[](1); + updates[0] = + abi.encodeCall(ISessionKeyPermissionsUpdates.setERC20SpendLimit, (address(token1), 1 ether, 1 days)); + vm.prank(owner1); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + // Run a user op that spends 0.5 ether, should succeed + Call[] memory calls = new Call[](1); + calls[0] = Call({ + target: address(token1), + data: abi.encodeCall(token1.transfer, (recipient1, 0.5 ether)), + value: 0 + }); + vm.expectCall(address(token1), 0 wei, calls[0].data, 1); + _runSessionKeyUserOp(calls, sessionKey1Private, ""); + + // Assert that the limit and last used time is updated + ISessionKeyPermissionsPlugin.SpendLimitInfo memory spendLimitInfo = + sessionKeyPermissionsPlugin.getERC20SpendLimitInfo(address(account1), sessionKey1, address(token1)); + assertEq(spendLimitInfo.limit, 1 ether); + assertEq(spendLimitInfo.limitUsed, 0.5 ether); + assertEq(spendLimitInfo.refreshInterval, 1 days); + assertEq(spendLimitInfo.lastUsedTime, time0); + + // Skip forward to the next interval + skip(1 days + 1 minutes); + + // Attempt a user operation that spends 1.5 ether (Exceeds limit, should fail). + calls[0] = Call({ + target: address(token1), + data: abi.encodeCall(token1.transfer, (recipient1, 1.5 ether)), + value: 0 + }); + vm.expectCall(address(token1), 0 wei, calls[0].data, 0); + _runSessionKeyUserOp(calls, sessionKey1Private, ""); + + // Assert that limits are not updated + spendLimitInfo = + sessionKeyPermissionsPlugin.getERC20SpendLimitInfo(address(account1), sessionKey1, address(token1)); + assertEq(spendLimitInfo.limit, 1 ether); + assertEq(spendLimitInfo.limitUsed, 0.5 ether); + assertEq(spendLimitInfo.refreshInterval, 1 days); + assertEq(spendLimitInfo.lastUsedTime, time0); + } + function _runSessionKeyUserOp(Call[] memory calls, uint256 sessionKeyPrivate, bytes memory expectedError) internal { diff --git a/test/plugin/session/permissions/SessionKeyGasLimits.t.sol b/test/plugin/session/permissions/SessionKeyGasLimits.t.sol index 1ad64178..39b741c6 100644 --- a/test/plugin/session/permissions/SessionKeyGasLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyGasLimits.t.sol @@ -222,6 +222,23 @@ contract SessionKeyGasLimitsTest is Test { ); } + function testFuzz_sessionKeyGasLimits_requireNonceAsAddress(uint192 nonceKey) public { + vm.assume(uint192(uint160(sessionKey1)) != nonceKey); + + bytes[] memory updates = new bytes[](1); + updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setGasSpendLimit, (1 ether, 0 days)); + vm.prank(owner1); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + UserOperation[] memory userOps = new UserOperation[](2); + userOps[0] = _generateAndSignUserOp( + 100_000, 300_000, 1_000 gwei, 0.4 ether, sessionKey1Private, uint256(nonceKey << 64) + ); + + vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error")); + entryPoint.handleOps(userOps, beneficiary); + } + function test_sessionKeyGasLimits_exceedLimit_single() public { bytes[] memory updates = new bytes[](1); updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setGasSpendLimit, (1 ether, 0 days)); diff --git a/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol b/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol index fc26e473..a267d90e 100644 --- a/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol @@ -244,6 +244,52 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { ); } + function test_sessionKeyNativeTokenSpendLimits_overflowState() public { + // Set the limit to 1 ether + bytes[] memory updates = new bytes[](1); + updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setNativeTokenSpendLimit, (1 ether, 0)); + vm.prank(owner1); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + // Run a user op that spends 1 wei, should succeed + Call[] memory calls = new Call[](1); + calls[0] = Call({target: recipient1, value: 1 wei, data: ""}); + _runSessionKeyUserOp(calls, sessionKey1Private, ""); + + // Attempt to run a user op that spends type(uint256).max (causing an overflow). Should fail. + calls[0] = Call({target: recipient1, value: type(uint256).max, data: ""}); + + // Assert that the failure happens during validation. + // We would like to assert here that the pre user op validation hook itself does not revert, but since the + // account catches those reverts internally and returns SIG_FAIL, we can't assert it from here easily. + _runSessionKeyUserOp( + calls, + sessionKey1Private, + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + ); + } + + function test_sessionKeyNativeTokenSpendLimits_overflowBatch() public { + // NOTE: overflows that happen within a single batch are not gracefully caught and returned as `SIG_FAIL`, + // instead the pre exec hook reverts. However, the modular account will catch that revert. + + // Set the limit to 1 ether + bytes[] memory updates = new bytes[](1); + updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setNativeTokenSpendLimit, (1 ether, 0)); + vm.prank(owner1); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + // Attempt to run a user op that spends > type(uint256).max (causing an overflow). Should fail. + Call[] memory calls = new Call[](2); + calls[0] = Call({target: recipient1, value: 1 wei, data: ""}); + calls[1] = Call({target: recipient1, value: type(uint256).max, data: ""}); + _runSessionKeyUserOp( + calls, + sessionKey1Private, + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + ); + } + function test_sessionKeyNativeTokenSpendLimits_exceedLimit_single() public { // Set the limit to 1 ether bytes[] memory updates = new bytes[](1); @@ -629,6 +675,53 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { assertEq(spendLimitInfo.lastUsedTime, 0); } + function test_sessionKeyNativeTokenSpendLimits_refreshInterval_exceedNewInterval() public { + // Set the time to the current unix timestamp as of writing + uint256 time0 = 1698708080; + vm.warp(time0); + + // Set the limit to 1 ether + bytes[] memory updates = new bytes[](1); + updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.setNativeTokenSpendLimit, (1 ether, 1 days)); + vm.prank(owner1); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + // Run a user op that spends 1 wei, should succeed + Call[] memory calls = new Call[](1); + calls[0] = Call({target: recipient1, value: 1 wei, data: ""}); + + _runSessionKeyUserOp(calls, sessionKey1Private, ""); + + // Assert that the limit is now updated and the last used timestamp is set. + ISessionKeyPermissionsPlugin.SpendLimitInfo memory spendLimitInfo = + sessionKeyPermissionsPlugin.getNativeTokenSpendLimitInfo(address(account1), sessionKey1); + + assertTrue(spendLimitInfo.hasLimit); + assertEq(spendLimitInfo.limit, 1 ether); + assertEq(spendLimitInfo.limitUsed, 1 wei); + assertEq(spendLimitInfo.refreshInterval, 1 days); + assertEq(spendLimitInfo.lastUsedTime, time0); + + skip(1 days + 1 minutes); + + // Attempt to run a user op that spends 1.1 ether, should fail. + calls[0] = Call({target: recipient1, value: 1.1 ether, data: ""}); + _runSessionKeyUserOp( + calls, + sessionKey1Private, + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + ); + + // Assert that limits are NOT updated + spendLimitInfo = sessionKeyPermissionsPlugin.getNativeTokenSpendLimitInfo(address(account1), sessionKey1); + + assertTrue(spendLimitInfo.hasLimit); + assertEq(spendLimitInfo.limit, 1 ether); + assertEq(spendLimitInfo.limitUsed, 1 wei); + assertEq(spendLimitInfo.refreshInterval, 1 days); + assertEq(spendLimitInfo.lastUsedTime, time0); + } + // There's an additional pre exec revert that I haven't been able to trigger in a real example, when a // usage of a session key with a native token spend limit reaches the "new time interval" section, but the // amount being spent exceeds the new limit. This seems to be impossible to reach because any prior usage would @@ -639,6 +732,9 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { // want to remove it for now, since it's a fairly cheap check that may prevent a dangerous issue. // // function test_sessionKeyNativeTokenSpendLimits_multiUserOpBundle_check_interval() + // + // Updated note since the addition of ERC-20 spend limits: this code branch is reachable from there, and was + // handled correctly. However, I still cannot reach it from the native token spend limit checking. function _runSessionKeyUserOp(Call[] memory calls, uint256 sessionKeyPrivate, bytes memory expectedError) internal diff --git a/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol b/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol index a7c699d7..e9f28a2f 100644 --- a/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol +++ b/test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol @@ -134,6 +134,14 @@ contract SessionKeyPermissionsPluginTest is Test { ); } + function test_sessionPerms_duplicateRegister() public { + vm.prank(owner1); + vm.expectRevert( + abi.encodeWithSelector(ISessionKeyPermissionsPlugin.KeyAlreadyRegistered.selector, sessionKey1) + ); + SessionKeyPermissionsPlugin(address(account1)).registerKey(sessionKey1, 0); + } + function test_sessionPerms_contractDefaultAllowList() public { _runSessionKeyExecUserOp( address(counter1), @@ -165,6 +173,13 @@ contract SessionKeyPermissionsPluginTest is Test { } function test_sessionPerms_contractAllowList() public { + // Assert the contracts to be added are not already on the allowlist + (bool isOnList, bool checkSelectors) = + sessionKeyPermissionsPlugin.getAccessControlEntry(address(account1), sessionKey1, address(counter1)); + + assertFalse(isOnList, "Address should not start on the list"); + assertFalse(checkSelectors, "Address should not start with selectors checked"); + // Add the allowlist bytes[] memory updates = new bytes[](1); updates[0] = abi.encodeCall( @@ -173,6 +188,12 @@ contract SessionKeyPermissionsPluginTest is Test { vm.prank(owner1); SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + // Plugin should report the entry as on the list + (isOnList, checkSelectors) = + sessionKeyPermissionsPlugin.getAccessControlEntry(address(account1), sessionKey1, address(counter1)); + assertTrue(isOnList); + assertFalse(checkSelectors); + // Call should succeed after adding the allowlist _runSessionKeyExecUserOp( address(counter1), sessionKey1, sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, "" @@ -227,6 +248,16 @@ contract SessionKeyPermissionsPluginTest is Test { } function test_sessionPerms_selectorAllowList() public { + // Validate that the address and the selector do not start out enabled. + (bool addressOnList, bool checkSelectors) = + sessionKeyPermissionsPlugin.getAccessControlEntry(address(account1), sessionKey1, address(counter1)); + assertFalse(addressOnList); + assertFalse(checkSelectors); + bool selectorOnList = sessionKeyPermissionsPlugin.isSelectorOnAccessControlList( + address(account1), sessionKey1, address(counter1), Counter.increment.selector + ); + assertFalse(selectorOnList); + // Add the allowlist bytes[] memory updates = new bytes[](2); updates[0] = abi.encodeCall( @@ -240,6 +271,16 @@ contract SessionKeyPermissionsPluginTest is Test { vm.prank(owner1); SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + // Validate that the address and the selector are now enabled + (addressOnList, checkSelectors) = + sessionKeyPermissionsPlugin.getAccessControlEntry(address(account1), sessionKey1, address(counter1)); + assertTrue(addressOnList); + assertTrue(checkSelectors); + selectorOnList = sessionKeyPermissionsPlugin.isSelectorOnAccessControlList( + address(account1), sessionKey1, address(counter1), Counter.increment.selector + ); + assertTrue(selectorOnList); + // Call should succeed after adding the allowlist _runSessionKeyExecUserOp( address(counter1), sessionKey1, sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, "" @@ -533,6 +574,18 @@ contract SessionKeyPermissionsPluginTest is Test { entryPoint.handleOps(userOps, beneficiary); } + function test_sessionKeyPerms_updatePermissions_invalidUpdates() public { + vm.startPrank(owner1); + bytes[] memory updates = new bytes[](1); + updates[0] = hex"112233"; // < 4 byte update + vm.expectRevert(abi.encodeWithSelector(ISessionKeyPermissionsPlugin.InvalidPermissionsUpdate.selector)); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + + updates[0] = hex"11223344"; // Invalid selector + vm.expectRevert(abi.encodeWithSelector(ISessionKeyPermissionsPlugin.InvalidPermissionsUpdate.selector)); + SessionKeyPermissionsPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); + } + function _runSessionKeyExecUserOp( address target, address sessionKey,