Skip to content

Commit

Permalink
fix: improve session key gas estimation
Browse files Browse the repository at this point in the history
  • Loading branch information
howydev committed Feb 9, 2024
1 parent 078f7c3 commit 7d9db25
Show file tree
Hide file tree
Showing 6 changed files with 40 additions and 28 deletions.
1 change: 1 addition & 0 deletions src/plugins/session/ISessionKeyPlugin.sol
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ interface ISessionKeyPlugin {
error LengthMismatch();
error NativeTokenSpendLimitExceeded(address account, address sessionKey);
error SessionKeyNotFound(address sessionKey);
error PermissionsCheckFailed();

// ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
// ┃ Execution functions ┃
Expand Down
15 changes: 10 additions & 5 deletions src/plugins/session/SessionKeyPlugin.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)")
);
}

Expand Down
12 changes: 6 additions & 6 deletions test/plugin/session/permissions/SessionKeyGasLimits.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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)")
);
}

Expand Down Expand Up @@ -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)")
);
}

Expand All @@ -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);
}

Expand All @@ -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)")
);
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)")
);
}

Expand All @@ -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)")
);
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
26 changes: 16 additions & 10 deletions test/plugin/session/permissions/SessionKeyPermissions.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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);

Expand Down Expand Up @@ -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);
}
}

Expand Down

0 comments on commit 7d9db25

Please sign in to comment.