fix: shield coverage result polling

packages/shield-controller/CHANGELOG.md

@@ -10,6 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Bump `@metamask/transaction-controller` from `^60.7.0` to `^60.8.0` ([#6883](https://github.com/MetaMask/core/pull/6883))
+- Updated internal coverage result polling and log logic. ([#6847](https://github.com/MetaMask/core/pull/6847))
+  - Added cancellation logic to the polling.
+  - Updated implementation of timeout.
+  - Cancel any pending requests before starting new polling or logging.
+- Updated TransactionMeta comparison in `TransactionController:stateChange` subscriber to avoid triggering multiple check coverage result unnecessarily. ([#6847](https://github.com/MetaMask/core/pull/6847))
+- Removed `Personal Sign` check from the check signature coverage result. ([#6847](https://github.com/MetaMask/core/pull/6847))
 
 ## [0.3.2]
 

packages/shield-controller/package.json

@@ -48,7 +48,9 @@
   },
   "dependencies": {
     "@metamask/base-controller": "^8.4.1",
-    "@metamask/utils": "^11.8.1"
+    "@metamask/controller-utils": "^11.14.1",
+    "@metamask/utils": "^11.8.1",
+    "cockatiel": "^3.1.2"
   },
   "devDependencies": {
     "@babel/runtime": "^7.23.9",

packages/shield-controller/src/ShieldController.test.ts

@@ -11,6 +11,7 @@ import {
 } from '@metamask/transaction-controller';
 
 import { ShieldController } from './ShieldController';
+import { TX_META_SIMULATION_DATA_MOCKS } from '../tests/data';
 import { createMockBackend, MOCK_COVERAGE_ID } from '../tests/mocks/backend';
 import { createMockMessenger } from '../tests/mocks/messenger';
 import {
@@ -169,6 +170,47 @@ describe('ShieldController', () => {
       });
     });
 
+    TX_META_SIMULATION_DATA_MOCKS.forEach(
+      ({ description, previousSimulationData, newSimulationData }) => {
+        it(`should check coverage when ${description}`, async () => {
+          const { baseMessenger, backend } = setup();
+          const previousTxMeta = {
+            ...generateMockTxMeta(),
+            simulationData: previousSimulationData,
+          };
+          const coverageResultReceived =
+            setupCoverageResultReceived(baseMessenger);
+
+          // Add transaction.
+          baseMessenger.publish(
+            'TransactionController:stateChange',
+            { transactions: [previousTxMeta] } as TransactionControllerState,
+            undefined as never,
+          );
+          expect(await coverageResultReceived).toBeUndefined();
+          expect(backend.checkCoverage).toHaveBeenCalledWith({
+            txMeta: previousTxMeta,
+          });
+
+          // Simulate transaction.
+          const txMeta2 = { ...previousTxMeta };
+          txMeta2.simulationData = newSimulationData;
+          const coverageResultReceived2 =
+            setupCoverageResultReceived(baseMessenger);
+          baseMessenger.publish(
+            'TransactionController:stateChange',
+            { transactions: [txMeta2] } as TransactionControllerState,
+            undefined as never,
+          );
+          expect(await coverageResultReceived2).toBeUndefined();
+          expect(backend.checkCoverage).toHaveBeenCalledWith({
+            coverageId: MOCK_COVERAGE_ID,
+            txMeta: txMeta2,
+          });
+        });
+      },
+    );
+
     it('throws an error when the coverage ID has changed', async () => {
       const { controller, backend } = setup();
       backend.checkCoverage.mockResolvedValueOnce({

packages/shield-controller/src/ShieldController.ts

@@ -5,7 +5,6 @@ import type {
 } from '@metamask/base-controller';
 import {
   SignatureRequestStatus,
-  SignatureRequestType,
   type SignatureRequest,
   type SignatureStateChange,
 } from '@metamask/signature-controller';
@@ -236,10 +235,7 @@ export class ShieldController extends BaseController<
 
       // Check coverage if the signature request is new and has type
       // `personal_sign`.
-      if (
-        !previousSignatureRequest &&
-        signatureRequest.type === SignatureRequestType.PersonalSign
-      ) {
+      if (!previousSignatureRequest) {
         this.checkSignatureCoverage(signatureRequest).catch(
           // istanbul ignore next
           (error) => log('Error checking coverage:', error),
@@ -268,15 +264,15 @@ export class ShieldController extends BaseController<
     );
     for (const transaction of transactions) {
       const previousTransaction = previousTransactionsById.get(transaction.id);
+      // Check if the simulation data has changed.
+      const simulationDataChanged = this.#compareTransactionSimulationData(
+        transaction.simulationData,
+        previousTransaction?.simulationData,
+      );
 
       // Check coverage if the transaction is new or if the simulation data has
       // changed.
-      if (
-        !previousTransaction ||
-        // Checking reference equality is sufficient because this object is
-        // replaced if the simulation data has changed.
-        previousTransaction.simulationData !== transaction.simulationData
-      ) {
+      if (!previousTransaction || simulationDataChanged) {
         this.checkCoverage(transaction).catch(
           // istanbul ignore next
           (error) => log('Error checking coverage:', error),
@@ -443,4 +439,61 @@ export class ShieldController extends BaseController<
   #getLatestCoverageId(itemId: string): string | undefined {
     return this.state.coverageResults[itemId]?.results[0]?.coverageId;
   }
+
+  /**
+   * Compares the simulation data of a transaction.
+   *
+   * @param simulationData - The simulation data of the transaction.
+   * @param previousSimulationData - The previous simulation data of the transaction.
+   * @returns Whether the simulation data has changed.
+   */
+  #compareTransactionSimulationData(
+    simulationData?: TransactionMeta['simulationData'],
+    previousSimulationData?: TransactionMeta['simulationData'],
+  ) {
+    if (!simulationData && !previousSimulationData) {
+      return false;
+    }
+
+    // check the simulation error
+    if (
+      simulationData?.error?.code !== previousSimulationData?.error?.code ||
+      simulationData?.error?.message !== previousSimulationData?.error?.message
+    ) {
+      return true;
+    }
+
+    // check the native balance change
+    if (
+      simulationData?.nativeBalanceChange?.difference !==
+        previousSimulationData?.nativeBalanceChange?.difference ||
+      simulationData?.nativeBalanceChange?.newBalance !==
+        previousSimulationData?.nativeBalanceChange?.newBalance ||
+      simulationData?.nativeBalanceChange?.previousBalance !==
+        previousSimulationData?.nativeBalanceChange?.previousBalance ||
+      simulationData?.nativeBalanceChange?.isDecrease !==
+        previousSimulationData?.nativeBalanceChange?.isDecrease
+    ) {
+      return true;
+    }
+
+    // check the token balance changes
+    if (
+      simulationData?.tokenBalanceChanges?.length !==
+        previousSimulationData?.tokenBalanceChanges?.length ||
+      simulationData?.tokenBalanceChanges?.some(
+        (tokenBalanceChange, index) =>
+          tokenBalanceChange.difference !==
+          previousSimulationData?.tokenBalanceChanges?.[index]?.difference,
+      )
+    ) {
+      return true;
+    }
+
+    // check the isUpdatedAfterSecurityCheck
+    return (
+      simulationData?.isUpdatedAfterSecurityCheck !==
+      previousSimulationData?.isUpdatedAfterSecurityCheck
+    );
+  }
 }

packages/shield-controller/src/backend.test.ts

@@ -45,6 +45,11 @@ function setup({
 }
 
 describe('ShieldRemoteBackend', () => {
+  afterEach(() => {
+    // Clean up mocks after each test
+    jest.clearAllMocks();
+  });
+
   it('should check coverage', async () => {
     const { backend, fetchMock, getAccessToken } = setup();
 
@@ -143,7 +148,7 @@ describe('ShieldRemoteBackend', () => {
 
     const txMeta = generateMockTxMeta();
     await expect(backend.checkCoverage({ txMeta })).rejects.toThrow(
-      'Timeout waiting for coverage result',
+      'getCoverageResult: Request timed out',
     );
 
     // Waiting here ensures coverage of the unexpected error and lets us know

packages/shield-controller/src/backend.ts

@@ -1,6 +1,7 @@
 import type { SignatureRequest } from '@metamask/signature-controller';
 import type { TransactionMeta } from '@metamask/transaction-controller';
 
+import { PollingWithTimeoutAndAbort } from './polling-with-timeout-abort';
 import type {
   CheckCoverageRequest,
   CheckSignatureCoverageRequest,
@@ -58,6 +59,8 @@ export class ShieldRemoteBackend implements ShieldBackend {
 
   readonly #fetch: typeof globalThis.fetch;
 
+  readonly #pollingWithTimeout: PollingWithTimeoutAndAbort;
+
   constructor({
     getAccessToken,
     getCoverageResultTimeout = 5000, // milliseconds
@@ -76,6 +79,10 @@ export class ShieldRemoteBackend implements ShieldBackend {
     this.#getCoverageResultPollInterval = getCoverageResultPollInterval;
     this.#baseUrl = baseUrl;
     this.#fetch = fetchFn;
+    this.#pollingWithTimeout = new PollingWithTimeoutAndAbort({
+      timeout: getCoverageResultTimeout,
+      pollInterval: getCoverageResultPollInterval,
+    });
   }
 
   async checkCoverage(req: CheckCoverageRequest): Promise<CoverageResult> {
@@ -90,6 +97,7 @@ export class ShieldRemoteBackend implements ShieldBackend {
 
     const txCoverageResultUrl = `${this.#baseUrl}/v1/transaction/coverage/result`;
     const coverageResult = await this.#getCoverageResult(coverageId, {
+      requestId: req.txMeta.id,
       coverageResultUrl: txCoverageResultUrl,
     });
     return {
@@ -114,6 +122,7 @@ export class ShieldRemoteBackend implements ShieldBackend {
 
     const signatureCoverageResultUrl = `${this.#baseUrl}/v1/signature/coverage/result`;
     const coverageResult = await this.#getCoverageResult(coverageId, {
+      requestId: req.signatureRequest.id,
       coverageResultUrl: signatureCoverageResultUrl,
     });
     return {
@@ -132,6 +141,9 @@ export class ShieldRemoteBackend implements ShieldBackend {
       ...initBody,
     };
 
+    // clean up the pending coverage result polling
+    await this.#pollingWithTimeout.abortPendingRequest(req.signatureRequest.id);
+
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/signature/coverage/log`,
       {
@@ -153,6 +165,9 @@ export class ShieldRemoteBackend implements ShieldBackend {
       ...initBody,
     };
 
+    // clean up the pending coverage result polling
+    await this.#pollingWithTimeout.abortPendingRequest(req.txMeta.id);
+
     const res = await this.#fetch(
       `${this.#baseUrl}/v1/transaction/coverage/log`,
       {
@@ -183,7 +198,8 @@ export class ShieldRemoteBackend implements ShieldBackend {
 
   async #getCoverageResult(
     coverageId: string,
-    configs: {
+    config: {
+      requestId: string;
       coverageResultUrl: string;
       timeout?: number;
       pollInterval?: number;
@@ -192,41 +208,34 @@ export class ShieldRemoteBackend implements ShieldBackend {
     const reqBody: GetCoverageResultRequest = {
       coverageId,
     };
-
-    const timeout = configs?.timeout ?? this.#getCoverageResultTimeout;
-    const pollInterval =
-      configs?.pollInterval ?? this.#getCoverageResultPollInterval;
-
+    const pollingOptions = {
+      timeout: config.timeout ?? this.#getCoverageResultTimeout,
+      pollInterval: config.pollInterval ?? this.#getCoverageResultPollInterval,
+      fnName: 'getCoverageResult',
+    };
     const headers = await this.#createHeaders();
-    return await new Promise((resolve, reject) => {
-      let timeoutReached = false;
-      setTimeout(() => {
-        timeoutReached = true;
-        reject(new Error('Timeout waiting for coverage result'));
-      }, timeout);
 
-      const poll = async (): Promise<GetCoverageResultResponse> => {
-        // The timeoutReached variable is modified in the timeout callback.
-        // eslint-disable-next-line no-unmodified-loop-condition
-        while (!timeoutReached) {
-          const startTime = Date.now();
-          const res = await this.#fetch(configs.coverageResultUrl, {
-            method: 'POST',
-            headers,
-            body: JSON.stringify(reqBody),
-          });
-          if (res.status === 200) {
-            return (await res.json()) as GetCoverageResultResponse;
-          }
-          await sleep(pollInterval - (Date.now() - startTime));
-        }
-        // The following line will not have an effect as the upper level promise
-        // will already be rejected by now.
-        throw new Error('unexpected error');
-      };
+    const requestCoverageFn = async (
+      signal: AbortSignal,
+    ): Promise<GetCoverageResultResponse> => {
+      const res = await this.#fetch(config.coverageResultUrl, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(reqBody),
+        signal,
+      });
+      if (res.status === 200) {
+        return (await res.json()) as GetCoverageResultResponse;
+      }
+      throw new Error(`Failed to get coverage result: ${res.status}`);
+    };
 
-      poll().then(resolve).catch(reject);
-    });
+    const coverageResult = await this.#pollingWithTimeout.pollRequest(
+      config.requestId,
+      requestCoverageFn,
+      pollingOptions,
+    );
+    return coverageResult;
   }
 
   async #createHeaders() {
@@ -238,16 +247,6 @@ export class ShieldRemoteBackend implements ShieldBackend {
   }
 }
 
-/**
- * Sleep for a specified amount of time.
- *
- * @param ms - The number of milliseconds to sleep.
- * @returns A promise that resolves after the specified amount of time.
- */
-async function sleep(ms: number) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 /**
  * Make the body for the init coverage check request.
  *