fix: Add simulation retrigger

packages/transaction-controller/jest.config.js

@@ -18,9 +18,9 @@ module.exports = merge(baseConfig, {
   coverageThreshold: {
     global: {
       branches: 94.42,
-      functions: 97.46,
-      lines: 98.44,
-      statements: 98.46,
+      functions: 97.22,
+      lines: 98.38,
+      statements: 98.39,
     },
   },
 

packages/transaction-controller/src/TransactionController.test.ts

@@ -36,6 +36,7 @@ import { errorCodes, providerErrors, rpcErrors } from '@metamask/rpc-errors';
 import type { Hex } from '@metamask/utils';
 import { createDeferredPromise } from '@metamask/utils';
 import assert from 'assert';
+import { merge } from 'lodash';
 import * as uuidModule from 'uuid';
 
 import { FakeBlockTracker } from '../../../tests/fake-block-tracker';
@@ -73,6 +74,7 @@ import type {
   SubmitHistoryEntry,
 } from './types';
 import {
+  BlockaidResultType,
   GasFeeEstimateType,
   SimulationErrorCode,
   SimulationTokenStandard,
@@ -5030,6 +5032,230 @@ describe('TransactionController', () => {
         'Cannot update security alert response as no transaction metadata found',
       );
     });
+
+    describe('should not trigger simulation data update', () => {
+      it('when simulationData is not available', async () => {
+        const transactionMeta = TRANSACTION_META_MOCK;
+
+        const { controller } = setupController({
+          options: {
+            state: {
+              transactions: [transactionMeta],
+            },
+          },
+        });
+
+        controller.updateSecurityAlertResponse(transactionMeta.id, {
+          reason: 'NA',
+          // This is API specific hence naming convention is not followed.
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          result_type: BlockaidResultType.Malicious,
+        });
+
+        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+      });
+
+      it('when security alert response type is not malicious', async () => {
+        const transactionMetaWithSimulationData = merge(
+          {},
+          TRANSACTION_META_MOCK,
+          {
+            txParams: {
+              value: '0x0',
+            },
+            simulationData: {
+              tokenBalanceChanges: [],
+              nativeBalanceChange: {
+                difference: '0x0',
+              },
+            },
+          },
+        );
+
+        getSimulationDataMock.mockResolvedValueOnce(SIMULATION_DATA_MOCK);
+
+        const { controller } = setupController({
+          options: {
+            state: {
+              transactions: [transactionMetaWithSimulationData],
+            },
+          },
+        });
+
+        controller.updateSecurityAlertResponse(
+          transactionMetaWithSimulationData.id,
+          {
+            reason: 'NA',
+            // This is API specific hence naming convention is not followed.
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            result_type: BlockaidResultType.Warning,
+          },
+        );
+
+        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+      });
+
+      it('when transaction value is already different than simulated value', async () => {
+        const transactionMetaWithSimulationData = merge(
+          {},
+          TRANSACTION_META_MOCK,
+          {
+            txParams: {
+              value: '0x12345678',
+            },
+            simulationData: {
+              tokenBalanceChanges: [],
+              nativeBalanceChange: {
+                difference: '0x0',
+              },
+            },
+          },
+        );
+
+        getSimulationDataMock.mockResolvedValueOnce(SIMULATION_DATA_MOCK);
+
+        const { controller } = setupController({
+          options: {
+            state: {
+              transactions: [transactionMetaWithSimulationData],
+            },
+          },
+        });
+
+        controller.updateSecurityAlertResponse(
+          transactionMetaWithSimulationData.id,
+          {
+            reason: 'NA',
+            // This is API specific hence naming convention is not followed.
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            result_type: BlockaidResultType.Malicious,
+          },
+        );
+
+        expect(getSimulationDataMock).toHaveBeenCalledTimes(0);
+      });
+    });
+
+    it('should retrigger simulation data update when simulationData available and if result_type is malicious', async () => {
+      const transactionMetaWithSimulationData = merge(
+        {},
+        TRANSACTION_META_MOCK,
+        {
+          txParams: {
+            value: '0x0',
+          },
+          simulationData: {
+            tokenBalanceChanges: [],
+            nativeBalanceChange: {
+              difference: '0x0',
+            },
+          },
+        },
+      );
+
+      getSimulationDataMock.mockResolvedValueOnce(SIMULATION_DATA_MOCK);
+
+      const { controller } = setupController({
+        options: {
+          state: {
+            transactions: [transactionMetaWithSimulationData],
+          },
+        },
+      });
+
+      controller.updateSecurityAlertResponse(
+        transactionMetaWithSimulationData.id,
+        {
+          reason: 'NA',
+          // This is API specific hence naming convention is not followed.
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          result_type: BlockaidResultType.Malicious,
+        },
+      );
+
+      await flushPromises();
+
+      expect(getSimulationDataMock).toHaveBeenCalledTimes(1);
+    });
+
+    describe('should mark simulationDataChanged after security alert update', () => {
+      it('as true if simulationData updated', async () => {
+        // Assume that previous simulationData tells that transaction has no token balance changes
+        const transactionMetaWithSimulationData = merge(
+          {},
+          TRANSACTION_META_MOCK,
+          {
+            txParams: {
+              value: '0x0',
+            },
+            simulationData: {
+              tokenBalanceChanges: [],
+              nativeBalanceChange: {
+                difference: '0x0',
+              },
+            },
+          },
+        );
+
+        getSimulationDataMock.mockResolvedValueOnce(SIMULATION_DATA_MOCK);
+
+        const { controller } = setupController({
+          options: {
+            state: {
+              transactions: [transactionMetaWithSimulationData],
+            },
+          },
+        });
+
+        controller.updateSecurityAlertResponse(
+          transactionMetaWithSimulationData.id,
+          {
+            reason: 'NA',
+            // This is API specific hence naming convention is not followed.
+            // eslint-disable-next-line @typescript-eslint/naming-convention
+            result_type: BlockaidResultType.Malicious,
+          },
+        );
+
+        await flushPromises();
+
+        expect(
+          controller.state.transactions[0].simulationData
+            ?.changeInSimulationData,
+        ).toBe(true);
+      });
+
+      it('as undefined if current and new simulationData equality check is false', async () => {
+        // Assume that previous simulationData tells that transaction has some token balance changes
+        const transactionMeta = merge({}, TRANSACTION_META_MOCK, {
+          simulationData: SIMULATION_DATA_MOCK,
+        });
+
+        getSimulationDataMock.mockResolvedValueOnce(SIMULATION_DATA_MOCK);
+
+        const { controller } = setupController({
+          options: {
+            state: {
+              transactions: [transactionMeta],
+            },
+          },
+        });
+
+        controller.updateSecurityAlertResponse(transactionMeta.id, {
+          reason: 'NA',
+          // This is API specific hence naming convention is not followed.
+          // eslint-disable-next-line @typescript-eslint/naming-convention
+          result_type: BlockaidResultType.Malicious,
+        });
+
+        await flushPromises();
+
+        expect(
+          controller.state.transactions[0].simulationData
+            ?.changeInSimulationData,
+        ).toBeUndefined();
+      });
+    });
   });
 
   describe('updateCustodialTransaction', () => {

packages/transaction-controller/src/TransactionController.ts

@@ -84,6 +84,7 @@ import type {
   SubmitHistoryEntry,
 } from './types';
 import {
+  BlockaidResultType,
   TransactionEnvelopeType,
   TransactionType,
   TransactionStatus,
@@ -3599,9 +3600,42 @@ export class TransactionController extends BaseController<
       );
     }
 
+    if (this.#checkIfSimulationRetriggerNeeded(transactionMeta)) {
+      this.#updateSimulationData(transactionMeta).catch((error) => {
+        log('Error updating simulation data', error);
+        throw error;
+      });
+    }
+
     return transactionMeta;
   }
 
