From 40ebd1833e38abfaccfb96005f462fbe518d7328 Mon Sep 17 00:00:00 2001 From: Michele Esposito Date: Tue, 16 May 2023 15:59:18 +0200 Subject: [PATCH] refactor: migrate KeyringController to BaseControllerV2 --- packages/keyring-controller/package.json | 1 + .../src/KeyringController.test.ts | 69 +++++-- .../src/KeyringController.ts | 180 ++++++++++-------- yarn.lock | 8 + 4 files changed, 155 insertions(+), 103 deletions(-) diff --git a/packages/keyring-controller/package.json b/packages/keyring-controller/package.json index 0dd3491bfe5..9fb8dc41759 100644 --- a/packages/keyring-controller/package.json +++ b/packages/keyring-controller/package.json @@ -48,6 +48,7 @@ "@metamask/scure-bip39": "^2.1.0", "@types/jest": "^27.4.1", "deepmerge": "^4.2.2", + "immer": "^10.0.2", "jest": "^27.5.1", "sinon": "^9.2.4", "ts-jest": "^27.1.4", diff --git a/packages/keyring-controller/src/KeyringController.test.ts b/packages/keyring-controller/src/KeyringController.test.ts index 64f90e38f93..b24adbc9b63 100644 --- a/packages/keyring-controller/src/KeyringController.test.ts +++ b/packages/keyring-controller/src/KeyringController.test.ts @@ -13,12 +13,15 @@ import * as uuid from 'uuid'; import { isValidHexAddress, NetworkType } from '@metamask/controller-utils'; import { keyringBuilderFactory } from '@metamask/eth-keyring-controller'; import { wordlist } from '@metamask/scure-bip39/dist/wordlists/english'; +import { ControllerMessenger } from '@metamask/base-controller'; import MockEncryptor, { mockKey } from '../tests/mocks/mockEncryptor'; import { AccountImportStrategy, - KeyringConfig, KeyringController, - KeyringMemState, + KeyringControllerEvents, + KeyringControllerMessenger, + KeyringControllerOptions, + KeyringState, KeyringTypes, } from './KeyringController'; @@ -247,10 +250,12 @@ describe('KeyringController', () => { await withController( { cacheEncryptionKey }, async ({ controller, initialState, preferences, encryptor }) => { - const cleanKeyringController = new KeyringController( + const cleanKeyringController = new KeyringController({ preferences, - { cacheEncryptionKey, encryptor }, - ); + cacheEncryptionKey, + encryptor, + messenger: buildKeyringControllerMessenger(), + }); const initialSeedWord = await controller.exportSeedPhrase( password, ); @@ -1043,20 +1048,14 @@ describe('KeyringController', () => { }); describe('subscribe and unsubscribe', () => { - it('should subscribe and unsubscribe', async () => { - await withController(async ({ controller }) => { + it('should fire stateChange event', async () => { + await withController(async ({ controller, messenger }) => { const listener = sinon.stub(); - controller.subscribe(listener); - await controller.importAccountWithStrategy( - AccountImportStrategy.privateKey, - [privateKey], - ); + messenger.subscribe('KeyringController:stateChange', listener); + + await controller.fullUpdate(); + expect(listener.called).toBe(true); - controller.unsubscribe(listener); - await controller.removeAccount( - '0x51253087e6f8358b5f10c0a94315d69db3357859', - ); - expect(listener.calledTwice).toBe(false); }); }); }); @@ -1581,6 +1580,7 @@ type WithControllerCallback = ({ preferences, initialState, encryptor, + messenger, }: { controller: KeyringController; preferences: { @@ -1591,15 +1591,40 @@ type WithControllerCallback = ({ setSelectedAddress: sinon.SinonStub; }; encryptor: MockEncryptor; - initialState: KeyringMemState; + initialState: KeyringState; + messenger: KeyringControllerMessenger; }) => Promise | ReturnValue; -type WithControllerOptions = Partial; +type WithControllerOptions = Partial; type WithControllerArgs = | [WithControllerCallback] | [WithControllerOptions, WithControllerCallback]; +/** + * Build a controller messenger that includes all events used by the keyring + * controller. + * + * @returns The controller messenger. + */ +function buildMessenger() { + return new ControllerMessenger(); +} + +/** + * Build a restricted controller messenger for the keyring controller. + * + * @param messenger - A controller messenger. + * @returns The keyring controller restricted messenger. + */ +function buildKeyringControllerMessenger(messenger = buildMessenger()) { + return messenger.getRestricted({ + name: 'KeyringController', + allowedActions: [], + allowedEvents: ['KeyringController:stateChange'], + }); +} + /** * Builds a controller based on the given options and creates a new vault * and keychain, then calls the given function with that controller. @@ -1622,9 +1647,12 @@ async function withController( updateIdentities: sinon.stub(), setSelectedAddress: sinon.stub(), }; - const controller = new KeyringController(preferences, { + const messenger = buildKeyringControllerMessenger(); + const controller = new KeyringController({ + preferences, encryptor, cacheEncryptionKey: false, + messenger, ...rest, }); const initialState = await controller.createNewVaultAndKeychain(password); @@ -1633,5 +1661,6 @@ async function withController( preferences, encryptor, initialState, + messenger, }); } diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts index e84a63a63a5..4c131b251b2 100644 --- a/packages/keyring-controller/src/KeyringController.ts +++ b/packages/keyring-controller/src/KeyringController.ts @@ -22,16 +22,17 @@ import { IKeyringState as IQRKeyringState, } from '@keystonehq/metamask-airgapped-keyring'; import { - BaseController, - BaseConfig, - BaseState, - Listener, + BaseControllerV2, + RestrictedControllerMessenger, } from '@metamask/base-controller'; import { PreferencesController } from '@metamask/preferences-controller'; import { PersonalMessageParams, TypedMessageParams, } from '@metamask/message-manager'; +import type { Patch } from 'immer'; + +const name = 'KeyringController'; /** * Available keyring types @@ -61,40 +62,54 @@ export interface KeyringObject { * * Keyring controller state * @property vault - Encrypted string representing keyring data - * @property keyrings - Group of accounts - */ -export interface KeyringState extends BaseState { - vault?: string; - keyrings: Keyring[]; -} - -/** - * @type KeyringMemState - * - * Keyring mem controller state * @property isUnlocked - Whether vault is unlocked * @property keyringTypes - Account types * @property keyrings - Group of accounts */ -export interface KeyringMemState extends BaseState { +export type KeyringState = { + vault?: string; isUnlocked: boolean; keyringTypes: string[]; keyrings: Keyring[]; encryptionKey?: string; encryptionSalt?: string; -} +}; + +export type KeyringControllerStateChangeEvent = { + type: `KeyringController:stateChange`; + payload: [KeyringState, Patch[]]; +}; + +export type KeyringControllerEvents = KeyringControllerStateChangeEvent; + +export type KeyringControllerMessenger = RestrictedControllerMessenger< + typeof name, + never, + KeyringControllerEvents, + string, + string +>; /** - * @type KeyringConfig + * @type KeyringControllerOptions * * Keyring controller configuration * @property encryptor - Keyring encryptor */ -export interface KeyringConfig extends BaseConfig { +export type KeyringControllerOptions = { + preferences: { + removeIdentity: PreferencesController['removeIdentity']; + syncIdentities: PreferencesController['syncIdentities']; + updateIdentities: PreferencesController['updateIdentities']; + setSelectedAddress: PreferencesController['setSelectedAddress']; + setAccountLabel?: PreferencesController['setAccountLabel']; + }; + messenger: KeyringControllerMessenger; encryptor?: any; keyringBuilders?: any[]; cacheEncryptionKey?: boolean; -} + state?: Partial; +}; /** * @type Keyring @@ -104,11 +119,11 @@ export interface KeyringConfig extends BaseConfig { * @property accounts - Associated accounts * @property index - Associated index */ -export interface Keyring { +export type Keyring = { accounts: string[]; type: string; index?: number; -} +}; /** * A strategy for importing an account @@ -129,6 +144,12 @@ export enum SignTypedDataVersion { V4 = 'V4', } +const defaultState: KeyringState = { + isUnlocked: false, + keyringTypes: [], + keyrings: [], +}; + /** * Controller responsible for establishing and managing user identity. * @@ -138,17 +159,13 @@ export enum SignTypedDataVersion { * with the internal keyring controller and handling certain complex operations that involve the * keyrings. */ -export class KeyringController extends BaseController< - KeyringConfig, - KeyringState +export class KeyringController extends BaseControllerV2< + typeof name, + KeyringState, + KeyringControllerMessenger > { private mutex = new Mutex(); - /** - * Name of this controller used during composition - */ - override name = 'KeyringController'; - private removeIdentity: PreferencesController['removeIdentity']; private syncIdentities: PreferencesController['syncIdentities']; @@ -165,46 +182,59 @@ export class KeyringController extends BaseController< * Creates a KeyringController instance. * * @param options - The controller options. - * @param options.removeIdentity - Remove the identity with the given address. - * @param options.syncIdentities - Sync identities with the given list of addresses. - * @param options.updateIdentities - Generate an identity for each address given that doesn't already have an identity. - * @param options.setSelectedAddress - Set the selected address. - * @param options.setAccountLabel - Set a new name for account. - * @param config - Initial options used to configure this controller. - * @param state - Initial state to set on this controller. + * @param options.preferences - The preferences controller. + * @param options.preferences.removeIdentity - Remove the identity with the given address. + * @param options.preferences.syncIdentities - Sync identities with the given list of addresses. + * @param options.preferences.updateIdentities - Generate an identity for each address given that doesn't already have an identity. + * @param options.preferences.setSelectedAddress - Set the selected address. + * @param options.preferences.setAccountLabel - Set a new name for account. + * @param options.encryptor - The encryptor for encrypting/decrypting keyrings. + * @param options.keyringBuilders - The keyring builders for initializing the keyring controller. + * @param options.cacheEncryptionKey - Whether to cache the encryption key in memory. + * @param options.state - Initial state to set on this controller. + * @param options.messenger - A restricted controller messenger. */ - constructor( - { + constructor({ + preferences: { removeIdentity, syncIdentities, updateIdentities, setSelectedAddress, setAccountLabel, - }: { - removeIdentity: PreferencesController['removeIdentity']; - syncIdentities: PreferencesController['syncIdentities']; - updateIdentities: PreferencesController['updateIdentities']; - setSelectedAddress: PreferencesController['setSelectedAddress']; - setAccountLabel?: PreferencesController['setAccountLabel']; }, - config?: Partial, - state?: Partial, - ) { - super(config, state); + encryptor, + keyringBuilders, + cacheEncryptionKey, + state, + messenger, + }: KeyringControllerOptions) { + super({ + name, + metadata: { + vault: { persist: true, anonymous: false }, + isUnlocked: { persist: false, anonymous: false }, + keyringTypes: { persist: false, anonymous: false }, + keyrings: { persist: false, anonymous: false }, + }, + messenger, + state: { + ...defaultState, + ...state, + }, + }); + this.#keyring = new EthKeyringController( - Object.assign({ initState: state }, config), + Object.assign( + { initState: state }, + { encryptor, keyringBuilders, cacheEncryptionKey }, + ), ); - this.defaultState = { - ...this.#keyring.store.getState(), - keyrings: [], - }; this.removeIdentity = removeIdentity; this.syncIdentities = syncIdentities; this.updateIdentities = updateIdentities; this.setSelectedAddress = setSelectedAddress; this.setAccountLabel = setAccountLabel; - this.initialize(); this.fullUpdate(); } @@ -217,7 +247,7 @@ export class KeyringController extends BaseController< * address. */ async addNewAccount(accountCount?: number): Promise<{ - keyringState: KeyringMemState; + keyringState: KeyringState; addedAccountAddress: string; }> { const primaryKeyring = this.#keyring.getKeyringsByType('HD Key Tree')[0]; @@ -259,7 +289,7 @@ export class KeyringController extends BaseController< * * @returns Promise resolving to current state when the account is added. */ - async addNewAccountWithoutUpdate(): Promise { + async addNewAccountWithoutUpdate(): Promise { const primaryKeyring = this.#keyring.getKeyringsByType('HD Key Tree')[0]; /* istanbul ignore if */ if (!primaryKeyring) { @@ -388,7 +418,7 @@ export class KeyringController extends BaseController< strategy: AccountImportStrategy, args: any[], ): Promise<{ - keyringState: KeyringMemState; + keyringState: KeyringState; importedAccountAddress: string; }> { let privateKey; @@ -449,7 +479,7 @@ export class KeyringController extends BaseController< * @param address - Address of the account to remove. * @returns Promise resolving current state when this account removal completes. */ - async removeAccount(address: string): Promise { + async removeAccount(address: string): Promise { this.removeIdentity(address); await this.#keyring.removeAccount(address); return this.fullUpdate(); @@ -460,7 +490,7 @@ export class KeyringController extends BaseController< * * @returns Promise resolving to current state. */ - setLocked(): Promise { + setLocked(): Promise { return this.#keyring.setLocked(); } @@ -576,7 +606,7 @@ export class KeyringController extends BaseController< async submitEncryptionKey( encryptionKey: string, encryptionSalt: string, - ): Promise { + ): Promise { return this.#keyring.submitEncryptionKey(encryptionKey, encryptionSalt); } @@ -587,32 +617,13 @@ export class KeyringController extends BaseController< * @param password - Password to unlock the keychain. * @returns Promise resolving to the current state. */ - async submitPassword(password: string): Promise { + async submitPassword(password: string): Promise { await this.#keyring.submitPassword(password); const accounts = await this.#keyring.getAccounts(); await this.syncIdentities(accounts); return this.fullUpdate(); } - /** - * Adds new listener to be notified of state changes. - * - * @param listener - Callback triggered when state changes. - */ - override subscribe(listener: Listener) { - this.#keyring.store.subscribe(listener); - } - - /** - * Removes existing listener from receiving state changes. - * - * @param listener - Callback to remove. - * @returns True if a listener is found and unsubscribed. - */ - override unsubscribe(listener: Listener) { - return this.#keyring.store.unsubscribe(listener); - } - /** * Adds new listener to be notified when the wallet is locked. * @@ -681,7 +692,7 @@ export class KeyringController extends BaseController< * * @returns The current state. */ - async fullUpdate(): Promise { + async fullUpdate(): Promise { const keyrings: Keyring[] = await Promise.all( this.#keyring.keyrings.map( async (keyring: KeyringObject, index: number): Promise => { @@ -697,7 +708,10 @@ export class KeyringController extends BaseController< }, ), ); - this.update({ keyrings: [...keyrings] }); + this.update((state) => { + state.vault = this.#keyring.vault; + state.keyrings = [...keyrings]; + }); return this.#keyring.fullUpdate(); } diff --git a/yarn.lock b/yarn.lock index 56dad60ed87..5b9703468b5 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1721,6 +1721,7 @@ __metadata: deepmerge: ^4.2.2 ethereumjs-util: ^7.0.10 ethereumjs-wallet: ^1.0.1 + immer: ^10.0.2 jest: ^27.5.1 sinon: ^9.2.4 ts-jest: ^27.1.4 @@ -5777,6 +5778,13 @@ __metadata: languageName: node linkType: hard +"immer@npm:^10.0.2": + version: 10.0.2 + resolution: "immer@npm:10.0.2" + checksum: 525a3b14210d02ae420c3b9f6ca14f7e9bcf625611d1356e773e7739f14c7c8de50dac442e6c7de3a6e24a782f7b792b6b8666bc0b3f00269d21a95f8f68ca84 + languageName: node + linkType: hard + "immer@npm:^9.0.6": version: 9.0.6 resolution: "immer@npm:9.0.6"