feat: add session key extensions for managing keys

OMGWINNING · Feb 6, 2024 · ac86a97 · ac86a97
1 parent 91e5778
commit ac86a97
Show file tree

Hide file tree

Showing 6 changed files with 131 additions and 35 deletions.
diff --git a/packages/accounts/src/msca/plugins/session-key/extension.ts b/packages/accounts/src/msca/plugins/session-key/extension.ts
@@ -5,7 +5,7 @@ import {
   type SmartContractAccount,
   type UserOperationOverrides,
 } from "@alchemy/aa-core";
-import type { Address, Chain, Client, Transport } from "viem";
+import type { Address, Chain, Client, Hex, Transport } from "viem";
 import {
   SessionKeyPlugin,
   sessionKeyPluginActions as sessionKeyPluginActions_,
@@ -17,7 +17,13 @@ export type SessionKeyPluginActions<
   TAccount extends SmartContractAccount | undefined =
     | SmartContractAccount
     | undefined
-> = Omit<SessionKeyPluginActions_<TAccount>, "removeSessionKey"> & {
+> = Omit<
+  SessionKeyPluginActions_<TAccount>,
+  | "removeSessionKey"
+  | "addSessionKey"
+  | "rotateSessionKey"
+  | "updateKeyPermissions"
+> & {
   isAccountSessionKey: (
     args: {
       key: Address;
@@ -38,6 +44,32 @@ export type SessionKeyPluginActions<
       overrides?: UserOperationOverrides;
     } & GetAccountParameter<TAccount>
   ) => Promise<SendUserOperationResult>;
+
+  addSessionKey: (
+    args: {
+      key: Address;
+      permissions: Hex[];
+      tag: Hex;
+      overrides?: UserOperationOverrides;
+    } & GetAccountParameter<TAccount>
+  ) => Promise<SendUserOperationResult>;
+
+  rotateSessionKey: (
+    args: {
+      oldKey: Address;
+      newKey: Address;
+      pluginAddress?: Address;
+      overrides?: UserOperationOverrides;
+    } & GetAccountParameter<TAccount>
+  ) => Promise<SendUserOperationResult>;
+
+  updateSessionKeyPermissions: (
+    args: {
+      key: Address;
+      permissions: Hex[];
+      overrides?: UserOperationOverrides;
+    } & GetAccountParameter<TAccount>
+  ) => Promise<SendUserOperationResult>;
 };
 
 export const sessionKeyPluginActions: <
@@ -57,7 +89,13 @@ export const sessionKeyPluginActions: <
 >(
   client: Client<TTransport, TChain, TAccount>
 ) => {
-  const { removeSessionKey, ...og } = sessionKeyPluginActions_(client);
+  const {
+    removeSessionKey,
+    addSessionKey,
+    rotateSessionKey,
+    updateKeyPermissions,
+    ...og
+  } = sessionKeyPluginActions_(client);
 
   return {
     ...og,
@@ -84,12 +122,18 @@ export const sessionKeyPluginActions: <
       return await contract.read.sessionKeysOf([account.address]);
     },
 
-    removeSessionKey: async ({ key, overrides, account = client.account }) => {
+    removeSessionKey: async ({
+      key,
+      overrides,
+      account = client.account,
+      pluginAddress,
+    }) => {
       if (!account) throw new AccountNotFoundError();
 
       const sessionKeysToRemove = await buildSessionKeysToRemoveStruct(client, {
         keys: [key],
         account,
+        pluginAddress,
       });
 
       return removeSessionKey({
@@ -98,5 +142,59 @@ export const sessionKeyPluginActions: <
         account,
       });
     },
+
+    addSessionKey: async ({
+      key,
+      tag,
+      permissions,
+      overrides,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      return addSessionKey({
+        args: [key, tag, permissions],
+        overrides,
+        account,
+      });
+    },
+
+    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,
+        account,
+      });
+    },
+
+    updateSessionKeyPermissions: async ({
+      key,
+      permissions,
+      overrides,
+      account = client.account,
+    }) => {
+      if (!account) throw new AccountNotFoundError();
+
+      return updateKeyPermissions({
+        args: [key, permissions],
+        overrides,
+        account,
+      });
+    },
   };
 };
