@account-kit/smart-contracts 4.0.0-alpha.0

package/src/msca/plugins/session-key/extension.ts

@@ -0,0 +1,301 @@
+import type {
+  GetEntryPointFromAccount,
+  UserOperationOverridesParameter,
+} from "@aa-sdk/core";
+import {
+  AccountNotFoundError,
+  isSmartAccountClient,
+  type GetAccountParameter,
+  type IsUndefined,
+  type SendUserOperationResult,
+  type SmartContractAccount,
+  type UserOperationOverrides,
+} from "@aa-sdk/core";
+import type { Address, Chain, Client, Hex, Transport } from "viem";
+import type { GetPluginAddressParameter } from "../types.js";
+import {
+  SessionKeyPlugin,
+  sessionKeyPluginActions as sessionKeyPluginActions_,
+  type SessionKeyPluginActions as SessionKeyPluginActions_,
+} from "./plugin.js";
+import { buildSessionKeysToRemoveStruct } from "./utils.js";
+
+export type SessionKeyPluginActions<
+  TAccount extends SmartContractAccount | undefined =
+    | SmartContractAccount
+    | undefined,
+  TEntryPointVersion extends GetEntryPointFromAccount<TAccount> = GetEntryPointFromAccount<TAccount>
+> = Omit<
+  SessionKeyPluginActions_<TAccount>,
+  | "removeSessionKey"
+  | "addSessionKey"
+  | "rotateSessionKey"
+  | "updateKeyPermissions"
+> & {
+  isAccountSessionKey: (
+    args: { key: Address } & GetPluginAddressParameter &
+      GetAccountParameter<TAccount>
+  ) => Promise<boolean>;
+
+  getAccountSessionKeys: (
+    args: GetPluginAddressParameter & GetAccountParameter<TAccount>
+  ) => Promise<ReadonlyArray<Address>>;
+
+  removeSessionKey: (
+    args: { key: Address } & GetPluginAddressParameter &
+      GetAccountParameter<TAccount> &
+      UserOperationOverridesParameter<TEntryPointVersion>
+  ) => Promise<SendUserOperationResult<TEntryPointVersion>>;
+
+  addSessionKey: (
+    args: {
+      key: Address;
+      permissions: Hex[];
+      tag: Hex;
+    } & GetPluginAddressParameter &
+      GetAccountParameter<TAccount> &
+      UserOperationOverridesParameter<TEntryPointVersion>
+  ) => Promise<SendUserOperationResult<TEntryPointVersion>>;
+
+  rotateSessionKey: (
+    args: {
+      oldKey: Address;
+      newKey: Address;
+    } & GetPluginAddressParameter &
+      GetAccountParameter<TAccount> &
+      UserOperationOverridesParameter<TEntryPointVersion>
+  ) => Promise<SendUserOperationResult<TEntryPointVersion>>;
+
+  updateSessionKeyPermissions: (
+    args: {
+      key: Address;
+      permissions: Hex[];
+    } & GetPluginAddressParameter &
+      GetAccountParameter<TAccount> &
+      UserOperationOverridesParameter<TEntryPointVersion>
+  ) => Promise<SendUserOperationResult<TEntryPointVersion>>;
+} & (IsUndefined<TAccount> extends false
+    ? {
+        getAccountSessionKeys: (
+          args?: GetPluginAddressParameter & GetAccountParameter<TAccount>
+        ) => Promise<ReadonlyArray<Address>>;
+      }
+    : {});
+
+export const sessionKeyPluginActions: <
+  TTransport extends Transport = Transport,
+  TChain extends Chain | undefined = Chain | undefined,
+  TAccount extends SmartContractAccount | undefined =
+    | SmartContractAccount
+    | undefined
+>(
+  client: Client<TTransport, TChain, TAccount>
+) => SessionKeyPluginActions<TAccount> = <
+  TTransport extends Transport = Transport,
+  TChain extends Chain | undefined = Chain | undefined,
+  TAccount extends SmartContractAccount | undefined =
+    | SmartContractAccount
+    | undefined,
+  TEntryPointVersion extends GetEntryPointFromAccount<TAccount> = GetEntryPointFromAccount<TAccount>
+>(
+  client: Client<TTransport, TChain, TAccount>
+) => {
+  const {
+    removeSessionKey,
+    addSessionKey,
+    rotateSessionKey,
+    updateKeyPermissions,
+    executeWithSessionKey,
+    ...og
+  } = sessionKeyPluginActions_<TTransport, TChain, TAccount>(client);
+
+  // The session key plugin returns an opaque error when the permissions check
+  // fails. We want the SDK to return a helpful error instead. After validation
+  // fails, rerun the user op simulation but using state overrides to replace
+  // the code of the session key plugin with a version that produces descriptive
+  // errors.
+  const fixedExecuteWithSessionKey: typeof executeWithSessionKey = async (
+    ...originalArgs
+  ) => {
+    let initialError: unknown;
+    try {
+      return await executeWithSessionKey(...originalArgs);
+    } catch (error) {
+      initialError = error;
+    }
+    const details = getRpcErrorMessageFromViemError(initialError);
+    if (!details?.includes("AA23 reverted (or OOG)")) {
+      throw initialError;
+    }
+    if (!isSmartAccountClient(client) || !client.chain) {
+      throw initialError;
+    }
+    const {
+      args,
+      overrides,
+      context,
+      account = client.account,
+    } = originalArgs[0];
+    if (!account) {
+      throw initialError;
+    }
+    const data = og.encodeExecuteWithSessionKey({ args });
+    const sessionKeyPluginAddress =
+      SessionKeyPlugin.meta.addresses[client.chain.id];
+    // The bytecode is 10kb of hex. Only import it if needed.
+    const { DEBUG_SESSION_KEY_BYTECODE } = await import(
+      "./debug-session-key-bytecode.js"
+    );
+    const updatedOverrides: UserOperationOverrides<TEntryPointVersion> = {
+      ...(overrides as UserOperationOverrides<TEntryPointVersion>),
+      stateOverride: [
+        ...(overrides?.stateOverride ?? []),
+        {
+          address: sessionKeyPluginAddress,
+          code: DEBUG_SESSION_KEY_BYTECODE,
+        },
+      ],
+    };
+    try {
+      await client.buildUserOperation({
+        uo: data,
+        overrides: updatedOverrides,
+        context,
+        account,
+      });
+      // We expect this to fail because we just saw the same user op fail a
+      // moment ago. It could succeed if the state of the chain has changed but
+      // there's not much we can do with a successful result so return the
+      // original error.
+      throw initialError;
+    } catch (improvedError) {
+      const details = getRpcErrorMessageFromViemError(improvedError) ?? "";
+      const reason = details.match(/AA23 reverted: (.+)/)?.[1];
+      if (!reason) {
+        throw initialError;
+      }
+      throw new SessionKeyPermissionError(reason);
+    }
+  };
+
+  return {
+    ...og,
+
+    executeWithSessionKey: fixedExecuteWithSessionKey,
+
+    isAccountSessionKey: async ({
+      key,
+      pluginAddress,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      const contract = SessionKeyPlugin.getContract(client, pluginAddress);
+
+      return await contract.read.isSessionKeyOf([account.address, key]);
+    },
+
+    getAccountSessionKeys: async (
+      args: GetPluginAddressParameter & GetAccountParameter<TAccount>
+    ) => {
+      const account = args?.account ?? client.account;
+      if (!account) throw new AccountNotFoundError();
+
+      const contract = SessionKeyPlugin.getContract(
+        client,
+        args?.pluginAddress
+      );
+
+      return await contract.read.sessionKeysOf([account.address]);
+    },
+
+    removeSessionKey: async ({
+      key,
+      overrides,
+      account = client.account,
+      pluginAddress,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      const sessionKeysToRemove = await buildSessionKeysToRemoveStruct(client, {
+        keys: [key],
+        account,
+        pluginAddress,
+      });
+
+      return removeSessionKey({
+        args: [key, sessionKeysToRemove[0].predecessor],
+        overrides: overrides as UserOperationOverrides<TEntryPointVersion>,
+        account,
+      });
+    },
+
+    addSessionKey: async ({
+      key,
+      tag,
+      permissions,
+      overrides,
+      pluginAddress,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      return addSessionKey({
+        args: [key, tag, permissions],
+        overrides: overrides as UserOperationOverrides<TEntryPointVersion>,
+        account,
+        pluginAddress,
+      });
+    },
+
+    rotateSessionKey: async ({
+      newKey,
+      oldKey,
+      overrides,
+      pluginAddress,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      const contract = SessionKeyPlugin.getContract(client, pluginAddress);
+
+      const predecessor = await contract.read.findPredecessor([
+        account.address,
+        oldKey,
+      ]);
+
+      return rotateSessionKey({
+        args: [oldKey, predecessor, newKey],
+        overrides: overrides as UserOperationOverrides<TEntryPointVersion>,
+        account,
+        pluginAddress,
+      });
+    },
+
+    updateSessionKeyPermissions: async ({
+      key,
+      permissions,
+      overrides,
+      pluginAddress,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      return updateKeyPermissions({
+        args: [key, permissions],
+        overrides: overrides as UserOperationOverrides<
+          GetEntryPointFromAccount<TAccount>
+        >,
+        account,
+        pluginAddress,
+      });
+    },
+  };
+};
+
+function getRpcErrorMessageFromViemError(error: unknown): string | undefined {
+  const details = (error as any)?.details;
+  return typeof details === "string" ? details : undefined;
+}
+
+class SessionKeyPermissionError extends Error {}

package/src/msca/plugins/session-key/index.ts

@@ -0,0 +1,3 @@
+export type * from "./extension.js";
+export { sessionKeyPluginActions } from "./extension.js";
+export { SessionKeyPlugin } from "./plugin.js";

package/src/msca/plugins/session-key/permissions.ts

@@ -0,0 +1,189 @@
+import { encodeFunctionData, type Address, type Hex } from "viem";
+import { SessionKeyPermissionsUpdatesAbi } from "./SessionKeyPermissionsUpdatesAbi.js";
+
+export enum SessionKeyAccessListType {
+  ALLOWLIST = 0,
+  DENYLIST = 1,
+  ALLOW_ALL_ACCESS = 2,
+}
+
+export type ContractAccessEntry = {
+  // The contract address to add or remove.
+  contractAddress: Address;
+  // Whether the contract address should be on the list.
+  isOnList: boolean;
+  // Whether to check selectors for the contract address.
+  checkSelectors: boolean;
+};
+
+export type ContractMethodEntry = {
+  // The contract address to add or remove.
+  contractAddress: Address;
+  // The function selector to add or remove.
+  methodSelector: Hex;
+  // Whether the function selector should be on the list.
+  isOnList: boolean;
+};
+
+export type TimeRange = {
+  validFrom: number;
+  validUntil: number;
+};
+
+export type NativeTokenLimit = {
+  spendLimit: bigint;
+  // The time interval over which the spend limit is enforced. If unset, there is no time
+  /// interval by which the limit is refreshed.
+  refreshInterval?: number;
+};
+
+export type Erc20TokenLimit = {
+  tokenAddress: Address;
+  spendLimit: bigint;
+  // The time interval over which the spend limit is enforced. If unset, there is no time
+  /// interval by which the limit is refreshed.
+  refreshInterval?: number;
+};
+
+// uint256 spendLimit, uint48 refreshInterval
+export type GasSpendLimit = {
+  // The amount, in wei, of native tokens that a session key can spend on gas.
+  // Note that this is not the same as the gas limit for a user operation, which is measured in units of gas.
+  // This tracks gas units * gas price.
+  spendLimit: bigint;
+  // The time interval over which the spend limit is enforced. If unset, there is no time
+  /// interval by which the limit is refreshed.
+  refreshInterval?: number;
+};
+
+export class SessionKeyPermissionsBuilder {
+  private _contractAccessControlType: SessionKeyAccessListType =
+    SessionKeyAccessListType.ALLOWLIST;
+  private _contractAddressAccessEntrys: ContractAccessEntry[] = [];
+  private _contractMethodAccessEntrys: ContractMethodEntry[] = [];
+  private _timeRange?: TimeRange;
+  private _nativeTokenSpendLimit?: NativeTokenLimit;
+  private _erc20TokenSpendLimits: Erc20TokenLimit[] = [];
+  private _gasSpendLimit?: GasSpendLimit;
+  private _requiredPaymaster?: Address;
+
+  public setContractAccessControlType(aclType: SessionKeyAccessListType) {
+    this._contractAccessControlType = aclType;
+    return this;
+  }
+
+  public addContractAddressAccessEntry(entry: ContractAccessEntry) {
+    this._contractAddressAccessEntrys.push(entry);
+    return this;
+  }
+
+  public addContractFunctionAccessEntry(entry: ContractMethodEntry) {
+    this._contractMethodAccessEntrys.push(entry);
+    return this;
+  }
+
+  public setTimeRange(timeRange: TimeRange) {
+    this._timeRange = timeRange;
+    return this;
+  }
+
+  public setNativeTokenSpendLimit(limit: NativeTokenLimit) {
+    this._nativeTokenSpendLimit = limit;
+    return this;
+  }
+
+  public addErc20TokenSpendLimit(limit: Erc20TokenLimit) {
+    this._erc20TokenSpendLimits.push(limit);
+    return this;
+  }
+
+  public setGasSpendLimit(limit: GasSpendLimit) {
+    this._gasSpendLimit = limit;
+    return this;
+  }
+
+  public setRequiredPaymaster(paymaster: Address) {
+    this._requiredPaymaster = paymaster;
+    return this;
+  }
+
+  public encode(): Hex[] {
+    return [
+      encodeFunctionData({
+        abi: SessionKeyPermissionsUpdatesAbi,
+        functionName: "setAccessListType",
+        args: [this._contractAccessControlType],
+      }),
+      ...this._contractAddressAccessEntrys.map((entry) =>
+        encodeFunctionData({
+          abi: SessionKeyPermissionsUpdatesAbi,
+          functionName: "updateAccessListAddressEntry",
+          args: [entry.contractAddress, entry.isOnList, entry.checkSelectors],
+        })
+      ),
+      ...this._contractMethodAccessEntrys.map((entry) =>
+        encodeFunctionData({
+          abi: SessionKeyPermissionsUpdatesAbi,
+          functionName: "updateAccessListFunctionEntry",
+          args: [entry.contractAddress, entry.methodSelector, entry.isOnList],
+        })
+      ),
+      this.encodeIfDefined(
+        (timeRange) =>
+          encodeFunctionData({
+            abi: SessionKeyPermissionsUpdatesAbi,
+            functionName: "updateTimeRange",
+            args: [timeRange.validFrom, timeRange.validUntil],
+          }),
+        this._timeRange
+      ),
+      this.encodeIfDefined(
+        (nativeSpendLimit) =>
+          encodeFunctionData({
+            abi: SessionKeyPermissionsUpdatesAbi,
+            functionName: "setNativeTokenSpendLimit",
+            args: [
+              nativeSpendLimit.spendLimit,
+              nativeSpendLimit.refreshInterval ?? 0,
+            ],
+          }),
+        this._nativeTokenSpendLimit
+      ),
+      ...this._erc20TokenSpendLimits.map((erc20SpendLimit) =>
+        encodeFunctionData({
+          abi: SessionKeyPermissionsUpdatesAbi,
+          functionName: "setERC20SpendLimit",
+          args: [
+            erc20SpendLimit.tokenAddress,
+            erc20SpendLimit.spendLimit,
+            erc20SpendLimit.refreshInterval ?? 0,
+          ],
+        })
+      ),
+      this.encodeIfDefined(
+        (spendLimit) =>
+          encodeFunctionData({
+            abi: SessionKeyPermissionsUpdatesAbi,
+            functionName: "setGasSpendLimit",
+            args: [spendLimit.spendLimit, spendLimit.refreshInterval ?? 0],
+          }),
+        this._gasSpendLimit
+      ),
+      this.encodeIfDefined(
+        (paymaster) =>
+          encodeFunctionData({
+            abi: SessionKeyPermissionsUpdatesAbi,
+            functionName: "setRequiredPaymaster",
+            args: [paymaster],
+          }),
+        this._requiredPaymaster
+      ),
+    ].filter((x) => x !== "0x");
+  }
+
+  private encodeIfDefined<T>(encode: (param: T) => Hex, param?: T): Hex {
+    if (!param) return "0x";
+
+    return encode(param);
+  }
+}