add-enclave-functionality

Author: dirtycajunrice

src/server/routes/contract/write/write.ts

@@ -16,12 +16,12 @@ import {
 import { txOverridesWithValueSchema } from "../../../schemas/tx-overrides";
 import {
   maybeAddress,
-  requiredAddress,
-  walletWithAAHeaderSchema,
+  type walletWithAAHeaderSchema, walletWithAAOrEnclaveHeaderSchema,
 } from "../../../schemas/wallet";
 import { sanitizeAbi, sanitizeFunctionName } from "../../../utils/abi";
 import { getChainIdFromChain } from "../../../utils/chain";
 import { parseTransactionOverrides } from "../../../utils/transaction-overrides";
+import { parseEnclaveHeaders } from "../../../utils/convertor";
 
 // INPUT
 const writeRequestBodySchema = Type.Object({
@@ -60,7 +60,7 @@ export async function writeToContract(fastify: FastifyInstance) {
       tags: ["Contract"],
       operationId: "write",
       params: contractParamSchema,
-      headers: walletWithAAHeaderSchema,
+      headers: walletWithAAOrEnclaveHeaderSchema,
       querystring: requestQuerystringSchema,
       body: writeRequestBodySchema,
       response: {
@@ -72,12 +72,13 @@ export async function writeToContract(fastify: FastifyInstance) {
       const { simulateTx } = request.query;
       const { functionName, args, txOverrides, abi } = request.body;
       const {
-        "x-backend-wallet-address": fromAddress,
+        "x-backend-wallet-address": backendWalletAddress,
         "x-account-address": accountAddress,
         "x-idempotency-key": idempotencyKey,
         "x-account-factory-address": accountFactoryAddress,
         "x-account-salt": accountSalt,
       } = request.headers as Static<typeof walletWithAAHeaderSchema>;
+      const enclave = parseEnclaveHeaders(request.headers);
 
       const chainId = await getChainIdFromChain(chain);
       const contract = await getContractV5({
@@ -118,14 +119,15 @@ export async function writeToContract(fastify: FastifyInstance) {
       const queueId = await queueTransaction({
         functionName,
         transaction,
-        fromAddress: requiredAddress(fromAddress, "x-backend-wallet-address"),
+        fromAddress: maybeAddress(backendWalletAddress, "x-backend-wallet-address"),
         toAddress: maybeAddress(contractAddress, "to"),
         accountAddress: maybeAddress(accountAddress, "x-account-address"),
         accountFactoryAddress: maybeAddress(
           accountFactoryAddress,
           "x-account-factory-address",
         ),
         accountSalt,
+        enclave,
         txOverrides,
         idempotencyKey,
         shouldSimulate: simulateTx,

src/server/schemas/wallet/index.ts

@@ -5,18 +5,48 @@ import { badAddressError } from "../../middleware/error";
 import { AddressSchema } from "../address";
 import { chainIdOrSlugSchema } from "../chain";
 
-export const walletHeaderSchema = Type.Object({
-  "x-backend-wallet-address": {
-    ...AddressSchema,
-    description: "Backend wallet address",
-  },
+export const idempotentHeaderSchema = Type.Object({
   "x-idempotency-key": Type.Optional(
     Type.String({
       maxLength: 200,
       description: `Transactions submitted with the same idempotency key will be de-duplicated. Only the last ${env.TRANSACTION_HISTORY_COUNT} transactions are compared.`,
     }),
   ),
 });
+export const enclaveWalletHeaderSchema = Type.Object({
+  "x-enclave-wallet-auth-token": Type.Optional(
+    Type.String({
+      description:
+        "Auth token of an enclave wallet. mutually exclusive with other wallet headers",
+    }),
+  ),
+  "x-client-id": Type.Optional(
+    Type.String({
+      description:
+        "Client id of an enclave wallet. mutually exclusive with other wallet headers",
+    }),
+  ),
+  "x-ecosystem-id": Type.Optional(
+    Type.RegExp(/^ecosystem\.[a-zA-Z0-9_-]+]$/, {
+      description:
+        "Ecosystem id of an enclave wallet. mutually exclusive with other wallet headers",
+    }),
+  ),
+  "x-ecosystem-partner-id": Type.Optional(
+    Type.String({
+      description:
+        "Ecosystem partner id of an enclave wallet. mutually exclusive with other wallet headers",
+    }),
+  )
+});
+
+export const walletHeaderSchema = Type.Object({
+  "x-backend-wallet-address": {
+    ...AddressSchema,
+    description: "Backend wallet address",
+  },
+  ...idempotentHeaderSchema.properties,
+});
 
 export const walletWithAAHeaderSchema = Type.Object({
   ...walletHeaderSchema.properties,
@@ -39,6 +69,11 @@ export const walletWithAAHeaderSchema = Type.Object({
   ),
 });
 
+export const walletWithAAOrEnclaveHeaderSchema = Type.Object({
+  ...walletWithAAHeaderSchema.properties,
+  ...enclaveWalletHeaderSchema.properties,
+});
+
 /**
  * Helper function to parse an address string.
  * Returns undefined if the address is undefined.

src/server/utils/convertor.ts

@@ -1,4 +1,10 @@
 import { BigNumber } from "ethers";
+import type { FastifyRequest} from "fastify"
+import type { Static } from "@sinclair/typebox";
+import type { enclaveWalletHeaderSchema, } from "../schemas/wallet";
+import type { EnclaveWalletParams } from "../../shared/utils/cache/get-enclave-wallet";
+import type { Ecosystem } from "thirdweb/dist/types/wallets/in-app/core/wallet/types";
+import type { EcosystemWalletId } from "thirdweb/dist/types/wallets/wallet-types";
 
 const isHexBigNumber = (value: unknown) => {
   const isNonNullObject = typeof value === "object" && value !== null;
@@ -17,3 +23,21 @@ export const bigNumberReplacer = (value: unknown): unknown => {
 
   return value;
 };
+
+export const parseEnclaveHeaders = (headers: FastifyRequest['headers']): EnclaveWalletParams => {
+  const {
+    "x-enclave-wallet-auth-token": authToken = "",
+    "x-client-id": clientId = "",
+    "x-ecosystem-id": id,
+    "x-ecosystem-partner-id": partnerId,
+  } = headers as Static<typeof enclaveWalletHeaderSchema>;
+  let ecosystem: Ecosystem | undefined;
+  if (id) {
+    ecosystem = { id: (id as EcosystemWalletId), partnerId }
+  }
+  return {
+    authToken,
+    clientId,
+    ecosystem,
+  }
+}

src/shared/utils/account.ts

@@ -18,17 +18,23 @@ import {
 import { getSmartWalletV5 } from "./cache/get-smart-wallet-v5";
 import { getChain } from "./chain";
 import { thirdwebClient } from "./sdk";
+import { type EnclaveWalletParams, getEnclaveWalletAccount } from "./cache/get-enclave-wallet";
 
 export const _accountsCache = new LRUMap<string, Account>(2048);
 
-export const getAccount = async (args: {
+export type GetAccountArgs = {
   chainId: number;
-  from: Address;
+  from: Address | undefined;
   accountAddress?: Address;
-}): Promise<Account> => {
-  const { chainId, from, accountAddress } = args;
+  enclave?: EnclaveWalletParams
+}
+
+export const getAccount = async (args: GetAccountArgs): Promise<Account> => {
+  const { chainId, from, accountAddress, enclave } = args;
   const chain = await getChain(chainId);
 
+  if (enclave) return getEnclaveWalletAccount(enclave);
+  if (!from) throw new Error("from is required");
   if (accountAddress) return getSmartWalletV5({ chain, accountAddress, from });
 
   // Get from cache.
@@ -48,9 +54,9 @@ export const getAccount = async (args: {
 };
 
 export const walletDetailsToAccount = async ({
-  walletDetails,
-  chain,
-}: {
+                                               walletDetails,
+                                               chain,
+                                             }: {
   walletDetails: ParsedWalletDetails;
   chain: Chain;
 }) => {
@@ -164,9 +170,9 @@ export const _adminAccountsCache = new LRUMap<string, Account>(2048);
  * Will throw if the wallet is not a smart backend wallet
  */
 export const getSmartBackendWalletAdminAccount = async ({
-  chainId,
-  accountAddress,
-}: {
+                                                          chainId,
+                                                          accountAddress,
+                                                        }: {
   chainId: number;
   accountAddress: Address;
 }) => {

src/shared/utils/cache/get-enclave-wallet.ts

@@ -0,0 +1,61 @@
+import { thirdwebClient as client } from "../sdk";
+import { EnclaveWallet } from "thirdweb/dist/types/wallets/in-app/core/wallet/enclave-wallet";
+import { getUserStatus } from "thirdweb/dist/types/wallets/in-app/core/actions/get-enclave-user-status";
+import type { Ecosystem } from "thirdweb/dist/types/wallets/in-app/core/wallet/types";
+import { ClientScopedStorage } from "thirdweb/dist/types/wallets/in-app/core/authentication/client-scoped-storage";
+import { MemoryStorage } from "./memory-storage";
+import type { Address } from "thirdweb";
+
+export interface EnclaveWalletParams {
+  authToken: string;
+  clientId: string;
+  ecosystem?: Ecosystem;
+}
+
+const getEnclaveUserWallet = async ({ authToken, ecosystem }: EnclaveWalletParams) => {
+  const user = await getUserStatus({ authToken, client, ecosystem });
+  if (!user) {
+    throw new Error("Cannot initialize wallet, no user logged in");
+  }
+  if (user.wallets.length === 0) {
+    throw new Error(
+      "Cannot initialize wallet, this user does not have a wallet generated yet",
+    );
+  }
+  const wallet = user.wallets[0];
+  if (wallet.type !== "enclave") {
+    throw new Error(
+      "Cannot initialize wallet, this user does not have an enclave wallet",
+    );
+  }
+  return wallet;
+}
+export const getEnclaveWallet = async (
+  {
+    authToken,
+    clientId,
+    ecosystem,
+  }: EnclaveWalletParams,
+) => {
+  const wallet = await getEnclaveUserWallet({ authToken, clientId, ecosystem });
+  return new EnclaveWallet({
+    client,
+    ecosystem,
+    address: wallet.address,
+    storage: new ClientScopedStorage({
+      storage: new MemoryStorage(),
+      clientId,
+      ecosystem,
+    }),
+  });
+};
+
+export const getEnclaveWalletAccount = async (params: EnclaveWalletParams) => {
+  const wallet = await getEnclaveWallet(params);
+  return wallet.getAccount();
+};
+
+export const getEnclaveWalletAddress = async (params: EnclaveWalletParams) => {
+  const wallet = await getEnclaveUserWallet(params)
+  return wallet.address as Address;
+};

src/shared/utils/cache/memory-storage.ts

@@ -0,0 +1,17 @@
+import type { AsyncStorage } from "@thirdweb-dev/wallets";
+
+export class MemoryStorage implements AsyncStorage {
+  data: Map<string, string> = new Map();
+
+  async getItem(key: string) {
+    return this.data.get(key) || null;
+  }
+
+  async setItem(key: string, value: string) {
+    this.data.set(key, value);
+  }
+
+  async removeItem(key: string) {
+    this.data.delete(key);
+  }
+}

src/shared/utils/transaction/insert-transaction.ts

@@ -1,11 +1,11 @@
 import { StatusCodes } from "http-status-codes";
 import { randomUUID } from "node:crypto";
-import { TransactionDB } from "../../../shared/db/transactions/db";
+import { TransactionDB } from "../../db/transactions/db";
 import {
   getWalletDetails,
   isSmartBackendWallet,
   type ParsedWalletDetails,
-} from "../../../shared/db/wallets/get-wallet-details";
+} from "../../db/wallets/get-wallet-details";
 import { doesChainSupportService } from "../../lib/chain/chain-capabilities";
 import { createCustomError } from "../../../server/middleware/error";
 import { SendTransactionQueue } from "../../../worker/queues/send-transaction-queue";
@@ -14,11 +14,13 @@ import { recordMetrics } from "../prometheus";
 import { reportUsage } from "../usage";
 import { doSimulateTransaction } from "./simulate-queued-transaction";
 import type { InsertedTransaction, QueuedTransaction } from "./types";
+import type { EnclaveWalletParams } from "../cache/get-enclave-wallet";
 
 interface InsertTransactionData {
   insertedTransaction: InsertedTransaction;
   idempotencyKey?: string;
   shouldSimulate?: boolean;
+  enclave?: EnclaveWalletParams;
 }
 
 /**
@@ -30,13 +32,17 @@ interface InsertTransactionData {
 export const insertTransaction = async (
   args: InsertTransactionData,
 ): Promise<string> => {
-  const { insertedTransaction, idempotencyKey, shouldSimulate = false } = args;
+  const {
+    insertedTransaction,
+    idempotencyKey,
+    shouldSimulate = false,
+    enclave
+  } = args;
 
   // The queueId uniquely represents an enqueued transaction.
   // It's also used as the idempotency key (default = no idempotence).
-  let queueId: string = randomUUID();
-  if (idempotencyKey) {
-    queueId = idempotencyKey;
+  const queueId: string = idempotencyKey ?? randomUUID();
+  if (queueId === idempotencyKey) {
     if (await TransactionDB.exists(queueId)) {
       // No-op. Return the existing queueId.
       return queueId;
@@ -151,7 +157,7 @@ export const insertTransaction = async (
 
   // Simulate the transaction.
   if (shouldSimulate) {
-    const error = await doSimulateTransaction(queuedTransaction);
+    const error = await doSimulateTransaction(queuedTransaction, enclave);
     if (error) {
       throw createCustomError(
         `Simulation failed: ${error.replace(/[\r\n]+/g, " --- ")}`,
@@ -165,6 +171,7 @@ export const insertTransaction = async (
   await SendTransactionQueue.add({
     queueId: queuedTransaction.queueId,
     resendCount: 0,
+    enclave,
   });
   reportUsage([{ action: "queue_tx", input: queuedTransaction }]);