alchemyplatform · adam-alchemy · Nov 20, 2023 · Nov 20, 2023
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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,