Merge pull request #2094 from MetaMask/AddBalanceController

Calculate balance based on pending txs

app/scripts/controllers/balance.js

@@ -0,0 +1,61 @@
+const ObservableStore = require('obs-store')
+const PendingBalanceCalculator = require('../lib/pending-balance-calculator')
+const BN = require('ethereumjs-util').BN
+
+class BalanceController {
+
+  constructor (opts = {}) {
+    const { address, accountTracker, txController, blockTracker } = opts
+    this.address = address
+    this.accountTracker = accountTracker
+    this.txController = txController
+    this.blockTracker = blockTracker
+
+    const initState = {
+      ethBalance: undefined,
+    }
+    this.store = new ObservableStore(initState)
+
+    this.balanceCalc = new PendingBalanceCalculator({
+      getBalance: () => this._getBalance(),
+      getPendingTransactions: this._getPendingTransactions.bind(this),
+    })
+
+    this._registerUpdates()
+  }
+
+  async updateBalance () {
+    const balance = await this.balanceCalc.getBalance()
+    this.store.updateState({
+      ethBalance: balance,
+    })
+  }
+
+  _registerUpdates () {
+    const update = this.updateBalance.bind(this)
+    this.txController.on('submitted', update)
+    this.txController.on('confirmed', update)
+    this.txController.on('failed', update)
+    this.accountTracker.subscribe(update)
+    this.blockTracker.on('block', update)
+  }
+
+  async _getBalance () {
+    const { accounts } = this.accountTracker.getState()
+    const entry = accounts[this.address]
+    const balance = entry.balance
+    return balance ? new BN(balance.substring(2), 16) : undefined
+  }
+
+  async _getPendingTransactions () {
+    const pending = this.txController.getFilteredTxList({
+      from: this.address,
+      status: 'submitted',
+      err: undefined,
+    })
+    return pending
+  }
+
+}
+
+module.exports = BalanceController

app/scripts/controllers/computed-balances.js

@@ -0,0 +1,66 @@
+const ObservableStore = require('obs-store')
+const extend = require('xtend')
+const BalanceController = require('./balance')
+
+class ComputedbalancesController {
+
+  constructor (opts = {}) {
+    const { accountTracker, txController, blockTracker } = opts
+    this.accountTracker = accountTracker
+    this.txController = txController
+    this.blockTracker = blockTracker
+
+    const initState = extend({
+      computedBalances: {},
+    }, opts.initState)
+    this.store = new ObservableStore(initState)
+    this.balances = {}
+
+    this._initBalanceUpdating()
+  }
+
+  updateAllBalances () {
+    for (let address in this.balances) {
+      this.balances[address].updateBalance()
+    }
+  }
+
+  _initBalanceUpdating () {
+    const store = this.accountTracker.getState()
+    this.addAnyAccountsFromStore(store)
+    this.accountTracker.subscribe(this.addAnyAccountsFromStore.bind(this))
+  }
+
+  addAnyAccountsFromStore(store) {
+    const balances = store.accounts
+
+    for (let address in balances) {
+      this.trackAddressIfNotAlready(address)
+    }
+  }
+
+  trackAddressIfNotAlready (address) {
+    const state = this.store.getState()
+    if (!(address in state.computedBalances)) {
+      this.trackAddress(address)
+    }
+  }
+
+  trackAddress (address) {
+    let updater = new BalanceController({
+      address,
+      accountTracker: this.accountTracker,
+      txController: this.txController,
+      blockTracker: this.blockTracker,
+    })
+    updater.store.subscribe((accountBalance) => {
+      let newState = this.store.getState()
+      newState.computedBalances[address] = accountBalance
+      this.store.updateState(newState)
+    })
+    this.balances[address] = updater
+    updater.updateBalance()
+  }
+}
+
+module.exports = ComputedbalancesController

app/scripts/controllers/transactions.js

