integrate shield gateway in transaction controller

Author: matthiasgeihs

packages/transaction-controller/src/TransactionController.ts

@@ -122,6 +122,7 @@ import type {
   BeforeSignHook,
   TransactionContainerType,
   NestedTransactionMetadata,
+  GetSimulationConfig,
 } from './types';
 import {
   GasFeeEstimateLevel,
@@ -450,6 +451,11 @@ export type TransactionControllerOptions = {
     ) => Promise<{ transactionHash: string }>;
     publishBatch?: PublishBatchHook;
   };
+
+  /**
+   * Optional parameters for the transaction simulation.
+   */
+  getSimulationConfig?: GetSimulationConfig;
 };
 
 /**
@@ -798,6 +804,8 @@ export class TransactionController extends BaseController<
 
   readonly #securityProviderRequest?: SecurityProviderRequest;
 
+  readonly #getSimulationConfig?: GetSimulationConfig;
+
   readonly #sign?: (
     transaction: TypedTransaction,
     from: string,
@@ -843,6 +851,7 @@ export class TransactionController extends BaseController<
       publicKeyEIP7702,
       securityProviderRequest,
       sign,
+      getSimulationConfig,
       state,
       testGasFeeFlows,
       trace,
@@ -899,6 +908,7 @@ export class TransactionController extends BaseController<
       hooks?.publish ?? (() => Promise.resolve({ transactionHash: undefined }));
     this.#publishBatchHook = hooks?.publishBatch;
     this.#securityProviderRequest = securityProviderRequest;
+    this.#getSimulationConfig = getSimulationConfig;
     this.#sign = sign;
     this.#testGasFeeFlows = testGasFeeFlows === true;
     this.#trace = trace ?? (((_request, fn) => fn?.()) as TraceCallback);
@@ -4171,6 +4181,7 @@ export class TransactionController extends BaseController<
             ethQuery: this.#getEthQuery({ networkClientId }),
             nestedTransactions,
             txParams,
+            getSimulationConfig: this.#getSimulationConfig,
           }),
       );
 
@@ -4191,6 +4202,7 @@ export class TransactionController extends BaseController<
         messenger: this.messagingSystem,
         publicKeyEIP7702: this.#publicKeyEIP7702,
         transactionMeta,
+        getSimulationConfig: this.#getSimulationConfig,
       });
       gasFeeTokens = gasFeeTokensResponse?.gasFeeTokens ?? [];
       isGasFeeSponsored = gasFeeTokensResponse?.isGasFeeSponsored ?? false;

packages/transaction-controller/src/api/simulation-api.ts

@@ -8,6 +8,7 @@ import {
 } from '../constants';
 import { SimulationChainNotSupportedError, SimulationError } from '../errors';
 import { projectLogger } from '../logger';
+import type { GetSimulationConfig } from '../types';
 
 const log = createModuleLogger(projectLogger, 'simulation-api');
 
@@ -111,6 +112,11 @@ export type SimulationRequest = {
    * Defaults to false.
    */
   withLogs?: boolean;
+
+  /**
+   * Optional function to modify the simulation request.
+   */
+  getSimulationConfig?: GetSimulationConfig;
 };
 
 /** Raw event log emitted by a simulated transaction. */
@@ -274,7 +280,13 @@ export async function simulateTransactions(
   chainId: Hex,
   request: SimulationRequest,
 ): Promise<SimulationResponse> {
-  const url = await getSimulationUrl(chainId);
+  let url = await getSimulationUrl(chainId);
+
+  const { newUrl, authorization } =
+    (await request.getSimulationConfig?.(url)) || {};
+  if (newUrl) {
+    url = newUrl;
+  }
 
   const requestId = requestIdCounter;
   requestIdCounter += 1;
@@ -283,8 +295,18 @@ export async function simulateTransactions(
 
   log('Sending request', url, request);
 
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+  };
+
+  // Add optional authorization header, if provided.
+  if (authorization) {
+    headers.Authorization = authorization;
+  }
+
   const response = await fetch(url, {
     method: 'POST',
+    headers,
     body: JSON.stringify({
       id: String(requestId),
       jsonrpc: '2.0',

packages/transaction-controller/src/types.ts

@@ -1910,3 +1910,11 @@ export type AssetsFiatValues = {
    */
   sending?: string;
 };
+
+/**
+ * Parameters for the transaction simulation API.
+ */
+export type GetSimulationConfig = (url: string) => Promise<{
+  newUrl?: string;
+  authorization?: string;
+}>;

packages/transaction-controller/src/utils/balance-changes.ts

@@ -32,6 +32,7 @@ import type {
   SimulationToken,
   TransactionParams,
   NestedTransactionMetadata,
+  GetSimulationConfig,
 } from '../types';
 import { SimulationTokenStandard } from '../types';
 
@@ -51,6 +52,7 @@ export type GetBalanceChangesRequest = {
   ethQuery: EthQuery;
   nestedTransactions?: NestedTransactionMetadata[];
   txParams: TransactionParams;
+  getSimulationConfig?: GetSimulationConfig;
 };
 
 type ParsedEvent = {
@@ -107,6 +109,7 @@ type BalanceTransactionMap = Map<SimulationToken, SimulationRequestTransaction>;
  * @param request.to - The recipient of the transaction.
  * @param request.value - The value of the transaction.
  * @param request.data - The data of the transaction.
+ * @param request.getSimulationConfig - Optional transaction simulation parameters.
  * @returns The simulation data.
  */
 export async function getBalanceChanges(
@@ -729,6 +732,8 @@ async function baseRequest({
 
   const isInsufficientBalance = currentBalanceBN.lt(requiredBalanceBN);
 
+  const { getSimulationConfig } = request;
+
   return await simulateTransactions(chainId, {
     ...params,
     transactions,
@@ -749,6 +754,7 @@ async function baseRequest({
         },
       },
     }),
+    getSimulationConfig,
   });
 }
 

packages/transaction-controller/src/utils/gas-fee-tokens.ts

@@ -17,6 +17,7 @@ import {
   type SimulationResponseTransaction,
 } from '../api/simulation-api';
 import { projectLogger } from '../logger';
+import type { GetSimulationConfig } from '../types';
 
 const log = createModuleLogger(projectLogger, 'gas-fee-tokens');
 
@@ -28,6 +29,7 @@ export type GetGasFeeTokensRequest = {
   messenger: TransactionControllerMessenger;
   publicKeyEIP7702?: Hex;
   transactionMeta: TransactionMeta;
+  getSimulationConfig?: GetSimulationConfig;
 };
 
 /**
@@ -39,6 +41,7 @@ export type GetGasFeeTokensRequest = {
  * @param request.messenger - The messenger instance.
  * @param request.publicKeyEIP7702 - Public key to validate EIP-7702 contract signatures.
  * @param request.transactionMeta - The transaction metadata.
+ * @param request.getSimulationConfig - Optional transaction simulation parameters.
  * @returns An array of gas fee tokens.
  */
 export async function getGasFeeTokens({
@@ -47,6 +50,7 @@ export async function getGasFeeTokens({
   messenger,
   publicKeyEIP7702,
   transactionMeta,
+  getSimulationConfig,
 }: GetGasFeeTokensRequest) {
   const { delegationAddress, txParams } = transactionMeta;
   const { authorizationList: authorizationListRequest } = txParams;
@@ -95,6 +99,7 @@ export async function getGasFeeTokens({
         withFeeTransfer: true,
         with7702,
       },
+      getSimulationConfig,
     });
 
     log('Response', response);