implement the missing pieces for smart accounts

Author: pcarranzav

packages/hypergraph/src/identity/abis.ts

@@ -44,3 +44,28 @@ export const smartSessionsAbi = [
     type: 'function',
   },
 ];
+
+// ABI for the Safe7579 module, only with the functions we need
+export const safe7579Abi = [
+  {
+    type: 'function',
+    name: 'isModuleInstalled',
+    inputs: [
+      {
+        name: 'moduleType',
+        type: 'uint256',
+        internalType: 'uint256',
+      },
+      { name: 'module', type: 'address', internalType: 'address' },
+      { name: 'additionalContext', type: 'bytes', internalType: 'bytes' },
+    ],
+    outputs: [
+      {
+        name: '',
+        type: 'bool',
+        internalType: 'bool',
+      },
+    ],
+    stateMutability: 'view',
+  },
+];

packages/hypergraph/src/identity/smart-account.ts

@@ -13,8 +13,12 @@ import {
   getOwnableValidatorMockSignature,
   getPermissionId,
   getSmartSessionsValidator,
+  getSpendingLimitsPolicy,
   getSudoPolicy,
+  getTimeFramePolicy,
   getUniversalActionPolicy,
+  getUsageLimitPolicy,
+  getValueLimitPolicy,
 } from '@rhinestone/module-sdk';
 import { privateKeyToAccount } from 'viem/accounts';
 
@@ -28,17 +32,24 @@ import {
   type AbiFunction,
   type Address,
   type Chain,
+  ContractFunctionExecutionError,
   type Hex,
+  type SignableMessage,
   type WalletClient,
   createPublicClient,
   encodeFunctionData,
   toBytes,
   toFunctionSelector,
   toHex,
 } from 'viem';
-import { type UserOperation, entryPoint07Address, getUserOperationHash } from 'viem/account-abstraction';
+import {
+  type UserOperation,
+  type WaitForUserOperationReceiptReturnType,
+  entryPoint07Address,
+  getUserOperationHash,
+} from 'viem/account-abstraction';
 import { bytesToHex } from '../utils/hexBytesAddressUtils.js';
