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}`);
  }