feat: shield-controller: check signature coverage

packages/shield-controller/CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Add two new controller state metadata properties: `includeInStateLogs` and `usedInUi` ([#6504](https://github.com/MetaMask/core/pull/6504))
+- Add signature coverage checking ([#6501](https://github.com/MetaMask/core/pull/6501))
 
 ### Changed
 

packages/shield-controller/package.json

@@ -55,6 +55,7 @@
     "@lavamoat/allow-scripts": "^3.0.4",
     "@lavamoat/preinstall-always-fail": "^2.1.0",
     "@metamask/auto-changelog": "^3.4.4",
+    "@metamask/signature-controller": "^33.0.0",
     "@metamask/transaction-controller": "^60.3.0",
     "@ts-bridge/cli": "^0.6.1",
     "@types/jest": "^27.4.1",
@@ -67,6 +68,7 @@
     "uuid": "^8.3.2"
   },
   "peerDependencies": {
+    "@metamask/signature-controller": "^33.0.0",
     "@metamask/transaction-controller": "^60.0.0"
   },
   "engines": {

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

@@ -1,10 +1,14 @@
 import { deriveStateFromMetadata } from '@metamask/base-controller';
+import type { SignatureControllerState } from '@metamask/signature-controller';
 import type { TransactionControllerState } from '@metamask/transaction-controller';
 
 import { ShieldController } from './ShieldController';
 import { createMockBackend } from '../tests/mocks/backend';
 import { createMockMessenger } from '../tests/mocks/messenger';
-import { generateMockTxMeta } from '../tests/utils';
+import {
+  generateMockSignatureRequest,
+  generateMockTxMeta,
+} from '../tests/utils';
 
 /**
  * Sets up a ShieldController for testing.
@@ -144,6 +148,76 @@ describe('ShieldController', () => {
     });
   });
 
+  describe('checkSignatureCoverage', () => {
+    it('should check signature coverage', async () => {
+      const { baseMessenger, backend } = setup();
+      const signatureRequest = generateMockSignatureRequest();
+      const coverageResultReceived = new Promise<void>((resolve) => {
+        baseMessenger.subscribe(
+          'ShieldController:coverageResultReceived',
+          (_coverageResult) => resolve(),
+        );
+      });
+      baseMessenger.publish(
+        'SignatureController:stateChange',
+        {
+          signatureRequests: { [signatureRequest.id]: signatureRequest },
+        } as SignatureControllerState,
+        undefined as never,
+      );
+      expect(await coverageResultReceived).toBeUndefined();
+      expect(backend.checkSignatureCoverage).toHaveBeenCalledWith(
+        signatureRequest,
+      );
+    });
+  });
+
+  it('should check coverage for multiple signature request', async () => {
+    const { baseMessenger, backend } = setup();
+    const signatureRequest1 = generateMockSignatureRequest();
+    const coverageResultReceived1 = new Promise<void>((resolve) => {
+      baseMessenger.subscribe(
+        'ShieldController:coverageResultReceived',
+        (_coverageResult) => resolve(),
+      );
+    });
+    baseMessenger.publish(
+      'SignatureController:stateChange',
+      {
+        signatureRequests: {
+          [signatureRequest1.id]: signatureRequest1,
+        },
+      } as SignatureControllerState,
+      undefined as never,
+    );
+    expect(await coverageResultReceived1).toBeUndefined();
+    expect(backend.checkSignatureCoverage).toHaveBeenCalledWith(
+      signatureRequest1,
+    );
+
+    const signatureRequest2 = generateMockSignatureRequest();
+    const coverageResultReceived2 = new Promise<void>((resolve) => {
+      baseMessenger.subscribe(
+        'ShieldController:coverageResultReceived',
+        (_coverageResult) => resolve(),
+      );
+    });
+    baseMessenger.publish(
+      'SignatureController:stateChange',
+      {
+        signatureRequests: {
+          [signatureRequest2.id]: signatureRequest2,
+        },
+      } as SignatureControllerState,
+      undefined as never,
+    );
+
+    expect(await coverageResultReceived2).toBeUndefined();
+    expect(backend.checkSignatureCoverage).toHaveBeenCalledWith(
+      signatureRequest2,
+    );
+  });
+
   describe('metadata', () => {
     it('includes expected state in debug snapshots', () => {
       const { controller } = setup();

packages/shield-controller/src/ShieldController.ts

@@ -3,6 +3,11 @@ import type {
   ControllerStateChangeEvent,
   RestrictedMessenger,
 } from '@metamask/base-controller';
+import {
+  SignatureRequestType,
+  type SignatureRequest,
+  type SignatureStateChange,
+} from '@metamask/signature-controller';
 import type {
   TransactionControllerStateChangeEvent,
   TransactionMeta,
@@ -82,7 +87,9 @@ type AllowedActions = never;
 /**
  * The external events available to the ShieldController.
  */
-type AllowedEvents = TransactionControllerStateChangeEvent;
+type AllowedEvents =
+  | SignatureStateChange
+  | TransactionControllerStateChangeEvent;
 
 /**
  * The messenger of the {@link ShieldController}.
@@ -138,6 +145,11 @@ export class ShieldController extends BaseController<
     previousTransactions: TransactionMeta[] | undefined,
   ) => void;
 
+  readonly #signatureControllerStateChangeHandler: (
+    signatureRequests: Record<string, SignatureRequest>,
+    previousSignatureRequests: Record<string, SignatureRequest> | undefined,
+  ) => void;
+
   constructor(options: ShieldControllerOptions) {
     const {
       messenger,
@@ -161,6 +173,8 @@ export class ShieldController extends BaseController<
     this.#transactionHistoryLimit = transactionHistoryLimit;
     this.#transactionControllerStateChangeHandler =
       this.#handleTransactionControllerStateChange.bind(this);
+    this.#signatureControllerStateChangeHandler =
+      this.#handleSignatureControllerStateChange.bind(this);
   }
 
   start() {
@@ -169,13 +183,54 @@ export class ShieldController extends BaseController<
       this.#transactionControllerStateChangeHandler,
       (state) => state.transactions,
     );
+
+    this.messagingSystem.subscribe(
+      'SignatureController:stateChange',
+      this.#signatureControllerStateChangeHandler,
+      (state) => state.signatureRequests,
+    );
   }
 
   stop() {
     this.messagingSystem.unsubscribe(
       'TransactionController:stateChange',
       this.#transactionControllerStateChangeHandler,
     );
+
+    this.messagingSystem.unsubscribe(
+      'SignatureController:stateChange',
+      this.#signatureControllerStateChangeHandler,
+    );
+  }
+
+  #handleSignatureControllerStateChange(
+    signatureRequests: Record<string, SignatureRequest>,
+    previousSignatureRequests: Record<string, SignatureRequest> | undefined,
+  ) {
+    const signatureRequestsArray = Object.values(signatureRequests);
+    const previousSignatureRequestsArray = Object.values(
+      previousSignatureRequests ?? {},
+    );
+    const previousSignatureRequestsById = new Map<string, SignatureRequest>(
+      previousSignatureRequestsArray.map((request) => [request.id, request]),
+    );
+    for (const signatureRequest of signatureRequestsArray) {
+      const previousSignatureRequest = previousSignatureRequestsById.get(
+        signatureRequest.id,
+      );
+
+      // Check coverage if the signature request is new and has type
+      // `personal_sign`.
+      if (
+        !previousSignatureRequest &&
+        signatureRequest.type === SignatureRequestType.PersonalSign
+      ) {
+        this.checkSignatureCoverage(signatureRequest).catch(
+          // istanbul ignore next
+          (error) => log('Error checking coverage:', error),
+        );
+      }
+    }
   }
 
   #handleTransactionControllerStateChange(
@@ -212,7 +267,7 @@ export class ShieldController extends BaseController<
    */
   async checkCoverage(txMeta: TransactionMeta): Promise<CoverageResult> {
     // Check coverage
-    const coverageResult = await this.#fetchCoverageResult(txMeta);
+    const coverageResult = await this.#backend.checkCoverage(txMeta);
 
     // Publish coverage result
     this.messagingSystem.publish(
@@ -226,8 +281,29 @@ export class ShieldController extends BaseController<
     return coverageResult;
   }
 
-  async #fetchCoverageResult(txMeta: TransactionMeta): Promise<CoverageResult> {
-    return this.#backend.checkCoverage(txMeta);
+  /**
+   * Checks the coverage of a signature request.
+   *
+   * @param signatureRequest - The signature request to check coverage for.
+   * @returns The coverage result.
+   */
+  async checkSignatureCoverage(
+    signatureRequest: SignatureRequest,
+  ): Promise<CoverageResult> {
+    // Check coverage
+    const coverageResult =
+      await this.#backend.checkSignatureCoverage(signatureRequest);
+
+    // Publish coverage result
+    this.messagingSystem.publish(
+      `${controllerName}:coverageResultReceived`,
+      coverageResult,
+    );
+
+    // Update state
+    this.#addCoverageResult(signatureRequest.id, coverageResult);
+
+    return coverageResult;
   }
 
   #addCoverageResult(txId: string, coverageResult: CoverageResult) {

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

@@ -1,5 +1,9 @@
 import { ShieldRemoteBackend } from './backend';
-import { generateMockTxMeta, getRandomCoverageStatus } from '../tests/utils';
+import {
+  generateMockSignatureRequest,
+  generateMockTxMeta,
+  getRandomCoverageStatus,
+} from '../tests/utils';
 
 /**
  * Setup the test environment.
@@ -141,4 +145,40 @@ describe('ShieldRemoteBackend', () => {
     // that the polling loop is exited as expected.
     await new Promise((resolve) => setTimeout(resolve, 10));
   });
+
+  describe('checkSignatureCoverage', () => {
+    it('should check signature coverage', async () => {
+      const { backend, fetchMock, getAccessToken } = setup();
+
+      // Mock init coverage check.
+      fetchMock.mockResolvedValueOnce({
+        status: 200,
+        json: jest.fn().mockResolvedValue({ coverageId: 'coverageId' }),
+      } as unknown as Response);
+
+      // Mock get coverage result.
+      const status = getRandomCoverageStatus();
+      fetchMock.mockResolvedValueOnce({
+        status: 200,
+        json: jest.fn().mockResolvedValue({ status }),
+      } as unknown as Response);
+
+      const signatureRequest = generateMockSignatureRequest();
+      const coverageResult =
+        await backend.checkSignatureCoverage(signatureRequest);
+      expect(coverageResult).toStrictEqual({ status });
+      expect(fetchMock).toHaveBeenCalledTimes(2);
+      expect(getAccessToken).toHaveBeenCalledTimes(2);
+    });
+
+    it('throws with invalid data', async () => {
+      const { backend } = setup();
+
+      const signatureRequest = generateMockSignatureRequest();
+      signatureRequest.messageParams.data = [];
+      await expect(
+        backend.checkSignatureCoverage(signatureRequest),
+      ).rejects.toThrow('Signature data must be a string');
+    });
+  });
 });

packages/shield-controller/src/backend.ts

@@ -1,3 +1,4 @@
+import type { SignatureRequest } from '@metamask/signature-controller';
 import type { TransactionMeta } from '@metamask/transaction-controller';
 
 import type { CoverageResult, CoverageStatus, ShieldBackend } from './types';
@@ -16,6 +17,14 @@ export type InitCoverageCheckRequest = {
   origin?: string;
 };
 
+export type InitSignatureCoverageCheckRequest = {
+  chainId: string;
+  data: string;
+  from: string;
+  method: string;
+  origin?: string;
+};
+
 export type InitCoverageCheckResponse = {
   coverageId: string;
 };
@@ -59,9 +68,7 @@ export class ShieldRemoteBackend implements ShieldBackend {
     this.#fetch = fetchFn;
   }
 
-  checkCoverage: (txMeta: TransactionMeta) => Promise<CoverageResult> = async (
-    txMeta,
-  ) => {
+  async checkCoverage(txMeta: TransactionMeta): Promise<CoverageResult> {
     const reqBody: InitCoverageCheckRequest = {
       txParams: [
         {
@@ -76,22 +83,46 @@ export class ShieldRemoteBackend implements ShieldBackend {
       origin: txMeta.origin,
     };
 
-    const { coverageId } = await this.#initCoverageCheck(reqBody);
+    const { coverageId } = await this.#initCoverageCheck(
+      'v1/transaction/coverage/init',
+      reqBody,
+    );
 
     return this.#getCoverageResult(coverageId);
-  };
+  }
+
+  async checkSignatureCoverage(
+    signatureRequest: SignatureRequest,
+  ): Promise<CoverageResult> {
+    if (typeof signatureRequest.messageParams.data !== 'string') {
+      throw new Error('Signature data must be a string');
+    }
+
+    const reqBody: InitSignatureCoverageCheckRequest = {
+      chainId: signatureRequest.chainId,
+      data: signatureRequest.messageParams.data,
+      from: signatureRequest.messageParams.from,
+      method: signatureRequest.type,
+      origin: signatureRequest.messageParams.origin,
+    };
+
+    const { coverageId } = await this.#initCoverageCheck(
+      'v1/signature/coverage/init',
+      reqBody,
+    );
+
+    return this.#getCoverageResult(coverageId);
+  }
 
   async #initCoverageCheck(
-    reqBody: InitCoverageCheckRequest,
+    path: string,
+    reqBody: unknown,
   ): Promise<InitCoverageCheckResponse> {
-    const res = await this.#fetch(
-      `${this.#baseUrl}/v1/transaction/coverage/init`,
-      {
-        method: 'POST',
-        headers: await this.#createHeaders(),
-        body: JSON.stringify(reqBody),
-      },
-    );
+    const res = await this.#fetch(`${this.#baseUrl}/${path}`, {
+      method: 'POST',
+      headers: await this.#createHeaders(),
+      body: JSON.stringify(reqBody),
+    });
     if (res.status !== 200) {
       throw new Error(`Failed to init coverage check: ${res.status}`);
     }