From df799d7fd6c49888e3eb659adbfca18030b85975 Mon Sep 17 00:00:00 2001 From: Dan Matthews Date: Thu, 26 Jul 2018 23:40:11 -0400 Subject: [PATCH] Restores accounts until one with a zero balance is found --- CHANGELOG.md | 2 + app/scripts/metamask-controller.js | 49 ++++++++++++++- .../controllers/metamask-controller-test.js | 62 +++++++++++++++++++ 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)