Skip to content

Commit

Permalink
feat: Implement Allowlist Module And Native Token Limit Module (#1232)
Browse files Browse the repository at this point in the history 
* feat: allowlist module installation and uninstallation

* feat: Add Native Token Limit Module Support (#1233)

* feat: (WIP) native token limit module support

* feat: adapt native token limit module to new arch and add tests

* chore: autogenerated file

* chore: cleanup

* chore: remove only test specifier

* test: test allowlist module usage

* test: flesh out native token limit module test

* chore: rebase fixes

* fix: await useroperation results correctly

* chore: remove redundant post hook

* feat: allow hooks on the fallback validation

* fix: fix error name mismatch

* fix: remove unnecessary await wrapping

* chore: removed unused module components, cleanup

* chore: rearrange tests to handle polluted end states
  • Loading branch information
@Zer0dot @howydev
Zer0dot authored and howydev committed Jan 2, 2025
1 parent f7e0e0b commit 4668c9c
Show file tree
Hide file tree
Showing 6 changed files with 342 additions and 19 deletions.
2 changes: 1 addition & 1 deletion aa-sdk/core/src/errors/client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ 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";
override name = "EntityIdOverrideError";

/**
* Initializes a new instance of the error message with a default message indicating that the nonce key is invalid.
Expand Down
23 changes: 16 additions & 7 deletions account-kit/smart-contracts/src/ma-v2/actions/install-validation/installValidation.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,19 @@ import {
AccountNotFoundError,
IncompatibleClientError,
isSmartAccountClient,
// EntityIdOverrideError,
EntityIdOverrideError,
type GetEntryPointFromAccount,
type SendUserOperationResult,
type UserOperationOverridesParameter,
type SmartAccountSigner,
} from "@aa-sdk/core";
import { type Address, type Hex, encodeFunctionData, concatHex } from "viem";
import {
type Address,
type Hex,
encodeFunctionData,
concatHex,
zeroAddress,
} from "viem";

import { semiModularAccountBytecodeAbi } from "../../abis/semiModularAccountBytecodeAbi.js";
import type { HookConfig, ValidationConfig } from "../common/types.js";
Expand All @@ -20,7 +26,7 @@ import {

import { type SMAV2AccountClient } from "../../client/client.js";
import { type SMAV2Account } from "../../account/semiModularAccountV2.js";
// import { DEFAULT_OWNER_ENTITY_ID } from "../../utils.js";
import { DEFAULT_OWNER_ENTITY_ID } from "../../utils.js";

export type InstallValidationParams<
TSigner extends SmartAccountSigner = SmartAccountSigner
Expand Down Expand Up @@ -85,10 +91,13 @@ export const installValidationActions: <
);
}

// TO DO: handle installing on fallback validation (entityId == 0) with non-zero address
// if (validationConfig.entityId === DEFAULT_OWNER_ENTITY_ID) {
// throw new EntityIdOverrideError();
// }
// an entityId of zero is only allowed if we're installing or uninstalling hooks on the fallback validation
if (
validationConfig.entityId === DEFAULT_OWNER_ENTITY_ID &&
validationConfig.moduleAddress !== zeroAddress
) {
throw new EntityIdOverrideError();
}

const { encodeCallData } = account;

Expand Down
225 changes: 216 additions & 9 deletions account-kit/smart-contracts/src/ma-v2/client/client.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,19 @@ import {
getDefaultPaymasterGuardModuleAddress,
getDefaultSingleSignerValidationModuleAddress,
getDefaultTimeRangeModuleAddress,
getDefaultAllowlistModuleAddress,
getDefaultNativeTokenLimitModuleAddress,
} 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";
import { TimeRangeModule } from "../modules/time-range-module/module.js";
import { allowlistModule } from "../modules/allowlist-module/module.js";
import { nativeTokenLimitModule } from "../modules/native-token-limit-module/module.js";

// TODO: Include a snapshot to reset to in afterEach.
describe("MA v2 Tests", async () => {
const instance = local070Instance;
let client: ReturnType<typeof instance.getClient> &
Expand Down Expand Up @@ -98,8 +103,7 @@ describe("MA v2 Tests", async () => {
hooks: [],
});

let txnHash = provider.waitForUserOperationTransaction(result);
await expect(txnHash).resolves.not.toThrowError();
await provider.waitForUserOperationTransaction(result);

const startingAddressBalance = await getTargetBalance();

Expand All @@ -120,8 +124,8 @@ describe("MA v2 Tests", async () => {
},
});

txnHash = sessionKeyClient.waitForUserOperationTransaction(result);
await expect(txnHash).resolves.not.toThrowError();
await sessionKeyClient.waitForUserOperationTransaction(result);

await expect(getTargetBalance()).resolves.toEqual(
startingAddressBalance + sendAmount
);
Expand Down Expand Up @@ -155,8 +159,7 @@ describe("MA v2 Tests", async () => {
hooks: [],
});

let txnHash = provider.waitForUserOperationTransaction(result);
await expect(txnHash).resolves.not.toThrowError();
await provider.waitForUserOperationTransaction(result);

result = await provider.uninstallValidation({
moduleAddress: getDefaultSingleSignerValidationModuleAddress(
Expand All @@ -169,8 +172,7 @@ describe("MA v2 Tests", async () => {
hookUninstallDatas: [],
});

txnHash = provider.waitForUserOperationTransaction(result);
await expect(txnHash).resolves.not.toThrowError();
await provider.waitForUserOperationTransaction(result);

// connect session key and send tx with session key
let sessionKeyClient = await createSMAV2AccountClient({
Expand Down Expand Up @@ -349,6 +351,211 @@ describe("MA v2 Tests", async () => {
).resolves.not.toThrowError();
});

it("installs allowlist module, uses, then uninstalls", async () => {
let provider = (await givenConnectedProvider({ signer })).extend(
installValidationActions
);

await setBalance(client, {
address: provider.getAddress(),
value: parseEther("2"),
});

const hookInstallData = allowlistModule.encodeOnInstallData({
entityId: 0,
inputs: [
{
target,
hasSelectorAllowlist: false,
hasERC20SpendLimit: false,
erc20SpendLimit: 0n,
selectors: [],
},
],
});

const installResult = await provider.installValidation({
validationConfig: {
moduleAddress: zeroAddress,
entityId: 0,
isGlobal: true,
isSignatureValidation: true,
isUserOpValidation: true,
},
selectors: [],
installData: "0x",
hooks: [
{
hookConfig: {
address: getDefaultAllowlistModuleAddress(provider.chain),
entityId: 0, // uint32
hookType: HookType.VALIDATION,
hasPreHooks: true,
hasPostHooks: false,
},
initData: hookInstallData,
},
],
});

await provider.waitForUserOperationTransaction(installResult);

// Test that the allowlist is active.
// We should *only* be able to call into the target address, as it's the only address we passed to onInstall.
const sendResult = await provider.sendUserOperation({
uo: {
target: target,
value: 0n,
data: "0x",
},
});

await provider.waitForUserOperationTransaction(sendResult);

// This should revert as we're calling an address separate fom the allowlisted target.
await expect(
provider.sendUserOperation({
uo: {
target: zeroAddress,
value: 0n,
data: "0x",
},
})
).rejects.toThrowError();

const hookUninstallData = allowlistModule.encodeOnUninstallData({
entityId: 0,
inputs: [
{
target,
hasSelectorAllowlist: false,
hasERC20SpendLimit: false,
erc20SpendLimit: 0n,
selectors: [],
},
],
});

const uninstallResult = await provider.uninstallValidation({
moduleAddress: zeroAddress,
entityId: 0,
uninstallData: "0x",
hookUninstallDatas: [hookUninstallData],
});

await provider.waitForUserOperationTransaction(uninstallResult);

// Post-uninstallation, we should now be able to call into any address successfully.
const postUninstallSendResult = await provider.sendUserOperation({
uo: {
target: zeroAddress,
value: 0n,
data: "0x",
},
});

await provider.waitForUserOperationTransaction(postUninstallSendResult);
});

it("installs native token limit module, uses, then uninstalls", async () => {
let provider = (await givenConnectedProvider({ signer })).extend(
installValidationActions
);

await setBalance(client, {
address: provider.getAddress(),
value: parseEther("2"),
});

const spendLimit = parseEther("0.5");

// Let's verify the module's limit is set correctly after installation
const hookInstallData = nativeTokenLimitModule.encodeOnInstallData({
entityId: 0,
spendLimit,
});

const installResult = await provider.installValidation({
validationConfig: {
moduleAddress: zeroAddress,
entityId: 0,
isGlobal: true,
isSignatureValidation: true,
isUserOpValidation: true,
},
selectors: [],
installData: "0x",
hooks: [
{
hookConfig: {
address: getDefaultNativeTokenLimitModuleAddress(provider.chain),
entityId: 0,
hookType: HookType.VALIDATION,
hasPreHooks: true,
hasPostHooks: false,
},
initData: hookInstallData,
},
{
hookConfig: {
address: getDefaultNativeTokenLimitModuleAddress(provider.chain),
entityId: 0,
hookType: HookType.EXECUTION,
hasPreHooks: true,
hasPostHooks: false,
},
initData: "0x",
},
],
});

await provider.waitForUserOperationTransaction(installResult);

// Try to send less than the limit - should pass
const passingSendResult = await provider.sendUserOperation({
uo: {
target: target,
value: parseEther("0.05"), // below the 0.5 limit
data: "0x",
},
});
await provider.waitForUserOperationTransaction(passingSendResult);

// Try to send more than the limit - should fail
await expect(
provider.sendUserOperation({
uo: {
target: target,
value: parseEther("0.6"), // passing the 0.5 limit
data: "0x",
},
})
).rejects.toThrowError();

const hookUninstallData = nativeTokenLimitModule.encodeOnUninstallData({
entityId: 0,
});

const uninstallResult = await provider.uninstallValidation({
moduleAddress: zeroAddress,
entityId: 0,
uninstallData: "0x",
hookUninstallDatas: [hookUninstallData, "0x"],
});

await provider.waitForUserOperationTransaction(uninstallResult);

// Sending over the limit should now pass
const postUninstallSendResult = await provider.sendUserOperation({
uo: {
target: target,
value: parseEther("0.6"),
data: "0x",
},
});
await provider.waitForUserOperationTransaction(postUninstallSendResult);
});

it("installs time range module, then uninstalls module within valid time range", async () => {
let provider = (
await givenConnectedProvider({
Expand Down Expand Up @@ -453,7 +660,7 @@ describe("MA v2 Tests", async () => {
],
});

// verify hook installtion succeeded
// verify hook installation succeeded
await provider.waitForUserOperationTransaction(installResult);

const hookUninstallData = TimeRangeModule.encodeOnUninstallData({
Expand Down
Loading

0 comments on commit 4668c9c

Please sign in to comment.