Added walletUtils and address in signMessage

5 files changed, 110 insertions, 54 deletions

diff --git a/packages/subproviders/src/subproviders/base_wallet_subprovider.ts b/packages/subproviders/src/subproviders/base_wallet_subprovider.ts
index 034f83e7f..47b45a126 100644
--- a/packages/subproviders/src/subproviders/base_wallet_subprovider.ts
+++ b/packages/subproviders/src/subproviders/base_wallet_subprovider.ts
@@ -21,7 +21,7 @@ export abstract class BaseWalletSubprovider extends Subprovider {
 
     public abstract async getAccountsAsync(): Promise<string[]>;
     public abstract async signTransactionAsync(txParams: PartialTxParams): Promise<string>;
-    public abstract async signPersonalMessageAsync(data: string): Promise<string>;
+    public abstract async signPersonalMessageAsync(data: string, address?: string): Promise<string>;
 
     /**
      * This method conforms to the web3-provider-engine interface.
@@ -85,8 +85,9 @@ export abstract class BaseWalletSubprovider extends Subprovider {
             case 'eth_sign':
             case 'personal_sign':
                 const data = payload.method === 'eth_sign' ? payload.params[1] : payload.params[0];
+                const address = payload.method === 'eth_sign' ? payload.params[0] : payload.params[1];
                 try {
-                    const ecSignatureHex = await this.signPersonalMessageAsync(data);
+                    const ecSignatureHex = await this.signPersonalMessageAsync(data, address);
                     end(null, ecSignatureHex);
                 } catch (err) {
                     end(err);
diff --git a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
index fdb497776..456bde05c 100644
--- a/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
+++ b/packages/subproviders/src/subproviders/mnemonic_wallet_subprovider.ts
@@ -4,14 +4,15 @@ import ethUtil = require('ethereumjs-util');
 import HDNode = require('hdkey');
 import * as _ from 'lodash';
 
-import { MnemonicSubproviderErrors, PartialTxParams } from '../types';
+import { DerivedHDKey, PartialTxParams, WalletSubproviderErrors } from '../types';
+import { walletUtils } from '../walletUtils';
 
 import { BaseWalletSubprovider } from './base_wallet_subprovider';
 import { PrivateKeyWalletSubprovider } from './private_key_wallet_subprovider';
 
 const DEFAULT_DERIVATION_PATH = `44'/60'/0'`;
 const DEFAULT_NUM_ADDRESSES_TO_FETCH = 10;
-const DEFAULT_ADDRESS_SEARCH_LIMIT = 100;
+const DEFAULT_ADDRESS_SEARCH_LIMIT = 1000;
 
 /**
  * This class implements the [web3-provider-engine](https://github.com/MetaMask/provider-engine) subprovider interface.
@@ -19,15 +20,22 @@ const DEFAULT_ADDRESS_SEARCH_LIMIT = 100;
  * all requests with accounts derived from the supplied mnemonic.
  */
 export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
+    private _addressSearchLimit: number;
     private _derivationPath: string;
     private _hdKey: HDNode;
-    private _derivationPathIndex: number;
-    constructor(mnemonic: string, derivationPath: string = DEFAULT_DERIVATION_PATH) {
+
+    constructor(
+        mnemonic: string,
+        derivationPath: string = DEFAULT_DERIVATION_PATH,
+        addressSearchLimit: number = DEFAULT_ADDRESS_SEARCH_LIMIT,
+    ) {
         assert.isString('mnemonic', mnemonic);
+        assert.isString('derivationPath', derivationPath);
+        assert.isNumber('addressSearchLimit', addressSearchLimit);
         super();
         this._hdKey = HDNode.fromMasterSeed(bip39.mnemonicToSeed(mnemonic));
-        this._derivationPathIndex = 0;
         this._derivationPath = derivationPath;
+        this._addressSearchLimit = addressSearchLimit;
     }
     /**
      * Retrieve the set derivation path
@@ -44,32 +52,14 @@ export class MnemonicWalletSubprovider 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 the account associated with the supplied private key.
+     * Retrieve the accounts associated with the mnemonic.
      * This method is implicitly called when issuing a `eth_accounts` JSON RPC request
      * via your providerEngine instance.
      * @return An array of accounts
      */
     public async getAccountsAsync(numberOfAccounts: number = DEFAULT_NUM_ADDRESSES_TO_FETCH): Promise<string[]> {
-        const accounts: string[] = [];
-        for (let i = 0; i < numberOfAccounts; i++) {
-            const derivedHDNode = this._hdKey.derive(`m/${this._derivationPath}/${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 derivedKeys = walletUtils._calculateDerivedHDKeys(this._hdKey, this._derivationPath, numberOfAccounts);
+        const accounts = _.map(derivedKeys, 'address');
         return accounts;
     }
 
@@ -82,9 +72,10 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
      * @return Signed transaction hex string
      */
     public async signTransactionAsync(txParams: PartialTxParams): Promise<string> {
-        const accounts = await this.getAccountsAsync();
-        const hdKey = this._findHDKeyByPublicAddress(txParams.from || accounts[0]);
-        const privateKeyWallet = new PrivateKeyWalletSubprovider(hdKey.privateKey.toString('hex'));
+        const derivedKey = _.isUndefined(txParams.from)
+            ? walletUtils._firstDerivedKey(this._hdKey, this._derivationPath)
+            : this._findDerivedKeyByPublicAddress(txParams.from);
+        const privateKeyWallet = new PrivateKeyWalletSubprovider(derivedKey.hdKey.privateKey.toString('hex'));
         const signedTx = privateKeyWallet.signTransactionAsync(txParams);
         return signedTx;
     }
@@ -95,29 +86,27 @@ export class MnemonicWalletSubprovider extends BaseWalletSubprovider {
      * 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> {
-        const accounts = await this.getAccountsAsync();
-        const hdKey = this._findHDKeyByPublicAddress(accounts[0]);
-        const privateKeyWallet = new PrivateKeyWalletSubprovider(hdKey.privateKey.toString('hex'));
-        const sig = await privateKeyWallet.signPersonalMessageAsync(data);
+    public async signPersonalMessageAsync(data: string, address?: string): Promise<string> {
+        const derivedKey = _.isUndefined(address)
+            ? walletUtils._firstDerivedKey(this._hdKey, this._derivationPath)
+            : this._findDerivedKeyByPublicAddress(address);
+        const privateKeyWallet = new PrivateKeyWalletSubprovider(derivedKey.hdKey.privateKey.toString('hex'));
+        const sig = await privateKeyWallet.signPersonalMessageAsync(data, derivedKey.address);
         return sig;
     }
-
-    private _findHDKeyByPublicAddress(address: string, searchLimit: number = DEFAULT_ADDRESS_SEARCH_LIMIT): HDNode {
-        for (let i = 0; i < searchLimit; i++) {
-            const derivedHDNode = this._hdKey.derive(`m/${this._derivationPath}/${i + this._derivationPathIndex}`);
-            const derivedPublicKey = derivedHDNode.publicKey;
-            const shouldSanitizePublicKey = true;
-            const ethereumAddressUnprefixed = ethUtil
-                .publicToAddress(derivedPublicKey, shouldSanitizePublicKey)
-                .toString('hex');
-            const ethereumAddressPrefixed = ethUtil.addHexPrefix(ethereumAddressUnprefixed);
-            if (ethereumAddressPrefixed === address) {
-                return derivedHDNode;
-            }
+    private _findDerivedKeyByPublicAddress(address: string): DerivedHDKey {
+        const matchedDerivedKey = walletUtils._findDerivedKeyByAddress(
+            address,
+            this._hdKey,
+            this._derivationPath,
+            this._addressSearchLimit,
+        );
+        if (_.isUndefined(matchedDerivedKey)) {
+            throw new Error(`${WalletSubproviderErrors.AddressNotFound}: ${address}`);
         }
-        throw new Error(MnemonicSubproviderErrors.AddressSearchExhausted);
+        return matchedDerivedKey;
     }
 }
diff --git a/packages/subproviders/src/subproviders/private_key_wallet_subprovider.ts b/packages/subproviders/src/subproviders/private_key_wallet_subprovider.ts
index 0aa2fb590..f6906bab6 100644
--- a/packages/subproviders/src/subproviders/private_key_wallet_subprovider.ts
+++ b/packages/subproviders/src/subproviders/private_key_wallet_subprovider.ts
@@ -52,12 +52,16 @@ export class PrivateKeyWalletSubprovider extends BaseWalletSubprovider {
      * 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(dataIfExists: string): Promise<string> {
+    public async signPersonalMessageAsync(dataIfExists: string, address?: string): Promise<string> {
         if (_.isUndefined(dataIfExists)) {
             throw new Error(WalletSubproviderErrors.DataMissingForSignPersonalMessage);
         }
+        if (!_.isUndefined(address) && address !== this._address) {
+            throw new Error(`${WalletSubproviderErrors.AddressNotFound}: ${address}`);
+        }
         assert.isHexString('data', dataIfExists);
         const dataBuff = ethUtil.toBuffer(dataIfExists);
         const msgHashBuff = ethUtil.hashPersonalMessage(dataBuff);
diff --git a/packages/subproviders/src/types.ts b/packages/subproviders/src/types.ts
index de04499ce..105ffa7cc 100644
--- a/packages/subproviders/src/types.ts
+++ b/packages/subproviders/src/types.ts
@@ -1,4 +1,5 @@
 import { ECSignature, JSONRPCRequestPayload } from '@0xproject/types';
+import HDNode = require('hdkey');
 import * as _ from 'lodash';
 
 export interface LedgerCommunicationClient {
@@ -95,10 +96,8 @@ export interface ResponseWithTxParams {
     tx: PartialTxParams;
 }
 
-export enum MnemonicSubproviderErrors {
-    AddressSearchExhausted = 'ADDRESS_SEARCH_EXHAUSTED',
-}
 export enum WalletSubproviderErrors {
+    AddressNotFound = 'ADDRESS_NOT_FOUND',
     DataMissingForSignPersonalMessage = 'DATA_MISSING_FOR_SIGN_PERSONAL_MESSAGE',
     SenderInvalidOrNotSupplied = 'SENDER_INVALID_OR_NOT_SUPPLIED',
 }
@@ -112,6 +111,11 @@ export enum NonceSubproviderErrors {
     EmptyParametersFound = 'EMPTY_PARAMETERS_FOUND',
     CannotDetermineAddressFromPayload = 'CANNOT_DETERMINE_ADDRESS_FROM_PAYLOAD',
 }
+export interface DerivedHDKey {
+    address: string;
+    derivationPath: string;
+    hdKey: HDNode;
+}
 
 export type ErrorCallback = (err: Error | null, data?: any) => void;
 export type Callback = () => void;
diff --git a/packages/subproviders/src/walletUtils.ts b/packages/subproviders/src/walletUtils.ts
new file mode 100644
index 000000000..631636a71
--- /dev/null
+++ b/packages/subproviders/src/walletUtils.ts
@@ -0,0 +1,58 @@
+import ethUtil = require('ethereumjs-util');
+import HDNode = require('hdkey');
+import * as _ from 'lodash';
+
+import { DerivedHDKey, WalletSubproviderErrors } from './types';
+
+const DEFAULT_ADDRESS_SEARCH_OFFSET = 0;
+const BATCH_SIZE = 10;
+export const walletUtils = {
+    _calculateDerivedHDKeys(
+        initialHDKey: HDNode,
+        derivationPath: string,
+        searchLimit: number,
+        offset: number = DEFAULT_ADDRESS_SEARCH_OFFSET,
+    ): DerivedHDKey[] {
+        const derivedKeys: DerivedHDKey[] = [];
+        _.times(searchLimit, i => {
+            const path = `m/${derivationPath}/${i + offset}`;
+            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);
+            const derivedKey: DerivedHDKey = {
+                derivationPath: path,
+                hdKey,
+                address,
+            };
+            derivedKeys.push(derivedKey);
+        });
+        return derivedKeys;
+    },
+
+    _findDerivedKeyByAddress(
+        address: string,
+        initialHDKey: HDNode,
+        derivationPath: string,
+        searchLimit: number,
+    ): 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);
+            matchedKey = _.find(derivedKeys, derivedKey => derivedKey.address === address);
+            if (matchedKey) {
+                break;
+            }
+        }
+        return matchedKey;
+    },
+
+    _firstDerivedKey(initialHDKey: HDNode, derivationPath: string): DerivedHDKey {
+        const derivedKeys = walletUtils._calculateDerivedHDKeys(initialHDKey, derivationPath, 1);
+        const firstDerivedKey = derivedKeys[0];
+        return firstDerivedKey;
+    },
+};