From 792a27fbce3c78d3e456fadee3bc8f08255553fe Mon Sep 17 00:00:00 2001 From: linna Date: Fri, 20 Dec 2024 14:31:02 -0500 Subject: [PATCH] feat: adds tests for paymaster guard module --- .../install-validation/installValidation.ts | 7 +- .../src/ma-v2/client/client.test.ts | 183 +++++++++++++++++- 2 files changed, 183 insertions(+), 7 deletions(-) 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 index 7d9edf589c..dc30cf6870 100644 --- 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 @@ -85,9 +85,10 @@ export const installValidationActions: < ); } - if (validationConfig.entityId === DEFAULT_OWNER_ENTITY_ID) { - throw new EntityIdOverrideError(); - } + // TO DO: handle installing on fallback validation (entityId == 0) with non-zero address + // if (validationConfig.entityId === DEFAULT_OWNER_ENTITY_ID) { + // throw new EntityIdOverrideError(); + // } const { encodeCallData } = account; 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 31e59d228c..4f766a6d51 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 @@ -1,12 +1,22 @@ -import { custom, parseEther, publicActions } from "viem"; -import { LocalAccountSigner, type SmartAccountSigner } from "@aa-sdk/core"; -import { createSMAV2AccountClient, type SMAV2AccountClient } from "./client.js"; +import { custom, parseEther, publicActions, zeroAddress } from "viem"; +import { + erc7677Middleware, + LocalAccountSigner, + type SmartAccountSigner, +} from "@aa-sdk/core"; +import { createSMAV2AccountClient } from "./client.js"; import { local070Instance } from "~test/instances.js"; import { setBalance } from "viem/actions"; import { accounts } from "~test/constants.js"; -import { getDefaultSingleSignerValidationModuleAddress } from "../modules/utils.js"; +import { + getDefaultPaymasterGuardModuleAddress, + getDefaultSingleSignerValidationModuleAddress, +} from "../modules/utils.js"; import { SingleSignerValidationModule } from "../modules/single-signer-validation/module.js"; import { installValidationActions } from "../actions/install-validation/installValidation.js"; +import { paymaster070 } from "~test/paymaster/paymaster070.js"; +import { PaymasterGuardModule } from "../modules/paymaster-guard-module/module.js"; +import { HookType } from "../actions/common/types.js"; describe("MA v2 Tests", async () => { const instance = local070Instance; @@ -182,17 +192,182 @@ describe("MA v2 Tests", async () => { ).rejects.toThrowError(); }); + it("installs paymaster guard module, verifies use of valid paymaster, then uninstalls module", async () => { + let provider = ( + await givenConnectedProvider({ + signer, + usePaymaster: true, + }) + ).extend(installValidationActions); + + await setBalance(client, { + address: provider.getAddress(), + value: parseEther("2"), + }); + + const paymaster = paymaster070.getPaymasterStubData(); + + const hookInstallData = PaymasterGuardModule.encodeOnInstallData({ + entityId: 0, + paymaster: "paymaster" in paymaster ? paymaster.paymaster : "0x0", // dummy value for paymaster address if it DNE + }); + + const installResult = await provider.installValidation({ + validationConfig: { + moduleAddress: zeroAddress, + entityId: 0, + isGlobal: true, + isSignatureValidation: true, + isUserOpValidation: true, + }, + selectors: [], + installData: "0x", + hooks: [ + { + hookConfig: { + address: getDefaultPaymasterGuardModuleAddress(provider.chain), + entityId: 0, // uint32 + hookType: HookType.VALIDATION, + hasPreHooks: true, + hasPostHooks: true, + }, + initData: hookInstallData, + }, + ], + }); + + // verify hook installtion succeeded + await expect( + provider.waitForUserOperationTransaction(installResult) + ).resolves.not.toThrowError(); + + // happy path: with correct paymaster + const result = await provider.sendUserOperation({ + uo: { + target: target, + value: sendAmount, + data: "0x", + }, + }); + + // verify if correct paymaster is used + const txnHash1 = provider.waitForUserOperationTransaction(result); + await expect(txnHash1).resolves.not.toThrowError(); + + const hookUninstallData = PaymasterGuardModule.encodeOnUninstallData({ + entityId: 0, + }); + + const uninstallResult = await provider.uninstallValidation({ + moduleAddress: zeroAddress, + entityId: 0, + uninstallData: "0x", + hookUninstallDatas: [hookUninstallData], + }); + + // verify uninstall + await expect( + provider.waitForUserOperationTransaction(uninstallResult) + ).resolves.not.toThrowError(); + }); + + it("installs paymaster guard module, verifies use of invalid paymaster, then uninstalls module", async () => { + let provider = ( + await givenConnectedProvider({ + signer, + usePaymaster: true, + }) + ).extend(installValidationActions); + + await setBalance(client, { + address: provider.getAddress(), + value: parseEther("2"), + }); + + const paymaster = paymaster070.getPaymasterStubData(); + + const hookInstallData = PaymasterGuardModule.encodeOnInstallData({ + entityId: 0, + paymaster: "paymaster" in paymaster ? paymaster.paymaster : "0x0", // dummy value for paymaster address if it DNE + }); + + const installResult = await provider.installValidation({ + validationConfig: { + moduleAddress: zeroAddress, + entityId: 0, + isGlobal: true, + isSignatureValidation: true, + isUserOpValidation: true, + }, + selectors: [], + installData: "0x", + hooks: [ + { + hookConfig: { + address: getDefaultPaymasterGuardModuleAddress(provider.chain), + entityId: 0, // uint32 + hookType: HookType.VALIDATION, + hasPreHooks: true, + hasPostHooks: true, + }, + initData: hookInstallData, + }, + ], + }); + + // verify hook installtion succeeded + await expect( + provider.waitForUserOperationTransaction(installResult) + ).resolves.not.toThrowError(); + + // sad path: with wrong paymaster + let providerNoPaymaster = await givenConnectedProvider({ + signer, + usePaymaster: false, + }); + + // TO DO: verify if correct paymaster is used + await expect( + providerNoPaymaster.sendUserOperation({ + uo: { + target: target, + value: sendAmount, + data: "0x", + }, + }) + ).rejects.toThrowError(); + + const hookUninstallData = PaymasterGuardModule.encodeOnUninstallData({ + entityId: 0, + }); + + const uninstallResult = await provider.uninstallValidation({ + moduleAddress: zeroAddress, + entityId: 0, + uninstallData: "0x", + hookUninstallDatas: [hookUninstallData], + }); + + // verify uninstall + await expect( + provider.waitForUserOperationTransaction(uninstallResult) + ).resolves.not.toThrowError(); + }); + const givenConnectedProvider = async ({ signer, accountAddress, + usePaymaster = false, }: { signer: SmartAccountSigner; accountAddress?: `0x${string}`; + usePaymaster?: boolean; }) => createSMAV2AccountClient({ chain: instance.chain, signer, accountAddress, transport: custom(instance.getClient()), + ...(usePaymaster ? erc7677Middleware() : {}), }); });