-import { safeModuleManagerAbi, smartSessionsAbi } from './abis.js';
+import { safe7579Abi, safeModuleManagerAbi, smartSessionsAbi } from './abis.js';
 
 const DEFAULT_RPC_URL = 'https://rpc-geo-genesis-h0q2s21xx8.t.conduit.xyz';
 /**
@@ -55,6 +66,8 @@ const ERC7579_LAUNCHPAD_ADDRESS = '0x7579011aB74c46090561ea277Ba79D510c6C00ff';
 const SPACE_FACTORY_ADDRESS = '0x0000000000000000000000000000000000000000'; // TODO: add address
 const SPACE_FACTORY_CREATE_SPACE_SELECTOR = '0x00000000'; // TODO: add selector
 
+const MODULE_TYPE_VALIDATOR = 1;
+
 // TODO: add details for testnet too
 export const GEOGENESIS = {
   id: Number('80451'),
@@ -74,15 +87,31 @@ export const GEOGENESIS = {
   },
 };
 
-type Action = {
+export type Action = {
   actionTarget: Address;
   actionTargetSelector: Hex;
   actionPolicies: { policy: Address; address: Address; initData: Hex }[];
 };
 
+// We re-export these functions to allow creating sessions with policies for
+// additional actions without needing the Rhinestone module SDK.
+export {
+  getSudoPolicy,
+  getUniversalActionPolicy,
+  getSpendingLimitsPolicy,
+  getTimeFramePolicy,
+  getUsageLimitPolicy,
+  getValueLimitPolicy,
+};
+
+export type SmartSessionClient = {
+  sendTransaction: <const calls extends readonly unknown[]>({ calls }: { calls: calls }) => Promise<string>;
+  signMessage: ({ message }: { message: SignableMessage }) => Promise<Hex>;
+};
+
 // Gets the legacy Geo smart account wallet client. If the smart account returned
 // by this function is deployed, it means it might need to be updated to have the 7579 module installed
-export const getLegacySmartAccountWalletClient = async (
+const getLegacySmartAccountWalletClient = async (
   walletClient: WalletClient,
   chain: Chain = GEOGENESIS,
   rpcUrl: string = DEFAULT_RPC_URL,
@@ -129,9 +158,11 @@ export const getLegacySmartAccountWalletClient = async (
   return smartAccountClient;
 };
 
-export const get7579SmartAccountWalletClient = async (
+// Gets the 7579 smart account wallet client. This is the new type of smart account that
+// includes the session keys validator and the 7579 module.
+const get7579SmartAccountWalletClient = async (
   walletClient: WalletClient,
-  address: `0x${string}` | undefined,
+  address: Hex | undefined,
   chain: Chain = GEOGENESIS,
   rpcUrl: string = DEFAULT_RPC_URL,
   apiKey: string = DEFAULT_API_KEY,
@@ -152,16 +183,16 @@ export const get7579SmartAccountWalletClient = async (
   });
   const smartSessionsValidator = getSmartSessionsValidator({});
 
-  const safeAccountParams: ToSafeSmartAccountParameters<'0.7', `0x${string}`> = {
+  const safeAccountParams: ToSafeSmartAccountParameters<'0.7', Hex> = {
     client: publicClient,
     owners: [walletClient],
     version: '1.4.1' as const,
     entryPoint: {
       address: entryPoint07Address,
       version: '0.7' as const,
     },
-    safe4337ModuleAddress: SAFE_7579_MODULE_ADDRESS as `0x${string}`,
-    erc7579LaunchpadAddress: ERC7579_LAUNCHPAD_ADDRESS as `0x${string}`,
+    safe4337ModuleAddress: SAFE_7579_MODULE_ADDRESS as Hex,
+    erc7579LaunchpadAddress: ERC7579_LAUNCHPAD_ADDRESS as Hex,
     attesters: [],
     attestersThreshold: 0,
     validators: [
@@ -205,32 +236,105 @@ export const get7579SmartAccountWalletClient = async (
   return smartAccountClient as unknown as SmartAccountClient;
 };
 
+// Checks if the smart account is deployed.
 export const isSmartAccountDeployed = async (smartAccountClient: SmartAccountClient): Promise<boolean> => {
   if (!smartAccountClient.account) {
     throw new Error('Invalid smart account');
   }
   return smartAccountClient.account.isDeployed();
 };
 
-export const legacySmartAccountNeedsUpdate = async (
+// Gets the smart account wallet client. This is the main function to use to get a smart account wallet client.
+// It will return the 7579 smart account wallet client if the smart account is deployed, otherwise it will return the legacy smart account wallet client, that might need to be updated.
+// You can use smartAccountNeedsUpdate to check if the smart account needs to be updated, and then call updateLegacySmartAccount to update it,
+// which requires executing a user operation.
+export const getSmartAccountWalletClient = async (
+  walletClient: WalletClient,
+  address: Hex | undefined,
+  chain: Chain,
+  rpcUrl: string,
+  apiKey: string,
+): Promise<SmartAccountClient> => {
+  if (address) {
+    return get7579SmartAccountWalletClient(walletClient, address, chain, rpcUrl, apiKey);
+  }
+  const legacyClient = await getLegacySmartAccountWalletClient(walletClient, chain, rpcUrl, apiKey);
+  if (await isSmartAccountDeployed(legacyClient)) {
+    return legacyClient;
+  }
+  return get7579SmartAccountWalletClient(walletClient, undefined, chain, rpcUrl, apiKey);
+};
+
+// Checks if the smart account has the 7579 module installed, the smart sessions validator installed, and the ownable validator installed.
+export const legacySmartAccountUpdateStatus = async (
   smartAccountClient: SmartAccountClient,
   chain: Chain,
   rpcUrl: string,
-): Promise<boolean> => {
+): Promise<{ has7579Module: boolean; hasSmartSessionsValidator: boolean; hasOwnableValidator: boolean }> => {
   if (!smartAccountClient.account) {
     throw new Error('Invalid smart account');
   }
   // We assume the smart account is deployed, so we just need to check if it has the 7579 module and smart sesions validator installed
-  // TODO: call the isModuleInstalled function from the Safe7579 ABI on the
+  // TODO: call the isModuleInstalled function from the safe7579Abi on the
   // smart account, checking if the smart sessions validator is installed. This would fail
   // if the smart account doesn't have the 7579 module installed.
-  return false;
+  const transport = http(rpcUrl);
+  const publicClient = createPublicClient({
+    transport,
+    chain,
+  });
+  const smartSessionsValidator = getSmartSessionsValidator({});
+  let isSmartSessionsValidatorInstalled = false;
+  try {
+    isSmartSessionsValidatorInstalled = (await publicClient.readContract({
+      abi: safe7579Abi,
+      address: smartAccountClient.account.address,
+      functionName: 'isModuleInstalled',
+      args: [MODULE_TYPE_VALIDATOR, smartSessionsValidator.address, '0x'],
+    })) as boolean;
+  } catch (error) {
+    if (error instanceof ContractFunctionExecutionError && error.details.includes('execution reverted')) {
+      // If the smart account doesn't have the 7579 module installed, the isModuleInstalled function will revert
+      return { has7579Module: false, hasSmartSessionsValidator: false, hasOwnableValidator: false };
+    }
+    throw error;
+  }
+  const ownableValidator = getOwnableValidator({
+    owners: [smartAccountClient.account.address],
+    threshold: 1,
+  });
+  // This shouldn't throw because by now we know the smart account has the 7579 module installed
+  const isOwnableValidatorInstalled = (await publicClient.readContract({
+    abi: safe7579Abi,
+    address: smartAccountClient.account.address,
+    functionName: 'isModuleInstalled',
+    args: [MODULE_TYPE_VALIDATOR, ownableValidator.address, '0x'],
+  })) as boolean;
+  return {
+    has7579Module: true,
+    hasSmartSessionsValidator: isSmartSessionsValidatorInstalled,
+    hasOwnableValidator: isOwnableValidatorInstalled,
+  };
+};
+
+// Checks if the smart account needs to be updated from a legacy ERC-4337 smart account to an ERC-7579 smart account
+// with support for smart sessions.
+export const smartAccountNeedsUpdate = async (
+  smartAccountClient: SmartAccountClient,
+  chain: Chain,
+  rpcUrl: string,
+): Promise<boolean> => {
+  const updateStatus = await legacySmartAccountUpdateStatus(smartAccountClient, chain, rpcUrl);
+  return !updateStatus.has7579Module || !updateStatus.hasSmartSessionsValidator || !updateStatus.hasOwnableValidator;
 };
 
 // Legacy Geo smart accounts (i.e. the ones that don't have the 7579 module installed)
-// need to be updated to have the 7579 module installed
-// with the ownable and smart sessions validators.
-export const updateLegacySmartAccount = async (smartAccountClient: SmartAccountClient) => {
+// need to be updated to have the 7579 module installed with the ownable and smart sessions validators.
+export const updateLegacySmartAccount = async (
+  smartAccountClient: SmartAccountClient,
+  chain: Chain,
+  rpcUrl: string,
+): Promise<WaitForUserOperationReceiptReturnType | undefined> => {
   if (!smartAccountClient.account?.address) {
     throw new Error('Invalid smart account');
   }
@@ -255,46 +359,54 @@ export const updateLegacySmartAccount = async (smartAccountClient: SmartAccountC
     ],
   });
 
-  const calls = [
-    {
+  const updateStatus = await legacySmartAccountUpdateStatus(smartAccountClient, chain, rpcUrl);
+  const calls = [];
+  if (!updateStatus.has7579Module) {
+    calls.push({
       to: smartAccountClient.account.address,
       data: encodeFunctionData({
         abi: safeModuleManagerAbi,
         functionName: 'enableModule',
-        args: [SAFE_7579_MODULE_ADDRESS as `0x${string}`],
+        args: [SAFE_7579_MODULE_ADDRESS as Hex],
       }),
       value: BigInt(0),
-    },
-    {
+    });
+    calls.push({
       to: smartAccountClient.account.address,
       data: encodeFunctionData({
         abi: safeModuleManagerAbi,
         functionName: 'setFallbackHandler',
-        args: [SAFE_7579_MODULE_ADDRESS as `0x${string}`],
+        args: [SAFE_7579_MODULE_ADDRESS as Hex],
       }),
       value: BigInt(0),
-    },
-    {
+    });
+    calls.push({
       to: smartAccountClient.account.address,
       data: encodeFunctionData({
         abi: safeModuleManagerAbi,
         functionName: 'disableModule',
-        args: [SAFE_4337_MODULE_ADDRESS as `0x${string}`],
+        args: [SAFE_4337_MODULE_ADDRESS as Hex],
       }),
       value: BigInt(0),
-    },
-    {
+    });
+  }
+  if (!updateStatus.hasOwnableValidator) {
+    calls.push({
       to: installValidatorsTx[0].to,
       data: installValidatorsTx[0].data,
       value: installValidatorsTx[0].value,
-    },
-    {
+    });
+  }
+  if (!updateStatus.hasSmartSessionsValidator) {
+    calls.push({
       to: installValidatorsTx[1].to,
       data: installValidatorsTx[1].data,
       value: installValidatorsTx[1].value,
-    },
-  ];
-
+    });
+  }
+  if (calls.length === 0) {
+    return;
+  }
   const tx = await smartAccountClient.sendUserOperation({
     calls,
   });
@@ -311,11 +423,11 @@ export const updateLegacySmartAccount = async (smartAccountClient: SmartAccountC
 // enable it on the smart account.
 // It will prompt the user to sign the message to enable the session, and then
 // execute the transaction to enable the session.
-// It will return the permissionId.
+// It will return the permissionId that can be used to create a smart session client.
 export const createSmartSession = async (
   walletClient: WalletClient,
   smartAccountClient: SmartAccountClient,
-  sessionPrivateKey: `0x${string}`,
+  sessionPrivateKey: Hex,
   chain: Chain,
   rpcUrl: string,
   {
@@ -324,10 +436,10 @@ export const createSmartSession = async (
     additionalActions = [],
   }: {
     allowCreateSpace?: boolean;
-    spaceAddresses?: `0x${string}`[];
+    spaceAddresses?: Hex[];
     additionalActions?: Action[];
   } = {},
-) => {
+): Promise<Hex> => {
   if (!smartAccountClient.account) {
     throw new Error('Invalid smart account');
   }
@@ -395,7 +507,7 @@ export const createSmartSession = async (
       threshold: 1,
       owners: [sessionKeyAccount.address],
     }),
-    salt: bytesToHex(randomBytes(32)) as `0x${string}`,
+    salt: bytesToHex(randomBytes(32)) as Hex,
     userOpPolicies: [getSudoPolicy()],
     erc7739Policies: {
       allowedERC7739Content: [],
@@ -485,13 +597,14 @@ export const createSmartSession = async (
 
 // This is the function that we use on the end user app to create a smart session client that can send transactions to the smart account.
 // The session must have previously been created by the createSmartSession function.
-export const getSmartSessionClient = async (
+// The client also includes a signMessage function that can be used to sign messages with the session key.
+export const getSmartSessionClient = (
   smartAccountClient: SmartAccountClient,
-  sessionPrivateKey: `0x${string}`,
+  sessionPrivateKey: Hex,
   permissionId: Hex,
   chain: Chain,
   rpcUrl: string,
-) => {
+): SmartSessionClient => {
   if (!smartAccountClient.account) {
     throw new Error('Invalid smart account');
   }
@@ -547,5 +660,8 @@ export const getSmartSessionClient = async (
 
       return smartAccountClient.sendUserOperation(userOperation as UserOperation);
     },
+    signMessage: async ({ message }: { message: SignableMessage }) => {
+      return sessionKeyAccount.signMessage({ message });
+    },
   };
 };