From 7d9db25fb80c6472ad3dc90dffc251c5128bc9bf Mon Sep 17 00:00:00 2001 From: howydev <132113803+howydev@users.noreply.github.com> Date: Fri, 9 Feb 2024 12:58:07 -0500 Subject: [PATCH] fix: improve session key gas estimation --- src/plugins/session/ISessionKeyPlugin.sol | 1 + src/plugins/session/SessionKeyPlugin.sol | 15 +++++++---- .../SessionKeyERC20SpendLimits.t.sol | 2 +- .../permissions/SessionKeyGasLimits.t.sol | 12 ++++----- .../SessionKeyNativeTokenSpendLimits.t.sol | 12 ++++----- .../permissions/SessionKeyPermissions.t.sol | 26 ++++++++++++------- 6 files changed, 40 insertions(+), 28 deletions(-) diff --git a/src/plugins/session/ISessionKeyPlugin.sol b/src/plugins/session/ISessionKeyPlugin.sol index a5270dfda..b5cea3f39 100644 --- a/src/plugins/session/ISessionKeyPlugin.sol +++ b/src/plugins/session/ISessionKeyPlugin.sol @@ -77,6 +77,7 @@ interface ISessionKeyPlugin { error LengthMismatch(); error NativeTokenSpendLimitExceeded(address account, address sessionKey); error SessionKeyNotFound(address sessionKey); + error PermissionsCheckFailed(); // ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ // ┃ Execution functions ┃ diff --git a/src/plugins/session/SessionKeyPlugin.sol b/src/plugins/session/SessionKeyPlugin.sol index 0bdf77214..1c225c8e4 100644 --- a/src/plugins/session/SessionKeyPlugin.sol +++ b/src/plugins/session/SessionKeyPlugin.sol @@ -35,7 +35,9 @@ import {Call, IStandardExecutor} from "../../interfaces/IStandardExecutor.sol"; import { AssociatedLinkedListSet, AssociatedLinkedListSetLib } from "../../libraries/AssociatedLinkedListSetLib.sol"; -import {SetValue, SENTINEL_VALUE, SIG_VALIDATION_FAILED} from "../../libraries/Constants.sol"; +import { + SetValue, SENTINEL_VALUE, SIG_VALIDATION_PASSED, SIG_VALIDATION_FAILED +} from "../../libraries/Constants.sol"; import {BasePlugin} from "../BasePlugin.sol"; import {ISessionKeyPlugin} from "./ISessionKeyPlugin.sol"; import {SessionKeyPermissions} from "./permissions/SessionKeyPermissions.sol"; @@ -183,10 +185,13 @@ contract SessionKeyPlugin is ISessionKeyPlugin, SessionKeyPermissions, BasePlugi (address recoveredSig, ECDSA.RecoverError err) = hash.tryRecover(userOp.signature); if (err == ECDSA.RecoverError.NoError) { - if ( - _sessionKeys.contains(msg.sender, CastLib.toSetValue(sessionKey)) && sessionKey == recoveredSig - ) { - return _checkUserOpPermissions(userOp, calls, sessionKey); + if (_sessionKeys.contains(msg.sender, CastLib.toSetValue(sessionKey))) { + uint256 validation = _checkUserOpPermissions(userOp, calls, sessionKey); + if (uint160(validation) > 0) { + revert PermissionsCheckFailed(); + } + return + validation | (sessionKey == recoveredSig ? SIG_VALIDATION_PASSED : SIG_VALIDATION_FAILED); } return SIG_VALIDATION_FAILED; } diff --git a/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol b/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol index c090aeda6..8913667d3 100644 --- a/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyERC20SpendLimits.t.sol @@ -220,7 +220,7 @@ contract SessionKeyERC20SpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } diff --git a/test/plugin/session/permissions/SessionKeyGasLimits.t.sol b/test/plugin/session/permissions/SessionKeyGasLimits.t.sol index 1ead5e0a1..c1d0d95a9 100644 --- a/test/plugin/session/permissions/SessionKeyGasLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyGasLimits.t.sol @@ -166,7 +166,7 @@ contract SessionKeyGasLimitsTest is Test { 1, 200_000 wei, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } @@ -206,7 +206,7 @@ contract SessionKeyGasLimitsTest is Test { 1_000 gwei, 0.6 ether, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } @@ -223,7 +223,7 @@ contract SessionKeyGasLimitsTest is Test { 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")); + vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)")); entryPoint.handleOps(userOps, beneficiary); } @@ -240,7 +240,7 @@ contract SessionKeyGasLimitsTest is Test { 2_000 gwei, 1.2 ether, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } @@ -289,7 +289,7 @@ contract SessionKeyGasLimitsTest is Test { // Run the user ops // The second op (index 1) should be the one that fails signature validation. - vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 1, "AA24 signature error")); + vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 1, "AA23 reverted (or OOG)")); entryPoint.handleOps(userOps, beneficiary); // The lack of usage update should be reflected in the limits @@ -492,7 +492,7 @@ contract SessionKeyGasLimitsTest is Test { // actually usable. skip(1 days + 1 minutes); - vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 2, "AA24 signature error")); + vm.expectRevert(abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 2, "AA23 reverted (or OOG)")); entryPoint.handleOps(userOps, beneficiary); } diff --git a/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol b/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol index 2c16762f7..c0579cb9c 100644 --- a/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol +++ b/test/plugin/session/permissions/SessionKeyNativeTokenSpendLimits.t.sol @@ -187,7 +187,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); // Run a user op that spends 0 wei, should succeed @@ -233,7 +233,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } @@ -258,7 +258,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); } @@ -302,7 +302,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); // Assert that the limit is NOT updated @@ -360,7 +360,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); // Assert that the limit is NOT updated @@ -702,7 +702,7 @@ contract SessionKeyNativeTokenSpendLimitsTest is Test { _runSessionKeyUserOp( calls, sessionKey1Private, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); // Assert that limits are NOT updated diff --git a/test/plugin/session/permissions/SessionKeyPermissions.t.sol b/test/plugin/session/permissions/SessionKeyPermissions.t.sol index a6aebfdab..b0f0d0054 100644 --- a/test/plugin/session/permissions/SessionKeyPermissions.t.sol +++ b/test/plugin/session/permissions/SessionKeyPermissions.t.sol @@ -17,7 +17,7 @@ pragma solidity ^0.8.22; -import {Test} from "forge-std/Test.sol"; +import {Test, console} from "forge-std/Test.sol"; import {EntryPoint} from "@eth-infinitism/account-abstraction/core/EntryPoint.sol"; import {ECDSA} from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol"; @@ -139,7 +139,7 @@ contract SessionKeyPermissionsTest is Test { sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); // Call should fail before removing the allowlist @@ -198,7 +198,7 @@ contract SessionKeyPermissionsTest is Test { sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); assertEq(counter2.number(), 1); @@ -223,7 +223,7 @@ contract SessionKeyPermissionsTest is Test { sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); assertEq(counter1.number(), 1); @@ -284,7 +284,7 @@ contract SessionKeyPermissionsTest is Test { sessionKey1Private, abi.encodeCall(Counter.setNumber, (5)), 0 wei, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); assertEq(counter1.number(), 2); @@ -314,7 +314,7 @@ contract SessionKeyPermissionsTest is Test { sessionKey1Private, abi.encodeCall(Counter.increment, ()), 0 wei, - abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA24 signature error") + abi.encodeWithSelector(IEntryPoint.FailedOp.selector, 0, "AA23 reverted (or OOG)") ); assertEq(counter1.number(), 1); @@ -328,8 +328,12 @@ contract SessionKeyPermissionsTest is Test { } function testFuzz_sessionKeyTimeRange(uint48 startTime, uint48 endTime) public { - bytes[] memory updates = new bytes[](1); + bytes[] memory updates = new bytes[](2); updates[0] = abi.encodeCall(ISessionKeyPermissionsUpdates.updateTimeRange, (startTime, endTime)); + updates[1] = abi.encodeCall( + ISessionKeyPermissionsUpdates.setAccessListType, + (ISessionKeyPlugin.ContractAccessControlType.ALLOW_ALL_ACCESS) + ); vm.prank(owner1); SessionKeyPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates); @@ -352,6 +356,7 @@ contract SessionKeyPermissionsTest is Test { }); bytes32 userOpHash = entryPoint.getUserOpHash(userOp); + console.log(sessionKey1); (uint8 v, bytes32 r, bytes32 s) = vm.sign(sessionKey1Private, userOpHash.toEthSignedMessageHash()); userOp.signature = abi.encodePacked(r, s, v); @@ -504,15 +509,16 @@ contract SessionKeyPermissionsTest is Test { vm.prank(owner1); SessionKeyPlugin(address(account1)).updateKeyPermissions(sessionKey1, updates2); - vm.prank(address(entryPoint)); - validationData = account1.validateUserOp(userOp, userOpHash, 0); + vm.startPrank(address(entryPoint)); if (requiredPaymaster == providedPaymaster || requiredPaymaster == address(0)) { // Assert that validation passes + validationData = account1.validateUserOp(userOp, userOpHash, 0); assertEq(uint160(validationData), 0); } else { // Assert that validation fails - assertEq(uint160(validationData), 1); + vm.expectRevert(ISessionKeyPlugin.PermissionsCheckFailed.selector); + validationData = account1.validateUserOp(userOp, userOpHash, 0); } }