From b576c5245ccd2811de5ef9c9953555ea6a78fe99 Mon Sep 17 00:00:00 2001 From: cryptodev-2s <109512101+cryptodev-2s@users.noreply.github.com> Date: Fri, 28 Jul 2023 20:09:14 +0100 Subject: [PATCH] Use getKeyringForAccount from core KeyringController (#20202) --- .../controllers/encryption-public-key.test.ts | 43 ++++++------------- .../controllers/encryption-public-key.ts | 23 ++++++---- app/scripts/metamask-controller.js | 28 ++++++++---- app/scripts/metamask-controller.test.js | 9 ++-- 4 files changed, 54 insertions(+), 49 deletions(-) diff --git a/app/scripts/controllers/encryption-public-key.test.ts b/app/scripts/controllers/encryption-public-key.test.ts index cb5581831..c36418abf 100644 --- a/app/scripts/controllers/encryption-public-key.test.ts +++ b/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 = () => }, } as any as jest.Mocked); -const createKeyringControllerMock = () => ({ - getKeyringForAccount: jest.fn(), - getEncryptionPublicKey: jest.fn(), -}); - describe('EncryptionPublicKeyController', () => { let encryptionPublicKeyController: EncryptionPublicKeyController; @@ -88,7 +83,8 @@ describe('EncryptionPublicKeyController', () => { const encryptionPublicKeyManagerMock = createEncryptionPublicKeyManagerMock(); 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), diff --git a/app/scripts/controllers/encryption-public-key.ts b/app/scripts/controllers/encryption-public-key.ts index da50f1afe..8448fba8a 100644 --- a/app/scripts/controllers/encryption-public-key.ts +++ b/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; + getAccountKeyringType: (account: string) => Promise; 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; + + private _getAccountKeyringType: (account: string) => Promise; 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 { - 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, ); diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js index 7ff6b5d26..9efbb8447 100644 --- a/app/scripts/metamask-controller.js +++ b/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() : {}; diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js index 32d7b3c2f..a8d360e10 100644 --- a/app/scripts/metamask-controller.test.js +++ b/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, ), );