diff --git a/packages/accounts/src/msca/plugins/session-key/utils.ts b/packages/accounts/src/msca/plugins/session-key/utils.ts
@@ -16,14 +16,17 @@ export const buildSessionKeysToRemoveStruct: <
     | undefined
 >(
   client: Client<TTransport, TChain, TAccount>,
-  args: { keys: ReadonlyArray<Address> } & GetAccountParameter<TAccount>
+  args: {
+    keys: ReadonlyArray<Address>;
+    pluginAddress?: Address;
+  } & GetAccountParameter<TAccount>
 ) => Promise<{ sessionKey: Address; predecessor: Address }[]> = async (
   client,
-  { keys, account = client.account }
+  { keys, pluginAddress, account = client.account }
 ) => {
   if (!account) throw new AccountNotFoundError();
 
-  const contract = SessionKeyPlugin.getContract(client);
+  const contract = SessionKeyPlugin.getContract(client, pluginAddress);
   return (
     await Promise.all(
       keys.map(async (key) => {

diff --git a/site/.vitepress/sidebar/index.ts b/site/.vitepress/sidebar/index.ts
@@ -105,15 +105,6 @@ export const sidebar: DefaultTheme.Sidebar = [
         ],
       },
       { text: "Simulate User Operations", link: "/simulate-user-operations" },
-      {
-        text: "Transfer Ownership",
-        base: "/using-smart-accounts/transfer-ownership",
-        collapsed: false,
-        items: [
-          { text: "Modular Account", link: "/modular-account" },
-          { text: "Light Account", link: "/light-account" },
-        ],
-      },
       {
         text: "Session Keys",
         base: "/smart-accounts/session-keys",
@@ -128,6 +119,15 @@ export const sidebar: DefaultTheme.Sidebar = [
           },
         ],
       },
+      {
+        text: "Transfer Ownership",
+        base: "/using-smart-accounts/transfer-ownership",
+        collapsed: false,
+        items: [
+          { text: "Modular Account", link: "/modular-account" },
+          { text: "Light Account", link: "/light-account" },
+        ],
+      },
       {
         text: "Alchemy Enhanced Apis",
         base: "/using-smart-accounts/enhanced-apis",

diff --git a/site/snippets/session-keys/add-session-key.ts b/site/snippets/session-keys/add-session-key.ts
@@ -2,10 +2,7 @@ import { SessionKeyPermissionsBuilder } from "@alchemy/aa-accounts";
 import { client } from "./base-client.js";
 
 const result = await client.addSessionKey({
-  args: [
-    "0xSessionKeyAddress",
-    "0xkeytag",
-    // you can pass other permissions here by building them up
-    new SessionKeyPermissionsBuilder().encode(),
-  ],
+  key: "0xSessionKeyAddress",
+  tag: "0xkeytag",
+  permissions: new SessionKeyPermissionsBuilder().encode(),
 });
diff --git a/site/snippets/session-keys/rotate-session-key.ts b/site/snippets/session-keys/rotate-session-key.ts
@@ -1,6 +1,6 @@
 import { client } from "./base-client.js";
 
-// TODO before GA, we should add a util like we do for remove session key to make this easier
 const result = await client.rotateSessionKey({
-  args: ["0xOldKey", "0xpredecessor", "0xNewKey"],
+  oldKey: "0xOldKey",
+  newKey: "0xNewKey",
 });
diff --git a/site/snippets/session-keys/update-session-key.ts b/site/snippets/session-keys/update-session-key.ts
@@ -1,16 +1,14 @@
 import { SessionKeyPermissionsBuilder } from "@alchemy/aa-accounts";
 import { client } from "./base-client.js";
 
-const result = await client.updateKeyPermissions({
+const result = await client.updateSessionKeyPermissions({
+  key: "0xSessionKeyAddress",
   // add other permissions to the builder
-  args: [
-    "0xSessionKeyAddress",
-    new SessionKeyPermissionsBuilder()
-      .setTimeRange({
-        validFrom: Date.now() / 1000,
-        // valid for 1 hour
-        validUntil: Date.now() / 1000 + 60 * 60,
-      })
-      .encode(),
-  ],
+  permissions: new SessionKeyPermissionsBuilder()
+    .setTimeRange({
+      validFrom: Date.now() / 1000,
+      // valid for 1 hour
+      validUntil: Date.now() / 1000 + 60 * 60,
+    })
+    .encode(),
 });