refactor: message-manager abstract addRequestToMessageParams and crea…

…teUnapprovedMessage
MetaMask · Aug 26, 2024 · 41878c9 · 41878c9
1 parent 9b12571
commit 41878c9
Show file tree

Hide file tree

Showing 5 changed files with 133 additions and 66 deletions.
diff --git a/packages/message-manager/src/AbstractMessageManager.ts b/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
@@ -96,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.
+ */
+type MessageParams =
+  | DecryptMessageParams
+  | EncryptionPublicKeyParams
+  | PersonalMessageParams
+  | TypedMessageParams;
+
 /**
  * A function for verifying a message, whether it is malicious or not
  */
@@ -130,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.
    *

diff --git a/packages/message-manager/src/DecryptMessageManager.ts b/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,23 +129,24 @@ export class DecryptMessageManager extends AbstractMessageManager<
     messageParams: DecryptMessageParams,
     req?: OriginalRequest,
   ) {
-    if (req) {
-      messageParams.requestId = req.id;
-      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;
   }

diff --git a/packages/message-manager/src/EncryptionPublicKeyManager.ts b/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,22 +120,23 @@ export class EncryptionPublicKeyManager extends AbstractMessageManager<
     messageParams: EncryptionPublicKeyParams,
     req?: OriginalRequest,
   ): Promise<string> {
-    if (req) {
-      messageParams.requestId = req.id;
-      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;
   }

diff --git a/packages/message-manager/src/PersonalMessageManager.ts b/packages/message-manager/src/PersonalMessageManager.ts
@@ -1,6 +1,5 @@
 import type { SIWEMessage } from '@metamask/controller-utils';
-import { detectSIWE } from '@metamask/controller-utils';
-import { v1 as random } from 'uuid';
+import { detectSIWE, ApprovalType } from '@metamask/controller-utils';
 
 import type {
   AbstractMessage,
@@ -91,28 +90,27 @@ export class PersonalMessageManager extends AbstractMessageManager<
     req?: OriginalRequest,
   ): Promise<string> {
     validateSignMessageData(messageParams);
-    if (req) {
-      messageParams.requestId = req.id;
-      messageParams.origin = req.origin;
-    }
-    messageParams.data = normalizeMessageData(messageParams.data);
 
     const ethereumSignInData = detectSIWE(messageParams);
-    const finalMsgParams = { ...messageParams, siwe: ethereumSignInData };
+    const updatedMessageParams = this.addRequestToMessageParams(
+      messageParams,
+      req,
+    ) as PersonalMessageParams;
 
-    const messageId = random();
-    const messageData: PersonalMessage = {
-      id: messageId,
-      messageParams: finalMsgParams,
-      securityAlertResponse: req?.securityAlertResponse,
-      status: 'unapproved',
-      time: Date.now(),
-      type: 'personal_sign',
-    };
+    updatedMessageParams.data = normalizeMessageData(messageParams.data);
+    updatedMessageParams.siwe = ethereumSignInData;
+
+    const messageData = this.createUnapprovedMessage(
+      updatedMessageParams,
+      ApprovalType.PersonalSign,
+      req,
+    ) as PersonalMessage;
+
+    const messageId = messageData.id;
     await this.addMessage(messageData);
     this.hub.emit(`unapprovedMessage`, {
-      ...finalMsgParams,
-      ...{ metamaskId: messageId },
+      ...updatedMessageParams,
+      metamaskId: messageId,
     });
     return messageId;
   }

diff --git a/packages/message-manager/src/TypedMessageManager.ts b/packages/message-manager/src/TypedMessageManager.ts
@@ -1,4 +1,4 @@
-import { v1 as random } from 'uuid';
+import { ApprovalType } from '@metamask/controller-utils';
 
 import type {
   AbstractMessage,
@@ -129,26 +129,31 @@ export class TypedMessageManager extends AbstractMessageManager<
       messageParams.data = JSON.stringify(messageParams.data);
     }
 
-    const messageId = random();
-    const messageParamsMetamask = {
+    const updatedMessageParams = this.addRequestToMessageParams(
+      messageParams,
+      req,
+    ) as TypedMessageParams;
+
+    const messageData = this.createUnapprovedMessage(
+      updatedMessageParams,
+      ApprovalType.EthSignTypedData,
+      req,
+    ) as TypedMessage;
+
+    const messageId = messageData.id;
+
+    await this.addMessage(messageData);
+
+    /**
+     * This intentionally splays messageParams rather than updatedMessageParams. I'm unsure if this
+     * is exactly what we want, but I am preserving existing logic.
+     */
+    this.hub.emit(`unapprovedMessage`, {
       ...messageParams,
       metamaskId: messageId,
       version,
-    };
-    if (req) {
-      messageParams.requestId = req.id;
-      messageParams.origin = req.origin;
-    }
-    const messageData: TypedMessage = {
-      id: messageId,
-      messageParams,
-      securityAlertResponse: req?.securityAlertResponse,
-      status: 'unapproved',
-      time: Date.now(),
-      type: 'eth_signTypedData',
-    };
-    await this.addMessage(messageData);
-    this.hub.emit(`unapprovedMessage`, messageParamsMetamask);
+    });
+
     return messageId;
   }