@@ -22,7 +22,7 @@ module.exports = class TransactionController extends EventEmitter {
     this.provider = opts.provider
     this.blockTracker = opts.blockTracker
     this.signEthTx = opts.signTransaction
-    this.ethStore = opts.ethStore
+    this.accountTracker = opts.accountTracker
 
     this.nonceTracker = new NonceTracker({
       provider: this.provider,
@@ -52,7 +52,7 @@ module.exports = class TransactionController extends EventEmitter {
       provider: this.provider,
       nonceTracker: this.nonceTracker,
       getBalance: (address) => {
-        const account = this.ethStore.getState().accounts[address]
+        const account = this.accountTracker.getState().accounts[address]
         if (!account) return
         return account.balance
       },
@@ -73,7 +73,7 @@ module.exports = class TransactionController extends EventEmitter {
     this.blockTracker.on('rawBlock', this.pendingTxTracker.checkForTxInBlock.bind(this.pendingTxTracker))
     // this is a little messy but until ethstore has been either
     // removed or redone this is to guard against the race condition
-    // where ethStore hasent been populated by the results yet
+    // where accountTracker hasent been populated by the results yet
     this.blockTracker.once('latest', () => {
       this.blockTracker.on('latest', this.pendingTxTracker.resubmitPendingTxs.bind(this.pendingTxTracker))
     })
@@ -434,6 +434,7 @@ module.exports = class TransactionController extends EventEmitter {
     const txMeta = this.getTx(txId)
     txMeta.status = status
     this.emit(`${txMeta.id}:${status}`, txId)
+    this.emit(`${status}`, txId)
     if (status === 'submitted' || status === 'rejected') {
       this.emit(`${txMeta.id}:finished`, txMeta)
     }

app/scripts/keyring-controller.js

@@ -35,7 +35,8 @@ class KeyringController extends EventEmitter {
       keyrings: [],
       identities: {},
     })
-    this.ethStore = opts.ethStore
+
+    this.accountTracker = opts.accountTracker
     this.encryptor = opts.encryptor || encryptor
     this.keyrings = []
     this.getNetwork = opts.getNetwork
@@ -338,7 +339,7 @@ class KeyringController extends EventEmitter {
   //
   // Initializes the provided account array
   // Gives them numerically incremented nicknames,
-  // and adds them to the ethStore for regular balance checking.
+  // and adds them to the accountTracker for regular balance checking.
   setupAccounts (accounts) {
     return this.getAccounts()
     .then((loadedAccounts) => {
@@ -361,7 +362,7 @@ class KeyringController extends EventEmitter {
       throw new Error('Problem loading account.')
     }
     const address = normalizeAddress(account)
-    this.ethStore.addAccount(address)
+    this.accountTracker.addAccount(address)
     return this.createNickname(address)
   }
 
@@ -567,12 +568,12 @@ class KeyringController extends EventEmitter {
   clearKeyrings () {
     let accounts
     try {
-      accounts = Object.keys(this.ethStore.getState())
+      accounts = Object.keys(this.accountTracker.getState())
     } catch (e) {
       accounts = []
     }
     accounts.forEach((address) => {
-      this.ethStore.removeAccount(address)
+      this.accountTracker.removeAccount(address)
     })
 
     // clear keyrings from memory

app/scripts/lib/account-tracker.js

@@ -1,4 +1,4 @@
-/* Ethereum Store
+/* Account Tracker
  *
  * This module is responsible for tracking any number of accounts
  * and caching their current balances & transaction counts.
@@ -18,10 +18,6 @@ class EthereumStore extends ObservableStore {
   constructor (opts = {}) {
     super({
       accounts: {},
-      transactions: {},
-      currentBlockNumber: '0',
-      currentBlockHash: '',
-      currentBlockGasLimit: '',
     })
     this._provider = opts.provider
     this._query = new EthQuery(this._provider)
@@ -50,21 +46,6 @@ class EthereumStore extends ObservableStore {
     this.updateState({ accounts })
   }
 
-  addTransaction (txHash) {
-    const transactions = this.getState().transactions
-    transactions[txHash] = {}
-    this.updateState({ transactions })
-    if (!this._currentBlockNumber) return
-    this._updateTransaction(this._currentBlockNumber, txHash, noop)
-  }
-
-  removeTransaction (txHash) {
-    const transactions = this.getState().transactions
-    delete transactions[txHash]
-    this.updateState({ transactions })
-  }
-
-
   //
   // private
   //
@@ -72,12 +53,9 @@ class EthereumStore extends ObservableStore {
   _updateForBlock (block) {
     const blockNumber = '0x' + block.number.toString('hex')
     this._currentBlockNumber = blockNumber
-    this.updateState({ currentBlockNumber: parseInt(blockNumber) })
-    this.updateState({ currentBlockHash: `0x${block.hash.toString('hex')}`})
-    this.updateState({ currentBlockGasLimit: `0x${block.gasLimit.toString('hex')}` })
+
     async.parallel([
       this._updateAccounts.bind(this),
-      this._updateTransactions.bind(this, blockNumber),
     ], (err) => {
       if (err) return console.error(err)
       this.emit('block', this.getState())
@@ -104,26 +82,6 @@ class EthereumStore extends ObservableStore {
     })
   }
 
-  _updateTransactions (block, cb = noop) {
-    const transactions = this.getState().transactions
-    const txHashes = Object.keys(transactions)
-    async.each(txHashes, this._updateTransaction.bind(this, block), cb)
-  }
-
-  _updateTransaction (block, txHash, cb = noop) {
-    // would use the block here to determine how many confirmations the tx has
-    const transactions = this.getState().transactions
-    this._query.getTransaction(txHash, (err, result) => {
-      if (err) return cb(err)
-      // only populate if the entry is still present
-      if (transactions[txHash]) {
-        transactions[txHash] = result
-        this.updateState({ transactions })
-      }
-      cb(null, result)
-    })
-  }
-
   _getAccount (address, cb = noop) {
     const query = this._query
     async.parallel({

app/scripts/lib/pending-balance-calculator.js

@@ -0,0 +1,51 @@
+const BN = require('ethereumjs-util').BN
+const normalize = require('eth-sig-util').normalize
+
+class PendingBalanceCalculator {
+
+  // Must be initialized with two functions:
+  // getBalance => Returns a promise of a BN of the current balance in Wei
+  // getPendingTransactions => Returns an array of TxMeta Objects,
+  // which have txParams properties, which include value, gasPrice, and gas,
+  // all in a base=16 hex format.
+  constructor ({ getBalance, getPendingTransactions }) {
+    this.getPendingTransactions = getPendingTransactions
+    this.getNetworkBalance = getBalance
+  }
+
+  async getBalance() {
+    const results = await Promise.all([
+      this.getNetworkBalance(),
+      this.getPendingTransactions(),
+    ])
+
+    const [ balance, pending ] = results
+    if (!balance) return undefined
+
+    const pendingValue = pending.reduce((total, tx) => {
+      return total.add(this.calculateMaxCost(tx))
+    }, new BN(0))
+
+    return `0x${balance.sub(pendingValue).toString(16)}`
+  }
+
+  calculateMaxCost (tx) {
+    const txValue = tx.txParams.value
+    const value = this.hexToBn(txValue)
+    const gasPrice = this.hexToBn(tx.txParams.gasPrice)
+
+    const gas = tx.txParams.gas
+    const gasLimit = tx.txParams.gasLimit
+    const gasLimitBn = this.hexToBn(gas || gasLimit)
+
+    const gasCost = gasPrice.mul(gasLimitBn)
+    return value.add(gasCost)
+  }
+
+  hexToBn (hex) {
+    return new BN(normalize(hex).substring(2), 16)
+  }
+
+}
+
+module.exports = PendingBalanceCalculator

app/scripts/metamask-controller.js

@@ -4,7 +4,7 @@ const promiseToCallback = require('promise-to-callback')
 const pump = require('pump')
 const Dnode = require('dnode')
 const ObservableStore = require('obs-store')
-const EthStore = require('./lib/eth-store')
+const AccountTracker = require('./lib/account-tracker')
 const EthQuery = require('eth-query')
 const RpcEngine = require('json-rpc-engine')
 const debounce = require('debounce')
@@ -26,6 +26,7 @@ const BlacklistController = require('./controllers/blacklist')
 const MessageManager = require('./lib/message-manager')
 const PersonalMessageManager = require('./lib/personal-message-manager')
 const TransactionController = require('./controllers/transactions')
+const BalancesController = require('./controllers/computed-balances')
 const ConfigManager = require('./lib/config-manager')
 const nodeify = require('./lib/nodeify')
 const accountImporter = require('./account-import-strategies')
@@ -85,7 +86,7 @@ module.exports = class MetamaskController extends EventEmitter {
 
     // eth data query tools
     this.ethQuery = new EthQuery(this.provider)
-    this.ethStore = new EthStore({
+    this.accountTracker = new AccountTracker({
       provider: this.provider,
       blockTracker: this.blockTracker,
     })
@@ -93,12 +94,22 @@ module.exports = class MetamaskController extends EventEmitter {
     // key mgmt
     this.keyringController = new KeyringController({
       initState: initState.KeyringController,
-      ethStore: this.ethStore,
+      accountTracker: this.accountTracker,
       getNetwork: this.networkController.getNetworkState.bind(this.networkController),
       encryptor: opts.encryptor || undefined,
     })
+
+    // account tracker watches balances, nonces, and any code at their address.
+    this.accountTracker = new AccountTracker({
+      provider: this.provider,
+      blockTracker: this.provider,
+    })
     this.keyringController.on('newAccount', (address) => {
       this.preferencesController.setSelectedAddress(address)
+      this.accountTracker.addAccount(address)
+    })
+    this.keyringController.on('removedAccount', (address) => {
+      this.accountTracker.removeAccount(address)
     })
 
     // address book controller
@@ -117,10 +128,21 @@ module.exports = class MetamaskController extends EventEmitter {
       provider: this.provider,
       blockTracker: this.blockTracker,
       ethQuery: this.ethQuery,
-      ethStore: this.ethStore,
+      accountTracker: this.accountTracker,
     })
     this.txController.on('newUnaprovedTx', opts.showUnapprovedTx.bind(opts))
 
+    // computed balances (accounting for pending transactions)
+    this.balancesController = new BalancesController({
+      accountTracker: this.accountTracker,
+      txController: this.txController,
+      blockTracker: this.blockTracker,
+    })
+    this.networkController.on('networkDidChange', () => {
+      this.balancesController.updateAllBalances()
+    })
+    this.balancesController.updateAllBalances()
+
     // notices
     this.noticeController = new NoticeController({
       initState: initState.NoticeController,
@@ -172,8 +194,9 @@ module.exports = class MetamaskController extends EventEmitter {
 
     // manual mem state subscriptions
     this.networkController.store.subscribe(this.sendUpdate.bind(this))
-    this.ethStore.subscribe(this.sendUpdate.bind(this))
+    this.accountTracker.subscribe(this.sendUpdate.bind(this))
     this.txController.memStore.subscribe(this.sendUpdate.bind(this))
+    this.balancesController.store.subscribe(this.sendUpdate.bind(this))
     this.messageManager.memStore.subscribe(this.sendUpdate.bind(this))
     this.personalMessageManager.memStore.subscribe(this.sendUpdate.bind(this))
     this.keyringController.memStore.subscribe(this.sendUpdate.bind(this))
@@ -248,16 +271,18 @@ module.exports = class MetamaskController extends EventEmitter {
     const wallet = this.configManager.getWallet()
     const vault = this.keyringController.store.getState().vault
     const isInitialized = (!!wallet || !!vault)
+
     return extend(
       {
         isInitialized,
       },
       this.networkController.store.getState(),
-      this.ethStore.getState(),
+      this.accountTracker.getState(),
       this.txController.memStore.getState(),
       this.messageManager.memStore.getState(),
       this.personalMessageManager.memStore.getState(),
       this.keyringController.memStore.getState(),
+      this.balancesController.store.getState(),
       this.preferencesController.store.getState(),
       this.addressBookController.store.getState(),
       this.currencyController.store.getState(),

development/states/first-time.json

@@ -4,6 +4,7 @@
     "isUnlocked": false,
     "rpcTarget": "https://rawtestrpc.metamask.io/",
     "identities": {},
+    "computedBalances": {},
     "frequentRpcList": [],
     "unapprovedTxs": {},
     "currentCurrency": "USD",
@@ -48,5 +49,6 @@
     "isLoading": false,
     "warning": null
   },
-  "identities": {}
+  "identities": {},
+  "computedBalances": {}
 }

test/lib/mock-encryptor.js

@@ -29,4 +29,8 @@ module.exports = {
     return 'WHADDASALT!'
   },
 
+  getRandomValues () {
+    return 'SOO RANDO!!!1'
+  }
+
 }

test/unit/components/pending-tx-test.js

@@ -34,10 +34,15 @@ describe('PendingTx', function () {
     const renderer = ReactTestUtils.createRenderer()
     const newGasPrice = '0x77359400'
 
+    const computedBalances = {}
+    computedBalances[Object.keys(identities)[0]] = {
+      ethBalance: '0x00000000000000056bc75e2d63100000',
+    }
     const props = {
       identities,
       accounts: identities,
       txData,
+      computedBalances,
       sendTransaction: (txMeta, event) => {
         // Assert changes:
         const result = ethUtil.addHexPrefix(txMeta.txParams.gasPrice)

test/unit/keyring-controller-test.js

@@ -24,7 +24,7 @@ describe('KeyringController', function () {
         getTxList: () => [],
         getUnapprovedTxList: () => [],
       },
-      ethStore: {
+      accountTracker: {
         addAccount (acct) { accounts.push(ethUtil.addHexPrefix(acct)) },
       },
       encryptor: mockEncryptor,

test/unit/pending-balance-test.js

@@ -0,0 +1,93 @@
+const assert = require('assert')
+const PendingBalanceCalculator = require('../../app/scripts/lib/pending-balance-calculator')
+const MockTxGen = require('../lib/mock-tx-gen')
+const BN = require('ethereumjs-util').BN
+let providerResultStub = {}
+
+const zeroBn = new BN(0)
+const etherBn = new BN(String(1e18))
+const ether = '0x' + etherBn.toString(16)
+
+describe('PendingBalanceCalculator', function () {
+  let balanceCalculator
+
+  describe('#calculateMaxCost(tx)', function () {
+    it('returns a BN for a given tx value', function () {
+      const txGen = new MockTxGen()
+      pendingTxs = txGen.generate({
+        status: 'submitted',
+        txParams: {
+          value: ether,
+          gasPrice: '0x0',
+          gas: '0x0',
+        }
+      }, { count: 1 })
+
+      const balanceCalculator = generateBalanceCalcWith([], zeroBn)
+      const result = balanceCalculator.calculateMaxCost(pendingTxs[0])
+      assert.equal(result.toString(), etherBn.toString(), 'computes one ether')
+    })
+
+    it('calculates gas costs as well', function () {
+      const txGen = new MockTxGen()
+      pendingTxs = txGen.generate({
+        status: 'submitted',
+        txParams: {
+          value: '0x0',
+          gasPrice: '0x2',
+          gas: '0x3',
+        }
+      }, { count: 1 })
+
+      const balanceCalculator = generateBalanceCalcWith([], zeroBn)
+      const result = balanceCalculator.calculateMaxCost(pendingTxs[0])
+      assert.equal(result.toString(), '6', 'computes 6 wei of gas')
+    })
+  })
+
+  describe('if you have no pending txs and one ether', function () {
+
+    beforeEach(function () {
+      balanceCalculator = generateBalanceCalcWith([], etherBn)
+    })
+
+    it('returns the network balance', async function () {
+      const result = await balanceCalculator.getBalance()
+      assert.equal(result, ether, `gave ${result} needed ${ether}`)
+    })
+  })
+
+  describe('if you have a one ether pending tx and one ether', function () {
+    beforeEach(function () {
+      const txGen = new MockTxGen()
+      pendingTxs = txGen.generate({
+        status: 'submitted',
+        txParams: {
+          value: ether,
+          gasPrice: '0x0',
+          gas: '0x0',
+        }
+      }, { count: 1 })
+
+      balanceCalculator = generateBalanceCalcWith(pendingTxs, etherBn)
+    })
+
+    it('returns the subtracted result', async function () {
+      const result = await balanceCalculator.getBalance()
+      assert.equal(result, '0x0', `gave ${result} needed '0x0'`)
+      return true
+    })
+
+  })
+})
+
+function generateBalanceCalcWith (transactions, providerStub = zeroBn) {
+  const getPendingTransactions = async () => transactions
+  const getBalance = async () => providerStub
+
+  return new PendingBalanceCalculator({
+    getBalance,
+    getPendingTransactions,
+  })
+}
+

test/unit/tx-controller-test.js

@@ -27,7 +27,7 @@ describe('Transaction Controller', function () {
       networkStore: new ObservableStore(currentNetworkId),
       txHistoryLimit: 10,
       blockTracker: { getCurrentBlock: noop, on: noop, once: noop },
-      ethStore: { getState: noop },
+      accountTracker: { getState: noop },
       signTransaction: (ethTx) => new Promise((resolve) => {
         ethTx.sign(privKey)
         resolve()
@@ -431,4 +431,4 @@ describe('Transaction Controller', function () {
       }).catch(done)
     })
   })
-})
+})

ui/app/account-detail.js

@@ -32,6 +32,7 @@ function mapStateToProps (state) {
     currentCurrency: state.metamask.currentCurrency,
     currentAccountTab: state.metamask.currentAccountTab,
     tokens: state.metamask.tokens,
+    computedBalances: state.metamask.computedBalances,
   }
 }
 
@@ -45,7 +46,7 @@ AccountDetailScreen.prototype.render = function () {
   var selected = props.address || Object.keys(props.accounts)[0]
   var checksumAddress = selected && ethUtil.toChecksumAddress(selected)
   var identity = props.identities[selected]
-  var account = props.accounts[selected]
+  var account = props.computedBalances[selected]
   const { network, conversionRate, currentCurrency } = props
 
   return (
@@ -180,7 +181,7 @@ AccountDetailScreen.prototype.render = function () {
         }, [
 
           h(EthBalance, {
-            value: account && account.balance,
+            value: account && account.ethBalance,
             conversionRate,
             currentCurrency,
             style: {

ui/app/components/pending-tx.js

@@ -33,7 +33,7 @@ function PendingTx () {
 
 PendingTx.prototype.render = function () {
   const props = this.props
-  const { currentCurrency, blockGasLimit } = props
+  const { currentCurrency, blockGasLimit, computedBalances } = props
 
   const conversionRate = props.conversionRate
   const txMeta = this.gatherTxMeta()
@@ -42,8 +42,8 @@ PendingTx.prototype.render = function () {
   // Account Details
   const address = txParams.from || props.selectedAddress
   const identity = props.identities[address] || { address: address }
-  const account = props.accounts[address]
-  const balance = account ? account.balance : '0x0'
+  const account = computedBalances[address]
+  const balance = account ? account.ethBalance : '0x0'
 
   // recipient check
   const isValidAddress = !txParams.to || util.isValidAddress(txParams.to)

ui/app/conf-tx.js

@@ -29,6 +29,7 @@ function mapStateToProps (state) {
     conversionRate: state.metamask.conversionRate,
     currentCurrency: state.metamask.currentCurrency,
     blockGasLimit: state.metamask.currentBlockGasLimit,
+    computedBalances: state.metamask.computedBalances,
   }
 }
 
@@ -39,7 +40,7 @@ function ConfirmTxScreen () {
 
 ConfirmTxScreen.prototype.render = function () {
   const props = this.props
-  const { network, provider, unapprovedTxs, currentCurrency,
+  const { network, provider, unapprovedTxs, currentCurrency, computedBalances,
     unapprovedMsgs, unapprovedPersonalMsgs, conversionRate, blockGasLimit } = props
 
   var unconfTxList = txHelper(unapprovedTxs, unapprovedMsgs, unapprovedPersonalMsgs, network)
@@ -48,7 +49,6 @@ ConfirmTxScreen.prototype.render = function () {
   var txParams = txData.params || {}
   var isNotification = isPopupOrNotification() === 'notification'
 
-
   log.info(`rendering a combined ${unconfTxList.length} unconf msg & txs`)
   if (unconfTxList.length === 0) return h(Loading, { isLoading: true })
 
@@ -104,6 +104,7 @@ ConfirmTxScreen.prototype.render = function () {
         currentCurrency,
         blockGasLimit,
         unconfTxListLength,
+        computedBalances,
         // Actions
         buyEth: this.buyEth.bind(this, txParams.from || props.selectedAddress),
         sendTransaction: this.sendTransaction.bind(this),