Restores accounts until one with a zero balance is found

3 files changed, 111 insertions, 2 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c61d566b4..ddaa496dd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,6 +33,8 @@
 - [#4691](https://github.com/MetaMask/metamask-extension/pull/4691): Redesign of the Confirm Transaction Screen.
 - [#4840](https://github.com/MetaMask/metamask-extension/pull/4840): Now shows notifications when transactions are completed.
 - [#4855](https://github.com/MetaMask/metamask-extension/pull/4855): Allow the use of HTTP prefix for custom rpc urls.
+- [#4855](https://github.com/MetaMask/metamask-extension/pull/4855): network.js: convert rpc protocol to lower case.
+- [#4898](https://github.com/MetaMask/metamask-extension/pull/4898): Restore multiple consecutive accounts with balances.
 
 ## 4.8.0 Thur Jun 14 2018
 
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 81bb080ab..166be720f 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -48,6 +48,7 @@ const percentile = require('percentile')
 const seedPhraseVerifier = require('./lib/seed-phrase-verifier')
 const log = require('loglevel')
 const TrezorKeyring = require('eth-trezor-keyring')
+const EthQuery = require('eth-query')
 
 module.exports = class MetamaskController extends EventEmitter {
 
@@ -475,12 +476,32 @@ module.exports = class MetamaskController extends EventEmitter {
   async createNewVaultAndRestore (password, seed) {
     const releaseLock = await this.createVaultMutex.acquire()
     try {
+      let accounts, lastBalance
+
+      const keyringController = this.keyringController
+
       // clear known identities
       this.preferencesController.setAddresses([])
       // create new vault
-      const vault = await this.keyringController.createNewVaultAndRestore(password, seed)
+      const vault = await keyringController.createNewVaultAndRestore(password, seed)
+
+      const ethQuery = new EthQuery(this.provider)
+      accounts = await keyringController.getAccounts()
+      lastBalance = await this.getBalance(accounts[accounts.length - 1], ethQuery)
+
+      const primaryKeyring = keyringController.getKeyringsByType('HD Key Tree')[0]
+      if (!primaryKeyring) {
+        throw new Error('MetamaskController - No HD Key Tree found')
+      }
+
+      // seek out the first zero balance
+      while (lastBalance !== '0x0') {
+        await keyringController.addNewAccount(primaryKeyring)
+        accounts = await keyringController.getAccounts()
+        lastBalance = await this.getBalance(accounts[accounts.length - 1], ethQuery)
+      }
+
       // set new identities
-      const accounts = await this.keyringController.getAccounts()
       this.preferencesController.setAddresses(accounts)
       this.selectFirstIdentity()
       releaseLock()
@@ -491,6 +512,30 @@ module.exports = class MetamaskController extends EventEmitter {
     }
   }
 
+  /**
+   * Get an account balance from the AccountTracker or request it directly from the network.
+   * @param {string} address - The account address
+   * @param {EthQuery} ethQuery - The EthQuery instance to use when asking the network
+   */
+  getBalance (address, ethQuery) {
+    return new Promise((resolve, reject) => {
+      const cached = this.accountTracker.store.getState().accounts[address]
+
+      if (cached && cached.balance) {
+        resolve(cached.balance)
+      } else {
+        ethQuery.getBalance(address, (error, balance) => {
+          if (error) {
+            reject(error)
+            log.error(error)
+          } else {
+            resolve(balance || '0x0')
+          }
+        })
+      }
+    })
+  }
+
   /*
    * Submits the user's password and attempts to unlock the vault.
    * Also synchronizes the preferencesController, to ensure its schema
diff --git a/test/unit/app/controllers/metamask-controller-test.js b/test/unit/app/controllers/metamask-controller-test.js
index 471321b9b..e51cd66aa 100644
--- a/test/unit/app/controllers/metamask-controller-test.js
+++ b/test/unit/app/controllers/metamask-controller-test.js
@@ -7,11 +7,15 @@ const blacklistJSON = require('eth-phishing-detect/src/config')
 const MetaMaskController = require('../../../../app/scripts/metamask-controller')
 const firstTimeState = require('../../../unit/localhostState')
 const createTxMeta = require('../../../lib/createTxMeta')
+const EthQuery = require('eth-query')
 
 const currentNetworkId = 42
 const DEFAULT_LABEL = 'Account 1'
+const DEFAULT_LABEL_2 = 'Account 2'
 const TEST_SEED = 'debris dizzy just program just float decrease vacant alarm reduce speak stadium'
 const TEST_ADDRESS = '0x0dcd5d886577d5081b0c52e242ef29e70be3e7bc'
+const TEST_ADDRESS_2 = '0xec1adf982415d2ef5ec55899b9bfb8bc0f29251b'
+const TEST_ADDRESS_3 = '0xeb9e64b93097bc15f01f13eae97015c57ab64823'
 const TEST_SEED_ALT = 'setup olympic issue mobile velvet surge alcohol burger horse view reopen gentle'
 const TEST_ADDRESS_ALT = '0xc42edfcc21ed14dda456aa0756c153f7985d8813'
 const CUSTOM_RPC_URL = 'http://localhost:8545'
@@ -136,6 +140,9 @@ describe('MetaMaskController', function () {
   describe('#createNewVaultAndRestore', function () {
     it('should be able to call newVaultAndRestore despite a mistake.', async function () {
       const password = 'what-what-what'
+      sandbox.stub(metamaskController, 'getBalance')
+      metamaskController.getBalance.callsFake(() => { return Promise.resolve('0x0') })
+
       await metamaskController.createNewVaultAndRestore(password, TEST_SEED.slice(0, -1)).catch((e) => null)
       await metamaskController.createNewVaultAndRestore(password, TEST_SEED)
 
@@ -143,6 +150,9 @@ describe('MetaMaskController', function () {
     })
 
     it('should clear previous identities after vault restoration', async () => {
+      sandbox.stub(metamaskController, 'getBalance')
+      metamaskController.getBalance.callsFake(() => { return Promise.resolve('0x0') })
+
       await metamaskController.createNewVaultAndRestore('foobar1337', TEST_SEED)
       assert.deepEqual(metamaskController.getState().identities, {
         [TEST_ADDRESS]: { address: TEST_ADDRESS, name: DEFAULT_LABEL },
@@ -158,6 +168,54 @@ describe('MetaMaskController', function () {
         [TEST_ADDRESS_ALT]: { address: TEST_ADDRESS_ALT, name: DEFAULT_LABEL },
       })
     })
+
+    it('should restore any consecutive accounts with balances', async () => {
+      sandbox.stub(metamaskController, 'getBalance')
+      metamaskController.getBalance.withArgs(TEST_ADDRESS).callsFake(() => {
+        return Promise.resolve('0x14ced5122ce0a000')
+      })
+      metamaskController.getBalance.withArgs(TEST_ADDRESS_2).callsFake(() => {
+        return Promise.resolve('0x0')
+      })
+      metamaskController.getBalance.withArgs(TEST_ADDRESS_3).callsFake(() => {
+        return Promise.resolve('0x14ced5122ce0a000')
+      })
+
+      await metamaskController.createNewVaultAndRestore('foobar1337', TEST_SEED)
+      assert.deepEqual(metamaskController.getState().identities, {
+        [TEST_ADDRESS]: { address: TEST_ADDRESS, name: DEFAULT_LABEL },
+        [TEST_ADDRESS_2]: { address: TEST_ADDRESS_2, name: DEFAULT_LABEL_2 },
+      })
+    })
+  })
+
+  describe('#getBalance', () => {
+    it('should return the balance known by accountTracker', async () => {
+      const accounts = {}
+      const balance = '0x14ced5122ce0a000'
+      accounts[TEST_ADDRESS] = { balance: balance }
+
+      metamaskController.accountTracker.store.putState({ accounts: accounts })
+
+      const gotten = await metamaskController.getBalance(TEST_ADDRESS)
+
+      assert.equal(balance, gotten)
+    })
+
+    it('should ask the network for a balance when not known by accountTracker', async () => {
+      const accounts = {}
+      const balance = '0x14ced5122ce0a000'
+      const ethQuery = new EthQuery()
+      sinon.stub(ethQuery, 'getBalance').callsFake((account, callback) => {
+        callback(undefined, balance)
+      })
+
+      metamaskController.accountTracker.store.putState({ accounts: accounts })
+
+      const gotten = await metamaskController.getBalance(TEST_ADDRESS, ethQuery)
+
+      assert.equal(balance, gotten)
+    })
   })
 
   describe('#getApi', function () {
@@ -553,6 +611,8 @@ describe('MetaMaskController', function () {
     const data = '0x43727970746f6b697474696573'
 
     beforeEach(async () => {
+      sandbox.stub(metamaskController, 'getBalance')
+      metamaskController.getBalance.callsFake(() => { return Promise.resolve('0x0') })
 
       await metamaskController.createNewVaultAndRestore('foobar1337', TEST_SEED_ALT)
 
@@ -622,6 +682,8 @@ describe('MetaMaskController', function () {
     const data = '0x43727970746f6b697474696573'
 
     beforeEach(async function () {
+      sandbox.stub(metamaskController, 'getBalance')
+      metamaskController.getBalance.callsFake(() => { return Promise.resolve('0x0') })
 
       await metamaskController.createNewVaultAndRestore('foobar1337', TEST_SEED_ALT)