From cbf1d53ac1a5a5dbccb7abf0b73ff75f6765ce6a Mon Sep 17 00:00:00 2001 From: howydev <132113803+howydev@users.noreply.github.com> Date: Tue, 17 Dec 2024 12:58:19 -0500 Subject: [PATCH] feat: add install validation, add tests for session key --- aa-sdk/core/src/errors/client.ts | 16 ++ aa-sdk/core/src/index.ts | 1 + .../src/ma-v2/account/semiModularAccountV2.ts | 10 +- .../src/ma-v2/actions/common/types.ts | 37 ++++ .../src/ma-v2/actions/common/utils.ts | 32 ++++ .../install-validation/installValidation.ts | 160 ++++++++++++++++++ .../src/ma-v2/client/client.test.ts | 151 ++++++++++++++++- .../EntityIdOverrideError/constructor.mdx | 18 ++ 8 files changed, 415 insertions(+), 10 deletions(-) create mode 100644 account-kit/smart-contracts/src/ma-v2/actions/common/types.ts create mode 100644 account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts create mode 100644 account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts create mode 100644 site/pages/reference/aa-sdk/core/classes/EntityIdOverrideError/constructor.mdx diff --git a/aa-sdk/core/src/errors/client.ts b/aa-sdk/core/src/errors/client.ts index b3040fe785..f90784ca6d 100644 --- a/aa-sdk/core/src/errors/client.ts +++ b/aa-sdk/core/src/errors/client.ts @@ -88,3 +88,19 @@ export class InvalidNonceKeyError extends BaseError { ); } } + +/** + * Error class denoting that the provided entity id is invalid because it's overriding the native entity id. + */ +export class EntityIdOverrideError extends BaseError { + override name = "InvalidNonceKeyError"; + + /** + * Initializes a new instance of the error message with a default message indicating that the nonce key is invalid. + */ + constructor() { + super( + `Installing entityId of 0 overrides the owner's entity id in the account` + ); + } +} diff --git a/aa-sdk/core/src/index.ts b/aa-sdk/core/src/index.ts index c5d33dea24..5f66aae5cc 100644 --- a/aa-sdk/core/src/index.ts +++ b/aa-sdk/core/src/index.ts @@ -75,6 +75,7 @@ export { InvalidRpcUrlError, InvalidEntityIdError, InvalidNonceKeyError, + EntityIdOverrideError, } from "./errors/client.js"; export { EntryPointNotFoundError, diff --git a/account-kit/smart-contracts/src/ma-v2/account/semiModularAccountV2.ts b/account-kit/smart-contracts/src/ma-v2/account/semiModularAccountV2.ts index b8977309cf..a3a917c0cc 100644 --- a/account-kit/smart-contracts/src/ma-v2/account/semiModularAccountV2.ts +++ b/account-kit/smart-contracts/src/ma-v2/account/semiModularAccountV2.ts @@ -8,6 +8,7 @@ import { createBundlerClient, getEntryPoint, toSmartContractAccount, + getAccountAddress, } from "@aa-sdk/core"; import { concatHex, @@ -119,6 +120,13 @@ export async function createSMAV2Account( ...singleSignerMessageSigner(signer), }); + const accountAddress_ = await getAccountAddress({ + client, + entryPoint, + accountAddress, + getAccountInitCode, + }); + // TODO: add deferred action flag const getAccountNonce = async (nonceKey?: bigint): Promise => { const nonceKeySuffix: Hex = `${toHex(entityId, { size: 4 })}${ @@ -136,7 +144,7 @@ export async function createSMAV2Account( }); return entryPointContract.read.getNonce([ - baseAccount.address, + accountAddress_, nonceKey ?? hexToBigInt(nonceKeySuffix), ]) as Promise; }; diff --git a/account-kit/smart-contracts/src/ma-v2/actions/common/types.ts b/account-kit/smart-contracts/src/ma-v2/actions/common/types.ts new file mode 100644 index 0000000000..14f5739e6e --- /dev/null +++ b/account-kit/smart-contracts/src/ma-v2/actions/common/types.ts @@ -0,0 +1,37 @@ +import type { Address, Hex } from "viem"; + +export type ModuleEntity = { + moduleAddress: Address; + entityId: number; +}; + +export type ValidationConfig = { + moduleAddress: Address; + entityId: number; // uint32 + isGlobal: boolean; + isSignatureValidation: boolean; + isUserOpValidation: boolean; +}; + +export enum HookType { + EXECUTION = "0x00", + VALIDATION = "0x01", +} + +export type HookConfig = { + address: Address; + entityId: number; // uint32 + hookType: HookType; + hasPreHooks: boolean; + hasPostHooks: boolean; +}; + +// maps to type ValidationStorage in MAv2 implementation +export type ValidationData = { + isGlobal: boolean; // validation flag + isSignatureValidation: boolean; // validation flag + isUserOpValidation: boolean; + validationHooks: HookConfig[]; + executionHooks: Hex[]; + selectors: Hex[]; +}; diff --git a/account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts b/account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts new file mode 100644 index 0000000000..8bd366ab76 --- /dev/null +++ b/account-kit/smart-contracts/src/ma-v2/actions/common/utils.ts @@ -0,0 +1,32 @@ +import { type Hex, toHex, concatHex } from "viem"; +import type { ValidationConfig, HookConfig, ModuleEntity } from "./types"; +import { HookType } from "./types.js"; + +export function serializeValidationConfig(config: ValidationConfig): Hex { + const isUserOpValidationBit = config.isUserOpValidation ? 1 : 0; + const isSignatureValidationBit = config.isSignatureValidation ? 2 : 0; + const isGlobalBit = config.isGlobal ? 4 : 0; + return concatHex([ + serializeModuleEntity(config), + toHex(isUserOpValidationBit + isSignatureValidationBit + isGlobalBit, { + size: 1, + }), + ]); +} + +export function serializeHookConfig(config: HookConfig): Hex { + const hookTypeBit = config.hookType === HookType.VALIDATION ? 1 : 0; + const hasPostHooksBit = config.hasPostHooks ? 2 : 0; + const hasPreHooksBit = config.hasPreHooks ? 4 : 0; + return concatHex([ + config.address, + toHex(config.entityId, { size: 4 }), + toHex(hookTypeBit + hasPostHooksBit + hasPreHooksBit, { + size: 1, + }), + ]); +} + +export function serializeModuleEntity(config: ModuleEntity): Hex { + return concatHex([config.moduleAddress, toHex(config.entityId, { size: 4 })]); +} diff --git a/account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts b/account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts new file mode 100644 index 0000000000..f2d77d386c --- /dev/null +++ b/account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts @@ -0,0 +1,160 @@ +import { + AccountNotFoundError, + IncompatibleClientError, + isSmartAccountClient, + EntityIdOverrideError, + type GetAccountParameter, + type GetEntryPointFromAccount, + type SendUserOperationResult, + type SmartContractAccount, + type SmartAccountClient, + type UserOperationOverridesParameter, +} from "@aa-sdk/core"; +import { + type Address, + type Hex, + type Chain, + type Transport, + encodeFunctionData, + concatHex, +} from "viem"; + +import { semiModularAccountBytecodeAbi } from "../../abis/semiModularAccountBytecodeAbi.js"; +import type { HookConfig, ValidationConfig } from "../common/types.js"; +import { + serializeValidationConfig, + serializeHookConfig, + serializeModuleEntity, +} from "../common/utils.js"; + +export type InstallValidationParams< + TAccount extends SmartContractAccount | undefined = + | SmartContractAccount + | undefined +> = { + validationConfig: ValidationConfig; + selectors: Hex[]; + installData: Hex; + hooks: { + hookConfig: HookConfig; + initData: Hex; + }[]; +} & UserOperationOverridesParameter> & + GetAccountParameter; + +export type UninstallValidationParams< + TAccount extends SmartContractAccount | undefined = + | SmartContractAccount + | undefined +> = { + moduleAddress: Address; + entityId: number; + uninstallData: Hex; + hookUninstallDatas: Hex[]; +} & UserOperationOverridesParameter> & + GetAccountParameter; + +export type InstallValidationActions< + TAccount extends SmartContractAccount | undefined = + | SmartContractAccount + | undefined +> = { + installValidation: ( + args: InstallValidationParams + ) => Promise; + uninstallValidation: ( + args: UninstallValidationParams + ) => Promise; +}; + +export const installValidationActions: < + TTransport extends Transport = Transport, + TChain extends Chain | undefined = Chain | undefined, + TAccount extends SmartContractAccount = SmartContractAccount +>( + client: SmartAccountClient +) => InstallValidationActions = (client) => ({ + installValidation: async ({ + validationConfig, + selectors, + installData, + hooks, + account = client.account, + overrides, + }) => { + if (!account) { + throw new AccountNotFoundError(); + } + + if (!isSmartAccountClient(client)) { + throw new IncompatibleClientError( + "SmartAccountClient", + "installValidation", + client + ); + } + + if (validationConfig.entityId === 0) { + throw new EntityIdOverrideError(); + } + + const callData = encodeFunctionData({ + abi: semiModularAccountBytecodeAbi, + functionName: "installValidation", + args: [ + serializeValidationConfig(validationConfig), + selectors, + installData, + hooks.map((hook: { hookConfig: HookConfig; initData: Hex }) => + concatHex([serializeHookConfig(hook.hookConfig), hook.initData]) + ), + ], + }); + + return client.sendUserOperation({ + uo: callData, + account, + overrides, + }); + }, + + uninstallValidation: async ({ + moduleAddress, + entityId, + uninstallData, + hookUninstallDatas, + account = client.account, + overrides, + }) => { + if (!account) { + throw new AccountNotFoundError(); + } + + if (!isSmartAccountClient(client)) { + throw new IncompatibleClientError( + "SmartAccountClient", + "uninstallValidation", + client + ); + } + + const callData = encodeFunctionData({ + abi: semiModularAccountBytecodeAbi, + functionName: "uninstallValidation", + args: [ + serializeModuleEntity({ + moduleAddress, + entityId, + }), + uninstallData, + hookUninstallDatas, + ], + }); + + return client.sendUserOperation({ + uo: callData, + account, + overrides, + }); + }, +}); diff --git a/account-kit/smart-contracts/src/ma-v2/client/client.test.ts b/account-kit/smart-contracts/src/ma-v2/client/client.test.ts index 8ead56e5fa..ee6af7f58c 100644 --- a/account-kit/smart-contracts/src/ma-v2/client/client.test.ts +++ b/account-kit/smart-contracts/src/ma-v2/client/client.test.ts @@ -7,6 +7,9 @@ import { createSMAV2AccountClient } from "./client.js"; import { local070Instance } from "~test/instances.js"; import { setBalance } from "viem/actions"; import { accounts } from "~test/constants.js"; +import { installValidationActions } from "../actions/install-validation/installValidation.js"; +import { getDefaultSingleSignerValidationModuleAddress } from "../modules/utils.js"; +import { SingleSignerValidationModule } from "../modules/single-signer-validation/module.js"; describe("MA v2 Tests", async () => { const instance = local070Instance; @@ -21,6 +24,18 @@ describe("MA v2 Tests", async () => { accounts.fundedAccountOwner ); + const sessionKey: SmartAccountSigner = new LocalAccountSigner( + accounts.unfundedAccountOwner + ); + + const target = "0x000000000000000000000000000000000000dEaD"; + const sendAmount = parseEther("1"); + + const getTargetBalance = async (): Promise => + client.getBalance({ + address: target, + }); + it("sends a simple UO", async () => { const provider = await givenConnectedProvider({ signer }); @@ -29,12 +44,7 @@ describe("MA v2 Tests", async () => { value: parseEther("2"), }); - const target = "0x000000000000000000000000000000000000dEaD"; - const sendAmount = parseEther("1"); - - const startingAddressBalance = await client.getBalance({ - address: target, - }); + const startingAddressBalance = await getTargetBalance(); const result = await provider.sendUserOperation({ uo: { @@ -47,23 +57,146 @@ describe("MA v2 Tests", async () => { const txnHash1 = provider.waitForUserOperationTransaction(result); await expect(txnHash1).resolves.not.toThrowError(); - const newAddressBalance = await client.getBalance({ - address: target, + await expect(await getTargetBalance()).toEqual( + startingAddressBalance + sendAmount + ); + }); + + it("adds a session key with no permissions", async () => { + let provider = (await givenConnectedProvider({ signer })).extend( + installValidationActions + ); + + await setBalance(client, { + address: provider.getAddress(), + value: parseEther("2"), + }); + + let result = await provider.installValidation({ + validationConfig: { + moduleAddress: getDefaultSingleSignerValidationModuleAddress( + provider.chain + ), + entityId: 1, + isGlobal: true, + isSignatureValidation: true, + isUserOpValidation: true, + }, + selectors: [], + installData: SingleSignerValidationModule.encodeOnInstallData({ + entityId: 1, + signer: await sessionKey.getAddress(), + }), + hooks: [], + }); + + let txnHash = provider.waitForUserOperationTransaction(result); + await expect(txnHash).resolves.not.toThrowError(); + + const startingAddressBalance = await getTargetBalance(); + + // connect session key and send tx with session key + let sessionKeyClient = await createSMAV2AccountClient({ + chain: instance.chain, + signer: sessionKey, + transport: custom(instance.getClient()), + accountAddress: provider.getAddress(), + entityId: 1, + isGlobalValidation: true, }); - await expect(newAddressBalance).toEqual( + result = await sessionKeyClient.sendUserOperation({ + uo: { + target: target, + value: sendAmount, + data: "0x", + }, + }); + + txnHash = sessionKeyClient.waitForUserOperationTransaction(result); + await expect(txnHash).resolves.not.toThrowError(); + await expect(await getTargetBalance()).toEqual( startingAddressBalance + sendAmount ); }); + it("uninstalls a session key", async () => { + let provider = (await givenConnectedProvider({ signer })).extend( + installValidationActions + ); + + await setBalance(client, { + address: provider.getAddress(), + value: parseEther("2"), + }); + + let result = await provider.installValidation({ + validationConfig: { + moduleAddress: getDefaultSingleSignerValidationModuleAddress( + provider.chain + ), + entityId: 1, + isGlobal: true, + isSignatureValidation: true, + isUserOpValidation: true, + }, + selectors: [], + installData: SingleSignerValidationModule.encodeOnInstallData({ + entityId: 1, + signer: await sessionKey.getAddress(), + }), + hooks: [], + }); + + let txnHash = provider.waitForUserOperationTransaction(result); + await expect(txnHash).resolves.not.toThrowError(); + + result = await provider.uninstallValidation({ + moduleAddress: getDefaultSingleSignerValidationModuleAddress( + provider.chain + ), + entityId: 1, + uninstallData: SingleSignerValidationModule.encodeOnUninstallData({ + entityId: 1, + }), + hookUninstallDatas: [], + }); + + txnHash = provider.waitForUserOperationTransaction(result); + await expect(txnHash).resolves.not.toThrowError(); + + // connect session key and send tx with session key + let sessionKeyClient = await createSMAV2AccountClient({ + chain: instance.chain, + signer: sessionKey, + transport: custom(instance.getClient()), + accountAddress: provider.getAddress(), + entityId: 1, + isGlobalValidation: true, + }); + + await expect( + sessionKeyClient.sendUserOperation({ + uo: { + target: target, + value: sendAmount, + data: "0x", + }, + }) + ).rejects.toThrowError(); + }); + const givenConnectedProvider = async ({ signer, + accountAddress, }: { signer: SmartAccountSigner; + accountAddress?: `0x${string}`; }) => createSMAV2AccountClient({ chain: instance.chain, signer, + accountAddress, transport: custom(instance.getClient()), }); }); diff --git a/site/pages/reference/aa-sdk/core/classes/EntityIdOverrideError/constructor.mdx b/site/pages/reference/aa-sdk/core/classes/EntityIdOverrideError/constructor.mdx new file mode 100644 index 0000000000..ef8e82cc95 --- /dev/null +++ b/site/pages/reference/aa-sdk/core/classes/EntityIdOverrideError/constructor.mdx @@ -0,0 +1,18 @@ +--- +# This file is autogenerated +title: EntityIdOverrideError +description: Overview of the EntityIdOverrideError method +--- + +# EntityIdOverrideError + +Initializes a new instance of the error message with a default message indicating that the nonce key is invalid. +:::note +`EntityIdOverrideError` extends `BaseError`, see the docs for BaseError for all supported methods. +::: + +## Import + +```ts +import { EntityIdOverrideError } from "@aa-sdk/core"; +```