From cfccfc7b43c34ad6c3c88f4743c0048ca9b489dd Mon Sep 17 00:00:00 2001 From: "Nicholas St. Germain" Date: Thu, 19 Dec 2024 07:01:22 -0600 Subject: [PATCH] add-enclave-functionality --- src/server/routes/contract/write/write.ts | 20 +++++--- src/server/schemas/wallet/index.ts | 47 ++++++++++++++--- src/server/utils/convertor.ts | 32 +++++++++++- src/shared/utils/account.ts | 26 ++++++---- src/shared/utils/cache/get-enclave-wallet.ts | 51 +++++++++++++++++++ src/shared/utils/cache/memory-storage.ts | 17 +++++++ .../utils/transaction/insert-transaction.ts | 21 +++++--- .../utils/transaction/queue-transation.ts | 33 +++++++----- .../simulate-queued-transaction.ts | 15 +++++- src/worker/queues/send-transaction-queue.ts | 4 +- src/worker/tasks/send-transaction-worker.ts | 11 ++-- 11 files changed, 228 insertions(+), 49 deletions(-) create mode 100644 src/shared/utils/cache/get-enclave-wallet.ts create mode 100644 src/shared/utils/cache/memory-storage.ts diff --git a/src/server/routes/contract/write/write.ts b/src/server/routes/contract/write/write.ts index fd5d6e99e..c4bac1557 100644 --- a/src/server/routes/contract/write/write.ts +++ b/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; + const enclave = await parseEnclaveHeaders(request.headers, chain); const chainId = await getChainIdFromChain(chain); const contract = await getContractV5({ @@ -118,7 +120,10 @@ 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( @@ -126,6 +131,7 @@ export async function writeToContract(fastify: FastifyInstance) { "x-account-factory-address", ), accountSalt, + enclave, txOverrides, idempotencyKey, shouldSimulate: simulateTx, diff --git a/src/server/schemas/wallet/index.ts b/src/server/schemas/wallet/index.ts index 3f4d940f4..5b0b89ba1 100644 --- a/src/server/schemas/wallet/index.ts +++ b/src/server/schemas/wallet/index.ts @@ -1,15 +1,11 @@ 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, @@ -17,6 +13,40 @@ export const walletHeaderSchema = Type.Object({ }), ), }); +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. diff --git a/src/server/utils/convertor.ts b/src/server/utils/convertor.ts index 5f0334679..4c68797e3 100644 --- a/src/server/utils/convertor.ts +++ b/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 => { + const { + "x-enclave-wallet-auth-token": authToken = "", + "x-client-id": clientId = "", + "x-ecosystem-id": id, + "x-ecosystem-partner-id": partnerId, + } = headers as Static; + let ecosystem: EnclaveWalletParams["ecosystem"]; + if (id) { + ecosystem = { id: id as EcosystemWalletId, partnerId }; + } + + return { + authToken, + clientId, + ecosystem, + chain, + }; +}; diff --git a/src/shared/utils/account.ts b/src/shared/utils/account.ts index 496071e3d..f63990ba2 100644 --- a/src/shared/utils/account.ts +++ b/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(2048); -export const getAccount = async (args: { +export type GetAccountArgs = { chainId: number; - from: Address; + from: Address | undefined; accountAddress?: Address; -}): Promise => { - const { chainId, from, accountAddress } = args; + enclave?: EnclaveWalletParams +} + +export const getAccount = async (args: GetAccountArgs): Promise => { + 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(2048); * Will throw if the wallet is not a smart backend wallet */ export const getSmartBackendWalletAdminAccount = async ({ - chainId, - accountAddress, -}: { + chainId, + accountAddress, + }: { chainId: number; accountAddress: Address; }) => { diff --git a/src/shared/utils/cache/get-enclave-wallet.ts b/src/shared/utils/cache/get-enclave-wallet.ts new file mode 100644 index 000000000..458876d1c --- /dev/null +++ b/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; +}; diff --git a/src/shared/utils/cache/memory-storage.ts b/src/shared/utils/cache/memory-storage.ts new file mode 100644 index 000000000..398b817c8 --- /dev/null +++ b/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 = 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); + } +} \ No newline at end of file diff --git a/src/shared/utils/transaction/insert-transaction.ts b/src/shared/utils/transaction/insert-transaction.ts index 625088a30..60de9a3ae 100644 --- a/src/shared/utils/transaction/insert-transaction.ts +++ b/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 => { - 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 }]); diff --git a/src/shared/utils/transaction/queue-transation.ts b/src/shared/utils/transaction/queue-transation.ts index b84a44515..76a044ace 100644 --- a/src/shared/utils/transaction/queue-transation.ts +++ b/src/shared/utils/transaction/queue-transation.ts @@ -13,14 +13,16 @@ import { parseTransactionOverrides } from "../../../server/utils/transaction-ove import { prettifyError } from "../error"; import { insertTransaction } from "./insert-transaction"; import type { InsertedTransaction } from "./types"; +import { type EnclaveWalletParams, getEnclaveWalletAddress } from "../cache/get-enclave-wallet"; export type QueuedTransactionParams = { transaction: PreparedTransaction; - fromAddress: Address; + fromAddress: Address | undefined; toAddress: Address | undefined; accountAddress: Address | undefined; accountFactoryAddress: Address | undefined; accountSalt: string | undefined; + enclave?: EnclaveWalletParams; txOverrides?: Static< typeof txOverridesWithValueSchema.properties.txOverrides >; @@ -36,7 +38,7 @@ export type QueuedTransactionParams = { * Encodes a transaction to generate data, and inserts it into the transaction queue using the insertTransaction() * * Note: - * - functionName must be be provided to populate the functionName field in the queued transaction + * - functionName must be provided to populate the functionName field in the queued transaction * - value and chain details are resolved from the transaction */ export async function queueTransaction(args: QueuedTransactionParams) { @@ -47,6 +49,7 @@ export async function queueTransaction(args: QueuedTransactionParams) { accountAddress, accountFactoryAddress, accountSalt, + enclave, txOverrides, idempotencyKey, shouldSimulate, @@ -64,10 +67,16 @@ export async function queueTransaction(args: QueuedTransactionParams) { "BAD_REQUEST", ); } - + let from = fromAddress; + if (!from) { + if (!enclave) { + throw new Error("Enclave wallet not provided"); + } + from = await getEnclaveWalletAddress(enclave) + } const insertedTransaction: InsertedTransaction = { chainId: transaction.chain.id, - from: fromAddress, + from, to: toAddress, data, value: await resolvePromisedValue(transaction.value), @@ -76,19 +85,19 @@ export async function queueTransaction(args: QueuedTransactionParams) { ...parseTransactionOverrides(txOverrides), ...(accountAddress ? { - isUserOp: true, - accountAddress: accountAddress, - signerAddress: fromAddress, - target: toAddress, - accountFactoryAddress, - accountSalt, - } - : { isUserOp: false }), + isUserOp: true, + accountAddress: accountAddress, + signerAddress: from, + target: toAddress, + accountFactoryAddress, + accountSalt, + } : { isUserOp: false }), }; return insertTransaction({ insertedTransaction, shouldSimulate, idempotencyKey, + enclave, }); } diff --git a/src/shared/utils/transaction/simulate-queued-transaction.ts b/src/shared/utils/transaction/simulate-queued-transaction.ts index 7fdbd5156..7972e8f87 100644 --- a/src/shared/utils/transaction/simulate-queued-transaction.ts +++ b/src/shared/utils/transaction/simulate-queued-transaction.ts @@ -11,14 +11,22 @@ import { getSmartWalletV5 } from "../cache/get-smart-wallet-v5"; import { getChain } from "../chain"; import { thirdwebClient } from "../sdk"; import type { AnyTransaction } from "./types"; +import type { Ecosystem } from "thirdweb/dist/types/wallets/in-app/core/wallet/types"; + +export type SimulateQueuedTransactionEnclaveParams = { + clientId: string; + authToken: string; + ecosystem?: Ecosystem; +} /** * Simulate the queued transaction. * @param transaction + * @param enclave * @returns string - The simulation error, or null if no error. */ export const doSimulateTransaction = async ( - transaction: AnyTransaction, + transaction: AnyTransaction, enclave: SimulateQueuedTransactionEnclaveParams | undefined = undefined, ): Promise => { const { chainId, @@ -61,6 +69,11 @@ export const doSimulateTransaction = async ( account = await getAccount({ chainId, from, + ...(enclave ? { + ecosystem: enclave.ecosystem, + clientId: enclave.clientId, + authToken: enclave.authToken, + } : {}) }); } diff --git a/src/worker/queues/send-transaction-queue.ts b/src/worker/queues/send-transaction-queue.ts index 0361da622..94e4a658b 100644 --- a/src/worker/queues/send-transaction-queue.ts +++ b/src/worker/queues/send-transaction-queue.ts @@ -2,10 +2,12 @@ import { Queue } from "bullmq"; import superjson from "superjson"; import { redis } from "../../shared/utils/redis/redis"; import { defaultJobOptions } from "./queues"; +import type { EnclaveWalletParams } from "../../shared/utils/cache/get-enclave-wallet"; export type SendTransactionData = { queueId: string; resendCount: number; + enclave?: EnclaveWalletParams; }; export class SendTransactionQueue { @@ -19,7 +21,7 @@ export class SendTransactionQueue { }, }); - // Allow enqueing the same queueId for multiple retries. + // Allow enqueuing the same queueId for multiple retries. static jobId = (data: SendTransactionData) => `${data.queueId}.${data.resendCount}`; diff --git a/src/worker/tasks/send-transaction-worker.ts b/src/worker/tasks/send-transaction-worker.ts index feee2b500..b1aa823f2 100644 --- a/src/worker/tasks/send-transaction-worker.ts +++ b/src/worker/tasks/send-transaction-worker.ts @@ -58,6 +58,7 @@ import { SendTransactionQueue, type SendTransactionData, } from "../queues/send-transaction-queue"; +import type { EnclaveWalletParams } from "../../shared/utils/cache/get-enclave-wallet"; /** * Submit a transaction to RPC (EOA transactions) or bundler (userOps). @@ -65,7 +66,7 @@ import { * This worker also handles retried EOA transactions. */ const handler: Processor = async (job: Job) => { - const { queueId, resendCount } = superjson.parse( + const { queueId, resendCount, enclave } = superjson.parse( job.data, ); @@ -85,7 +86,7 @@ const handler: Processor = async (job: Job) => { if (transaction.isUserOp) { resultTransaction = await _sendUserOp(job, transaction); } else { - resultTransaction = await _sendTransaction(job, transaction); + resultTransaction = await _sendTransaction(job, transaction, enclave); } } else if (transaction.status === "sent") { resultTransaction = await _resendTransaction(job, transaction, resendCount); @@ -266,6 +267,7 @@ const _sendUserOp = async ( const _sendTransaction = async ( job: Job, queuedTransaction: QueuedTransaction, + enclave?: EnclaveWalletParams ): Promise => { assert(!queuedTransaction.isUserOp); @@ -283,6 +285,7 @@ const _sendTransaction = async ( const account = await getAccount({ chainId: chainId, from: from, + enclave }); // Populate the transaction to resolve gas values. @@ -413,6 +416,7 @@ const _resendTransaction = async ( job: Job, sentTransaction: SentTransaction, resendCount: number, + enclave?: EnclaveWalletParams ): Promise => { assert(!sentTransaction.isUserOp); @@ -450,7 +454,7 @@ const _resendTransaction = async ( // This call throws if the RPC rejects the transaction. let transactionHash: Hex; try { - const account = await getAccount({ chainId, from }); + const account = await getAccount({ chainId, from, enclave }); const result = await account.sendTransaction(populatedTransaction); transactionHash = result.transactionHash; } catch (error) { @@ -568,6 +572,7 @@ const _minutesFromNow = (minutes: number) => * * @param populatedTransaction The transaction with estimated gas from RPC. * @param resendCount The resend attempt #. Example: 2 = the transaction was initially sent, then resent once. This is the second resend attempt. + * @param overrides */ export function _updateGasFees( populatedTransaction: PopulatedTransaction,