fix: Use latest DecryptMessageManager EncryptMessageManager to ex…

…tend controllers (#29237)

<!--
Please submit this PR as a draft initially.
Do not mark it as "Ready for review" until the template has been
completely filled out, and PR status checks have passed at least once.
-->

## **Description**

<!--
Write a short description of the changes included in this pull request,
also include relevant motivation and context. Have in mind the following
questions:
1. What is the reason for the change?
2. What is the improvement/solution?
-->

This PR aims to adapt latest changes in the `@metamask/message-manager`
`DecryptMessageManager` and `EncryptionPublicKeyManager` classes. The
latest change in core mainly tries to adapt `BaseControllerV2`. Hence in
the extension, minimum changes made in the wrapper classes to keep
functionality as is.

[![Open in GitHub
Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/MetaMask/metamask-extension/pull/29237?quickstart=1)

## **Related issues**

Fixes: MetaMask/MetaMask-planning#3747
Related core PR: MetaMask/core#5103

## **Manual testing steps**

1. Go to test-dapp and `Encrypt / Decrypt` section
2. "Get Encryption Key", "Encrypt" and "Decrypt" must be functional and
work without issue

## **Screenshots/Recordings**



https://github.com/user-attachments/assets/3857cb34-2542-41c3-ac1e-6d95e98b0dfb



### **Before**

<!-- [screenshots/recordings] -->

### **After**

<!-- [screenshots/recordings] -->

## **Pre-merge author checklist**

- [X] I've followed [MetaMask Contributor
Docs](https://github.com/MetaMask/contributor-docs) and [MetaMask
Extension Coding
Standards](https://github.com/MetaMask/metamask-extension/blob/main/.github/guidelines/CODING_GUIDELINES.md).
- [X] I've completed the PR template to the best of my ability
- [X] I’ve included tests if applicable
- [X] I’ve documented my code using [JSDoc](https://jsdoc.app/) format
if applicable
- [X] I’ve applied the right labels on the PR (see [labeling
guidelines](https://github.com/MetaMask/metamask-extension/blob/main/.github/guidelines/LABELING_GUIDELINES.md)).
Not required for external contributors.

## **Pre-merge reviewer checklist**

- [ ] I've manually tested the PR (e.g. pull and build branch, run the
app, test code being changed).
- [ ] I confirm that this PR addresses all acceptance criteria described
in the ticket it closes and includes the necessary testing evidence such
as recordings and or screenshots.

---------

Co-authored-by: MetaMask Bot <[email protected]>

app/scripts/background.js

@@ -1045,12 +1045,12 @@ export function setupController(
   //
   updateBadge();
 
-  controller.decryptMessageController.hub.on(
-    METAMASK_CONTROLLER_EVENTS.UPDATE_BADGE,
+  controller.controllerMessenger.subscribe(
+    METAMASK_CONTROLLER_EVENTS.DECRYPT_MESSAGE_MANAGER_UPDATE_BADGE,
     updateBadge,
   );
-  controller.encryptionPublicKeyController.hub.on(
-    METAMASK_CONTROLLER_EVENTS.UPDATE_BADGE,
+  controller.controllerMessenger.subscribe(
+    METAMASK_CONTROLLER_EVENTS.ENCRYPTION_PUBLIC_KEY_MANAGER_UPDATE_BADGE,
     updateBadge,
   );
   controller.signatureController.hub.on(

app/scripts/controllers/decrypt-message.test.ts

@@ -2,6 +2,7 @@ import {
   DecryptMessageManager,
   DecryptMessageParams,
 } from '@metamask/message-manager';
+import type { DecryptMessageManagerMessenger } from '@metamask/message-manager';
 import { MetaMetricsEventCategory } from '../../../shared/constants/metametrics';
 import DecryptMessageController, {
   DecryptMessageControllerMessenger,
@@ -36,11 +37,15 @@ const createMessengerMock = () =>
   ({
     registerActionHandler: jest.fn(),
     registerInitialEventPayload: jest.fn(),
+    subscribe: jest.fn(),
     publish: jest.fn(),
     call: jest.fn(),
-    // TODO: Replace `any` with type
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  } as any as jest.Mocked<DecryptMessageControllerMessenger>);
+  } as unknown as jest.Mocked<DecryptMessageControllerMessenger>);
+
+const createManagerMessengerMock = () =>
+  ({
+    subscribe: jest.fn(),
+  } as unknown as jest.Mocked<DecryptMessageManagerMessenger>);
 
 const createDecryptMessageManagerMock = <T>() =>
   ({
@@ -64,20 +69,14 @@ const createDecryptMessageManagerMock = <T>() =>
   } as any as jest.Mocked<T>);
 
 describe('DecryptMessageController', () => {
-  class MockDecryptMessageController extends DecryptMessageController {
-    // update is protected, so we expose it for typechecking here
-    public update(callback: Parameters<DecryptMessageController['update']>[0]) {
-      return super.update(callback);
-    }
-  }
-
-  let decryptMessageController: MockDecryptMessageController;
+  let decryptMessageController: DecryptMessageController;
 
   const decryptMessageManagerConstructorMock =
     DecryptMessageManager as jest.MockedClass<typeof DecryptMessageManager>;
   const getStateMock = jest.fn();
   const keyringControllerMock = createKeyringControllerMock();
   const messengerMock = createMessengerMock();
+  const managerMessengerMock = createManagerMessengerMock();
   const metricsEventMock = jest.fn();
 
   const decryptMessageManagerMock =
@@ -105,7 +104,7 @@ describe('DecryptMessageController', () => {
       decryptMessageManagerMock,
     );
 
-    decryptMessageController = new MockDecryptMessageController({
+    decryptMessageController = new DecryptMessageController({
       // TODO: Replace `any` with type
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       getState: getStateMock as any,
@@ -118,6 +117,7 @@ describe('DecryptMessageController', () => {
       // TODO: Replace `any` with type
       // eslint-disable-next-line @typescript-eslint/no-explicit-any
       metricsEvent: metricsEventMock as any,
+      managerMessenger: managerMessengerMock,
     } as DecryptMessageControllerOptions);
   });
 
@@ -127,23 +127,10 @@ describe('DecryptMessageController', () => {
   });
 
   it('should reset state', () => {
-    decryptMessageController.update(() => ({
-      unapprovedDecryptMsgs: {
-        [messageIdMock]: messageMock,
-        // TODO: Replace `any` with type
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      } as any,
-      unapprovedDecryptMsgCount: 1,
-    }));
     decryptMessageController.resetState();
     expect(decryptMessageController.state).toStrictEqual(getDefaultState());
   });
 
-  it('should clear unapproved messages', () => {
-    decryptMessageController.clearUnapproved();
-    expect(decryptMessageController.state).toStrictEqual(getDefaultState());
-    expect(decryptMessageManagerMock.update).toBeCalledTimes(1);
-  });
   it('should add unapproved messages', async () => {
     await decryptMessageController.newRequestDecryptMessage(
       messageMock,

app/scripts/controllers/decrypt-message.ts

@@ -1,8 +1,6 @@
-import EventEmitter from 'events';
 import log from 'loglevel';
 import {
   AbstractMessage,
-  AbstractMessageManager,
   AbstractMessageParams,
   AbstractMessageParamsMetamask,
   MessageManagerState,
@@ -11,6 +9,11 @@ import {
   DecryptMessageParams,
   DecryptMessageParamsMetamask,
 } from '@metamask/message-manager';
+import type {
+  DecryptMessageManagerMessenger,
+  DecryptMessageManagerState,
+  DecryptMessageManagerUnapprovedMessageAddedEvent,
+} from '@metamask/message-manager';
 import {
   BaseController,
   RestrictedControllerMessenger,
@@ -34,6 +37,8 @@ const stateMetadata = {
   unapprovedDecryptMsgCount: { persist: false, anonymous: false },
 };
 
+export const managerName = 'DecryptMessageManager';
+
 /**
  * Type guard that checks for the presence of the required properties
  * for EIP-1024 encrypted data.
@@ -86,38 +91,49 @@ export type DecryptMessageControllerState = {
   unapprovedDecryptMsgCount: number;
 };
 
-export type GetDecryptMessageState = {
+export type GetDecryptMessageControllerState = {
   type: `${typeof controllerName}:getState`;
   handler: () => DecryptMessageControllerState;
 };
 
-export type DecryptMessageStateChange = {
+export type DecryptMessageControllerStateChange = {
   type: `${typeof controllerName}:stateChange`;
   payload: [DecryptMessageControllerState, Patch[]];
 };
 
-export type DecryptMessageControllerActions = GetDecryptMessageState;
+export type DecryptMessageControllerActions = GetDecryptMessageControllerState;
 
-export type DecryptMessageControllerEvents = DecryptMessageStateChange;
+export type DecryptMessageControllerEvents =
+  DecryptMessageControllerStateChange;
 
 type AllowedActions =
   | AddApprovalRequest
   | AcceptRequest
   | RejectRequest
   | KeyringControllerDecryptMessageAction;
 
+type DecryptMessageManagerStateChangeEvent = {
+  type: `DecryptMessageManager:stateChange`;
+  payload: [DecryptMessageManagerState, Patch[]];
+};
+
+type AllowedEvents =
+  | DecryptMessageManagerStateChangeEvent
+  | DecryptMessageManagerUnapprovedMessageAddedEvent;
+
 export type DecryptMessageControllerMessenger = RestrictedControllerMessenger<
   typeof controllerName,
   DecryptMessageControllerActions | AllowedActions,
-  DecryptMessageControllerEvents,
+  DecryptMessageControllerEvents | AllowedEvents,
   AllowedActions['type'],
-  never
+  AllowedEvents['type']
 >;
 
 export type DecryptMessageControllerOptions = {
   // TODO: Replace `any` with type
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   getState: () => any;
+  managerMessenger: DecryptMessageManagerMessenger;
   messenger: DecryptMessageControllerMessenger;
   // TODO: Replace `any` with type
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -132,8 +148,6 @@ export default class DecryptMessageController extends BaseController<
   DecryptMessageControllerState,
   DecryptMessageControllerMessenger
 > {
-  hub: EventEmitter;
-
   // TODO: Replace `any` with type
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   private _getState: () => any;
@@ -151,11 +165,13 @@ export default class DecryptMessageController extends BaseController<
    * @param options.getState - Callback to retrieve all user state.
    * @param options.messenger - A reference to the messaging system.
    * @param options.metricsEvent - A function for emitting a metric event.
+   * @param options.managerMessenger - A reference to the messenger need by the message manager.
    */
   constructor({
     getState,
     metricsEvent,
     messenger,
+    managerMessenger,
   }: DecryptMessageControllerOptions) {
     super({
       metadata: stateMetadata,
@@ -166,28 +182,18 @@ export default class DecryptMessageController extends BaseController<
     this._getState = getState;
     this._metricsEvent = metricsEvent;
 
-    this.hub = new EventEmitter();
-
-    this._decryptMessageManager = new DecryptMessageManager(
-      undefined,
-      undefined,
-      undefined,
-      ['decrypted'],
-    );
-
-    this._decryptMessageManager.hub.on('updateBadge', () => {
-      this.hub.emit('updateBadge');
+    this._decryptMessageManager = new DecryptMessageManager({
+      additionalFinishStatuses: ['decrypted'],
+      messenger: managerMessenger,
     });
 
-    this._decryptMessageManager.hub.on(
-      'unapprovedMessage',
-      (messageParams: AbstractMessageParamsMetamask) => {
-        this._requestApproval(messageParams);
-      },
+    messenger.subscribe(
+      `${managerName}:unapprovedMessage`,
+      this._requestApproval.bind(this),
     );
 
     this._subscribeToMessageState(
-      this._decryptMessageManager,
+      messenger,
       (state, newMessages, messageCount) => {
         state.unapprovedDecryptMsgs = newMessages;
         state.unapprovedDecryptMsgCount = messageCount;
@@ -215,10 +221,7 @@ export default class DecryptMessageController extends BaseController<
    * Clears all unapproved messages from memory.
    */
   clearUnapproved() {
-    this._decryptMessageManager.update({
-      unapprovedMessages: {},
-      unapprovedMessagesCount: 0,
-    });
+    this._decryptMessageManager.clearUnapprovedMessages();
   }
 
   /**
@@ -331,11 +334,7 @@ export default class DecryptMessageController extends BaseController<
   }
 
   private _cancelAbstractMessage(
-    messageManager: AbstractMessageManager<
-      AbstractMessage,
-      AbstractMessageParams,
-      AbstractMessageParamsMetamask
-    >,
+    messageManager: DecryptMessageManager,
     messageId: string,
     reason?: string,
   ) {
@@ -356,27 +355,26 @@ export default class DecryptMessageController extends BaseController<
   }
 
   private _subscribeToMessageState(
-    messageManager: AbstractMessageManager<
-      AbstractMessage,
-      AbstractMessageParams,
-      AbstractMessageParamsMetamask
-    >,
+    controllerMessenger: DecryptMessageControllerMessenger,
     updateState: (
       state: DecryptMessageControllerState,
       newMessages: Record<string, StateMessage>,
       messageCount: number,
     ) => void,
   ) {
-    messageManager.subscribe((state: MessageManagerState<AbstractMessage>) => {
-      const newMessages = this._migrateMessages(
-        // TODO: Replace `any` with type
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        state.unapprovedMessages as any,
-      );
-      this.update((draftState) => {
-        updateState(draftState, newMessages, state.unapprovedMessagesCount);
-      });
-    });
+    controllerMessenger.subscribe(
+      `${managerName}:stateChange`,
+      (state: MessageManagerState<AbstractMessage>) => {
+        const newMessages = this._migrateMessages(
+          // TODO: Replace `any` with type
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+          state.unapprovedMessages as any,
+        );
+        this.update((draftState) => {
+          updateState(draftState, newMessages, state.unapprovedMessagesCount);
+        });
+      },
+    );
   }
 
   private _migrateMessages(