4 files changed, 146 insertions, 2 deletions
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dd2770a53..8383e104d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 - Remove certain non-essential permissions from certain builds.
 - Add a check for when a tx is included in a block.
 - Implement replay attack protections allowed by EIP 155.
+- Fix bug where sometimes loading account data would fail by querying a future block.
 
 ## 2.14.1 2016-12-20
 
diff --git a/app/scripts/lib/eth-store.js b/app/scripts/lib/eth-store.js
new file mode 100644
index 000000000..a42b2417f
--- /dev/null
+++ b/app/scripts/lib/eth-store.js
@@ -0,0 +1,144 @@
+/* Ethereum Store
+ *
+ * This module is responsible for tracking any number of accounts
+ * and caching their current balances & transaction counts.
+ *
+ * It also tracks transaction hashes, and checks their inclusion status
+ * on each new block.
+ */
+
+const EventEmitter = require('events').EventEmitter
+const inherits = require('util').inherits
+const async = require('async')
+const clone = require('clone')
+const EthQuery = require('eth-query')
+
+module.exports = EthereumStore
+
+
+inherits(EthereumStore, EventEmitter)
+function EthereumStore(engine) {
+  const self = this
+  EventEmitter.call(self)
+  self._currentState = {
+    accounts: {},
+    transactions: {},
+  }
+  self._query = new EthQuery(engine)
+
+  engine.on('block', self._updateForBlock.bind(self))
+}
+
+//
+// public
+//
+
+EthereumStore.prototype.getState = function () {
+  const self = this
+  return clone(self._currentState)
+}
+
+EthereumStore.prototype.addAccount = function (address) {
+  const self = this
+  self._currentState.accounts[address] = {}
+  self._didUpdate()
+  if (!self.currentBlockNumber) return
+  self._updateAccount(address, noop)
+}
+
+EthereumStore.prototype.removeAccount = function (address) {
+  const self = this
+  delete self._currentState.accounts[address]
+  self._didUpdate()
+}
+
+EthereumStore.prototype.addTransaction = function (txHash) {
+  const self = this
+  self._currentState.transactions[txHash] = {}
+  self._didUpdate()
+  if (!self.currentBlockNumber) return
+  self._updateTransaction(self.currentBlockNumber, txHash, noop)
+}
+
+EthereumStore.prototype.removeTransaction = function (address) {
+  const self = this
+  delete self._currentState.transactions[address]
+  self._didUpdate()
+}
+
+
+//
+// private
+//
+
+EthereumStore.prototype._didUpdate = function () {
+  const self = this
+  var state = self.getState()
+  self.emit('update', state)
+}
+
+EthereumStore.prototype._updateForBlock = function (block) {
+  const self = this
+  var blockNumber = '0x' + block.number.toString('hex')
+  self.currentBlockNumber = blockNumber
+  async.parallel([
+    self._updateAccounts.bind(self),
+    self._updateTransactions.bind(self, blockNumber),
+  ], function (err) {
+    if (err) return console.error(err)
+    self.emit('block', self.getState())
+    self._didUpdate()
+  })
+}
+
+EthereumStore.prototype._updateAccounts = function (cb) {
+  var accountsState = this._currentState.accounts
+  var addresses = Object.keys(accountsState)
+  async.each(addresses, this._updateAccount.bind(this), cb)
+}
+
+EthereumStore.prototype._updateAccount = function (address, cb) {
+  var accountsState = this._currentState.accounts
+  this.getAccount(address, function (err, result) {
+    if (err) return cb(err)
+    result.address = address
+    // only populate if the entry is still present
+    if (accountsState[address]) {
+      accountsState[address] = result
+    }
+    cb(null, result)
+  })
+}
+
+EthereumStore.prototype.getAccount = function (address, cb) {
+  const query = this._query
+  async.parallel({
+    balance: query.getBalance.bind(query, address),
+    nonce: query.getTransactionCount.bind(query, address),
+    code: query.getCode.bind(query, address),
+  }, cb)
+}
+
+EthereumStore.prototype._updateTransactions = function (block, cb) {
+  const self = this
+  var transactionsState = self._currentState.transactions
+  var txHashes = Object.keys(transactionsState)
+  async.each(txHashes, self._updateTransaction.bind(self, block), cb)
+}
+
+EthereumStore.prototype._updateTransaction = function (block, txHash, cb) {
+  const self = this
+  // would use the block here to determine how many confirmations the tx has
+  var transactionsState = self._currentState.transactions
+  self._query.getTransaction(txHash, function (err, result) {
+    if (err) return cb(err)
+    // only populate if the entry is still present
+    if (transactionsState[txHash]) {
+      transactionsState[txHash] = result
+      self._didUpdate()
+    }
+    cb(null, result)
+  })
+}
+
+function noop() {}
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 02e65bf10..961130dfd 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -1,5 +1,5 @@
 const extend = require('xtend')
-const EthStore = require('eth-store')
+const EthStore = require('./lib/eth-store')
 const MetaMaskProvider = require('web3-provider-engine/zero.js')
 const KeyringController = require('./keyring-controller')
 const NoticeController = require('./notice-controller')
diff --git a/package.json b/package.json
index 4c33ad9ab..6f0fe00c7 100644
--- a/package.json
+++ b/package.json
@@ -49,7 +49,6 @@
     "eth-bin-to-ops": "^1.0.1",
     "eth-lightwallet": "^2.3.3",
     "eth-query": "^1.0.3",
-    "eth-store": "^1.1.0",
     "ethereumjs-tx": "^1.0.0",
     "ethereumjs-util": "^4.4.0",
     "ethereumjs-wallet": "^0.6.0",