+  #checkIfSimulationRetriggerNeeded(transactionMeta: TransactionMeta) {
+    const {
+      securityAlertResponse,
+      simulationData,
+      txParams: { value },
+    } = transactionMeta;
+
+    const isSimulationDataAvailable = Boolean(simulationData);
+    const isMaliciousTransfer =
+      securityAlertResponse?.result_type === BlockaidResultType.Malicious;
+    const simulationNativeBalanceDifference =
+      simulationData?.nativeBalanceChange?.difference;
+    const isBalanceDifferenceEqual =
+      simulationNativeBalanceDifference === value;
+
+    if (
+      !isSimulationDataAvailable ||
+      !isMaliciousTransfer ||
+      !isBalanceDifferenceEqual
+    ) {
+      return false;
+    }
+
+    return true;
+  }
+
   #checkIfTransactionParamsUpdated(newTransactionMeta: TransactionMeta) {
     const { id: transactionId, txParams: newParams } = newTransactionMeta;
 
@@ -3649,7 +3683,12 @@ export class TransactionController extends BaseController<
     transactionMeta: TransactionMeta,
     { traceContext }: { traceContext?: TraceContext } = {},
   ) {
-    const { id: transactionId, chainId, txParams } = transactionMeta;
+    const {
+      id: transactionId,
+      chainId,
+      txParams,
+      simulationData: prevSimulationData,
+    } = transactionMeta;
     const { from, to, value, data } = txParams;
 
     let simulationData: SimulationData = {
@@ -3679,6 +3718,13 @@ export class TransactionController extends BaseController<
             data: data as Hex,
           }),
       );
+
+      if (prevSimulationData && !isEqual(simulationData, prevSimulationData)) {
+        simulationData = {
+          ...simulationData,
+          changeInSimulationData: true,
+        };
+      }
     }
 
     const finalTransactionMeta = this.getTransaction(transactionId);

packages/transaction-controller/src/types.ts

@@ -1119,9 +1119,9 @@ export type TransactionError = {
 export type SecurityAlertResponse = {
   reason: string;
   features?: string[];
-  // TODO: Either fix this lint violation or explain why it's necessary to ignore.
+  // This is API specific hence naming convention is not followed.
   // eslint-disable-next-line @typescript-eslint/naming-convention
-  result_type: string;
+  result_type: BlockaidResultType | string;
   providerRequestsCount?: Record<string, number>;
 };
 
@@ -1307,6 +1307,9 @@ export type SimulationError = {
 
 /** Simulation data for a transaction. */
 export type SimulationData = {
+  /** Set to true if simulationData changed after updating security alert */
+  changeInSimulationData?: boolean;
+
   /** Error data if the simulation failed or the transaction reverted. */
   error?: SimulationError;
 
@@ -1368,3 +1371,14 @@ export type SubmitHistoryEntry = {
   /** The transaction parameters that were submitted. */
   transaction: TransactionParams;
 };
+
+export enum BlockaidResultType {
+  Malicious = 'Malicious',
+  Warning = 'Warning',
+  Benign = 'Benign',
+  Errored = 'Error',
+
+  // MetaMask defined result types
+  NotApplicable = 'NotApplicable',
+  Loading = 'loading',
+}