Use getKeyringForAccount from core KeyringController (#20202)

app/scripts/controllers/encryption-public-key.test.ts

@@ -19,7 +19,7 @@ const messageIdMock2 = '456';
 const stateMock = { test: 123 };
 const addressMock = '0xc38bf1ad06ef69f0c04e29dbeb4152b4175f0a8d';
 const publicKeyMock = '32762347862378feb87123781623a=';
-const keyringMock = { type: KeyringType.hdKeyTree };
+const keyringTypeMock = KeyringType.hdKeyTree;
 
 const messageParamsMock = {
   from: addressMock,
@@ -73,11 +73,6 @@ const createEncryptionPublicKeyManagerMock = <T>() =>
     },
   } as any as jest.Mocked<T>);
 
-const createKeyringControllerMock = () => ({
-  getKeyringForAccount: jest.fn(),
-  getEncryptionPublicKey: jest.fn(),
-});
-
 describe('EncryptionPublicKeyController', () => {
   let encryptionPublicKeyController: EncryptionPublicKeyController;
 
@@ -88,7 +83,8 @@ describe('EncryptionPublicKeyController', () => {
   const encryptionPublicKeyManagerMock =
     createEncryptionPublicKeyManagerMock<EncryptionPublicKeyManager>();
   const messengerMock = createMessengerMock();
-  const keyringControllerMock = createKeyringControllerMock();
+  const getEncryptionPublicKeyMock = jest.fn();
+  const getAccountKeyringTypeMock = jest.fn();
   const getStateMock = jest.fn();
   const metricsEventMock = jest.fn();
 
@@ -101,7 +97,8 @@ describe('EncryptionPublicKeyController', () => {
 
     encryptionPublicKeyController = new EncryptionPublicKeyController({
       messenger: messengerMock as any,
-      keyringController: keyringControllerMock as any,
+      getEncryptionPublicKey: getEncryptionPublicKeyMock as any,
+      getAccountKeyringType: getAccountKeyringTypeMock as any,
       getState: getStateMock as any,
       metricsEvent: metricsEventMock as any,
     } as EncryptionPublicKeyControllerOptions);
@@ -203,9 +200,7 @@ describe('EncryptionPublicKeyController', () => {
     ])(
       'throws if keyring is not supported',
       async (keyringName, keyringType) => {
-        keyringControllerMock.getKeyringForAccount.mockResolvedValueOnce({
-          type: keyringType,
-        });
+        getAccountKeyringTypeMock.mockResolvedValueOnce(keyringType);
 
         await expect(
           encryptionPublicKeyController.newRequestEncryptionPublicKey(
@@ -219,9 +214,7 @@ describe('EncryptionPublicKeyController', () => {
     );
 
     it('adds message to message manager', async () => {
-      keyringControllerMock.getKeyringForAccount.mockResolvedValueOnce(
-        keyringMock,
-      );
+      getAccountKeyringTypeMock.mockResolvedValueOnce(keyringTypeMock);
 
       await encryptionPublicKeyController.newRequestEncryptionPublicKey(
         addressMock,
@@ -243,9 +236,7 @@ describe('EncryptionPublicKeyController', () => {
         from: messageParamsMock.data,
       });
 
-      keyringControllerMock.getEncryptionPublicKey.mockResolvedValueOnce(
-        publicKeyMock,
-      );
+      getEncryptionPublicKeyMock.mockResolvedValueOnce(publicKeyMock);
     });
 
     it('approves message and signs', async () => {
@@ -253,10 +244,8 @@ describe('EncryptionPublicKeyController', () => {
         messageParamsMock,
       );
 
-      expect(
-        keyringControllerMock.getEncryptionPublicKey,
-      ).toHaveBeenCalledTimes(1);
-      expect(keyringControllerMock.getEncryptionPublicKey).toHaveBeenCalledWith(
+      expect(getEncryptionPublicKeyMock).toHaveBeenCalledTimes(1);
+      expect(getEncryptionPublicKeyMock).toHaveBeenCalledWith(
         messageParamsMock.data,
       );
 
@@ -294,10 +283,8 @@ describe('EncryptionPublicKeyController', () => {
     });
 
     it('rejects message on error', async () => {
-      keyringControllerMock.getEncryptionPublicKey.mockReset();
-      keyringControllerMock.getEncryptionPublicKey.mockRejectedValue(
-        new Error('Test Error'),
-      );
+      getEncryptionPublicKeyMock.mockReset();
+      getEncryptionPublicKeyMock.mockRejectedValue(new Error('Test Error'));
 
       await expect(
         encryptionPublicKeyController.encryptionPublicKey(messageParamsMock),
@@ -312,10 +299,8 @@ describe('EncryptionPublicKeyController', () => {
     });
 
     it('rejects approval on error', async () => {
-      keyringControllerMock.getEncryptionPublicKey.mockReset();
-      keyringControllerMock.getEncryptionPublicKey.mockRejectedValue(
-        new Error('Test Error'),
-      );
+      getEncryptionPublicKeyMock.mockReset();
+      getEncryptionPublicKeyMock.mockRejectedValue(new Error('Test Error'));
 
       await expect(
         encryptionPublicKeyController.encryptionPublicKey(messageParamsMock),

app/scripts/controllers/encryption-public-key.ts

@@ -4,7 +4,6 @@ import {
   EncryptionPublicKeyManager,
   EncryptionPublicKeyParamsMetamask,
 } from '@metamask/message-manager';
-import { KeyringController } from '@metamask/eth-keyring-controller';
 import {
   AbstractMessageManager,
   AbstractMessage,
@@ -83,7 +82,8 @@ export type EncryptionPublicKeyControllerMessenger =
 
 export type EncryptionPublicKeyControllerOptions = {
   messenger: EncryptionPublicKeyControllerMessenger;
-  keyringController: KeyringController;
+  getEncryptionPublicKey: (address: string) => Promise<string>;
+  getAccountKeyringType: (account: string) => Promise<string>;
   getState: () => any;
   metricsEvent: (payload: any, options?: any) => void;
 };
@@ -98,7 +98,9 @@ export default class EncryptionPublicKeyController extends BaseControllerV2<
 > {
   hub: EventEmitter;
 
-  private _keyringController: KeyringController;
+  private _getEncryptionPublicKey: (address: string) => Promise<string>;
+
+  private _getAccountKeyringType: (account: string) => Promise<string>;
 
   private _getState: () => any;
 
@@ -111,13 +113,15 @@ export default class EncryptionPublicKeyController extends BaseControllerV2<
    *
    * @param options - The controller options.
    * @param options.messenger - The restricted controller messenger for the EncryptionPublicKey controller.
-   * @param options.keyringController - An instance of a keyring controller used to extract the encryption public key.
+   * @param options.getEncryptionPublicKey - Callback to get the keyring encryption public key.
+   * @param options.getAccountKeyringType - Callback to get the keyring type.
    * @param options.getState - Callback to retrieve all user state.
    * @param options.metricsEvent - A function for emitting a metric event.
    */
   constructor({
     messenger,
-    keyringController,
+    getEncryptionPublicKey,
+    getAccountKeyringType,
     getState,
     metricsEvent,
   }: EncryptionPublicKeyControllerOptions) {
@@ -128,7 +132,8 @@ export default class EncryptionPublicKeyController extends BaseControllerV2<
       state: getDefaultState(),
     });
 
-    this._keyringController = keyringController;
+    this._getEncryptionPublicKey = getEncryptionPublicKey;
+    this._getAccountKeyringType = getAccountKeyringType;
     this._getState = getState;
     this._metricsEvent = metricsEvent;
 
@@ -186,9 +191,9 @@ export default class EncryptionPublicKeyController extends BaseControllerV2<
     address: string,
     req: OriginalRequest,
   ): Promise<string> {
-    const keyring = await this._keyringController.getKeyringForAccount(address);
+    const keyringType = await this._getAccountKeyringType(address);
 
-    switch (keyring.type) {
+    switch (keyringType) {
       case KeyringType.ledger: {
         return new Promise((_, reject) => {
           reject(
@@ -244,7 +249,7 @@ export default class EncryptionPublicKeyController extends BaseControllerV2<
         await this._encryptionPublicKeyManager.approveMessage(msgParams);
 
       // EncryptionPublicKey message
-      const publicKey = await this._keyringController.getEncryptionPublicKey(
+      const publicKey = await this._getEncryptionPublicKey(
         cleanMessageParams.from,
       );
 

app/scripts/metamask-controller.js

@@ -904,9 +904,8 @@ export default class MetamaskController extends EventEmitter {
               (address) => !identities[address],
             );
             const keyringTypesWithMissingIdentities =
-              accountsMissingIdentities.map(
-                (address) =>
-                  this.keyringController.getKeyringForAccount(address)?.type,
+              accountsMissingIdentities.map((address) =>
+                this.coreKeyringController.getAccountKeyringType(address),
               );
 
             const identitiesCount = Object.keys(identities || {}).length;
@@ -1349,7 +1348,14 @@ export default class MetamaskController extends EventEmitter {
           `${this.approvalController.name}:rejectRequest`,
         ],
       }),
-      keyringController: this.keyringController,
+      getEncryptionPublicKey:
+        this.keyringController.getEncryptionPublicKey.bind(
+          this.keyringController,
+        ),
+      getAccountKeyringType:
+        this.coreKeyringController.getAccountKeyringType.bind(
+          this.coreKeyringController,
+        ),
       getState: this.getState.bind(this),
       metricsEvent: this.metaMetricsController.trackEvent.bind(
         this.metaMetricsController,
@@ -3265,8 +3271,10 @@ export default class MetamaskController extends EventEmitter {
    * @returns {'hardware' | 'imported' | 'MetaMask'}
    */
   async getAccountType(address) {
-    const keyring = await this.keyringController.getKeyringForAccount(address);
-    switch (keyring.type) {
+    const keyringType = await this.coreKeyringController.getAccountKeyringType(
+      address,
+    );
+    switch (keyringType) {
       case KeyringType.trezor:
       case KeyringType.lattice:
       case KeyringType.qr:
@@ -3288,7 +3296,9 @@ export default class MetamaskController extends EventEmitter {
    * @returns {'ledger' | 'lattice' | 'N/A' | string}
    */
   async getDeviceModel(address) {
-    const keyring = await this.keyringController.getKeyringForAccount(address);
+    const keyring = await this.coreKeyringController.getKeyringForAccount(
+      address,
+    );
     switch (keyring.type) {
       case KeyringType.trezor:
         return keyring.getModel();
@@ -3527,7 +3537,9 @@ export default class MetamaskController extends EventEmitter {
     this.custodyController.removeAccount(address);
     ///: END:ONLY_INCLUDE_IN(build-mmi)
 
-    const keyring = await this.keyringController.getKeyringForAccount(address);
+    const keyring = await this.coreKeyringController.getKeyringForAccount(
+      address,
+    );
     // Remove account from the keyring
     await this.keyringController.removeAccount(address);
     const updatedKeyringAccounts = keyring ? await keyring.getAccounts() : {};

app/scripts/metamask-controller.test.js

@@ -859,7 +859,10 @@ describe('MetaMaskController', function () {
         sinon.stub(metamaskController.keyringController, 'removeAccount');
         sinon.stub(metamaskController, 'removeAllAccountPermissions');
         sinon
-          .stub(metamaskController.keyringController, 'getKeyringForAccount')
+          .stub(
+            metamaskController.coreKeyringController,
+            'getKeyringForAccount',
+          )
           .returns(Promise.resolve(mockKeyring));
 
         ret = await metamaskController.removeAccount(addressToRemove);
@@ -906,9 +909,9 @@ describe('MetaMaskController', function () {
       it('should return address', async function () {
         assert.equal(ret, '0x1');
       });
-      it('should call keyringController.getKeyringForAccount', async function () {
+      it('should call coreKeyringController.getKeyringForAccount', async function () {
         assert(
-          metamaskController.keyringController.getKeyringForAccount.calledWith(
+          metamaskController.coreKeyringController.getKeyringForAccount.calledWith(
             addressToRemove,
           ),
         );