chore: track walletType usage

arcoraven · arcoraven · commit c8d3c13ce546 · 2024-11-18T12:40:59.000+08:00
diff --git a/src/db/wallets/getWalletDetails.ts b/src/db/wallets/getWalletDetails.ts
@@ -1,3 +1,4 @@
+import LRUMap from "mnemonist/lru-map";
 import { getAddress } from "thirdweb";
 import { z } from "zod";
 import type { PrismaTransaction } from "../../schema/prisma";
@@ -8,7 +9,7 @@ import { getPrismaWithPostgresTx } from "../client";
 
 interface GetWalletDetailsParams {
   pgtx?: PrismaTransaction;
-  address: string;
+  walletAddress: string;
 }
 
 export class WalletDetailsError extends Error {
@@ -130,6 +131,8 @@ export type SmartBackendWalletType = (typeof SmartBackendWalletTypes)[number];
 export type BackendWalletType = (typeof BackendWalletTypes)[number];
 export type ParsedWalletDetails = z.infer<typeof walletDetailsSchema>;
 
+export const walletDetailsCache = new LRUMap<string, ParsedWalletDetails>(2048);
+
 /**
  * Return the wallet details for the given address.
  *
@@ -143,20 +146,28 @@ export type ParsedWalletDetails = z.infer<typeof walletDetailsSchema>;
  */
 export const getWalletDetails = async ({
   pgtx,
-  address,
+  walletAddress: _walletAddress,
 }: GetWalletDetailsParams) => {
+  // Wallet details are stored in lowercase.
+  const walletAddress = _walletAddress.toLowerCase();
+
+  const cachedDetails = walletDetailsCache.get(walletAddress);
+  if (cachedDetails) {
+    return cachedDetails;
+  }
+
   const prisma = getPrismaWithPostgresTx(pgtx);
   const config = await getConfig();
 
   const walletDetails = await prisma.walletDetails.findUnique({
     where: {
-      address: address.toLowerCase(),
+      address: walletAddress,
     },
   });
 
   if (!walletDetails) {
     throw new WalletDetailsError(
-      `No wallet details found for address ${address}`,
+      `No wallet details found for address ${walletAddress}`,
     );
   }
 
@@ -167,7 +178,7 @@ export const getWalletDetails = async ({
   ) {
     if (!walletDetails.awsKmsArn) {
       throw new WalletDetailsError(
-        `AWS KMS ARN is missing for the wallet with address ${address}`,
+        `AWS KMS ARN is missing for the wallet with address ${walletAddress}`,
       );
     }
 
@@ -188,7 +199,7 @@ export const getWalletDetails = async ({
   ) {
     if (!walletDetails.gcpKmsResourcePath) {
       throw new WalletDetailsError(
-        `GCP KMS resource path is missing for the wallet with address ${address}`,
+        `GCP KMS resource path is missing for the wallet with address ${walletAddress}`,
       );
     }
 
@@ -209,14 +220,17 @@ export const getWalletDetails = async ({
 
   // zod schema can validate all necessary fields are populated after decryption
   try {
-    return walletDetailsSchema.parse(walletDetails, {
+    const result = walletDetailsSchema.parse(walletDetails, {
       errorMap: (issue) => {
         const fieldName = issue.path.join(".");
         return {
-          message: `${fieldName} is necessary for wallet ${address} of type ${walletDetails.type}, but not found in wallet details or configuration`,
+          message: `${fieldName} is necessary for wallet ${walletAddress} of type ${walletDetails.type}, but not found in wallet details or configuration`,
         };
       },
     });
+
+    walletDetailsCache.set(walletAddress, result);
+    return result;
   } catch (e) {
     if (e instanceof z.ZodError) {
       throw new WalletDetailsError(
diff --git a/src/server/routes/backend-wallet/signMessage.ts b/src/server/routes/backend-wallet/signMessage.ts
@@ -56,7 +56,7 @@ export async function signMessageRoute(fastify: FastifyInstance) {
       }
 
       const walletDetails = await getWalletDetails({
-        address: walletAddress,
+        walletAddress,
       });
 
       if (isSmartBackendWallet(walletDetails) && !chainId) {
diff --git a/src/server/utils/wallets/getLocalWallet.ts b/src/server/utils/wallets/getLocalWallet.ts
@@ -49,7 +49,9 @@ export const getLocalWallet = async ({
     });
 
     // If that works, save the wallet using the encryption password for the future
-    const walletDetails = await getWalletDetails({ address: walletAddress });
+    const walletDetails = await getWalletDetails({
+      walletAddress,
+    });
 
     logger({
       service: "worker",
@@ -73,7 +75,9 @@ export const getLocalWallet = async ({
 export const getLocalWalletAccount = async (
   walletAddress: Address,
 ): Promise<Account> => {
-  const walletDetails = await getWalletDetails({ address: walletAddress });
+  const walletDetails = await getWalletDetails({
+    walletAddress,
+  });
 
   if (walletDetails.type !== "local") {
     throw new Error(`Local Wallet not found for address ${walletAddress}`);
diff --git a/src/utils/account.ts b/src/utils/account.ts
@@ -39,7 +39,7 @@ export const getAccount = async (args: {
   }
 
   const walletDetails = await getWalletDetails({
-    address: from,
+    walletAddress: from,
   });
 
   const { account } = await walletDetailsToAccount({ walletDetails, chain });
@@ -180,7 +180,7 @@ export const getSmartBackendWalletAdminAccount = async ({
   }
 
   const walletDetails = await getWalletDetails({
-    address: accountAddress,
+    walletAddress: accountAddress,
   });
 
   if (!isSmartBackendWallet(walletDetails)) {
diff --git a/src/utils/cache/clearCache.ts b/src/utils/cache/clearCache.ts
@@ -1,7 +1,9 @@
+import { walletDetailsCache } from "../../db/wallets/getWalletDetails";
 import type { env } from "../env";
 import { accessTokenCache } from "./accessToken";
 import { invalidateConfig } from "./getConfig";
 import { sdkCache } from "./getSdk";
+import { smartWalletsCache } from "./getSmartWalletV5";
 import { walletsCache } from "./getWallet";
 import { webhookCache } from "./getWebhook";
 import { keypairCache } from "./keypair";
@@ -15,4 +17,6 @@ export const clearCache = async (
   walletsCache.clear();
   accessTokenCache.clear();
   keypairCache.clear();
+  smartWalletsCache.clear();
+  walletDetailsCache.clear();
 };
diff --git a/src/utils/cache/getWallet.ts b/src/utils/cache/getWallet.ts
@@ -44,7 +44,7 @@ export const getWallet = async <TWallet extends EVMWallet>({
   try {
     walletDetails = await getWalletDetails({
       pgtx,
-      address: walletAddress,
+      walletAddress,
     });
   } catch (e) {
     if (e instanceof WalletDetailsError) {
diff --git a/src/utils/transaction/insertTransaction.ts b/src/utils/transaction/insertTransaction.ts
@@ -2,9 +2,9 @@ import { StatusCodes } from "http-status-codes";
 import { randomUUID } from "node:crypto";
 import { TransactionDB } from "../../db/transactions/db";
 import {
+  ParsedWalletDetails,
   getWalletDetails,
   isSmartBackendWallet,
-  type ParsedWalletDetails,
 } from "../../db/wallets/getWalletDetails";
 import { doesChainSupportService } from "../../lib/chain/chain-capabilities";
 import { createCustomError } from "../../server/middleware/error";
@@ -43,13 +43,40 @@ export const insertTransaction = async (
     }
   }
 
+  // Get wallet details. For EOA and SBW (v5 endpoints), `from` should return a valid backend wallet.
+  // For SBW (v4 endpoints), `accountAddress` should return a valid backend wallet.
+  // Else the provided details are incorrect (user error).
+  let walletDetails: ParsedWalletDetails | undefined;
+  let isSmartBackendWalletV4 = false;
+  try {
+    walletDetails = await getWalletDetails({
+      walletAddress: insertedTransaction.from,
+    });
+  } catch {}
+  if (!walletDetails && insertedTransaction.accountAddress) {
+    try {
+      walletDetails = await getWalletDetails({
+        walletAddress: insertedTransaction.accountAddress,
+      });
+      isSmartBackendWalletV4 = true;
+    } catch {}
+  }
+  if (!walletDetails) {
+    throw createCustomError(
+      "Account not found",
+      StatusCodes.BAD_REQUEST,
+      "ACCOUNT_NOT_FOUND",
+    );
+  }
+
   let queuedTransaction: QueuedTransaction = {
     ...insertedTransaction,
     status: "queued",
     queueId,
     queuedAt: new Date(),
     resendCount: 0,
 
+    walletType: walletDetails.type,
     from: getChecksumAddress(insertedTransaction.from),
     to: getChecksumAddress(insertedTransaction.to),
     signerAddress: getChecksumAddress(insertedTransaction.signerAddress),
@@ -60,37 +87,34 @@ export const insertTransaction = async (
     value: insertedTransaction.value ?? 0n,
   };
 
-  let walletDetails: ParsedWalletDetails | undefined;
+  // Handle smart backend wallets details.
+  if (isSmartBackendWallet(walletDetails)) {
+    if (
+      !(await doesChainSupportService(
+        queuedTransaction.chainId,
+        "account-abstraction",
+      ))
+    ) {
+      throw createCustomError(
+        `Smart backend wallets do not support chain ${queuedTransaction.chainId}.`,
+        StatusCodes.BAD_REQUEST,
+        "INVALID_SMART_BACKEND_WALLET_TRANSACTION",
+      );
+    }
 
-  try {
-    walletDetails = await getWalletDetails({
-      address: queuedTransaction.from,
-    });
+    queuedTransaction = {
+      ...queuedTransaction,
+      accountFactoryAddress: walletDetails.accountFactoryAddress ?? undefined,
+      entrypointAddress: walletDetails.entrypointAddress ?? undefined,
+    };
 
-    // when using the v5 SDK with smart backend wallets, the following values are not set correctly:
-    // isUserOp is set to false
-    // account address is blank or the user provided value (this should be the SBW account address)
-    // from is set to the SBW account address (this should be the SBW signer address)
-    // these values need to be corrected so the worker can process the transaction
-    if (isSmartBackendWallet(walletDetails)) {
+    if (!isSmartBackendWalletV4) {
       if (queuedTransaction.accountAddress) {
+        // Disallow smart backend wallets from sending userOps.
         throw createCustomError(
-          "Smart backend wallets do not support interacting with other smart accounts",
+          "Smart backend wallets do not support sending transactions with other smart accounts",
           StatusCodes.BAD_REQUEST,
-          "INVALID_SMART_BACKEND_WALLET_INTERACTION",
-        );
-      }
-
-      if (
-        !(await doesChainSupportService(
-          queuedTransaction.chainId,
-          "account-abstraction",
-        ))
-      ) {
-        throw createCustomError(
-          "Chain does not support smart backend wallets",
-          StatusCodes.BAD_REQUEST,
-          "SBW_CHAIN_NOT_SUPPORTED",
+          "INVALID_SMART_BACKEND_WALLET_TRANSACTION",
         );
       }
 
@@ -101,52 +125,8 @@ export const insertTransaction = async (
         from: walletDetails.accountSignerAddress,
         accountAddress: queuedTransaction.from,
         target: queuedTransaction.to,
-        accountFactoryAddress: walletDetails.accountFactoryAddress ?? undefined,
-        entrypointAddress: walletDetails.entrypointAddress ?? undefined,
       };
     }
-  } catch {
-    // if wallet details are not found, this is a smart backend wallet using a v4 endpoint
-  }
-
-  if (!walletDetails && queuedTransaction.accountAddress) {
-    try {
-      walletDetails = await getWalletDetails({
-        address: queuedTransaction.accountAddress,
-      });
-
-      // when using v4 SDK with smart backend wallets, the following values are not set correctly:
-      // entrypointAddress is not set
-      // accountFactoryAddress is not set
-      if (walletDetails && isSmartBackendWallet(walletDetails)) {
-        if (
-          !(await doesChainSupportService(
-            queuedTransaction.chainId,
-            "account-abstraction",
-          ))
-        ) {
-          throw createCustomError(
-            "Chain does not support smart backend wallets",
-            StatusCodes.BAD_REQUEST,
-            "SBW_CHAIN_NOT_SUPPORTED",
-          );
-        }
-
-        queuedTransaction = {
-          ...queuedTransaction,
-          entrypointAddress: walletDetails.entrypointAddress ?? undefined,
-          accountFactoryAddress:
-            walletDetails.accountFactoryAddress ?? undefined,
-        };
-      }
-    } catch {
-      // if wallet details are not found for this either, this backend wallet does not exist at all
-      throw createCustomError(
-        "Account not found",
-        StatusCodes.BAD_REQUEST,
-        "ACCOUNT_NOT_FOUND",
-      );
-    }
   }
 
   // Simulate the transaction.
diff --git a/src/utils/transaction/types.ts b/src/utils/transaction/types.ts
@@ -1,5 +1,6 @@
 import type { Address, Hex, toSerializableTransaction } from "thirdweb";
 import type { TransactionType } from "viem";
+import { BackendWalletType } from "../../db/wallets/getWalletDetails";
 
 // TODO: Replace with thirdweb SDK exported type when available.
 export type PopulatedTransaction = Awaited<
@@ -52,6 +53,7 @@ export type InsertedTransaction = {
 export type QueuedTransaction = InsertedTransaction & {
   status: "queued";
 
+  walletType: BackendWalletType;
   resendCount: number;
   queueId: string;
   queuedAt: Date;
diff --git a/src/utils/usage.ts b/src/utils/usage.ts
@@ -2,6 +2,7 @@ import { Static } from "@sinclair/typebox";
 import { UsageEvent } from "@thirdweb-dev/service-utils/cf-worker";
 import { FastifyInstance } from "fastify";
 import { Address, Hex } from "thirdweb";
+import { BackendWalletType } from "../db/wallets/getWalletDetails";
 import { ADMIN_QUEUES_BASEPATH } from "../server/middleware/adminRoutes";
 import { OPENAPI_ROUTES } from "../server/middleware/open-api";
 import { contractParamSchema } from "../server/schemas/sharedApiSchemas";
@@ -22,6 +23,7 @@ export interface ReportUsageParams {
     | "api_request";
   input: {
     chainId?: number;
+    walletType?: BackendWalletType;
     from?: Address;
     to?: Address;
     value?: bigint;
@@ -113,6 +115,7 @@ export const reportUsage = (usageEvents: ReportUsageParams[]) => {
         action,
         clientId: thirdwebClientId,
         chainId: input.chainId,
+        walletType: input.walletType,
         walletAddress: input.from,
         contractAddress: input.to,
         transactionValue: input.value?.toString(),
@@ -136,7 +139,7 @@ export const reportUsage = (usageEvents: ReportUsageParams[]) => {
       logger({
         service: "worker",
         level: "error",
-        message: `Error:`,
+        message: "Error reporting usage event:",
         error: e,
       });
     }

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,7 @@ export async function signMessageRoute(fastify: FastifyInstance) {`
`56`	`56`	`}`
`57`	`57`
`58`	`58`	`const walletDetails = await getWalletDetails({`
`59`		`- address: walletAddress,`
	`59`	`+ walletAddress,`
`60`	`60`	`});`
`61`	`61`
`62`	`62`	`if (isSmartBackendWallet(walletDetails) && !chainId) {`
Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@ export const getAccount = async (args: {`
`39`	`39`	`}`
`40`	`40`
`41`	`41`	`const walletDetails = await getWalletDetails({`
`42`		`- address: from,`
	`42`	`+ walletAddress: from,`
`43`	`43`	`});`
`44`	`44`
`45`	`45`	`const { account } = await walletDetailsToAccount({ walletDetails, chain });`
`@@ -180,7 +180,7 @@ export const getSmartBackendWalletAdminAccount = async ({`
`180`	`180`	`}`
`181`	`181`
`182`	`182`	`const walletDetails = await getWalletDetails({`
`183`		`- address: accountAddress,`
	`183`	`+ walletAddress: accountAddress,`
`184`	`184`	`});`
`185`	`185`
`186`	`186`	`if (!isSmartBackendWallet(walletDetails)) {`