Refactor ledger to support multiple accounts with from address

4 files changed, 91 insertions, 56 deletions

diff --git a/packages/subproviders/src/subproviders/ledger.ts b/packages/subproviders/src/subproviders/ledger.ts
index 23ed3c93e..a7b79c128 100644
--- a/packages/subproviders/src/subproviders/ledger.ts
+++ b/packages/subproviders/src/subproviders/ledger.ts
@@ -8,6 +8,7 @@ import { Lock } from 'semaphore-async-await';
 
 import {
     Callback,
+    DerivedHDKey,
     LedgerEthereumClient,
     LedgerEthereumClientFactoryAsync,
     LedgerSubproviderConfigs,
@@ -16,6 +17,7 @@ import {
     ResponseWithTxParams,
     WalletSubproviderErrors,
 } from '../types';
+import { walletUtils } from '../utils/wallet_utils';
 
 import { BaseWalletSubprovider } from './base_wallet_subprovider';
 
@@ -34,10 +36,11 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
     private _connectionLock = new Lock();
     private _networkId: number;
     private _derivationPath: string;
-    private _derivationPathIndex: number;
     private _ledgerEthereumClientFactoryAsync: LedgerEthereumClientFactoryAsync;
     private _ledgerClientIfExists?: LedgerEthereumClient;
     private _shouldAlwaysAskForConfirmation: boolean;
+    private _addressSearchLimit: number;
+    private _hardenedKey: boolean = true;
     /**
      * Instantiates a LedgerSubprovider. Defaults to derivationPath set to `44'/60'/0'`.
      * TestRPC/Ganache defaults to `m/44'/60'/0'/0`, so set this in the configs if desired.
@@ -54,7 +57,11 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
             !_.isUndefined(config.accountFetchingConfigs.shouldAskForOnDeviceConfirmation)
                 ? config.accountFetchingConfigs.shouldAskForOnDeviceConfirmation
                 : ASK_FOR_ON_DEVICE_CONFIRMATION;
-        this._derivationPathIndex = 0;
+        this._addressSearchLimit =
+            !_.isUndefined(config.accountFetchingConfigs) &&
+            !_.isUndefined(config.accountFetchingConfigs.numAddressesToReturn)
+                ? config.accountFetchingConfigs.numAddressesToReturn
+                : DEFAULT_NUM_ADDRESSES_TO_FETCH;
     }
     /**
      * Retrieve the set derivation path
@@ -71,15 +78,6 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
         this._derivationPath = derivationPath;
     }
     /**
-     * Set the final derivation path index. If a user wishes to sign a message with the
-     * 6th address in a derivation path, before calling `signPersonalMessageAsync`, you must
-     * call this method with pathIndex `6`.
-     * @param pathIndex Desired derivation path index
-     */
-    public setPathIndex(pathIndex: number) {
-        this._derivationPathIndex = pathIndex;
-    }
-    /**
      * Retrieve a users Ledger accounts. The accounts are derived from the derivationPath,
      * master public key and chainCode. Because of this, you can request as many accounts
      * as you wish and it only requires a single request to the Ledger device. This method
@@ -89,34 +87,15 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
      * @return An array of accounts
      */
     public async getAccountsAsync(numberOfAccounts: number = DEFAULT_NUM_ADDRESSES_TO_FETCH): Promise<string[]> {
-        this._ledgerClientIfExists = await this._createLedgerClientAsync();
-
-        let ledgerResponse;
-        try {
-            ledgerResponse = await this._ledgerClientIfExists.getAddress(
-                this._derivationPath,
-                this._shouldAlwaysAskForConfirmation,
-                SHOULD_GET_CHAIN_CODE,
-            );
-        } finally {
-            await this._destroyLedgerClientAsync();
-        }
-
-        const hdKey = new HDNode();
-        hdKey.publicKey = new Buffer(ledgerResponse.publicKey, 'hex');
-        hdKey.chainCode = new Buffer(ledgerResponse.chainCode, 'hex');
-
-        const accounts: string[] = [];
-        for (let i = 0; i < numberOfAccounts; i++) {
-            const derivedHDNode = hdKey.derive(`m/${i + this._derivationPathIndex}`);
-            const derivedPublicKey = derivedHDNode.publicKey;
-            const shouldSanitizePublicKey = true;
-            const ethereumAddressUnprefixed = ethUtil
-                .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
-                .toString('hex');
-            const ethereumAddressPrefixed = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
-            accounts.push(ethereumAddressPrefixed.toLowerCase());
-        }
+        const initialHDKey = await this._initialHDKeyAsync();
+        const derivedKeys = walletUtils._calculateDerivedHDKeys(
+            initialHDKey,
+            this._derivationPath,
+            numberOfAccounts,
+            0,
+            true,
+        );
+        const accounts = _.map(derivedKeys, 'address');
         return accounts;
     }
     /**
@@ -129,6 +108,11 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
      */
     public async signTransactionAsync(txParams: PartialTxParams): Promise<string> {
         LedgerSubprovider._validateTxParams(txParams);
+        const initialHDKey = await this._initialHDKeyAsync();
+        const derivedKey = _.isUndefined(txParams.from)
+            ? walletUtils._firstDerivedKey(initialHDKey, this._derivationPath, this._hardenedKey)
+            : this._findDerivedKeyByPublicAddress(initialHDKey, txParams.from);
+
         this._ledgerClientIfExists = await this._createLedgerClientAsync();
 
         const tx = new EthereumTx(txParams);
@@ -140,7 +124,7 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
 
         const txHex = tx.serialize().toString('hex');
         try {
-            const derivationPath = this._getDerivationPath();
+            const derivationPath = `${derivedKey.derivationPath}/${derivedKey.derivationIndex}`;
             const result = await this._ledgerClientIfExists.signTransaction(derivationPath, txHex);
             // Store signature in transaction
             tx.r = Buffer.from(result.r, 'hex');
@@ -165,22 +149,28 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
     }
     /**
      * Sign a personal Ethereum signed message. The signing address will be the one
-     * retrieved given a derivationPath and pathIndex set on the subprovider.
+     * either the provided address or the first address on the ledger device.
      * The Ledger adds the Ethereum signed message prefix on-device.  If you've added
      * the LedgerSubprovider to your app's provider, you can simply send an `eth_sign`
      * or `personal_sign` JSON RPC request, and this method will be called auto-magically.
      * If you are not using this via a ProviderEngine instance, you can call it directly.
      * @param data Message to sign
+     * @param address Address to sign with
      * @return Signature hex string (order: rsv)
      */
-    public async signPersonalMessageAsync(data: string): Promise<string> {
+    public async signPersonalMessageAsync(data: string, address?: string): Promise<string> {
         if (_.isUndefined(data)) {
             throw new Error(WalletSubproviderErrors.DataMissingForSignPersonalMessage);
         }
         assert.isHexString('data', data);
+        const initialHDKey = await this._initialHDKeyAsync();
+        const derivedKey = _.isUndefined(address)
+            ? walletUtils._firstDerivedKey(initialHDKey, this._derivationPath, this._hardenedKey)
+            : this._findDerivedKeyByPublicAddress(initialHDKey, address);
+
         this._ledgerClientIfExists = await this._createLedgerClientAsync();
         try {
-            const derivationPath = this._getDerivationPath();
+            const derivationPath = `${derivedKey.derivationPath}/${derivedKey.derivationIndex}`;
             const result = await this._ledgerClientIfExists.signPersonalMessage(
                 derivationPath,
                 ethUtil.stripHexPrefix(data),
@@ -198,10 +188,6 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
             throw err;
         }
     }
-    private _getDerivationPath() {
-        const derivationPath = `${this.getPath()}/${this._derivationPathIndex}`;
-        return derivationPath;
-    }
     private async _createLedgerClientAsync(): Promise<LedgerEthereumClient> {
         await this._connectionLock.acquire();
         if (!_.isUndefined(this._ledgerClientIfExists)) {
@@ -222,4 +208,35 @@ export class LedgerSubprovider extends BaseWalletSubprovider {
         this._ledgerClientIfExists = undefined;
         this._connectionLock.release();
     }
+    private async _initialHDKeyAsync(): Promise<HDNode> {
+        this._ledgerClientIfExists = await this._createLedgerClientAsync();
+
+        let ledgerResponse;
+        try {
+            ledgerResponse = await this._ledgerClientIfExists.getAddress(
+                this._derivationPath,
+                this._shouldAlwaysAskForConfirmation,
+                SHOULD_GET_CHAIN_CODE,
+            );
+        } finally {
+            await this._destroyLedgerClientAsync();
+        }
+        const hdKey = new HDNode();
+        hdKey.publicKey = new Buffer(ledgerResponse.publicKey, 'hex');
+        hdKey.chainCode = new Buffer(ledgerResponse.chainCode, 'hex');
+        return hdKey;
+    }
+    private _findDerivedKeyByPublicAddress(initalHDKey: HDNode, address: string): DerivedHDKey {
+        const matchedDerivedKey = walletUtils._findDerivedKeyByAddress(
+            address,
+            initalHDKey,
+            this._derivationPath,
+            this._addressSearchLimit,
+            this._hardenedKey,
+        );
+        if (_.isUndefined(matchedDerivedKey)) {
+            throw new Error(`${WalletSubproviderErrors.AddressNotFound}: ${address}`);
+        }
+        return matchedDerivedKey;
+    }
 }
diff --git a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
index 456bde05c..3ff437659 100644
--- a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
+++ b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
@@ -5,12 +5,12 @@ import HDNode = require('hdkey');
 import * as _ from 'lodash';
 
 import { DerivedHDKey, PartialTxParams, WalletSubproviderErrors } from '../types';
-import { walletUtils } from '../walletUtils';
+import { walletUtils } from '../utils/wallet_utils';
 
 import { BaseWalletSubprovider } from './base_wallet_subprovider';
 import { PrivateKeyWalletSubprovider } from './private_key_wallet_subprovider';
 
-const DEFAULT_DERIVATION_PATH = `44'/60'/0'`;
+const DEFAULT_DERIVATION_PATH = `44'/60'/0'/0`;
 const DEFAULT_NUM_ADDRESSES_TO_FETCH = 10;
 const DEFAULT_ADDRESS_SEARCH_LIMIT = 1000;
 
@@ -23,7 +23,6 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
     private _addressSearchLimit: number;
     private _derivationPath: string;
     private _hdKey: HDNode;
-
     constructor(
         mnemonic: string,
         derivationPath: string = DEFAULT_DERIVATION_PATH,
diff --git a/packages/subproviders/src/types.ts b/packages/subproviders/src/types.ts
index 105ffa7cc..fe7ae921e 100644
--- a/packages/subproviders/src/types.ts
+++ b/packages/subproviders/src/types.ts
@@ -114,6 +114,7 @@ export enum NonceSubproviderErrors {
 export interface DerivedHDKey {
     address: string;
     derivationPath: string;
+    derivationIndex: number;
     hdKey: HDNode;
 }
 
diff --git a/packages/subproviders/src/walletUtils.ts b/packages/subproviders/src/utils/wallet_utils.ts
index 631636a71..6c698a006 100644
--- a/packages/subproviders/src/walletUtils.ts
+++ b/packages/subproviders/src/utils/wallet_utils.ts
@@ -2,20 +2,30 @@ import ethUtil = require('ethereumjs-util');
 import HDNode = require('hdkey');
 import * as _ from 'lodash';
 
-import { DerivedHDKey, WalletSubproviderErrors } from './types';
+import { DerivedHDKey, WalletSubproviderErrors } from '../types';
 
 const DEFAULT_ADDRESS_SEARCH_OFFSET = 0;
 const BATCH_SIZE = 10;
+
+// Derivation Paths
+// BIP44 m / purpose' / coin_type' / account' / change / address_index
+// m/44'/60'/0'/0
+// m/44'/60'/0'/0/0
+// m/44'/60'/0'/0/{account_index} - testrpc
+// m/44'/60'/0'  - ledger
+
 export const walletUtils = {
     _calculateDerivedHDKeys(
         initialHDKey: HDNode,
         derivationPath: string,
         searchLimit: number,
         offset: number = DEFAULT_ADDRESS_SEARCH_OFFSET,
+        hardened: boolean = false,
     ): DerivedHDKey[] {
         const derivedKeys: DerivedHDKey[] = [];
         _.times(searchLimit, i => {
-            const path = `m/${derivationPath}/${i + offset}`;
+            const derivationIndex = offset + i;
+            const path = hardened ? `m/${derivationIndex}` : `m/${derivationPath}/${derivationIndex}`;
             const hdKey = initialHDKey.derive(path);
             const derivedPublicKey = hdKey.publicKey;
             const shouldSanitizePublicKey = true;
@@ -24,9 +34,10 @@ export const walletUtils = {
                 .toString('hex');
             const address = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
             const derivedKey: DerivedHDKey = {
-                derivationPath: path,
+                derivationPath,
                 hdKey,
                 address,
+                derivationIndex,
             };
             derivedKeys.push(derivedKey);
         });
@@ -38,10 +49,17 @@ export const walletUtils = {
         initialHDKey: HDNode,
         derivationPath: string,
         searchLimit: number,
+        hardened: 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);
+            const derivedKeys = walletUtils._calculateDerivedHDKeys(
+                initialHDKey,
+                derivationPath,
+                BATCH_SIZE,
+                index,
+                hardened,
+            );
             matchedKey = _.find(derivedKeys, derivedKey => derivedKey.address === address);
             if (matchedKey) {
                 break;
@@ -50,8 +68,8 @@ export const walletUtils = {
         return matchedKey;
     },
 
-    _firstDerivedKey(initialHDKey: HDNode, derivationPath: string): DerivedHDKey {
-        const derivedKeys = walletUtils._calculateDerivedHDKeys(initialHDKey, derivationPath, 1);
+    _firstDerivedKey(initialHDKey: HDNode, derivationPath: string, hardened: boolean = false): DerivedHDKey {
+        const derivedKeys = walletUtils._calculateDerivedHDKeys(initialHDKey, derivationPath, 1, 0, hardened);
         const firstDerivedKey = derivedKeys[0];
         return firstDerivedKey;
     },