From 5db6a424cf6b53f6576922ca9eba8d05903f4a6a Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Wed, 11 Dec 2024 22:27:14 +0800 Subject: [PATCH 1/2] feat: support multiple hd keyrings --- packages/keyring-controller/package.json | 9 +- .../src/KeyringController.ts | 114 +++++++++++++++--- yarn.lock | 10 ++ 3 files changed, 113 insertions(+), 20 deletions(-) diff --git a/packages/keyring-controller/package.json b/packages/keyring-controller/package.json index 73a80303ba..469fa7abd8 100644 --- a/packages/keyring-controller/package.json +++ b/packages/keyring-controller/package.json @@ -59,7 +59,8 @@ "@metamask/utils": "^10.0.0", "async-mutex": "^0.5.0", "ethereumjs-wallet": "^1.0.1", - "immer": "^9.0.6" + "immer": "^9.0.6", + "ulid": "^2.3.0" }, "devDependencies": { "@ethereumjs/common": "^3.2.0", @@ -94,6 +95,10 @@ "registry": "https://registry.npmjs.org/" }, "lavamoat": { - "allowScripts": {} + "allowScripts": { + "@lavamoat/preinstall-always-fail": false, + "ethereumjs-wallet>ethereum-cryptography>keccak": false, + "ethereumjs-wallet>ethereum-cryptography>secp256k1": false + } } } diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts index 375397fd60..9bbfd729f7 100644 --- a/packages/keyring-controller/src/KeyringController.ts +++ b/packages/keyring-controller/src/KeyringController.ts @@ -43,6 +43,7 @@ import { Mutex } from 'async-mutex'; import type { MutexInterface } from 'async-mutex'; import Wallet, { thirdparty as importers } from 'ethereumjs-wallet'; import type { Patch } from 'immer'; +import { ulid } from 'ulid'; import { KeyringControllerError } from './constants'; @@ -256,10 +257,14 @@ export type KeyringControllerOptions = { * Keyring object to return in fullUpdate * @property type - Keyring type * @property accounts - Associated accounts + * @property typeIndex - The index of the keyring in the array of keyrings of the same type + * @property id - The id of the keyring */ export type KeyringObject = { accounts: string[]; type: string; + typeIndex: number; + id: string; }; /** @@ -396,6 +401,9 @@ export type KeyringSelector = } | { address: Hex; + } + | { + id: string; }; /** @@ -523,16 +531,24 @@ function isSerializedKeyringsArray( * @param keyring - The keyring to display. * @returns A keyring display object, with type and accounts properties. */ -async function displayForKeyring( - keyring: EthKeyring, -): Promise<{ type: string; accounts: string[] }> { +async function displayForKeyring(keyring: EthKeyring): Promise<{ + type: string; + accounts: string[]; + typeIndex: number; + id: string; +}> { const accounts = await keyring.getAccounts(); + const { opts } = keyring as EthKeyring & { + opts: { typeIndex: number; id: string }; + }; return { type: keyring.type, // Cast to `string[]` here is safe here because `accounts` has no nullish // values, and `normalize` returns `string` unless given a nullish value accounts: accounts.map(normalize) as string[], + typeIndex: opts?.typeIndex, + id: opts?.id, }; } @@ -658,18 +674,27 @@ export class KeyringController extends BaseController< * Adds a new account to the default (first) HD seed phrase keyring. * * @param accountCount - Number of accounts before adding a new one, used to + * @param keyringId - The id of the keyring to add the account to. * make the method idempotent. * @returns Promise resolving to the added account address. */ - async addNewAccount(accountCount?: number): Promise { + async addNewAccount( + accountCount?: number, + keyringId?: string, + ): Promise { return this.#persistOrRollback(async () => { - const primaryKeyring = this.getKeyringsByType('HD Key Tree')[0] as - | EthKeyring - | undefined; - if (!primaryKeyring) { + let selectedKeyring: EthKeyring | undefined; + if (keyringId) { + selectedKeyring = this.getKeyringById(keyringId) as EthKeyring; + } else { + selectedKeyring = this.getKeyringsByType( + KeyringTypes.hd, + )[0] as EthKeyring; + } + if (!selectedKeyring) { throw new Error('No HD keyring found'); } - const oldAccounts = await primaryKeyring.getAccounts(); + const oldAccounts = await selectedKeyring.getAccounts(); if (accountCount && oldAccounts.length !== accountCount) { if (accountCount > oldAccounts.length) { @@ -685,7 +710,7 @@ export class KeyringController extends BaseController< return existingAccount; } - const [addedAccountAddress] = await primaryKeyring.addAccounts(1); + const [addedAccountAddress] = await selectedKeyring.addAccounts(1); await this.verifySeedPhrase(); return addedAccountAddress; @@ -758,6 +783,15 @@ export class KeyringController extends BaseController< }); } + async createKeyringFromMnemonic(mnemonic: string): Promise { + return this.#persistOrRollback(async () => { + return await this.#createKeyringWithFirstAccount(KeyringTypes.hd, { + mnemonic, + numberOfAccounts: 1, + }); + }); + } + /** * Create a new vault and primary keyring. * @@ -853,13 +887,16 @@ export class KeyringController extends BaseController< /** * Returns the public addresses of all accounts from every keyring. * + * @param keyringId - The id of the keyring to get the accounts from. * @returns A promise resolving to an array of addresses. */ - async getAccounts(): Promise { - return this.state.keyrings.reduce( - (accounts, keyring) => accounts.concat(keyring.accounts), - [], - ); + async getAccounts(keyringId?: string): Promise { + return this.state.keyrings + .filter((keyring) => (keyringId ? keyring.id === keyringId : true)) + .reduce( + (accounts, keyring) => accounts.concat(keyring.accounts), + [], + ); } /** @@ -908,6 +945,24 @@ export class KeyringController extends BaseController< return keyring.decryptMessage(address, messageParams.data); } + /** + * Returns the keyring with the given id. + * + * @param id - The id of the keyring to return. + * @returns The keyring with the given id. + */ + getKeyringById(id: string): unknown { + const keyring = this.#keyrings.find( + (item) => + (item as EthKeyring & { opts: { id: string } }).opts.id === id, + ); + if (!keyring) { + throw new Error(KeyringControllerError.KeyringNotFound); + } + + return keyring; + } + /** * Returns the currently initialized keyring that manages * the specified `address` if one exists. @@ -1472,7 +1527,7 @@ export class KeyringController extends BaseController< keyring = (await this.getKeyringForAccount(selector.address)) as | SelectedKeyring | undefined; - } else { + } else if ('type' in selector) { keyring = this.getKeyringsByType(selector.type)[selector.index || 0] as | SelectedKeyring | undefined; @@ -1483,6 +1538,8 @@ export class KeyringController extends BaseController< options.createWithData, )) as SelectedKeyring; } + } else if ('id' in selector) { + keyring = this.getKeyringById(selector.id) as SelectedKeyring; } if (!keyring) { @@ -2148,6 +2205,7 @@ export class KeyringController extends BaseController< if (!firstAccount) { throw new Error(KeyringControllerError.NoFirstAccount); } + return firstAccount; } /** @@ -2174,8 +2232,28 @@ export class KeyringController extends BaseController< const keyring = keyringBuilder(); - // @ts-expect-error Enforce data type after updating clients - await keyring.deserialize(data); + // find the last index of the type + const lastIndexOfType = this.#keyrings.reduce((maxIndex, item) => { + if (item.type === type) { + return Math.max( + maxIndex, + (item as EthKeyring & { opts: { typeIndex: number } }).opts + ?.typeIndex ?? 0, + ); + } + return maxIndex; + }, 0); + + if (type === KeyringTypes.hd) { + await keyring.deserialize({ + ...(data ?? {}), + typeIndex: lastIndexOfType, + id: ulid(), + }); + } else { + // @ts-expect-error Enforce data type after updating clients + await keyring.deserialize(data); + } if (keyring.init) { await keyring.init(); diff --git a/yarn.lock b/yarn.lock index 5662cf9c72..fd76183568 100644 --- a/yarn.lock +++ b/yarn.lock @@ -2994,6 +2994,7 @@ __metadata: typedoc: "npm:^0.24.8" typedoc-plugin-missing-exports: "npm:^2.0.0" typescript: "npm:~5.2.2" + ulid: "npm:^2.3.0" uuid: "npm:^8.3.2" webextension-polyfill: "npm:^0.12.0" peerDependencies: @@ -12277,6 +12278,15 @@ __metadata: languageName: node linkType: hard +"ulid@npm:^2.3.0": + version: 2.3.0 + resolution: "ulid@npm:2.3.0" + bin: + ulid: ./bin/cli.js + checksum: 10/11d7dd35072b863effb1249f66fb03070142a625610f00e5afd99af7e909b5de9cc7ebca6ede621a6bb1b7479b2489d6f064db6742b55c14bff6496ac60f290f + languageName: node + linkType: hard + "unbox-primitive@npm:^1.0.2": version: 1.0.2 resolution: "unbox-primitive@npm:1.0.2" From 953c11ad173369a051dde0a4b31943d786e1f0ed Mon Sep 17 00:00:00 2001 From: Monte Lai Date: Thu, 12 Dec 2024 00:41:12 +0800 Subject: [PATCH 2/2] fix: typeIndex increment and add arg for exportSeedPhrase --- .../src/KeyringController.ts | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/packages/keyring-controller/src/KeyringController.ts b/packages/keyring-controller/src/KeyringController.ts index 9bbfd729f7..dc67bcf85a 100644 --- a/packages/keyring-controller/src/KeyringController.ts +++ b/packages/keyring-controller/src/KeyringController.ts @@ -856,12 +856,27 @@ export class KeyringController extends BaseController< * Gets the seed phrase of the HD keyring. * * @param password - Password of the keyring. + * @param typeIndex - Hd keyring identifier * @returns Promise resolving to the seed phrase. */ - async exportSeedPhrase(password: string): Promise { + async exportSeedPhrase( + password: string, + typeIndex: number, + ): Promise { await this.verifyPassword(password); - assertHasUint8ArrayMnemonic(this.#keyrings[0]); - return this.#keyrings[0].mnemonic; + + const keyring = this.getKeyringsByType(KeyringTypes.hd).find( + (innerKeyring) => + (innerKeyring as EthKeyring & { opts: { typeIndex: number } }) + .opts.typeIndex === typeIndex, + ) as EthKeyring; + + if (!keyring) { + throw new Error(KeyringControllerError.KeyringNotFound); + } + + assertHasUint8ArrayMnemonic(keyring); + return keyring.mnemonic; } /** @@ -2247,7 +2262,7 @@ export class KeyringController extends BaseController< if (type === KeyringTypes.hd) { await keyring.deserialize({ ...(data ?? {}), - typeIndex: lastIndexOfType, + typeIndex: lastIndexOfType + 1, id: ulid(), }); } else {