Merge pull request #1617 from MetaMask/nonce-tracker

 transaction controller - use nonce-tracker

9 files changed, 205 insertions, 121 deletions

diff --git a/app/scripts/controllers/transactions.js b/app/scripts/controllers/transactions.js
index 4d037ce98..61e96ca13 100644
--- a/app/scripts/controllers/transactions.js
+++ b/app/scripts/controllers/transactions.js
@@ -1,12 +1,11 @@
 const EventEmitter = require('events')
 const async = require('async')
 const extend = require('xtend')
-const Semaphore = require('semaphore')
 const ObservableStore = require('obs-store')
 const ethUtil = require('ethereumjs-util')
 const TxProviderUtil = require('../lib/tx-utils')
 const createId = require('../lib/random-id')
-const denodeify = require('denodeify')
+const NonceTracker = require('../lib/nonce-tracker')
 
 module.exports = class TransactionController extends EventEmitter {
   constructor (opts) {
@@ -20,6 +19,17 @@ module.exports = class TransactionController extends EventEmitter {
     this.txHistoryLimit = opts.txHistoryLimit
     this.provider = opts.provider
     this.blockTracker = opts.blockTracker
+    this.nonceTracker = new NonceTracker({
+      provider: this.provider,
+      blockTracker: this.provider._blockTracker,
+      getPendingTransactions: (address) => {
+        return this.getFilteredTxList({
+          from: address,
+          status: 'submitted',
+          err: undefined,
+        })
+      },
+    })
     this.query = opts.ethQuery
     this.txProviderUtils = new TxProviderUtil(this.query)
     this.blockTracker.on('rawBlock', this.checkForTxInBlock.bind(this))
@@ -29,7 +39,6 @@ module.exports = class TransactionController extends EventEmitter {
     this.blockTracker.once('latest', () => this.blockTracker.on('latest', this.resubmitPendingTxs.bind(this)))
     this.blockTracker.on('sync', this.queryPendingTxs.bind(this))
     this.signEthTx = opts.signTransaction
-    this.nonceLock = Semaphore(1)
     this.ethStore = opts.ethStore
     // memstore is computed from a few different stores
     this._updateMemstore()
@@ -173,29 +182,32 @@ module.exports = class TransactionController extends EventEmitter {
     }, {})
   }
 
-  approveTransaction (txId, cb = warn) {
-    const self = this
-    // approve
-    self.setTxStatusApproved(txId)
-    // only allow one tx at a time for atomic nonce usage
-    self.nonceLock.take(() => {
-      // begin signature process
-      async.waterfall([
-        (cb) => self.fillInTxParams(txId, cb),
-        (cb) => self.signTransaction(txId, cb),
-        (rawTx, cb) => self.publishTransaction(txId, rawTx, cb),
-      ], (err) => {
-        self.nonceLock.leave()
-        if (err) {
-          this.setTxStatusFailed(txId, {
-            errCode: err.errCode || err,
-            message: err.message || 'Transaction failed during approval',
-          })
-          return cb(err)
-        }
-        cb()
+  async approveTransaction (txId) {
+    let nonceLock
+    try {
+      // approve
+      this.setTxStatusApproved(txId)
+      // get next nonce
+      const txMeta = this.getTx(txId)
+      const fromAddress = txMeta.txParams.from
+      nonceLock = await this.nonceTracker.getNonceLock(fromAddress)
+      txMeta.txParams.nonce = nonceLock.nextNonce
+      this.updateTx(txMeta)
+      // sign transaction
+      const rawTx = await this.signTransaction(txId)
+      await this.publishTransaction(txId, rawTx)
+      // must set transaction to submitted/failed before releasing lock
+      nonceLock.releaseLock()
+    } catch (err) {
+      this.setTxStatusFailed(txId, {
+        errCode: err.errCode || err,
+        message: err.message || 'Transaction failed during approval',
       })
-    })
+      // must set transaction to submitted/failed before releasing lock
+      if (nonceLock) nonceLock.releaseLock()
+      // continue with error chain
+      throw err
+    }
   }
 
   cancelTransaction (txId, cb = warn) {
@@ -203,13 +215,9 @@ module.exports = class TransactionController extends EventEmitter {
     cb()
   }
 
-  fillInTxParams (txId, cb) {
-    const txMeta = this.getTx(txId)
-    this.txProviderUtils.fillInTxParams(txMeta.txParams, (err) => {
-      if (err) return cb(err)
-      this.updateTx(txMeta)
-      cb()
-    })
+  async updateAndApproveTransaction (txMeta) {
+    this.updateTx(txMeta)
+    await this.approveTransaction(txMeta.id)
   }
 
   getChainId () {
@@ -222,31 +230,27 @@ module.exports = class TransactionController extends EventEmitter {
     }
   }
 
-  signTransaction (txId, cb) {
+  async signTransaction (txId) {
     const txMeta = this.getTx(txId)
     const txParams = txMeta.txParams
     const fromAddress = txParams.from
     // add network/chain id
     txParams.chainId = this.getChainId()
     const ethTx = this.txProviderUtils.buildEthTxFromParams(txParams)
-    this.signEthTx(ethTx, fromAddress).then(() => {
+    const rawTx = await this.signEthTx(ethTx, fromAddress).then(() => {
       this.setTxStatusSigned(txMeta.id)
-      cb(null, ethUtil.bufferToHex(ethTx.serialize()))
-    }).catch((err) => {
-      cb(err)
+      return ethUtil.bufferToHex(ethTx.serialize())
     })
+    return rawTx
   }
 
-  publishTransaction (txId, rawTx, cb = warn) {
+  async publishTransaction (txId, rawTx) {
     const txMeta = this.getTx(txId)
     txMeta.rawTx = rawTx
     this.updateTx(txMeta)
-
-    this.txProviderUtils.publishTransaction(rawTx, (err, txHash) => {
-      if (err) return cb(err)
+    await this.txProviderUtils.publishTransaction(rawTx).then((txHash) => {
       this.setTxHash(txId, txHash)
       this.setTxStatusSubmitted(txId)
-      cb()
     })
   }
 
@@ -264,10 +268,19 @@ module.exports = class TransactionController extends EventEmitter {
     to: '0x0..',
     from: '0x0..',
     status: 'signed',
+    err: undefined,
   }
   and returns a list of tx with all
   options matching
 
+  ****************HINT****************
+  | `err: undefined` is like looking |
+  | for a tx with no err             |
+  | so you can also search txs that  |
+  | dont have something as well by   |
+  | setting the value as undefined   |
+  ************************************
+
   this is for things like filtering a the tx list
   for only tx's from 1 account
   or for filltering for all txs from one account
@@ -416,8 +429,7 @@ module.exports = class TransactionController extends EventEmitter {
     const pending = this.getTxsByMetaData('status', 'submitted')
     // only try resubmitting if their are transactions to resubmit
     if (!pending.length) return
-    const resubmit = denodeify(this._resubmitTx.bind(this))
-    pending.forEach((txMeta) => resubmit(txMeta).catch((err) => {
+    pending.forEach((txMeta) => this._resubmitTx(txMeta).catch((err) => {
       /*
       Dont marked as failed if the error is a "known" transaction warning
       "there is already a transaction with the same sender-nonce
@@ -445,7 +457,7 @@ module.exports = class TransactionController extends EventEmitter {
     }))
   }
 
-  _resubmitTx (txMeta, cb) {
+  async _resubmitTx (txMeta, cb) {
     const address = txMeta.txParams.from
     const balance = this.ethStore.getState().accounts[address].balance
     if (!('retryCount' in txMeta)) txMeta.retryCount = 0
@@ -464,7 +476,7 @@ module.exports = class TransactionController extends EventEmitter {
     // Increment a try counter.
     txMeta.retryCount++
     const rawTx = txMeta.rawTx
-    this.txProviderUtils.publishTransaction(rawTx, cb)
+    return await this.txProviderUtils.publishTransaction(rawTx, cb)
   }
 
   // checks the network for signed txs and
diff --git a/app/scripts/lib/nodeify.js b/app/scripts/lib/nodeify.js
index 51d89a8fb..299bfe624 100644
--- a/app/scripts/lib/nodeify.js
+++ b/app/scripts/lib/nodeify.js
@@ -1,24 +1,9 @@
-module.exports = function (promiseFn) {
-  return function () {
-    var args = []
-    for (var i = 0; i < arguments.length - 1; i++) {
-      args.push(arguments[i])
-    }
-    var cb = arguments[arguments.length - 1]
+const promiseToCallback = require('promise-to-callback')
 
-    const nodeified = promiseFn.apply(this, args)
-
-    if (!nodeified) {
-      const methodName = String(promiseFn).split('(')[0]
-      throw new Error(`The ${methodName} did not return a Promise, but was nodeified.`)
-    }
-    nodeified.then(function (result) {
-      cb(null, result)
-    })
-    .catch(function (reason) {
-      cb(reason)
-    })
-
-    return nodeified
+module.exports = function(fn, context) {
+  return function(){
+    const args = [].slice.call(arguments)
+    const callback = args.pop()
+    promiseToCallback(fn.apply(context, args))(callback)
   }
 }
diff --git a/app/scripts/lib/nonce-tracker.js b/app/scripts/lib/nonce-tracker.js
new file mode 100644
index 000000000..ab2893b10
--- /dev/null
+++ b/app/scripts/lib/nonce-tracker.js
@@ -0,0 +1,59 @@
+const EthQuery = require('eth-query')
+
+class NonceTracker {
+
+  constructor ({ blockTracker, provider, getPendingTransactions }) {
+    this.blockTracker = blockTracker
+    this.ethQuery = new EthQuery(provider)
+    this.getPendingTransactions = getPendingTransactions
+    this.lockMap = {}
+  }
+
+  // releaseLock must be called
+  // releaseLock must be called after adding signed tx to pending transactions (or discarding)
+  async getNonceLock (address) {
+    // await lock free
+    await this.lockMap[address]
+    // take lock
+    const releaseLock = this._takeLock(address)
+    // calculate next nonce
+    // we need to make sure our base count
+    // and pending count are from the same block
+    const currentBlock = await this._getCurrentBlock()
+    const pendingTransactions = this.getPendingTransactions(address)
+    const baseCount = await this._getTxCount(address, currentBlock)
+    const nextNonce = parseInt(baseCount) + pendingTransactions.length
+    // return next nonce and release cb
+    return { nextNonce: nextNonce.toString(16), releaseLock }
+  }
+
+  async _getCurrentBlock () {
+    const currentBlock = this.blockTracker.getCurrentBlock()
+    if (currentBlock) return currentBlock
+    return await Promise((reject, resolve) => {
+      this.blockTracker.once('latest', resolve)
+    })
+  }
+
+  _takeLock (lockId) {
+    let releaseLock = null
+    // create and store lock
+    const lock = new Promise((resolve, reject) => { releaseLock = resolve })
+    this.lockMap[lockId] = lock
+    // setup lock teardown
+    lock.then(() => delete this.lockMap[lockId])
+    return releaseLock
+  }
+
+  async _getTxCount (address, currentBlock) {
+    const blockNumber = currentBlock.number
+    return new Promise((resolve, reject) => {
+      this.ethQuery.getTransactionCount(address, blockNumber, (err, result) => {
+        err ? reject(err) : resolve(result)
+      })
+    })
+  }
+
+}
+
+module.exports = NonceTracker
diff --git a/app/scripts/lib/tx-utils.js b/app/scripts/lib/tx-utils.js
index 4e780fcc0..8f6943937 100644
--- a/app/scripts/lib/tx-utils.js
+++ b/app/scripts/lib/tx-utils.js
@@ -106,8 +106,13 @@ module.exports = class txProviderUtils {
     return ethTx
   }
 
-  publishTransaction (rawTx, cb) {
-    this.query.sendRawTransaction(rawTx, cb)
+  publishTransaction (rawTx) {
+    return new Promise((resolve, reject) => {
+      this.query.sendRawTransaction(rawTx, (err, ress) => {
+        if (err) reject(err)
+        else resolve(ress)
+      })
+    })
   }
 
   validateTxParams (txParams, cb) {
@@ -118,11 +123,11 @@ module.exports = class txProviderUtils {
     }
   }
 
-  sufficientBalance (tx, hexBalance) {
+  sufficientBalance (txParams, hexBalance) {
     const balance = hexToBn(hexBalance)
-    const value = hexToBn(tx.value)
-    const gasLimit = hexToBn(tx.gas)
-    const gasPrice = hexToBn(tx.gasPrice)
+    const value = hexToBn(txParams.value)
+    const gasLimit = hexToBn(txParams.gas)
+    const gasPrice = hexToBn(txParams.gasPrice)
 
     const maxCost = value.add(gasLimit.mul(gasPrice))
     return balance.gte(maxCost)
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 0e7ccbd66..c6c3fde1e 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -294,34 +294,33 @@ module.exports = class MetamaskController extends EventEmitter {
       submitPassword: this.submitPassword.bind(this),
 
       // PreferencesController
-      setSelectedAddress: nodeify(preferencesController.setSelectedAddress).bind(preferencesController),
-      addToken: nodeify(preferencesController.addToken).bind(preferencesController),
-      setCurrentAccountTab: nodeify(preferencesController.setCurrentAccountTab).bind(preferencesController),
-      setDefaultRpc: nodeify(this.setDefaultRpc).bind(this),
-      setCustomRpc: nodeify(this.setCustomRpc).bind(this),
+      setSelectedAddress: nodeify(preferencesController.setSelectedAddress, preferencesController),
+      addToken: nodeify(preferencesController.addToken, preferencesController),
+      setCurrentAccountTab: nodeify(preferencesController.setCurrentAccountTab, preferencesController),
+      setDefaultRpc: nodeify(this.setDefaultRpc, this),
+      setCustomRpc: nodeify(this.setCustomRpc, this),
 
       // AddressController
-      setAddressBook: nodeify(addressBookController.setAddressBook).bind(addressBookController),
+      setAddressBook: nodeify(addressBookController.setAddressBook, addressBookController),
 
       // KeyringController
-      setLocked: nodeify(keyringController.setLocked).bind(keyringController),
-      createNewVaultAndKeychain: nodeify(keyringController.createNewVaultAndKeychain).bind(keyringController),
-      createNewVaultAndRestore: nodeify(keyringController.createNewVaultAndRestore).bind(keyringController),
-      addNewKeyring: nodeify(keyringController.addNewKeyring).bind(keyringController),
-      saveAccountLabel: nodeify(keyringController.saveAccountLabel).bind(keyringController),
-      exportAccount: nodeify(keyringController.exportAccount).bind(keyringController),
+      setLocked: nodeify(keyringController.setLocked, keyringController),
+      createNewVaultAndKeychain: nodeify(keyringController.createNewVaultAndKeychain, keyringController),
+      createNewVaultAndRestore: nodeify(keyringController.createNewVaultAndRestore, keyringController),
+      addNewKeyring: nodeify(keyringController.addNewKeyring, keyringController),
+      saveAccountLabel: nodeify(keyringController.saveAccountLabel, keyringController),
+      exportAccount: nodeify(keyringController.exportAccount, keyringController),
 
       // txController
-      approveTransaction: txController.approveTransaction.bind(txController),
       cancelTransaction: txController.cancelTransaction.bind(txController),
-      updateAndApproveTransaction: this.updateAndApproveTx.bind(this),
+      updateAndApproveTransaction: nodeify(txController.updateAndApproveTransaction, txController),
 
       // messageManager
-      signMessage: nodeify(this.signMessage).bind(this),
+      signMessage: nodeify(this.signMessage, this),
       cancelMessage: this.cancelMessage.bind(this),
 
       // personalMessageManager
-      signPersonalMessage: nodeify(this.signPersonalMessage).bind(this),
+      signPersonalMessage: nodeify(this.signPersonalMessage, this),
       cancelPersonalMessage: this.cancelPersonalMessage.bind(this),
 
       // notices
@@ -502,13 +501,6 @@ module.exports = class MetamaskController extends EventEmitter {
     })
   }
 
-  updateAndApproveTx (txMeta, cb) {
-    log.debug(`MetaMaskController - updateAndApproveTx: ${JSON.stringify(txMeta)}`)
-    const txController = this.txController
-    txController.updateTx(txMeta)
-    txController.approveTransaction(txMeta.id, cb)
-  }
-
   signMessage (msgParams, cb) {
     log.info('MetaMaskController - signMessage')
     const msgId = msgParams.metamaskId
diff --git a/package.json b/package.json
index d3de895c3..1f2d0d591 100644
--- a/package.json
+++ b/package.json
@@ -11,6 +11,7 @@
     "dist": "npm run clear && npm install && gulp dist",
     "test": "npm run lint && npm run test-unit && npm run test-integration",
     "test-unit": "METAMASK_ENV=test mocha --require test/helper.js --recursive \"test/unit/**/*.js\"",
+    "single-test": "METAMASK_ENV=test mocha --require test/helper.js",
     "test-integration": "npm run buildMock && npm run buildCiUnits && testem ci -P 2",
     "lint": "gulp lint",
     "buildCiUnits": "node test/integration/index.js",
@@ -56,7 +57,6 @@
     "copy-to-clipboard": "^2.0.0",
     "debounce": "^1.0.0",
     "deep-extend": "^0.4.1",
-    "denodeify": "^1.2.1",
     "detect-node": "^2.0.3",
     "disc": "^1.3.2",
     "dnode": "^1.2.2",
diff --git a/test/unit/nodeify-test.js b/test/unit/nodeify-test.js
index 5aed758fa..06241334d 100644
--- a/test/unit/nodeify-test.js
+++ b/test/unit/nodeify-test.js
@@ -11,7 +11,7 @@ describe('nodeify', function () {
   }
 
   it('should retain original context', function (done) {
-    var nodified = nodeify(obj.promiseFunc).bind(obj)
+    var nodified = nodeify(obj.promiseFunc, obj)
     nodified('baz', function (err, res) {
       assert.equal(res, 'barbaz')
       done()
diff --git a/test/unit/nonce-tracker-test.js b/test/unit/nonce-tracker-test.js
new file mode 100644
index 000000000..16cd6d008
--- /dev/null
+++ b/test/unit/nonce-tracker-test.js
@@ -0,0 +1,40 @@
+const assert = require('assert')
+const NonceTracker = require('../../app/scripts/lib/nonce-tracker')
+
+describe('Nonce Tracker', function () {
+  let nonceTracker, provider, getPendingTransactions, pendingTxs
+
+
+  beforeEach(function () {
+    pendingTxs = [{
+      'status': 'submitted',
+      'txParams': {
+        'from': '0x7d3517b0d011698406d6e0aed8453f0be2697926',
+        'gas': '0x30d40',
+        'value': '0x0',
+        'nonce': '0x0',
+      },
+    }]
+
+
+    getPendingTransactions = () => pendingTxs
+    provider = { sendAsync: (_, cb) => { cb(undefined, {result: '0x0'}) } }
+    nonceTracker = new NonceTracker({
+      blockTracker: {
+        getCurrentBlock: () => '0x11b568',
+      },
+      provider,
+      getPendingTransactions,
+    })
+  })
+
+  describe('#getNonceLock', function () {
+    it('should work', async function (done) {
+      this.timeout(15000)
+      const nonceLock = await nonceTracker.getNonceLock('0x7d3517b0d011698406d6e0aed8453f0be2697926')
+      assert.equal(nonceLock.nextNonce, '1', 'nonce should be 1')
+      await nonceLock.releaseLock()
+      done()
+    })
+  })
+})
diff --git a/test/unit/tx-controller-test.js b/test/unit/tx-controller-test.js
index 01a498820..7b86cfe14 100644
--- a/test/unit/tx-controller-test.js
+++ b/test/unit/tx-controller-test.js
@@ -1,5 +1,4 @@
 const assert = require('assert')
-const EventEmitter = require('events')
 const ethUtil = require('ethereumjs-util')
 const EthTx = require('ethereumjs-tx')
 const EthQuery = require('eth-query')
@@ -19,15 +18,16 @@ describe('Transaction Controller', function () {
     txController = new TransactionController({
       networkStore: new ObservableStore(currentNetworkId),
       txHistoryLimit: 10,
+      blockTracker: { getCurrentBlock: noop, on: noop, once: noop },
+      provider: { sendAsync: noop },
+      ethQuery: new EthQuery({ sendAsync: noop }),
       ethStore: { getState: noop },
-      provider: { _blockTracker: new EventEmitter()},
-      blockTracker: new EventEmitter(),
-      ethQuery: new EthQuery(new EventEmitter()),
       signTransaction: (ethTx) => new Promise((resolve) => {
         ethTx.sign(privKey)
         resolve()
       }),
     })
+    txController.nonceTracker.getNonceLock = () => Promise.resolve({ nextNonce: 0, releaseLock: noop })
   })
 
   describe('#validateTxParams', function () {
@@ -271,56 +271,47 @@ describe('Transaction Controller', function () {
 
 
     it('does not overwrite set values', function (done) {
+      this.timeout(15000)
       const wrongValue = '0x05'
 
       txController.addTx(txMeta)
 
       const estimateStub = sinon.stub(txController.txProviderUtils.query, 'estimateGas')
-      .callsArgWith(1, null, wrongValue)
+      .callsArgWithAsync(1, null, wrongValue)
 
       const priceStub = sinon.stub(txController.txProviderUtils.query, 'gasPrice')
-      .callsArgWith(0, null, wrongValue)
+      .callsArgWithAsync(0, null, wrongValue)
 
-      const nonceStub = sinon.stub(txController.txProviderUtils.query, 'getTransactionCount')
-      .callsArgWith(2, null, wrongValue)
 
-      const signStub = sinon.stub(txController, 'signTransaction')
-      .callsArgWith(1, null, noop)
+      const signStub = sinon.stub(txController, 'signTransaction', () => Promise.resolve())
 
-      const pubStub = sinon.stub(txController.txProviderUtils, 'publishTransaction')
-      .callsArgWith(1, null, originalValue)
-
-      txController.approveTransaction(txMeta.id, (err) => {
-        assert.ifError(err, 'should not error')
+      const pubStub = sinon.stub(txController.txProviderUtils, 'publishTransaction', () => Promise.resolve(originalValue))
 
+      txController.approveTransaction(txMeta.id).then(() => {
         const result = txController.getTx(txMeta.id)
         const params = result.txParams
 
         assert.equal(params.gas, originalValue, 'gas unmodified')
         assert.equal(params.gasPrice, originalValue, 'gas price unmodified')
-        assert.equal(params.nonce, originalValue, 'nonce unmodified')
-        assert.equal(result.hash, originalValue, 'hash was set')
+        assert.equal(result.hash, originalValue, `hash was set \n got: ${result.hash} \n expected: ${originalValue}`)
 
         estimateStub.restore()
         priceStub.restore()
         signStub.restore()
-        nonceStub.restore()
         pubStub.restore()
-
         done()
-      })
+      }).catch(done)
     })
   })
 
   describe('#sign replay-protected tx', function () {
     it('prepares a tx with the chainId set', function (done) {
       txController.addTx({ id: '1', status: 'unapproved', metamaskNetworkId: currentNetworkId, txParams: {} }, noop)
-      txController.signTransaction('1', (err, rawTx) => {
-        if (err) return done('it should not fail')
+      txController.signTransaction('1').then((rawTx) => {
         const ethTx = new EthTx(ethUtil.toBuffer(rawTx))
         assert.equal(ethTx.getChainId(), currentNetworkId)
         done()
-      })
+      }).catch(done)
     })
   })