feat: add request id to messageParams in @metamask/message-managers

packages/message-manager/src/AbstractMessageManager.test.ts

@@ -1,4 +1,10 @@
-import type { SecurityProviderRequest } from './AbstractMessageManager';
+import { ApprovalType } from '@metamask/controller-utils';
+
+import type {
+  MessageParams,
+  OriginalRequest,
+  SecurityProviderRequest,
+} from './AbstractMessageManager';
 import { AbstractMessageManager } from './AbstractMessageManager';
 import type {
   TypedMessage,
@@ -50,6 +56,32 @@ const rawSigMock = '0xsignaturemocked';
 const messageIdMock = 'message-id-mocked';
 const fromMock = '0xc38bf1ad06ef69f0c04e29dbeb4152b4175f0a8d';
 
+const mockSecurityProviderResponse = { flagAsDangerous: 2 };
+const mockRequest = {
+  origin: 'example.com',
+  id: 123,
+  securityAlertResponse: mockSecurityProviderResponse,
+};
+const mockMessageParams = { from, data: 'test' };
+
+/** Mock extended AbstractTestManager to test protected methods */
+class MockExtendedAbstractTestManager extends AbstractTestManager {
+  addRequestToMessageParams(
+    messageParams: MessageParams,
+    req?: OriginalRequest,
+  ) {
+    return super.addRequestToMessageParams(messageParams, req);
+  }
+
+  createUnapprovedMessage(
+    messageParams: MessageParams,
+    type: ApprovalType,
+    req?: OriginalRequest,
+  ) {
+    return super.createUnapprovedMessage(messageParams, type, req);
+  }
+}
+
 describe('AbstractTestManager', () => {
   it('should set default state', () => {
     const controller = new AbstractTestManager();
@@ -317,30 +349,41 @@ describe('AbstractTestManager', () => {
     expect(message.status).toBe('approved');
   });
 
-  describe('setMessageStatus', () => {
-    it('should set the given message status', async () => {
-      const controller = new AbstractTestManager();
-      await controller.addMessage({
-        id: messageId,
-        messageParams: { from: '0x1234', data: 'test' },
-        status: 'status',
-        time: 10,
-        type: 'type',
-      });
-      const messageBefore = controller.getMessage(messageId);
-      expect(messageBefore?.status).toBe('status');
+  describe('addRequestToMessageParams', () => {
+    it('adds original request id and origin to messageParams', () => {
+      const controller = new MockExtendedAbstractTestManager();
 
-      controller.setMessageStatus(messageId, 'newstatus');
-      const messageAfter = controller.getMessage(messageId);
-      expect(messageAfter?.status).toBe('newstatus');
+      const result = controller.addRequestToMessageParams(
+        mockMessageParams,
+        mockRequest,
+      );
+
+      expect(result).toStrictEqual({
+        ...mockMessageParams,
+        origin: mockRequest.origin,
+        requestId: mockRequest.id,
+      });
     });
+  });
 
-    it('should throw an error if message is not found', () => {
-      const controller = new AbstractTestManager();
+  describe('createUnapprovedMessage', () => {
+    it('creates a Message object with an unapproved status', () => {
+      const controller = new MockExtendedAbstractTestManager();
 
-      expect(() => controller.setMessageStatus(messageId, 'newstatus')).toThrow(
-        'AbstractMessageManager: Message not found for id: 1.',
+      const result = controller.createUnapprovedMessage(
+        mockMessageParams,
+        ApprovalType.PersonalSign,
+        mockRequest,
+      );
+
+      expect(result.messageParams).toBe(mockMessageParams);
+      expect(result.securityAlertResponse).toBe(
+        mockRequest.securityAlertResponse,
       );
+      expect(result.status).toBe('unapproved');
+      expect(result.type).toBe(ApprovalType.PersonalSign);
+      expect(typeof result.time).toBe('number');
+      expect(typeof result.id).toBe('string');
     });
   });
 
@@ -353,7 +396,7 @@ describe('AbstractTestManager', () => {
       const controller = new AbstractTestManager();
       await controller.addMessage({
         id: messageId,
-        messageParams: { from: '0x1234', data: 'test' },
+        messageParams: { ...mockMessageParams },
         status: 'status',
         time: 10,
         type: 'type',
@@ -374,7 +417,7 @@ describe('AbstractTestManager', () => {
       const controller = new AbstractTestManager();
       await controller.addMessage({
         id: messageId,
-        messageParams: { from: '0x1234', data: 'test' },
+        messageParams: { ...mockMessageParams },
         status: 'status',
         time: 10,
         type: 'type',

packages/message-manager/src/AbstractMessageManager.ts

@@ -1,7 +1,14 @@
 import type { BaseConfig, BaseState } from '@metamask/base-controller';
 import { BaseControllerV1 } from '@metamask/base-controller';
+import type { ApprovalType } from '@metamask/controller-utils';
 import type { Hex, Json } from '@metamask/utils';
 import { EventEmitter } from 'events';
+import { v1 as random } from 'uuid';
+
+import type { DecryptMessageParams } from './DecryptMessageManager';
+import type { EncryptionPublicKeyParams } from './EncryptionPublicKeyManager';
+import type { PersonalMessageParams } from './PersonalMessageManager';
+import type { TypedMessageParams } from './TypedMessageManager';
 
 /**
  * @type OriginalRequest
@@ -13,6 +20,7 @@ import { EventEmitter } from 'events';
 // Convert to a `type` in a future major version.
 // eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 export interface OriginalRequest {
+  id?: number;
   origin?: string;
   securityAlertResponse?: Record<string, Json>;
 }
@@ -49,6 +57,7 @@ export interface AbstractMessage {
  * Represents the parameters to pass to the signing method once the signature request is approved.
  * @property from - Address from which the message is processed
  * @property origin? - Added for request origin identification
+ * @property requestId? - Original request id
  * @property deferSetAsSigned? - Whether to defer setting the message as signed immediately after the keyring is told to sign it
  */
 // This interface was created before this ESLint rule was added.
@@ -57,6 +66,7 @@ export interface AbstractMessage {
 export interface AbstractMessageParams {
   from: string;
   origin?: string;
+  requestId?: number;
   deferSetAsSigned?: boolean;
 }
 
@@ -93,6 +103,15 @@ export interface MessageManagerState<M extends AbstractMessage>
   unapprovedMessagesCount: number;
 }
 
+/**
+ * Represents the parameters to pass to the signing method once the signature request is approved.
+ */
+export type MessageParams =
+  | DecryptMessageParams
+  | EncryptionPublicKeyParams
+  | PersonalMessageParams
+  | TypedMessageParams;
+
 /**
  * A function for verifying a message, whether it is malicious or not
  */
@@ -127,6 +146,52 @@ export abstract class AbstractMessageManager<
 
   private readonly additionalFinishStatuses: string[];
 
+  /**
+   * Adds request props to the messsage params and returns a new messageParams object.
+   * @param messageParams - The messageParams to add the request props to.
+   * @param req - The original request object.
+   * @returns The messageParams with the request props added.
+   */
+  protected addRequestToMessageParams(
+    messageParams: MessageParams,
+    req?: OriginalRequest,
+  ) {
+    const updatedMessageParams = {
+      ...messageParams,
+    };
+
+    if (req) {
+      updatedMessageParams.requestId = req.id;
+      updatedMessageParams.origin = req.origin;
+    }
+
+    return updatedMessageParams;
+  }
+
+  /**
+   * Creates a new Message with a random id and an 'unapproved' status.
+   * @param messageParams - The messageParams to add the request props to.
+   * @param type - The approval type of the message.
+   * @param req - The original request object.
+   * @returns The new unapproved message for a specified type.
+   */
+  protected createUnapprovedMessage(
+    messageParams: MessageParams,
+    type: ApprovalType,
+    req?: OriginalRequest,
+  ) {
+    const messageId = random();
+
+    return {
+      id: messageId,
+      messageParams,
+      securityAlertResponse: req?.securityAlertResponse,
+      status: 'unapproved',
+      time: Date.now(),
+      type,
+    };
+  }
+
   /**
    * Saves the unapproved messages, and their count to state.
    *

packages/message-manager/src/DecryptMessageManager.test.ts

@@ -140,7 +140,7 @@ describe('DecryptMessageManager', () => {
     const messageStatus = 'unapproved';
     const messageType = 'eth_decrypt';
     const messageParams = { from: fromMock, data: dataMock };
-    const originalRequest = { origin: 'origin' };
+    const originalRequest = { id: 111, origin: 'origin' };
     const messageId = await controller.addUnapprovedMessage(
       messageParams,
       originalRequest,
@@ -151,6 +151,7 @@ describe('DecryptMessageManager', () => {
       throw new Error('"message" is falsy');
     }
     expect(message.messageParams.from).toBe(messageParams.from);
+    expect(message.messageParams.requestId).toBe(originalRequest.id);
     expect(message.time).toBeDefined();
     expect(message.status).toBe(messageStatus);
     expect(message.type).toBe(messageType);

packages/message-manager/src/DecryptMessageManager.ts

@@ -1,4 +1,4 @@
-import { v1 as random } from 'uuid';
+import { ApprovalType } from '@metamask/controller-utils';
 
 import type {
   AbstractMessage,
@@ -129,22 +129,24 @@ export class DecryptMessageManager extends AbstractMessageManager<
     messageParams: DecryptMessageParams,
     req?: OriginalRequest,
   ) {
-    if (req) {
-      messageParams.origin = req.origin;
-    }
-    messageParams.data = normalizeMessageData(messageParams.data);
-    const messageId = random();
-    const messageData: DecryptMessage = {
-      id: messageId,
+    const updatedMessageParams = this.addRequestToMessageParams(
       messageParams,
-      status: 'unapproved',
-      time: Date.now(),
-      type: 'eth_decrypt',
-    };
+      req,
+    ) as DecryptMessageParams;
+    updatedMessageParams.data = normalizeMessageData(messageParams.data);
+
+    const messageData = this.createUnapprovedMessage(
+      updatedMessageParams,
+      ApprovalType.EthDecrypt,
+      req,
+    ) as DecryptMessage;
+
+    const messageId = messageData.id;
+
     await this.addMessage(messageData);
     this.hub.emit(`unapprovedMessage`, {
-      ...messageParams,
-      ...{ metamaskId: messageId },
+      ...updatedMessageParams,
+      metamaskId: messageId,
     });
     return messageId;
   }

packages/message-manager/src/EncryptionPublicKeyManager.test.ts

@@ -120,7 +120,7 @@ describe('EncryptionPublicKeyManager', () => {
     const messageParams = {
       from: fromMock,
     };
-    const originalRequest = { origin: 'origin' };
+    const originalRequest = { id: 111, origin: 'origin' };
     const messageId = await controller.addUnapprovedMessage(
       messageParams,
       originalRequest,
@@ -131,6 +131,7 @@ describe('EncryptionPublicKeyManager', () => {
       throw new Error('"message" is falsy');
     }
     expect(message.messageParams.from).toBe(messageParams.from);
+    expect(message.messageParams.requestId).toBe(originalRequest.id);
     expect(message.time).toBeDefined();
     expect(message.status).toBe(messageStatus);
     expect(message.type).toBe(messageType);

packages/message-manager/src/EncryptionPublicKeyManager.ts

@@ -1,4 +1,4 @@
-import { v1 as random } from 'uuid';
+import { ApprovalType } from '@metamask/controller-utils';
 
 import type {
   AbstractMessage,
@@ -120,21 +120,23 @@ export class EncryptionPublicKeyManager extends AbstractMessageManager<
     messageParams: EncryptionPublicKeyParams,
     req?: OriginalRequest,
   ): Promise<string> {
-    if (req) {
-      messageParams.origin = req.origin;
-    }
-    const messageId = random();
-    const messageData: EncryptionPublicKey = {
-      id: messageId,
+    const updatedMessageParams = this.addRequestToMessageParams(
       messageParams,
-      status: 'unapproved',
-      time: Date.now(),
-      type: 'eth_getEncryptionPublicKey',
-    };
+      req,
+    ) as EncryptionPublicKeyParams;
+
+    const messageData = this.createUnapprovedMessage(
+      updatedMessageParams,
+      ApprovalType.EthGetEncryptionPublicKey,
+      req,
+    ) as EncryptionPublicKey;
+
+    const messageId = messageData.id;
+
     await this.addMessage(messageData);
     this.hub.emit(`unapprovedMessage`, {
-      ...messageParams,
-      ...{ metamaskId: messageId },
+      ...updatedMessageParams,
+      metamaskId: messageId,
     });
     return messageId;
   }

packages/message-manager/src/PersonalMessageManager.test.ts

@@ -79,6 +79,7 @@ describe('PersonalMessageManager', () => {
       from: fromMock,
     };
     const originalRequest = {
+      id: 111,
       origin: 'origin',
       // TODO: Either fix this lint violation or explain why it's necessary to ignore.
       // eslint-disable-next-line @typescript-eslint/naming-convention
@@ -95,6 +96,7 @@ describe('PersonalMessageManager', () => {
     }
     expect(message.messageParams.from).toBe(messageParams.from);
     expect(message.messageParams.data).toBe(messageParams.data);
+    expect(message.messageParams.requestId).toBe(originalRequest.id);
     expect(message.time).toBeDefined();
     expect(message.status).toBe(messageStatus);
     expect(message.type).toBe(messageType);