Skip to content

Commit

Permalink
Expand session key test coverage
Browse files Browse the repository at this point in the history
  • Loading branch information
@adamegyed
adamegyed committed Nov 20, 2023
1 parent 91a281f commit 2331472
Show file tree
Hide file tree
Showing 5 changed files with 225 additions and 5 deletions.
8 changes: 4 additions & 4 deletions src/plugins/session/permissions/SessionKeyPermissionsPlugin.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
56 changes: 55 additions & 1 deletion test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down Expand Up @@ -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
{
Expand Down
17 changes: 17 additions & 0 deletions test/plugin/session/permissions/SessionKeyGasLimits.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
96 changes: 96 additions & 0 deletions test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
53 changes: 53 additions & 0 deletions test/plugin/session/permissions/SessionKeyPermissionsPlugin.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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(
Expand All @@ -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, ""
Expand Down Expand Up @@ -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(
Expand All @@ -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, ""
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 2331472

Please sign in to comment.