Correctly clear ethStore cache on new vault restore

2 files changed, 38 insertions, 13 deletions

diff --git a/app/scripts/lib/idStore.js b/app/scripts/lib/idStore.js
index d2d37b0f4..69ffd3f72 100644
--- a/app/scripts/lib/idStore.js
+++ b/app/scripts/lib/idStore.js
@@ -451,7 +451,11 @@ IdentityStore.prototype.tryPassword = function (password, cb) {
 }
 
 IdentityStore.prototype._createIdmgmt = function (password, seedPhrase, entropy, cb) {
-  const opts = { password }
+  const opts = {
+    password,
+    hdPathString: this.hdPathString,
+  }
+
   if (seedPhrase) {
     opts.seedPhrase = seedPhrase
   }
@@ -464,10 +468,7 @@ IdentityStore.prototype._createIdmgmt = function (password, seedPhrase, entropy,
     keyStore.keyFromPassword(password, (err, derivedKey) => {
       if (err) return cb(err)
 
-      this._ethStore._currentState = {
-        accounts: {},
-        transactions: {},
-      }
+      this.purgeCache()
 
       keyStore.addHdDerivationPath(this.hdPathString, derivedKey, {curve: 'secp256k1', purpose: 'sign'})
 
@@ -486,10 +487,16 @@ IdentityStore.prototype._createIdmgmt = function (password, seedPhrase, entropy,
   })
 }
 
+IdentityStore.prototype.purgeCache = function () {
+  this._getAddresses().forEach((address) => {
+    this._ethStore.del(address)
+  })
+}
+
 IdentityStore.prototype._createFirstWallet = function (derivedKey) {
   const keyStore = this._keyStore
   keyStore.setDefaultHdDerivationPath(this.hdPathString)
-  keyStore.generateNewAddress(derivedKey)
+  keyStore.generateNewAddress(derivedKey, 1)
   var addresses = keyStore.getAddresses()
   this._ethStore.addAccount(addresses[0])
   this.configManager.setWallet(keyStore.serialize())
@@ -497,7 +504,9 @@ IdentityStore.prototype._createFirstWallet = function (derivedKey) {
 
 // get addresses and normalize address hexString
 IdentityStore.prototype._getAddresses = function () {
-  return this._keyStore.getAddresses(this.hdPathString).map((address) => { return '0x' + address })
+  return this._keyStore.getAddresses(this.hdPathString).map((address) => {
+    return ethUtil.addHexPrefix(address)
+  })
 }
 
 IdentityStore.prototype._autoFaucet = function () {
diff --git a/test/unit/idStore-test.js b/test/unit/idStore-test.js
index 3d28ddf90..422db6f7b 100644
--- a/test/unit/idStore-test.js
+++ b/test/unit/idStore-test.js
@@ -1,6 +1,7 @@
 var assert = require('assert')
 var IdentityStore = require('../../app/scripts/lib/idStore')
 var configManagerGen = require('../lib/mock-config-manager')
+const ethUtil = require('ethereumjs-util')
 
 describe('IdentityStore', function() {
 
@@ -18,7 +19,7 @@ describe('IdentityStore', function() {
       idStore = new IdentityStore({
         configManager: configManagerGen(),
         ethStore: {
-          addAccount(acct) { accounts.push(acct) },
+          addAccount(acct) { accounts.push(ethUtil.addHexPrefix(acct)) },
         },
       })
 
@@ -39,7 +40,7 @@ describe('IdentityStore', function() {
         idStore = new IdentityStore({
           configManager: configManagerGen(),
           ethStore: {
-            addAccount(acct) { newAccounts.push(acct) },
+            addAccount(acct) { newAccounts.push(ethUtil.addHexPrefix(acct)) },
           },
         })
       })
@@ -62,6 +63,9 @@ describe('IdentityStore', function() {
     let firstAccount = '0x5d8de92c205279c10e5669f797b853ccef4f739a'
     const salt = 'lightwalletSalt'
 
+    const secondSeed = 'radar blur cabbage chef fix engine embark joy scheme fiction master release'
+    const secondAcct = '0xac39b311dceb2a4b2f5d8461c1cdaf756f4f7ae9'
+
     let password = 'secret!'
     let accounts = []
     let idStore
@@ -72,11 +76,15 @@ describe('IdentityStore', function() {
       idStore = new IdentityStore({
         configManager: configManagerGen(),
         ethStore: {
-          addAccount(acct) { accounts.push('0x' + acct) },
+          addAccount(acct) { accounts.push(ethUtil.addHexPrefix(acct)) },
         },
       })
     })
 
+    beforeEach(function() {
+      accounts = []
+    })
+
     it('should return the expected first account', function (done) {
 
       idStore.recoverFromSeed(password, seedWords, (err) => {
@@ -87,9 +95,6 @@ describe('IdentityStore', function() {
         assert.equal(accounts[0], firstAccount)
 
         accounts = []
-        const secondSeed = 'radar blur cabbage chef fix engine embark joy scheme fiction master release'
-        const secondAcct = '0xac39b311dceb2a4b2f5d8461c1cdaf756f4f7ae9'
-
         idStore.recoverFromSeed(password, secondSeed, (err) => {
 
           let accounts = idStore._getAddresses()
@@ -98,5 +103,16 @@ describe('IdentityStore', function() {
         })
       })
     })
+
+    it('should return the expected second account', function (done) {
+      idStore.recoverFromSeed(password, secondSeed, (err) => {
+        assert.ifError(err)
+
+        let newKeystore = idStore._idmgmt.keyStore
+
+        assert.equal(accounts[0], firstAccount)
+        done()
+      })
+    })
   })
 })