Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Session key module #99

Merged
merged 39 commits into from
Oct 31, 2024
Merged
Changes from 1 commit
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
9437387
feat: new session key module compiled
ly0va Oct 16, 2024
731073d
test: add some tests
ly0va Oct 16, 2024
1fd3731
test: forgot to push
ly0va Oct 16, 2024
9ea2a3a
fix: session key transaction
ly0va Oct 17, 2024
5a0c423
fix: install / uninstall
ly0va Oct 17, 2024
6feded2
fix: a bunch of bugs
ly0va Oct 17, 2024
bc05b52
feat: create mini test framework
ly0va Oct 18, 2024
c1fff2a
feat: big refactor and ERC20 tests
ly0va Oct 19, 2024
e4f922f
feat: more tests
ly0va Oct 20, 2024
d84713b
fix: reorder tests
ly0va Oct 20, 2024
0b84192
feat: add getters
ly0va Oct 21, 2024
b68ec0e
fix: ERC165 support
ly0va Oct 21, 2024
c135f4e
feat: properly separate calls and transfers
ly0va Oct 22, 2024
8acedfe
fix: better fee limits
ly0va Oct 22, 2024
ab9eb0b
feat: if -> require
ly0va Oct 23, 2024
42d54dd
feat: storage layout overhaul to comply with AA restrictions
ly0va Oct 23, 2024
75c7813
feat: store number of constraints
ly0va Oct 24, 2024
9e66f62
feat: make tests work using local node
ly0va Oct 25, 2024
4122801
test: revoke key and getters
ly0va Oct 25, 2024
1621a7a
chore: format and lint
ly0va Oct 28, 2024
8ed637b
Merge remote-tracking branch 'origin/main' into lyova-session-keys-mo…
ly0va Oct 28, 2024
adc1efd
feat: session state getter
ly0va Oct 28, 2024
8961895
feat: try to adapt sdk and gateway to the new interface
ly0va Oct 29, 2024
dfa3428
chore: format
ly0va Oct 29, 2024
3c18e57
fix: update addresses
ly0va Oct 29, 2024
b794dc0
fix: erc165
ly0va Oct 29, 2024
044e6f0
fix: pass hooks data
ly0va Oct 30, 2024
e30fac9
fix: serialize bigints in a message
ly0va Oct 30, 2024
2658c92
chore: split e2e and contract tests
ly0va Oct 30, 2024
ea6c590
fix: update addresses
ly0va Oct 30, 2024
7069b20
chore: fix lints
ly0va Oct 30, 2024
76bda4f
chore: cspell fix
ly0va Oct 30, 2024
6ab5a04
fix: update address
ly0va Oct 30, 2024
86d5df4
fix: erc165 again
ly0va Oct 30, 2024
ebb927a
fix: review
ly0va Oct 31, 2024
903acbd
docs: update docs
ly0va Oct 31, 2024
8a606f6
chore: fmt
ly0va Oct 31, 2024
84ec928
feat: add comments and refactor getters a bit
ly0va Oct 31, 2024
1a5d231
Merge branch 'main' into lyova-session-keys-module
MexicanAce Oct 31, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
feat: add comments and refactor getters a bit
ly0va committed Oct 31, 2024
commit 84ec928725971f2eb6547f1e6aed0217daa01546
118 changes: 57 additions & 61 deletions packages/contracts/src/validators/SessionKeyValidator.sol
Original file line number Diff line number Diff line change
@@ -22,30 +22,45 @@ library SessionLib {

uint256 constant MAX_CONSTRAINTS = 16;

// We do not permit session keys to be reused to open multiple sessions
// (after one expires or is closed, e.g.).
// For each session key, its session status can only be changed
// from NotInitialized to Active, and from Active to Closed.
enum Status {
NotInitialized,
Active,
Closed
}

// This struct is used to track usage information for each session.
// Along with `status`, this is considered the session state.
// While everything else is considered the session spec.
struct UsageTrackers {
UsageTracker fee;
// (target) => transfer value tracker
mapping(address => UsageTracker) transferValue;
// (target, selector) => call value tracker
mapping(address => mapping(bytes4 => UsageTracker)) callValue;
// (target, selector, index) => call parameter tracker
// index is the constraint index in callPolicy, not the parameter index
mapping(address => mapping(bytes4 => mapping(uint256 => UsageTracker))) params;
}

// This struct has weird layout because of the AA storage access restrictions for validation
// This is the main struct that holds information about all sessions and their state.
// This struct has weird layout because of the AA storage access restrictions for validation.
// Innermost mappings are all mapping(address account => ...) because of this.
struct SessionStorage {
// (target, selector) => call policy
mapping(address => mapping(bytes4 => mapping(address => CallPolicy))) callPolicy;
// (target) => empty (no selector) call policy
// (target) => transfer policy. Used for calls with calldata.length < 4.
mapping(address => mapping(address => TransferPolicy)) transferPolicy;
mapping(address => Status) status;
// Timestamp after which session is considered expired
mapping(address => uint256) expiry;
// Tracks gasLimit * maxFeePerGas of each transaction
mapping(address => UsageLimit) feeLimit;
UsageTrackers trackers;
// only used in getters / view functions, not used during validation
// These 2 mappings are only used in getters / view functions, not used during validation.
mapping(address => CallTarget[]) callTargets;
mapping(address => address[]) transferTargets;
}
@@ -56,11 +71,13 @@ library SessionLib {
bool isAllowed;
uint256 maxValuePerUse;
UsageLimit valueLimit;
// We restrain from using a dynamic array here, as it would mean further
// complications for the storage layout due to the AA storage access restrictions.
uint256 totalConstraints;
Constraint[MAX_CONSTRAINTS] constraints;
}

// for transfers, i.e. calls without a selector
// For transfers, i.e. calls without a selector
struct TransferPolicy {
bool isAllowed;
uint256 maxValuePerUse;
@@ -75,14 +92,16 @@ library SessionLib {
}

struct UsageTracker {
// Used for LimitType.Lifetime
mapping(address => uint256) lifetimeUsage;
// Used for LimitType.Allowance
// period => used that period
mapping(uint256 => mapping(address => uint256)) allowanceUsage;
}

struct UsageLimit {
LimitType limitType;
uint256 limit;
uint256 limit; // ignored if limitType == Unlimited
uint256 period; // ignored if limitType != Allowance
}

@@ -132,6 +151,8 @@ library SessionLib {
}

struct LimitState {
// this might also be limited by a constraint or `maxValuePerUse`,
// which is not reflected here
uint256 remaining;
address target;
// ignored for transfer value
@@ -140,8 +161,9 @@ library SessionLib {
uint256 index;
}

// Info about remaining session limits
// Info about remaining session limits and its status
struct SessionState {
Status status;
uint256 fee;
LimitState[] transferValue;
LimitState[] callValue;
@@ -162,10 +184,6 @@ library SessionLib {
}

function checkAndUpdate(Constraint storage constraint, UsageTracker storage tracker, bytes calldata data) internal {
if (constraint.condition == Condition.Unconstrained && constraint.limit.limitType == LimitType.Unlimited) {
return;
}

uint256 index = 4 + constraint.index * 32;
bytes32 param = bytes32(data[index:index + 32]);
Condition condition = constraint.condition;
@@ -246,7 +264,7 @@ library SessionLib {
}
}

function getSpec(SessionStorage storage session, address account) internal view returns (Status, SessionSpec memory) {
function getSpec(SessionStorage storage session, address account) internal view returns (SessionSpec memory) {
CallSpec[] memory callPolicies = new CallSpec[](session.callTargets[account].length);
TransferSpec[] memory transferPolicies = new TransferSpec[](session.transferTargets[account].length);
for (uint256 i = 0; i < session.callTargets[account].length; i++) {
@@ -273,25 +291,25 @@ library SessionLib {
valueLimit: transferPolicy.valueLimit
});
}
return (
session.status[account],
return
SessionSpec({
// Signer addresses are not stored in SessionStorage,
// and are filled in later in the `sessionSpec()` getter.
signer: address(0),
ly0va marked this conversation as resolved.
Show resolved Hide resolved
expiry: session.expiry[account],
feeLimit: session.feeLimit[account],
callPolicies: callPolicies,
transferPolicies: transferPolicies
})
);
});
}

function remainingUsage(
function remainingLimit(
UsageLimit memory limit,
UsageTracker storage tracker,
address account
) internal view returns (uint256) {
if (limit.limitType == LimitType.Unlimited) {
// this might be still limited by `maxValuePerUse`
// this might be still limited by `maxValuePerUse` or a constraint
return type(uint256).max;
}
if (limit.limitType == LimitType.Lifetime) {
@@ -303,11 +321,8 @@ library SessionLib {
}
}

function remainingLimits(
SessionStorage storage session,
address account
) internal view returns (SessionState memory) {
(, SessionSpec memory spec) = getSpec(session, account);
function getState(SessionStorage storage session, address account) internal view returns (SessionState memory) {
SessionSpec memory spec = getSpec(session, account);

LimitState[] memory transferValue = new LimitState[](spec.transferPolicies.length);
LimitState[] memory callValue = new LimitState[](spec.callPolicies.length);
@@ -316,20 +331,12 @@ library SessionLib {

for (uint256 i = 0; i < transferValue.length; i++) {
TransferSpec memory transferSpec = spec.transferPolicies[i];
uint256 remaining;

if (transferSpec.valueLimit.limitType == LimitType.Unlimited) {
remaining = transferSpec.maxValuePerUse;
} else {
remaining = remainingUsage(
transferValue[i] = LimitState({
remaining: remainingLimit(
transferSpec.valueLimit,
session.trackers.transferValue[transferSpec.target],
account
);
}

transferValue[i] = LimitState({
remaining: remaining,
),
target: spec.transferPolicies[i].target,
selector: bytes4(0),
index: 0
@@ -338,20 +345,12 @@ library SessionLib {

for (uint256 i = 0; i < callValue.length; i++) {
CallSpec memory callSpec = spec.callPolicies[i];
uint256 remaining;

if (callSpec.valueLimit.limitType == LimitType.Unlimited) {
remaining = callSpec.maxValuePerUse;
} else {
remaining = remainingUsage(
callValue[i] = LimitState({
remaining: remainingLimit(
callSpec.valueLimit,
session.trackers.callValue[callSpec.target][callSpec.selector],
account
);
}

callValue[i] = LimitState({
remaining: remaining,
),
target: callSpec.target,
selector: callSpec.selector,
index: 0
@@ -360,7 +359,7 @@ library SessionLib {
for (uint256 j = 0; j < callSpec.constraints.length; j++) {
if (callSpec.constraints[j].limit.limitType != LimitType.Unlimited) {
callParams[paramLimitIndex++] = LimitState({
remaining: remainingUsage(
remaining: remainingLimit(
callSpec.constraints[j].limit,
session.trackers.params[callSpec.target][callSpec.selector][j],
account
@@ -380,7 +379,8 @@ library SessionLib {

return
SessionState({
fee: remainingUsage(spec.feeLimit, session.trackers.fee, account),
status: session.status[account],
fee: remainingLimit(spec.feeLimit, session.trackers.fee, account),
transferValue: transferValue,
callValue: callValue,
callParams: callParams
@@ -399,19 +399,13 @@ contract SessionKeyValidator is IHook, IValidationHook, IModuleValidator, IModul
// account => owners
mapping(address => EnumerableSet.AddressSet) sessionOwners;

function session(
address account,
address signer
) public view returns (SessionLib.Status status, SessionLib.SessionSpec memory spec) {
(status, spec) = sessions[signer].getSpec(account);
function sessionSpec(address account, address signer) public view returns (SessionLib.SessionSpec memory spec) {
spec = sessions[signer].getSpec(account);
spec.signer = signer;
}

function remainingSessionLimits(
address account,
address signer
) external view returns (SessionLib.SessionState memory) {
return sessions[signer].remainingLimits(account);
function sessionState(address account, address signer) public view returns (SessionLib.SessionState memory) {
return sessions[signer].getState(account);
}

function activeSigners(address account) external view returns (address[] memory) {
@@ -420,12 +414,14 @@ contract SessionKeyValidator is IHook, IValidationHook, IModuleValidator, IModul

function sessionList(
address account
) external view returns (SessionLib.Status[] memory statuses, SessionLib.SessionSpec[] memory specs) {
specs = new SessionLib.SessionSpec[](sessionOwners[account].length());
statuses = new SessionLib.Status[](specs.length);
for (uint256 i = 0; i < specs.length; i++) {
) external view returns (SessionLib.SessionState[] memory states, SessionLib.SessionSpec[] memory specs) {
uint256 length = sessionOwners[account].length();
states = new SessionLib.SessionState[](length);
specs = new SessionLib.SessionSpec[](length);
for (uint256 i = 0; i < length; i++) {
address signer = sessionOwners[account].at(i);
(statuses[i], specs[i]) = session(account, signer);
specs[i] = sessionSpec(account, signer);
states[i] = sessionState(account, signer);
}
}

24 changes: 12 additions & 12 deletions packages/contracts/test/SessionKeyTest.ts
Original file line number Diff line number Diff line change
@@ -86,8 +86,8 @@ class SessionTester {
}, provider);

const [oldList] = await sessionKeyModuleContract.sessionList(this.proxyAccountAddress);
const [oldStatus] = await sessionKeyModuleContract.session(this.proxyAccountAddress, this.sessionOwner.address);
expect(oldStatus).to.equal(0, "session should not be initialized");
const oldState = await sessionKeyModuleContract.sessionState(this.proxyAccountAddress, this.sessionOwner.address);
expect(oldState.status).to.equal(0, "session should not be initialized");

this.session = this.getSession(newSession);

@@ -102,11 +102,11 @@ class SessionTester {
const tx = await provider.broadcastTransaction(signedTransaction);
await tx.wait();

const [statusList] = await sessionKeyModuleContract.sessionList(this.proxyAccountAddress);
expect(statusList).to.have.lengthOf(oldList.length + 1, "session should be created");
const [status, session] = await sessionKeyModuleContract.session(this.proxyAccountAddress, this.sessionOwner.address);
expect(status).to.equal(1, "session should be active");
this.assertSession(session);
const [newList] = await sessionKeyModuleContract.sessionList(this.proxyAccountAddress);
expect(newList).to.have.lengthOf(oldList.length + 1, "session should be created");
const newState = await sessionKeyModuleContract.sessionState(this.proxyAccountAddress, this.sessionOwner.address);
expect(newState.status).to.equal(1, "session should be active");
this.assertSession(await sessionKeyModuleContract.sessionSpec(this.proxyAccountAddress, this.sessionOwner.address));
}

assertSession(session: SessionLib.SessionSpecStruct) {
@@ -126,8 +126,8 @@ class SessionTester {

async revokeKey() {
const sessionKeyModuleContract = await fixtures.getSessionKeyContract();
let [status] = await sessionKeyModuleContract.session(this.proxyAccountAddress, this.sessionOwner.address);
expect(status).to.equal(1, "session should be active");
let state = await sessionKeyModuleContract.sessionState(this.proxyAccountAddress, this.sessionOwner.address);
expect(state.status).to.equal(1, "session should be active");

const smartAccount = new SmartAccount({
address: this.proxyAccountAddress,
@@ -144,8 +144,8 @@ class SessionTester {
const signedTransaction = await smartAccount.signTransaction(aaTx);
const tx = await provider.broadcastTransaction(signedTransaction);
await tx.wait();
[status] = await sessionKeyModuleContract.session(this.proxyAccountAddress, this.sessionOwner.address);
expect(status).to.equal(2, "session should be revoked");
state = await sessionKeyModuleContract.sessionState(this.proxyAccountAddress, this.sessionOwner.address);
expect(state.status).to.equal(2, "session should be revoked");
}

async sendTxSuccess(txRequest: ethers.TransactionRequest = {}) {
@@ -384,7 +384,7 @@ describe("SessionKeyModule tests", function () {

it("should reject a session key transaction that goes over total limit", async () => {
const sessionKeyModuleContract = await fixtures.getSessionKeyContract();
const remainingLimits = await sessionKeyModuleContract.remainingSessionLimits(proxyAccountAddress, tester.sessionOwner.address);
const remainingLimits = await sessionKeyModuleContract.sessionState(proxyAccountAddress, tester.sessionOwner.address);
expect(remainingLimits.callParams[0].remaining).to.equal(500n, "should have 500 tokens remaining in allowance");

await tester.sendTxFail({
2 changes: 1 addition & 1 deletion packages/gateway/stores/client.ts
Original file line number Diff line number Diff line change
@@ -29,7 +29,7 @@ export const contractsByChain: Record<SupportedChainId, ChainContracts> = {
accountImplementation: "0x",
},
[zksyncInMemoryNode.id]: {
session: "0x9B7FF4010a3d36a17d423Ef4240c86A8ECeaf105",
session: "0xC9c7F31CCf72daDFd18924e8111Fe90a35400734",
passkey: "0x0cc51Dc85E147B66271E34BCd92AA6Cf9458D2a2",
accountFactory: "0x23b13d016E973C9915c6252271fF06cCA2098885",
accountImplementation: "0x0fA8Ed8e24db620f5d80c2683D16d405a5357450",