Add isChildKey to derived key

4 files changed, 73 insertions, 72 deletions

diff --git a/packages/subproviders/src/subproviders/ledger.ts b/packages/subproviders/src/subproviders/ledger.ts
index fabff88cd..6f66e3018 100644
--- a/packages/subproviders/src/subproviders/ledger.ts
+++ b/packages/subproviders/src/subproviders/ledger.ts
@@ -22,7 +22,6 @@ import { walletUtils } from '../utils/wallet_utils';
 import { BaseWalletSubprovider } from './base_wallet_subprovider';
 
 const DEFAULT_DERIVATION_PATH = `44'/60'/0'`;
-const DEFAULT_NUM_ADDRESSES_TO_FETCH = 10;
 const ASK_FOR_ON_DEVICE_CONFIRMATION = false;
 const SHOULD_GET_CHAIN_CODE = true;
 const IS_CHILD_KEY = true;
@@ -59,9 +58,9 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
                 : ASK_FOR_ON_DEVICE_CONFIRMATION;
         this._addressSearchLimit =
             !_.isUndefined(config.accountFetchingConfigs) &&
-            !_.isUndefined(config.accountFetchingConfigs.numAddressesToReturn)
-                ? config.accountFetchingConfigs.numAddressesToReturn
-                : DEFAULT_NUM_ADDRESSES_TO_FETCH;
+            !_.isUndefined(config.accountFetchingConfigs.addressSearchLimit)
+                ? config.accountFetchingConfigs.addressSearchLimit
+                : walletUtils.DEFAULT_ADDRESS_SEARCH_LIMIT;
     }
     /**
      * Retrieve the set derivation path
@@ -86,15 +85,12 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
      * @param numberOfAccounts Number of accounts to retrieve (default: 10)
      * @return An array of accounts
      */
