add-enclave-functionality

Author: dirtycajunrice

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

@@ -1,8 +1,8 @@
-import { Type, type Static } from "@sinclair/typebox";
+import { type Static, Type } from "@sinclair/typebox";
 import type { FastifyInstance } from "fastify";
 import { StatusCodes } from "http-status-codes";
 import { prepareContractCall, resolveMethod } from "thirdweb";
-import { parseAbiParams, type AbiFunction } from "thirdweb/utils";
+import { type AbiFunction, parseAbiParams } from "thirdweb/utils";
 import { getContractV5 } from "../../../../shared/utils/cache/get-contractv5";
 import { prettifyError } from "../../../../shared/utils/error";
 import { queueTransaction } from "../../../../shared/utils/transaction/queue-transation";
@@ -16,11 +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 { parseEnclaveHeaders } from "../../../utils/convertor";
 import { parseTransactionOverrides } from "../../../utils/transaction-overrides";
 
 // INPUT
@@ -60,7 +61,7 @@ export async function writeToContract(fastify: FastifyInstance) {
       tags: ["Contract"],
       operationId: "write",
       params: contractParamSchema,
-      headers: walletWithAAHeaderSchema,
+      headers: walletWithAAOrEnclaveHeaderSchema,
       querystring: requestQuerystringSchema,
       body: writeRequestBodySchema,
       response: {
@@ -72,12 +73,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 = await parseEnclaveHeaders(request.headers, chain);
 
       const chainId = await getChainIdFromChain(chain);
       const contract = await getContractV5({
@@ -118,14 +120,18 @@ 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

@@ -1,22 +1,52 @@
 import { Type } from "@sinclair/typebox";
-import { getAddress, type Address } from "thirdweb";
+import { type Address, getAddress } from "thirdweb";
 import { env } from "../../../shared/utils/env";
 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,10 +1,15 @@
+import type { Static } from "@sinclair/typebox";
 import { BigNumber } from "ethers";
+import type { FastifyRequest } from "fastify";
+import type { EcosystemWalletId } from "thirdweb/dist/types/wallets/wallet-types";
+import type { EnclaveWalletParams } from "../../shared/utils/cache/get-enclave-wallet";
+import type { enclaveWalletHeaderSchema } from "../schemas/wallet";
 
 const isHexBigNumber = (value: unknown) => {
   const isNonNullObject = typeof value === "object" && value !== null;
   const hasType = isNonNullObject && "type" in value;
-  return hasType && value.type === "BigNumber" && "hex" in value
-}
+  return hasType && value.type === "BigNumber" && "hex" in value;
+};
 export const bigNumberReplacer = (value: unknown): unknown => {
   // if we find a BigNumber then make it into a string (since that is safe)
   if (BigNumber.isBigNumber(value) || isHexBigNumber(value)) {
@@ -17,3 +22,26 @@ export const bigNumberReplacer = (value: unknown): unknown => {
 
   return value;
 };
+
+export const parseEnclaveHeaders = async (
+  headers: FastifyRequest["headers"],
+  chain: string,
+): Promise<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: EnclaveWalletParams["ecosystem"];
+  if (id) {
+    ecosystem = { id: id as EcosystemWalletId, partnerId };
+  }
+
+  return {
+    authToken,
+    clientId,
+    ecosystem,
+    chain,
+  };
+};

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,51 @@
+import type { Address } from "thirdweb";
+import { ecosystemWallet, inAppWallet } from "thirdweb/wallets";
+import { getChainIdFromChain } from "../../../server/utils/chain";
+import { getChain } from "../chain";
+import { thirdwebClient as client } from "../sdk";
+
+export interface EnclaveWalletParams {
+  chain: string;
+  authToken: string;
+  clientId: string;
+  ecosystem?: {
+    id: `ecosystem.${string}`;
+    partnerId?: string;
+  };
+}
+
+export const getEnclaveWalletAccount = async ({
+  chain: chainSlug,
+  authToken,
+  clientId,
+  ecosystem,
+}: EnclaveWalletParams) => {
+  const wallet = ecosystem
+    ? ecosystemWallet(ecosystem.id, { partnerId: ecosystem.partnerId })
+    : inAppWallet();
+  const chainId = await getChainIdFromChain(chainSlug);
+  const chain = await getChain(chainId);
+  return wallet.autoConnect({
+    client,
+    authResult: {
+      storedToken: {
+        jwtToken: "",
+        authProvider: "Cognito",
+        authDetails: {
+          userWalletId: "",
+          recoveryShareManagement: "ENCLAVE",
+        },
+        developerClientId: clientId,
+        cookieString: authToken,
+        shouldStoreCookieString: false,
+        isNewUser: false,
+      },
+    },
+    chain,
+  });
+};
+
+export const getEnclaveWalletAddress = async (params: EnclaveWalletParams) => {
+  const account = await getEnclaveWalletAccount(params);
+  return account.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 }]);