From c9f299024d0d5c42c9032e9d03062667cdad619c Mon Sep 17 00:00:00 2001
From: zer0dot <zer0dot.dev@gmail.com>
Date: Tue, 17 Dec 2024 19:14:14 -0500
Subject: [PATCH] feat: (WIP) native token limit module support

---
 .../src/ma-v2/client/client.test.ts           | 100 ++++++++++++++++++
 .../native-token-limit-module/module.ts       |  34 ++++++
 2 files changed, 134 insertions(+)
 create mode 100644 account-kit/smart-contracts/src/ma-v2/modules/native-token-limit-module/module.ts

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 4323175511..646bfbcf11 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
@@ -12,6 +12,7 @@ import { getDefaultSingleSignerValidationModuleAddress } from "../modules/utils.
 import { SingleSignerValidationModule } from "../modules/single-signer-validation/module.js";
 import { allowlistModule } from "../modules/allowlist-module/module.js";
 import { HookType } from "../actions/common/types.js";
+import { nativeTokenLimitModule } from "../modules/native-token-limit-module/module.js";
 
 describe("MA v2 Tests", async () => {
   const instance = local070Instance;
@@ -263,6 +264,105 @@ describe("MA v2 Tests", async () => {
     ).resolves.not.toThrowError();
   });
 
+  it.only("installs native token limit module, then uninstalls", async () => {
+    let provider = (await givenConnectedProvider({ signer })).extend(
+      installValidationActions
+    );
+
+    await setBalance(client, {
+      address: provider.getAddress(),
+      value: parseEther("2"),
+    });
+
+    const spendLimit = parseEther("0.1"); // 0.1 ETH limit
+    const hookInstallData = nativeTokenLimitModule.encodeOnInstallData({
+      entityId: 1,
+      spendLimit,
+    });
+
+    const installResult = await provider.installValidation({
+      validationConfig: {
+        // TODO: Make this a constant "FALLBACK_VALIDATION_CONFIG or something"
+        moduleAddress: zeroAddress,
+        entityId: 0,
+        isGlobal: true,
+        isSignatureValidation: true,
+        isUserOpValidation: true,
+      },
+      selectors: [],
+      installData: "0x",
+      hooks: [
+        {
+          hookConfig: {
+            address: addresses.nativeTokenLimitModule,
+            entityId: 0,
+            hookType: HookType.VALIDATION,
+            hasPreHooks: true,
+            hasPostHooks: false,
+          },
+          initData: hookInstallData,
+        },
+      ],
+    });
+
+    await expect(
+      provider.waitForUserOperationTransaction(installResult)
+    ).resolves.not.toThrowError();
+
+    // Try to send more than the limit - should fail
+    await expect(
+      provider.sendUserOperation({
+        uo: {
+          target: target,
+          value: parseEther("0.2"), // Try to send 0.2 ETH, above the 0.1 limit
+          data: "0x",
+        },
+      })
+    ).rejects.toThrowError();
+
+    // Try to send less than the limit - should succeed
+    const sendResult = await provider.sendUserOperation({
+      uo: {
+        target: target,
+        value: parseEther("0.05"), // Send 0.05 ETH, below the 0.1 limit
+        data: "0x",
+      },
+    });
+
+    // await expect(
+    //   provider.waitForUserOperationTransaction(sendResult)
+    // ).resolves.not.toThrowError();
+
+    // // Now uninstall the module
+    // const hookUninstallData = nativeTokenLimitModule.encodeOnUninstallData({
+    //   entityId: 0,
+    // });
+
+    // const uninstallResult = await provider.uninstallValidation({
+    //   moduleAddress: zeroAddress,
+    //   entityId: 0,
+    //   uninstallData: "0x",
+    //   hookUninstallDatas: [hookUninstallData],
+    // });
+
+    // await expect(
+    //   provider.waitForUserOperationTransaction(uninstallResult)
+    // ).resolves.not.toThrowError();
+
+    // // After uninstall, should be able to send more than the previous limit
+    // const finalSendResult = await provider.sendUserOperation({
+    //   uo: {
+    //     target: target,
+    //     value: parseEther("0.2"), // Now we can send 0.2 ETH
+    //     data: "0x",
+    //   },
+    // });
+
+    // await expect(
+    //   provider.waitForUserOperationTransaction(finalSendResult)
+    // ).resolves.not.toThrowError();
+  });
+
   const givenConnectedProvider = async ({
     signer,
     accountAddress,
diff --git a/account-kit/smart-contracts/src/ma-v2/modules/native-token-limit-module/module.ts b/account-kit/smart-contracts/src/ma-v2/modules/native-token-limit-module/module.ts
new file mode 100644
index 0000000000..0d4ba6bb83
--- /dev/null
+++ b/account-kit/smart-contracts/src/ma-v2/modules/native-token-limit-module/module.ts
@@ -0,0 +1,34 @@
+import { encodeAbiParameters, type Address, type Hex } from "viem";
+
+import { nativeTokenLimitModuleAbi } from "./abis/nativeTokenLimitModuleAbi.js";
+
+const addresses: Record<number | "default", Address> = {
+  default: "0xEa6a05306315196f2A7CA2ec7eEA45290bae00A0",
+};
+
+const meta = {
+  name: "NativeTokenLimitModule",
+  version: "alpha.1",
+  addresses,
+};
+
+export const nativeTokenLimitModule = {
+  meta,
+  abi: nativeTokenLimitModuleAbi,
+
+  encodeOnInstallData: (args: {
+    entityId: number;
+    spendLimit: bigint;
+  }): Hex => {
+    const { entityId, spendLimit } = args;
+    return encodeAbiParameters(
+      [{ type: "uint32" }, { type: "uint256" }],
+      [entityId, spendLimit]
+    );
+  },
+
+  encodeOnUninstallData: (args: { entityId: number }): Hex => {
+    const { entityId } = args;
+    return encodeAbiParameters([{ type: "uint32" }], [entityId]);
+  },
+};