zksync-sso 0.2.0 → 0.3.0

package/src/utils/session.ts

@@ -1,4 +1,6 @@
-import { type Address, getAddress, type Hash, type Hex } from "viem";
+import { type Address, bytesToHex, getAddress, type Hash, type Hex, hexToBigInt, hexToBytes } from "viem";
+
+import { calculateMaxFee, findSmallestBigInt } from "./helpers.js";
 
 export enum LimitType {
   Unlimited = 0,
@@ -167,3 +169,356 @@ export const getPeriodIdsForTransaction = (args: {
   ];
   return periodIds;
 };
+
+export enum SessionErrorType {
+  SessionInactive = "session_inactive",
+  SessionExpired = "session_expired",
+  FeeLimitExceeded = "fee_limit_exceeded",
+  NoCallPolicy = "no_call_policy",
+  NoTransferPolicy = "no_transfer_policy",
+  MaxValuePerUseExceeded = "max_value_per_use_exceeded",
+  ValueLimitExceeded = "value_limit_exceeded",
+  ConstraintIndexOutOfBounds = "constraint_index_out_of_bounds",
+  NoLimitStateFound = "no_limit_state_found",
+  ParameterLimitExceeded = "parameter_limit_exceeded",
+  ConstraintEqualViolated = "constraint_equal_violated",
+  ConstraintGreaterViolated = "constraint_greater_violated",
+  ConstraintLessViolated = "constraint_less_violated",
+  ConstraintGreaterEqualViolated = "constraint_greater_equal_violated",
+  ConstraintLessEqualViolated = "constraint_less_equal_violated",
+  ConstraintNotEqualViolated = "constraint_not_equal_violated",
+}
+
+export enum SessionEventType {
+  Expired = "session_expired",
+  Revoked = "session_revoked",
+  Inactive = "session_inactive",
+}
+
+export type SessionStateEvent = {
+  type: SessionEventType;
+  message: string;
+};
+
+export type SessionStateEventCallback = (event: SessionStateEvent) => void;
+
+export type ValidationResult = {
+  valid: boolean;
+  error: null | {
+    type: SessionErrorType;
+    message: string;
+  };
+};
+
+export type TransactionValidationArgs = {
+  sessionState: SessionState;
+  sessionConfig: SessionConfig;
+  transaction: {
+    to: Address;
+    value?: bigint;
+    data?: Hex;
+    gas?: bigint;
+    gasPrice?: bigint;
+    maxFeePerGas?: bigint;
+    maxPriorityFeePerGas?: bigint;
+    paymaster?: Address;
+  };
+  currentTimestamp?: bigint;
+};
+
+export function validateSessionTransaction(args: TransactionValidationArgs): ValidationResult {
+  const { sessionState, sessionConfig, transaction, currentTimestamp } = args;
+  const timestamp = currentTimestamp || BigInt(Math.floor(Date.now() / 1000));
+
+  // 1. Check if session is active
+  if (sessionState.status !== SessionStatus.Active) {
+    return {
+      valid: false,
+      error: {
+        type: SessionErrorType.SessionInactive,
+        message: "Session is not active",
+      },
+    };
+  }
+
+  // 2. Check if session hasn't expired
+  if (sessionConfig.expiresAt <= timestamp) {
+    return {
+      valid: false,
+      error: {
+        type: SessionErrorType.SessionExpired,
+        message: "Session has expired",
+      },
+    };
+  }
+
+  // 3. Extract transaction data
+  const to = getAddress(transaction.to.toLowerCase());
+  const value = transaction.value || 0n;
+  const data = transaction.data || "0x";
+  const gas = transaction.gas || 0n;
+  const selector = data.length >= 10 ? data.slice(0, 10) as Hash : undefined;
+
+  if (!transaction.paymaster) {
+    const maxFee = calculateMaxFee({
+      gas,
+      gasPrice: transaction.gasPrice,
+      maxFeePerGas: transaction.maxFeePerGas,
+      maxPriorityFeePerGas: transaction.maxPriorityFeePerGas,
+    });
+
+    if (maxFee > sessionState.feesRemaining) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.FeeLimitExceeded,
+          message: `Transaction max fee ${maxFee} exceeds remaining fee limit ${sessionState.feesRemaining}`,
+        },
+      };
+    }
+  }
+
+  const isContractCall = !!selector;
+  if (isContractCall) {
+    const policies = sessionConfig.callPolicies.filter(
+      (policy) => policy.target === to && policy.selector === selector,
+    );
+
+    if (!policies.length) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.NoCallPolicy,
+          message: `No call policy found for target ${to} with selector ${selector}`,
+        },
+      };
+    }
+
+    // Verify max value per use
+    const lowestMaxValuePerUse = findSmallestBigInt(policies.map((policy) => policy.maxValuePerUse));
+    if (value > lowestMaxValuePerUse) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.MaxValuePerUseExceeded,
+          message: `Transaction value ${value} exceeds max value per use ${lowestMaxValuePerUse}`,
+        },
+      };
+    }
+
+    // Verify remaining value limit
+    const remainingValue = findLowestRemainingValue(sessionState.callValue, to, selector);
+    if (value > remainingValue) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.ValueLimitExceeded,
+          message: `Transaction value ${value} exceeds remaining value limit ${remainingValue || 0n}`,
+        },
+      };
+    }
+
+    for (const policy of policies) {
+      const constraintResult = validateConstraints(data, policy.constraints, sessionState.callParams);
+      if (!constraintResult.valid) return constraintResult;
+    }
+  } else {
+    // This is a simple transfer
+    const policies = sessionConfig.transferPolicies.filter((policy) => policy.target === to);
+    if (!policies.length) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.NoTransferPolicy,
+          message: `No transfer policy found for target ${to}`,
+        },
+      };
+    }
+
+    // Verify max value per use
+    const lowestMaxValuePerUse = findSmallestBigInt(policies.map((policy) => policy.maxValuePerUse));
+    if (value > lowestMaxValuePerUse) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.MaxValuePerUseExceeded,
+          message: `Transaction value ${value} exceeds max value per use ${lowestMaxValuePerUse}`,
+        },
+      };
+    }
+
+    // Verify remaining value limit
+    const remainingValue = findLowestRemainingValue(sessionState.transferValue, to);
+    if (value > remainingValue) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.ValueLimitExceeded,
+          message: `Transaction value ${value} exceeds remaining value limit ${remainingValue || 0n}`,
+        },
+      };
+    }
+  }
+
+  return {
+    valid: true,
+    error: null,
+  };
+}
+
+function findLowestRemainingValue(
+  valueArray: SessionState["transferValue"] | SessionState["callValue"],
+  target: Address,
+  selector?: Hash,
+): bigint {
+  if (selector) {
+    const filtered = valueArray
+      .filter((item) => item.target === target && item.selector === selector)
+      .map((item) => item.remaining);
+    if (!filtered.length) return 0n;
+    return findSmallestBigInt(filtered);
+  } else {
+    const filtered = valueArray
+      .filter((item) => item.target === target)
+      .map((item) => item.remaining);
+    if (!filtered.length) return 0n;
+    return findSmallestBigInt(filtered);
+  }
+}
+
+function validateConstraints(
+  data: Hex,
+  constraints: Constraint[],
+  callParams: SessionState["callParams"],
+): ValidationResult {
+  const dataBytes = hexToBytes(data);
+
+  for (const constraint of constraints) {
+    // Find proper index and extract data
+    const index = Number(constraint.index);
+    if (index + 32 > dataBytes.length) {
+      return {
+        valid: false,
+        error: {
+          type: SessionErrorType.ConstraintIndexOutOfBounds,
+          message: `Constraint index ${index} out of bounds for data length ${dataBytes.length}`,
+        },
+      };
+    }
+
+    const parameterBytes = dataBytes.slice(index, index + 32);
+    const parameterValue = bytesToHex(parameterBytes);
+    const refValue = constraint.refValue;
+
+    // Find remaining limit value if applicable
+    let remaining: bigint | undefined;
+    if (constraint.limit.limitType !== LimitType.Unlimited) {
+      const paramItem = callParams.find((item) => Number(item.index) === index);
+      remaining = paramItem?.remaining;
+
+      // If this is a limited constraint, but we don't have state for it, something is wrong
+      if (remaining === undefined) {
+        return {
+          valid: false,
+          error: {
+            type: SessionErrorType.NoLimitStateFound,
+            message: `No remaining limit state found for constraint at index ${index}`,
+          },
+        };
+      }
+
+      // Check if parameter value exceeds remaining limit
+      const parameterValueBigInt = hexToBigInt(parameterValue);
+      if (parameterValueBigInt > remaining) {
+        return {
+          valid: false,
+          error: {
+            type: SessionErrorType.ParameterLimitExceeded,
+            message: `Parameter value ${parameterValueBigInt} at index ${index} exceeds remaining limit ${remaining}`,
+          },
+        };
+      }
+    }
+
+    // Check condition
+    switch (constraint.condition) {
+      case ConstraintCondition.Equal:
+        if (parameterValue !== refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintEqualViolated,
+              message: `Parameter value ${parameterValue} must equal ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.Greater:
+        if (parameterValue <= refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintGreaterViolated,
+              message: `Parameter value ${parameterValue} must be greater than ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.Less:
+        if (parameterValue >= refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintLessViolated,
+              message: `Parameter value ${parameterValue} must be less than ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.GreaterEqual:
+        if (parameterValue < refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintGreaterEqualViolated,
+              message: `Parameter value ${parameterValue} must be greater than or equal to ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.LessEqual:
+        if (parameterValue > refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintLessEqualViolated,
+              message: `Parameter value ${parameterValue} must be less than or equal to ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.NotEqual:
+        if (parameterValue === refValue) {
+          return {
+            valid: false,
+            error: {
+              type: SessionErrorType.ConstraintNotEqualViolated,
+              message: `Parameter value ${parameterValue} must not equal ${refValue}`,
+            },
+          };
+        }
+        break;
+      case ConstraintCondition.Unconstrained:
+        // Unconstrained means no checks, so we skip
+        break;
+      default:
+        console.warn(`Unhandled constraint condition: ${constraint.condition}`);
+        break;
+    }
+  }
+
+  return {
+    valid: true,
+    error: null,
+  };
+}