4337 with session keys

src/4337/SessionLib.sol

@@ -0,0 +1,304 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.26;
+
+import {LibBytes} from "solady/utils/LibBytes.sol";
+
+library SessionLib {
+
+    using SessionLib for SessionLib.Constraint;
+    using SessionLib for SessionLib.UsageLimit;
+    using LibBytes for bytes;
+
+    enum LimitType {
+        Unlimited,
+        Lifetime,
+        Allowance
+    }
+
+    enum Condition {
+        Unconstrained,
+        Equal,
+        Greater,
+        Less,
+        GreaterOrEqual,
+        LessOrEqual,
+        NotEqual
+    }
+
+    struct UsageLimit {
+        LimitType limitType;
+        uint256 limit; // ignored if limitType == Unlimited
+        uint256 period; // ignored if limitType != Allowance
+    }
+
+    struct Constraint {
+        Condition condition;
+        uint64 index;
+        bytes32 refValue;
+        UsageLimit limit;
+    }
+
+    struct CallSpec {
+        address target;
+        bytes4 selector;
+        uint256 maxValuePerUse;
+        UsageLimit valueLimit;
+        Constraint[] constraints;
+    }
+
+    struct TransferSpec {
+        address target;
+        uint256 maxValuePerUse;
+        UsageLimit valueLimit;
+    }
+
+    struct SessionSpec {
+        address signer;
+        uint256 expiresAt;
+        CallSpec[] callPolicies;
+        TransferSpec[] transferPolicies;
+    }
+
+    struct StoredSessionSpec {
+        uint256 expiresAt;
+        CallSpec[] callPolicies;
+        TransferSpec[] transferPolicies;
+    }
+
+    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
+        bytes4 selector;
+        // ignored for transfer and call value
+        uint256 index;
+    }
+
+    // Info about remaining session limits and its status
+    struct SessionState {
+        LimitState[] transferValue;
+        LimitState[] callValue;
+        LimitState[] callParams;
+    }
+
+    struct UsageTracker {
+        // Used for LimitType.Lifetime
+        uint256 lifetimeUsage;
+        // Used for LimitType.Allowance
+        // period => used that period
+        mapping(uint64 => uint256) allowanceUsage;
+    }
+
+    struct SessionStorage {
+        // (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;
+    }
+
+    function checkAndUpdate(UsageLimit memory limit, UsageTracker storage tracker, uint256 value)
+        internal
+        returns (bool)
+    {
+        if (limit.limitType == LimitType.Lifetime) {
+            if (tracker.lifetimeUsage + value > limit.limit) {
+                // revert SessionLifetimeUsageExceeded(tracker.lifetimeUsage[msg.sender], limit.limit);
+                return false;
+            }
+            tracker.lifetimeUsage += value;
+        } else if (limit.limitType == LimitType.Allowance) {
+            uint64 period = uint64(block.timestamp / limit.period);
+
+            if (tracker.allowanceUsage[period] + value > limit.limit) {
+                // revert SessionAllowanceExceeded(tracker.allowanceUsage[period][msg.sender], limit.limit, period);
+                return false;
+            }
+
+            tracker.allowanceUsage[period] += value;
+        }
+
+        return true;
+    }
+
+    function checkAndUpdate(Constraint memory constraint, UsageTracker storage tracker, bytes memory data)
+        internal
+        returns (bool)
+    {
+        uint256 expectedLength = 4 + constraint.index * 32 + 32;
+
+        if (data.length < expectedLength) {
+            //   revert SessionInvalidDataLength(data.length, expectedLength);
+            return false;
+        }
+
+        bytes32 param = data.load(4 + constraint.index * 32);
+        Condition condition = constraint.condition;
+        bytes32 refValue = constraint.refValue;
+
+        if (
+            (condition == Condition.Equal && param != refValue) || (condition == Condition.Greater && param <= refValue)
+                || (condition == Condition.Less && param >= refValue)
+                || (condition == Condition.GreaterOrEqual && param < refValue)
+                || (condition == Condition.LessOrEqual && param > refValue)
+                || (condition == Condition.NotEqual && param == refValue)
+        ) {
+            //   revert SessionConditionFailed(param, refValue, uint8(condition));
+            return false;
+        }
+
+        bool check = constraint.limit.checkAndUpdate(tracker, uint256(param));
+
+        return check;
+    }
+
+    function checkCallPolicy(
+        SessionStorage storage state,
+        bytes memory data,
+        address target,
+        bytes4 selector,
+        CallSpec memory callPolicy
+    ) private returns (bool found) {
+        if (target != callPolicy.target || selector != callPolicy.selector) {
+            return false;
+        }
+
+        for (uint256 i = 0; i < callPolicy.constraints.length; i++) {
+            bool check = callPolicy.constraints[i].checkAndUpdate(state.params[target][selector][i], data);
+            if (!check) {
+                // return (false, callPolicy);
+                return false;
+            }
+        }
+        return true;
+    }
+
+    function validate(
+        SessionStorage storage state,
+        address target,
+        bytes4 selector,
+        uint256 value,
+        bytes memory callData,
+        uint256 expiresAt,
+        CallSpec memory callPolicy,
+        TransferSpec memory transferPolicy
+    ) internal returns (bool) {
+        if (expiresAt < block.timestamp) {
+            return false;
+        }
+
+        if (callData.length >= 4) {
+            bool callPolicyValid = checkCallPolicy(state, callData, target, selector, callPolicy);
+
+            if (!callPolicyValid) {
+                return false;
+            }
+
+            if (value > callPolicy.maxValuePerUse) {
+                // revert SessionMaxValueExceeded(value, callPolicy.maxValuePerUse);
+                return false;
+            }
+
+            return callPolicy.valueLimit.checkAndUpdate(state.callValue[target][selector], value);
+        } else {
+            if (target != transferPolicy.target) {
+                return false;
+            }
+
+            if (value > transferPolicy.maxValuePerUse) {
+                // revert SessionMaxValueExceeded(value, transferPolicy.maxValuePerUse);
+                return false;
+            }
+            return transferPolicy.valueLimit.checkAndUpdate(state.transferValue[target], value);
+        }
+    }
+
+    function remainingLimit(UsageLimit memory limit, UsageTracker storage tracker) private view returns (uint256) {
+        if (limit.limitType == LimitType.Unlimited) {
+            // this might be still limited by `maxValuePerUse` or a constraint
+            return type(uint256).max;
+        }
+        if (limit.limitType == LimitType.Lifetime) {
+            return limit.limit - tracker.lifetimeUsage;
+        }
+        if (limit.limitType == LimitType.Allowance) {
+            // this is not used during validation, so it's fine to use block.timestamp
+            uint64 period = uint64(block.timestamp / limit.period);
+            return limit.limit - tracker.allowanceUsage[period];
+        }
+    }
+
+    function getState(
+        SessionStorage storage session,
+        CallSpec[] memory callPolicies,
+        TransferSpec[] memory transferPolicies
+    ) internal view returns (SessionState memory) {
+        LimitState[] memory transferValue;
+        LimitState[] memory callValue;
+        LimitState[] memory callParams;
+
+        {
+            uint256 totalConstraints = 0;
+            for (uint256 i = 0; i < callPolicies.length; i++) {
+                totalConstraints += callPolicies[i].constraints.length;
+            }
+
+            transferValue = new LimitState[](transferPolicies.length);
+            callValue = new LimitState[](callPolicies.length);
+            callParams = new LimitState[](totalConstraints); // there will be empty ones at the end
+        }
+        uint256 paramLimitIndex = 0;
+
+        {
+            for (uint256 i = 0; i < transferValue.length; i++) {
+                TransferSpec memory transferSpec = transferPolicies[i];
+                transferValue[i] = LimitState({
+                    remaining: remainingLimit(transferSpec.valueLimit, session.transferValue[transferSpec.target]),
+                    target: transferSpec.target,
+                    selector: bytes4(0),
+                    index: 0
+                });
+            }
+        }
+
+        for (uint256 i = 0; i < callValue.length; i++) {
+            CallSpec memory callSpec = callPolicies[i];
+
+            {
+                callValue[i] = LimitState({
+                    remaining: remainingLimit(callSpec.valueLimit, session.callValue[callSpec.target][callSpec.selector]),
+                    target: callSpec.target,
+                    selector: callSpec.selector,
+                    index: 0
+                });
+            }
+
+            {
+                for (uint256 j = 0; j < callSpec.constraints.length; j++) {
+                    if (callSpec.constraints[j].limit.limitType != LimitType.Unlimited) {
+                        callParams[paramLimitIndex++] = LimitState({
+                            remaining: remainingLimit(
+                                callSpec.constraints[j].limit, session.params[callSpec.target][callSpec.selector][j]
+                            ),
+                            target: callSpec.target,
+                            selector: callSpec.selector,
+                            index: callSpec.constraints[j].index
+                        });
+                    }
+                }
+            }
+        }
+
+        // shrink array to actual size
+        assembly {
+            mstore(callParams, paramLimitIndex)
+        }
+
+        return SessionState({transferValue: transferValue, callValue: callValue, callParams: callParams});
+    }
+
+}