fix all the things

joaquim-verges · joaquim-verges · commit 5de54420574b · 2024-12-06T09:47:01.000+13:00
diff --git a/packages/thirdweb/src/exports/wallets/smart.ts b/packages/thirdweb/src/exports/wallets/smart.ts
@@ -35,4 +35,5 @@ export {
   ENTRYPOINT_ADDRESS_v0_7,
   DEFAULT_ACCOUNT_FACTORY_V0_6,
   DEFAULT_ACCOUNT_FACTORY_V0_7,
+  TokenPaymaster,
 } from "../../wallets/smart/lib/constants.js";
diff --git a/packages/thirdweb/src/wallets/smart/index.ts b/packages/thirdweb/src/wallets/smart/index.ts
@@ -46,6 +46,7 @@ import {
 } from "./lib/calls.js";
 import {
   ENTRYPOINT_ADDRESS_v0_6,
+  ENTRYPOINT_ADDRESS_v0_7,
   getDefaultAccountFactory,
   getEntryPointVersion,
 } from "./lib/constants.js";
@@ -120,6 +121,17 @@ export async function connectSmartWallet(
     }
   }
 
+  if (
+    options.overrides?.tokenPaymaster &&
+    !options.overrides?.entrypointAddress
+  ) {
+    // if token paymaster is set, but no entrypoint address, set the entrypoint address to v0.7
+    options.overrides = {
+      ...options.overrides,
+      entrypointAddress: ENTRYPOINT_ADDRESS_v0_7,
+    };
+  }
+
   const factoryAddress =
     options.factoryAddress ??
     getDefaultAccountFactory(options.overrides?.entrypointAddress);