-    public async getAccountsAsync(numberOfAccounts: number = DEFAULT_NUM_ADDRESSES_TO_FETCH): Promise<string[]> {
-        const initialHDKey = await this._initialHDKeyAsync();
-        const derivedKeys = walletUtils._calculateDerivedHDKeys(
-            initialHDKey,
-            this._derivationPath,
-            numberOfAccounts,
-            0,
-            true,
-        );
+    public async getAccountsAsync(
+        numberOfAccounts: number = walletUtils.DEFAULT_NUM_ADDRESSES_TO_FETCH,
+    ): Promise<string[]> {
+        const offset = 0;
+        const initialHDerivedKey = await this._initialDerivedKeyAsync();
+        const derivedKeys = walletUtils.calculateDerivedHDKeys(initialHDerivedKey, numberOfAccounts, offset);
         const accounts = _.map(derivedKeys, 'address');
         return accounts;
     }
@@ -108,10 +104,10 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
      */
     public async signTransactionAsync(txParams: PartialTxParams): Promise<string> {
         LedgerSubprovider._validateTxParams(txParams);
-        const initialHDKey = await this._initialHDKeyAsync();
+        const initialDerivedKey = await this._initialDerivedKeyAsync();
         const derivedKey = _.isUndefined(txParams.from)
-            ? walletUtils._firstDerivedKey(initialHDKey, this._derivationPath, IS_CHILD_KEY)
-            : this._findDerivedKeyByPublicAddress(initialHDKey, txParams.from);
+            ? walletUtils._firstDerivedKey(initialDerivedKey)
+            : this._findDerivedKeyByPublicAddress(initialDerivedKey, txParams.from);
 
         this._ledgerClientIfExists = await this._createLedgerClientAsync();
 
@@ -163,10 +159,10 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
             throw new Error(WalletSubproviderErrors.DataMissingForSignPersonalMessage);
         }
         assert.isHexString('data', data);
-        const initialHDKey = await this._initialHDKeyAsync();
+        const initialDerivedKey = await this._initialDerivedKeyAsync();
         const derivedKey = _.isUndefined(address)
-            ? walletUtils._firstDerivedKey(initialHDKey, this._derivationPath, IS_CHILD_KEY)
-            : this._findDerivedKeyByPublicAddress(initialHDKey, address);
+            ? walletUtils._firstDerivedKey(initialDerivedKey)
+            : this._findDerivedKeyByPublicAddress(initialDerivedKey, address);
 
         this._ledgerClientIfExists = await this._createLedgerClientAsync();
         try {
@@ -208,7 +204,7 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
         this._ledgerClientIfExists = undefined;
         this._connectionLock.release();
     }
-    private async _initialHDKeyAsync(): Promise<HDNode> {
+    private async _initialDerivedKeyAsync(): Promise<DerivedHDKey> {
         this._ledgerClientIfExists = await this._createLedgerClientAsync();
 
         let ledgerResponse;
@@ -224,16 +220,16 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
         const hdKey = new HDNode();
         hdKey.publicKey = new Buffer(ledgerResponse.publicKey, 'hex');
         hdKey.chainCode = new Buffer(ledgerResponse.chainCode, 'hex');
-        return hdKey;
+        return {
+            hdKey,
+            address: ledgerResponse.address,
+            isChildKey: true,
+            derivationPath: this._derivationPath,
+            derivationIndex: 0,
+        };
     }
-    private _findDerivedKeyByPublicAddress(initalHDKey: HDNode, address: string): DerivedHDKey {
-        const matchedDerivedKey = walletUtils._findDerivedKeyByAddress(
-            address,
-            initalHDKey,
-            this._derivationPath,
-            this._addressSearchLimit,
-            IS_CHILD_KEY,
-        );
+    private _findDerivedKeyByPublicAddress(initalHDKey: DerivedHDKey, address: string): DerivedHDKey {
+        const matchedDerivedKey = walletUtils.findDerivedKeyByAddress(address, initalHDKey, this._addressSearchLimit);
         if (_.isUndefined(matchedDerivedKey)) {
             throw new Error(`${WalletSubproviderErrors.AddressNotFound}: ${address}`);
         }
diff --git a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
index 02d06e2cd..53013c44c 100644
--- a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
+++ b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
@@ -11,8 +11,6 @@ import { BaseWalletSubprovider } from './base_wallet_subprovider';
 import { PrivateKeyWalletSubprovider } from './private_key_wallet_subprovider';
 
 const DEFAULT_DERIVATION_PATH = `44'/60'/0'/0`;
-const DEFAULT_NUM_ADDRESSES_TO_FETCH = 10;
-const DEFAULT_ADDRESS_SEARCH_LIMIT = 1000;
 
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
@@ -22,7 +20,7 @@ const DEFAULT_ADDRESS_SEARCH_LIMIT = 1000;
 export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
     private _addressSearchLimit: number;
     private _derivationPath: string;
-    private _hdKey: HDNode;
+    private _derivedKey: DerivedHDKey;
     /**
      * Instantiates a MnemonicWalletSubprovider. Defaults to derivationPath set to `44'/60'/0'/0`.
      * This is the default in TestRPC/Ganache, this can be overridden if desired.
@@ -34,15 +32,22 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
     constructor(
         mnemonic: string,
         derivationPath: string = DEFAULT_DERIVATION_PATH,
-        addressSearchLimit: number = DEFAULT_ADDRESS_SEARCH_LIMIT,
+        addressSearchLimit: number = walletUtils.DEFAULT_ADDRESS_SEARCH_LIMIT,
     ) {
         assert.isString('mnemonic', mnemonic);
         assert.isString('derivationPath', derivationPath);
         assert.isNumber('addressSearchLimit', addressSearchLimit);
         super();
         const seed = bip39.mnemonicToSeed(mnemonic);
-        this._hdKey = HDNode.fromMasterSeed(seed);
+        const hdKey = HDNode.fromMasterSeed(seed);
         this._derivationPath = derivationPath;
+        this._derivedKey = {
+            address: walletUtils.addressOfHDKey(hdKey),
+            derivationPath: this._derivationPath,
+            derivationIndex: 0,
+            hdKey,
+            isChildKey: false,
+        };
         this._addressSearchLimit = addressSearchLimit;
     }
     /**
@@ -66,8 +71,10 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
      * @param numberOfAccounts Number of accounts to retrieve (default: 10)
      * @return An array of accounts
      */
-    public async getAccountsAsync(numberOfAccounts: number = DEFAULT_NUM_ADDRESSES_TO_FETCH): Promise<string[]> {
-        const derivedKeys = walletUtils._calculateDerivedHDKeys(this._hdKey, this._derivationPath, numberOfAccounts);
+    public async getAccountsAsync(
+        numberOfAccounts: number = walletUtils.DEFAULT_NUM_ADDRESSES_TO_FETCH,
+    ): Promise<string[]> {
+        const derivedKeys = walletUtils.calculateDerivedHDKeys(this._derivedKey, numberOfAccounts);
         const accounts = _.map(derivedKeys, 'address');
         return accounts;
     }
@@ -82,7 +89,7 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
      */
     public async signTransactionAsync(txParams: PartialTxParams): Promise<string> {
         const derivedKey = _.isUndefined(txParams.from)
-            ? walletUtils._firstDerivedKey(this._hdKey, this._derivationPath)
+            ? walletUtils._firstDerivedKey(this._derivedKey)
             : this._findDerivedKeyByPublicAddress(txParams.from);
         const privateKeyWallet = new PrivateKeyWalletSubprovider(derivedKey.hdKey.privateKey.toString('hex'));
         const signedTx = privateKeyWallet.signTransactionAsync(txParams);
@@ -100,17 +107,16 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
      */
     public async signPersonalMessageAsync(data: string, address?: string): Promise<string> {
         const derivedKey = _.isUndefined(address)
-            ? walletUtils._firstDerivedKey(this._hdKey, this._derivationPath)
+            ? walletUtils._firstDerivedKey(this._derivedKey)
             : this._findDerivedKeyByPublicAddress(address);
         const privateKeyWallet = new PrivateKeyWalletSubprovider(derivedKey.hdKey.privateKey.toString('hex'));
         const sig = await privateKeyWallet.signPersonalMessageAsync(data, derivedKey.address);
         return sig;
     }
     private _findDerivedKeyByPublicAddress(address: string): DerivedHDKey {
-        const matchedDerivedKey = walletUtils._findDerivedKeyByAddress(
+        const matchedDerivedKey = walletUtils.findDerivedKeyByAddress(
             address,
-            this._hdKey,
-            this._derivationPath,
+            this._derivedKey,
             this._addressSearchLimit,
         );
         if (_.isUndefined(matchedDerivedKey)) {
diff --git a/packages/subproviders/src/types.ts b/packages/subproviders/src/types.ts
index fe7ae921e..9f557731f 100644
--- a/packages/subproviders/src/types.ts
+++ b/packages/subproviders/src/types.ts
@@ -51,6 +51,7 @@ export interface LedgerSubproviderConfigs {
  *                                   before fetching their addresses
  */
 export interface AccountFetchingConfigs {
+    addressSearchLimit?: number;
     numAddressesToReturn?: number;
     shouldAskForOnDeviceConfirmation?: boolean;
 }
@@ -116,6 +117,7 @@ export interface DerivedHDKey {
     derivationPath: string;
     derivationIndex: number;
     hdKey: HDNode;
+    isChildKey: boolean;
 }
 
 export type ErrorCallback = (err: Error | null, data?: any) => void;
diff --git a/packages/subproviders/src/utils/wallet_utils.ts b/packages/subproviders/src/utils/wallet_utils.ts
index 48d475559..0c3e895b6 100644
--- a/packages/subproviders/src/utils/wallet_utils.ts
+++ b/packages/subproviders/src/utils/wallet_utils.ts
@@ -8,54 +8,51 @@ const DEFAULT_ADDRESS_SEARCH_OFFSET = 0;
 const BATCH_SIZE = 10;
 
 export const walletUtils = {
-    _calculateDerivedHDKeys(
-        initialHDKey: HDNode,
-        derivationPath: string,
+    DEFAULT_NUM_ADDRESSES_TO_FETCH: 10,
+    DEFAULT_ADDRESS_SEARCH_LIMIT: 1000,
+    calculateDerivedHDKeys(
+        initialDerivedKey: DerivedHDKey,
         searchLimit: number,
         offset: number = DEFAULT_ADDRESS_SEARCH_OFFSET,
-        isChildKey: boolean = false,
     ): DerivedHDKey[] {
         const derivedKeys: DerivedHDKey[] = [];
         _.times(searchLimit, i => {
+            const derivationPath = initialDerivedKey.derivationPath;
             const derivationIndex = offset + i;
-            // Normally we need to set the full derivation path to walk the tree from the root
-            // as the initial key is at the root.
-            // But with ledger the initial key is a child so we walk the tree relative to that child
-            const path = isChildKey ? `m/${derivationIndex}` : `m/${derivationPath}/${derivationIndex}`;
-            const hdKey = initialHDKey.derive(path);
-            const derivedPublicKey = hdKey.publicKey;
-            const shouldSanitizePublicKey = true;
-            const ethereumAddressUnprefixed = ethUtil
-                .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
-                .toString('hex');
-            const address = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
+            // If the DerivedHDKey is a child then we walk relative, if not we walk the full derivation path
+            const path = initialDerivedKey.isChildKey
+                ? `m/${derivationIndex}`
+                : `m/${derivationPath}/${derivationIndex}`;
+            const hdKey = initialDerivedKey.hdKey.derive(path);
+            const address = walletUtils.addressOfHDKey(hdKey);
             const derivedKey: DerivedHDKey = {
-                derivationPath,
-                hdKey,
                 address,
+                hdKey,
+                derivationPath,
                 derivationIndex,
+                isChildKey: initialDerivedKey.isChildKey,
             };
             derivedKeys.push(derivedKey);
         });
         return derivedKeys;
     },
-
-    _findDerivedKeyByAddress(
+    addressOfHDKey(hdKey: HDNode): string {
+        const shouldSanitizePublicKey = true;
+        const derivedPublicKey = hdKey.publicKey;
+        const ethereumAddressUnprefixed = ethUtil
+            .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
+            .toString('hex');
+        const address = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
+        return address;
+    },
+    findDerivedKeyByAddress(
         address: string,
-        initialHDKey: HDNode,
-        derivationPath: string,
+        initialDerivedKey: DerivedHDKey,
         searchLimit: number,
-        isChild: boolean = false,
     ): DerivedHDKey | undefined {
         let matchedKey: DerivedHDKey | undefined;
         for (let index = 0; index < searchLimit; index = index + BATCH_SIZE) {
-            const derivedKeys = walletUtils._calculateDerivedHDKeys(
-                initialHDKey,
-                derivationPath,
-                BATCH_SIZE,
-                index,
-                isChild,
-            );
+            const derivedKeys = walletUtils.calculateDerivedHDKeys(initialDerivedKey, BATCH_SIZE, index);
             matchedKey = _.find(derivedKeys, derivedKey => derivedKey.address === address);
             if (matchedKey) {
                 break;
@@ -64,8 +61,8 @@ export const walletUtils = {
         return matchedKey;
     },
 
-    _firstDerivedKey(initialHDKey: HDNode, derivationPath: string, isChild: boolean = false): DerivedHDKey {
-        const derivedKeys = walletUtils._calculateDerivedHDKeys(initialHDKey, derivationPath, 1, 0, isChild);
+    _firstDerivedKey(initialDerivedKey: DerivedHDKey): DerivedHDKey {
+        const derivedKeys = walletUtils.calculateDerivedHDKeys(initialDerivedKey, 1, 0);
         const firstDerivedKey = derivedKeys[0];
         return firstDerivedKey;
     },