@@ -200,8 +212,8 @@ export async function disconnectSmartWallet(
 async function createSmartAccount(
   options: SmartAccountOptions,
 ): Promise<Account> {
-  let erc20Paymaster = options.overrides?.tokenPaymaster;
-  if (erc20Paymaster && typeof erc20Paymaster === "string") {
+  const erc20Paymaster = options.overrides?.tokenPaymaster;
+  if (erc20Paymaster) {
     if (
       getEntryPointVersion(
         options.overrides?.entrypointAddress || ENTRYPOINT_ADDRESS_v0_6,
@@ -211,32 +223,6 @@ async function createSmartAccount(
         "Token paymaster is only supported for entrypoint version v0.7",
       );
     }
-    switch (erc20Paymaster) {
-      case "BASE_USDC":
-        erc20Paymaster = {
-          chainId: 8453,
-          paymasterAddress: "0x2222f2738BE6bB7aA0Bfe4AEeAf2908172CF5539",
-          tokenAddress: "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913",
-          balanceStorageSlot: 9,
-        };
-        break;
-      case "CELO_CUSD":
-        erc20Paymaster = {
-          chainId: 42220,
-          paymasterAddress: "0x3feA3c5744D715ff46e91C4e5C9a94426DfF2aF9",
-          tokenAddress: "0x765DE816845861e75A25fCA122bb6898B8B1282a",
-          balanceStorageSlot: 9,
-        };
-        break;
-      case "LISK_LSK":
-        erc20Paymaster = {
-          chainId: 1135,
-          paymasterAddress: "0x9eb8cf7fBa5ed9EeDCC97a0d52254cc0e9B1AC25",
-          tokenAddress: "0xac485391EB2d7D88253a7F1eF18C37f4242D1A24",
-          balanceStorageSlot: 9,
-        };
-        break;
-    }
   }
 
   const { accountContract } = options;
diff --git a/packages/thirdweb/src/wallets/smart/lib/constants.ts b/packages/thirdweb/src/wallets/smart/lib/constants.ts
@@ -1,6 +1,7 @@
 import type { Chain } from "../../../chains/types.js";
 import { getAddress } from "../../../utils/address.js";
 import { getThirdwebDomains } from "../../../utils/domains.js";
+import type { TokenPaymasterConfig } from "../types.js";
 
 export const DUMMY_SIGNATURE =
   "0xfffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c";
@@ -17,6 +18,28 @@ export const ENTRYPOINT_ADDRESS_v0_7 =
 
 export const MANAGED_ACCOUNT_GAS_BUFFER = 50000n;
 
+export type PAYMASTERS = "BASE_USDC" | "CELO_CUSD" | "LISK_LSK";
+export const TokenPaymaster: Record<PAYMASTERS, TokenPaymasterConfig> = {
+  BASE_USDC: {
+    chainId: 8453,
+    paymasterAddress: "0x2222f2738BE6bB7aA0Bfe4AEeAf2908172CF5539",
+    tokenAddress: "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913",
+    balanceStorageSlot: 9n,
+  },
+  CELO_CUSD: {
+    chainId: 42220,
+    paymasterAddress: "0x3feA3c5744D715ff46e91C4e5C9a94426DfF2aF9",
+    tokenAddress: "0x765DE816845861e75A25fCA122bb6898B8B1282a",
+    balanceStorageSlot: 9n,
+  },
+  LISK_LSK: {
+    chainId: 1135,
+    paymasterAddress: "0x9eb8cf7fBa5ed9EeDCC97a0d52254cc0e9B1AC25",
+    tokenAddress: "0xac485391EB2d7D88253a7F1eF18C37f4242D1A24",
+    balanceStorageSlot: 9n,
+  },
+};
+
 /*
  * @internal
  */
diff --git a/packages/thirdweb/src/wallets/smart/lib/userop.ts b/packages/thirdweb/src/wallets/smart/lib/userop.ts
@@ -1,4 +1,5 @@
-import { concat, encodePacked, keccak256, toBytes, toHex } from "viem";
+import { maxUint96 } from "ox/Solidity";
+import { concat, keccak256, toHex } from "viem";
 import type { Chain } from "../../../chains/types.js";
 import type { ThirdwebClient } from "../../../client/client.js";
 import {
@@ -13,6 +14,7 @@ import { encode } from "../../../transaction/actions/encode.js";
 import { toSerializableTransaction } from "../../../transaction/actions/to-serializable-transaction.js";
 import type { PreparedTransaction } from "../../../transaction/prepare-transaction.js";
 import type { TransactionReceipt } from "../../../transaction/types.js";
+import { encodeAbiParameters } from "../../../utils/abi/encodeAbiParameters.js";
 import { isContractDeployed } from "../../../utils/bytecode/is-contract-deployed.js";
 import type { Hex } from "../../../utils/encoding/hex.js";
 import { hexToBytes } from "../../../utils/encoding/to-bytes.js";
@@ -23,7 +25,6 @@ import type {
   BundlerOptions,
   PaymasterResult,
   SmartWalletOptions,
-  TokenPaymasterConfig,
   UserOperationV06,
   UserOperationV07,
 } from "../types.js";
@@ -364,26 +365,21 @@ async function populateUserOp_v0_7(args: {
       // otherwise fallback to bundler for gas limits
       const stateOverrides = overrides?.tokenPaymaster
         ? {
-            [(overrides.tokenPaymaster as TokenPaymasterConfig).tokenAddress]: {
+            [overrides.tokenPaymaster.tokenAddress]: {
               stateDiff: {
                 [keccak256(
-                  encodePacked(
-                    ["address", "uint256"],
+                  encodeAbiParameters(
+                    [{ type: "address" }, { type: "uint256" }],
                     [
                       accountContract.address,
-                      BigInt(
-                        (overrides?.tokenPaymaster as TokenPaymasterConfig)
-                          .balanceStorageSlot,
-                      ),
+                      overrides.tokenPaymaster.balanceStorageSlot,
                     ],
                   ),
-                )]: toBytes(toHex(BigInt(2) * BigInt(96) - BigInt(1)), {
-                  size: 32,
-                }),
+                )]: toHex(maxUint96, { size: 32 }),
               },
             },
           }
-        : {};
+        : undefined;
       const estimates = await estimateUserOpGas(
         {
           userOp: partialOp,
@@ -395,10 +391,10 @@ async function populateUserOp_v0_7(args: {
       partialOp.verificationGasLimit = estimates.verificationGasLimit;
       partialOp.preVerificationGas = estimates.preVerificationGas;
       partialOp.paymasterPostOpGasLimit = overrides?.tokenPaymaster
-        ? 500_000n
-        : paymasterResult.paymasterPostOpGasLimit || 0n;
+        ? 500000n // TODO: estimate this better, needed if there's an extra swap needed in the paymaster
+        : estimates.paymasterPostOpGasLimit || 0n;
       partialOp.paymasterVerificationGasLimit =
-        paymasterResult.paymasterVerificationGasLimit || 0n;
+        estimates.paymasterVerificationGasLimit || 0n;
       // need paymaster to re-sign after estimates
       const paymasterResult2 = (await getPaymasterAndData({
         userOp: partialOp,
diff --git a/packages/thirdweb/src/wallets/smart/smart-wallet-tokenpaymaster.test.ts b/packages/thirdweb/src/wallets/smart/smart-wallet-tokenpaymaster.test.ts
@@ -1,55 +1,162 @@
-import { base } from "src/chains/chain-definitions/base.js";
-import { sendTransaction } from "src/transaction/actions/send-transaction.js";
-import { prepareTransaction } from "src/transaction/prepare-transaction.js";
-import { beforeAll, describe, expect, it } from "vitest";
+import { describe, expect, it } from "vitest";
 import { TEST_CLIENT } from "../../../test/src/test-clients.js";
-import { setThirdwebDomains } from "../../utils/domains.js";
-import type { Account } from "../interfaces/wallet.js";
+import { base } from "../../chains/chain-definitions/base.js";
+import { celo } from "../../chains/chain-definitions/celo.js";
+import { defineChain } from "../../chains/utils.js";
+import { sendTransaction } from "../../transaction/actions/send-transaction.js";
+import { prepareTransaction } from "../../transaction/prepare-transaction.js";
+import type { Address } from "../../utils/address.js";
 import { privateKeyToAccount } from "../private-key.js";
+import { getWalletBalance } from "../utils/getWalletBalance.js";
+import { TokenPaymaster } from "./lib/constants.js";
 import { smartWallet } from "./smart-wallet.js";
 
-let personalAccount: Account;
-
 const client = TEST_CLIENT;
 
 describe.runIf(process.env.TW_SECRET_KEY).skip.sequential(
-  "SmartWallet policy tests",
+  "SmartWallet token paymaster tests",
   {
     retry: 0,
     timeout: 240_000,
   },
   () => {
-    beforeAll(async () => {
-      setThirdwebDomains({
-        rpc: "rpc.thirdweb-dev.com",
-        storage: "storage.thirdweb-dev.com",
-        bundler: "bundler.thirdweb-dev.com",
-      });
-      personalAccount = await privateKeyToAccount({
+    it.skip("should send a transaction with base usdc", async () => {
+      const chain = base;
+      const tokenPaymaster = TokenPaymaster.BASE_USDC;
+      const personalAccount = privateKeyToAccount({
         client,
         privateKey:
           "edf401e8ddbb743f3353b055081cb220ce4c5c04e08da162d86e0dba7c6f0f01", // 0xa470E7c88611364f55B2d7912613e10AF2eA918D
       });
+      const wallet = smartWallet({
+        chain,
+        gasless: true,
+        overrides: {
+          tokenPaymaster,
+        },
+      });
+      const smartAccount = await wallet.connect({
+        client: TEST_CLIENT,
+        personalAccount,
+      });
+      const smartWalletAddress = smartAccount.address as Address;
+      console.log("smartWalletAddress", smartWalletAddress);
+      console.log(
+        "balance before",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
+      const tx = prepareTransaction({
+        client,
+        chain,
+        to: smartAccount.address,
+        value: 0n,
+      });
+      const receipt = await sendTransaction({
+        transaction: tx,
+        account: smartAccount,
+      });
+      expect(receipt.transactionHash).toBeDefined();
+      console.log(
+        "balance after",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
     });
 
-    it("can self transfer with BASE_USDC", async () => {
+    it.skip("should send a transaction with base celo", async () => {
+      const chain = celo;
+      const tokenPaymaster = TokenPaymaster.CELO_CUSD;
+      const personalAccount = privateKeyToAccount({
+        client,
+        privateKey:
+          "edf401e8ddbb743f3353b055081cb220ce4c5c04e08da162d86e0dba7c6f0f01", // 0xa470E7c88611364f55B2d7912613e10AF2eA918D
+      });
       const wallet = smartWallet({
-        chain: base,
+        chain,
         gasless: true,
         overrides: {
-          tokenPaymaster: "BASE_USDC",
+          tokenPaymaster,
         },
       });
       const smartAccount = await wallet.connect({
         client: TEST_CLIENT,
         personalAccount,
       });
+      const smartWalletAddress = smartAccount.address as Address;
+      console.log("smartWalletAddress", smartWalletAddress);
+      console.log(
+        "balance before",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
+      const tx = prepareTransaction({
+        client,
+        chain,
+        to: smartAccount.address,
+        value: 0n,
+      });
+      const receipt = await sendTransaction({
+        transaction: tx,
+        account: smartAccount,
+      });
+      expect(receipt.transactionHash).toBeDefined();
+      console.log(
+        "balance after",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
+    });
 
-      console.log("smartAccount", smartAccount.address);
-
+    it("should send a transaction with base lisk", async () => {
+      const chain = defineChain(1135);
+      const tokenPaymaster = TokenPaymaster.LISK_LSK;
+      const personalAccount = privateKeyToAccount({
+        client,
+        privateKey:
+          "edf401e8ddbb743f3353b055081cb220ce4c5c04e08da162d86e0dba7c6f0f01", // 0xa470E7c88611364f55B2d7912613e10AF2eA918D
+      });
+      const wallet = smartWallet({
+        chain,
+        gasless: true,
+        overrides: {
+          tokenPaymaster,
+        },
+      });
+      const smartAccount = await wallet.connect({
+        client: TEST_CLIENT,
+        personalAccount,
+      });
+      const smartWalletAddress = smartAccount.address as Address;
+      console.log("smartWalletAddress", smartWalletAddress);
+      console.log(
+        "balance before",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
       const tx = prepareTransaction({
         client,
-        chain: base,
+        chain,
         to: smartAccount.address,
         value: 0n,
       });
@@ -58,6 +165,15 @@ describe.runIf(process.env.TW_SECRET_KEY).skip.sequential(
         account: smartAccount,
       });
       expect(receipt.transactionHash).toBeDefined();
+      console.log(
+        "balance after",
+        await getWalletBalance({
+          address: smartAccount.address,
+          chain: chain,
+          client,
+          tokenAddress: tokenPaymaster.tokenAddress,
+        }),
+      );
     });
   },
 );
diff --git a/packages/thirdweb/src/wallets/smart/types.ts b/packages/thirdweb/src/wallets/smart/types.ts
