7 files changed, 166 insertions, 115 deletions
diff --git a/app/scripts/lib/account-tracker.js b/app/scripts/lib/account-tracker.js
index 0f7b3d865..b7e2c7cbe 100644
--- a/app/scripts/lib/account-tracker.js
+++ b/app/scripts/lib/account-tracker.js
@@ -7,14 +7,13 @@
  * on each new block.
  */
 
-const async = require('async')
 const EthQuery = require('eth-query')
 const ObservableStore = require('obs-store')
-const EventEmitter = require('events').EventEmitter
-function noop () {}
+const log = require('loglevel')
+const pify = require('pify')
 
 
-class AccountTracker extends EventEmitter {
+class AccountTracker {
 
   /**
    * This module is responsible for tracking any number of accounts and caching their current balances & transaction
@@ -35,8 +34,6 @@ class AccountTracker extends EventEmitter {
    *
    */
   constructor (opts = {}) {
-    super()
-
     const initState = {
       accounts: {},
       currentBlockGasLimit: '',
@@ -44,12 +41,12 @@ class AccountTracker extends EventEmitter {
     this.store = new ObservableStore(initState)
 
     this._provider = opts.provider
-    this._query = new EthQuery(this._provider)
+    this._query = pify(new EthQuery(this._provider))
     this._blockTracker = opts.blockTracker
     // subscribe to latest block
-    this._blockTracker.on('block', this._updateForBlock.bind(this))
+    this._blockTracker.on('latest', this._updateForBlock.bind(this))
     // blockTracker.currentBlock may be null
-    this._currentBlockNumber = this._blockTracker.currentBlock
+    this._currentBlockNumber = this._blockTracker.getCurrentBlock()
   }
 
   /**
@@ -67,49 +64,57 @@ class AccountTracker extends EventEmitter {
     const accounts = this.store.getState().accounts
     const locals = Object.keys(accounts)
 
-    const toAdd = []
+    const accountsToAdd = []
     addresses.forEach((upstream) => {
       if (!locals.includes(upstream)) {
-        toAdd.push(upstream)
+        accountsToAdd.push(upstream)
       }
     })
 
-    const toRemove = []
+    const accountsToRemove = []
     locals.forEach((local) => {
       if (!addresses.includes(local)) {
-        toRemove.push(local)
+        accountsToRemove.push(local)
       }
     })
 
-    toAdd.forEach(upstream => this.addAccount(upstream))
-    toRemove.forEach(local => this.removeAccount(local))
-    this._updateAccounts()
+    this.addAccounts(accountsToAdd)
+    this.removeAccount(accountsToRemove)
   }
 
   /**
-   * Adds a new address to this AccountTracker's accounts object, which points to an empty object. This object will be
+   * Adds new addresses to track the balances of
    * given a balance as long this._currentBlockNumber is defined.
    *
-   * @param {string} address A hex address of a new account to store in this AccountTracker's accounts object
+   * @param {array} addresses An array of hex addresses of new accounts to track
    *
    */
-  addAccount (address) {
+  addAccounts (addresses) {
     const accounts = this.store.getState().accounts
-    accounts[address] = {}
+    // add initial state for addresses
+    addresses.forEach(address => {
+      accounts[address] = {}
+    })
+    // save accounts state
     this.store.updateState({ accounts })
+    // fetch balances for the accounts if there is block number ready
     if (!this._currentBlockNumber) return
-    this._updateAccount(address)
+    addresses.forEach(address => this._updateAccount(address))
   }
 
   /**
-   * Removes an account from this AccountTracker's accounts object
+   * Removes accounts from being tracked
    *
-   * @param {string} address A hex address of a the account to remove
+   * @param {array} an array of hex addresses to stop tracking
    *
    */
-  removeAccount (address) {
+  removeAccount (addresses) {
     const accounts = this.store.getState().accounts
-    delete accounts[address]
+    // remove each state object
+    addresses.forEach(address => {
+      delete accounts[address]
+    })
+    // save accounts state
     this.store.updateState({ accounts })
   }
 
@@ -118,71 +123,56 @@ class AccountTracker extends EventEmitter {
    * via EthQuery
    *
    * @private
-   * @param {object} block Data about the block that contains the data to update to.
+   * @param {number} blockNumber the block number to update to.
    * @fires 'block' The updated state, if all account updates are successful
    *
    */
-  _updateForBlock (block) {
-    this._currentBlockNumber = block.number
-    const currentBlockGasLimit = block.gasLimit
+  async _updateForBlock (blockNumber) {
+    this._currentBlockNumber = blockNumber
 
+    // block gasLimit polling shouldn't be in account-tracker shouldn't be here...
+    const currentBlock = await this._query.getBlockByNumber(blockNumber, false)
+    if (!currentBlock) return
+    const currentBlockGasLimit = currentBlock.gasLimit
     this.store.updateState({ currentBlockGasLimit })
 
-    async.parallel([
-      this._updateAccounts.bind(this),
-    ], (err) => {
-      if (err) return console.error(err)
-      this.emit('block', this.store.getState())
-    })
+    try {
+      await this._updateAccounts()
+    } catch (err) {
+      log.error(err)
+    }
   }
 
   /**
    * Calls this._updateAccount for each account in this.store
    *
-   * @param {Function} cb A callback to pass to this._updateAccount, called after each account is successfully updated
+   * @returns {Promise} after all account balances updated
    *
    */
-  _updateAccounts (cb = noop) {
+  async _updateAccounts () {
     const accounts = this.store.getState().accounts
     const addresses = Object.keys(accounts)
-    async.each(addresses, this._updateAccount.bind(this), cb)
+    await Promise.all(addresses.map(this._updateAccount.bind(this)))
   }
 
   /**
-   * Updates the current balance of an account. Gets an updated balance via this._getAccount.
+   * Updates the current balance of an account.
    *
    * @private
    * @param {string} address A hex address of a the account to be updated
-   * @param {Function} cb A callback to call once the account at address is successfully update
+   * @returns {Promise} after the account balance is updated
    *
    */
-  _updateAccount (address, cb = noop) {
-    this._getAccount(address, (err, result) => {
-      if (err) return cb(err)
-      result.address = address
-      const accounts = this.store.getState().accounts
-      // only populate if the entry is still present
-      if (accounts[address]) {
-        accounts[address] = result
-        this.store.updateState({ accounts })
-      }
-      cb(null, result)
-    })
-  }
-
-  /**
-   * Gets the current balance of an account via EthQuery.
-   *
-   * @private
-   * @param {string} address A hex address of a the account to query
-   * @param {Function} cb A callback to call once the account at address is successfully update
-   *
-   */
-  _getAccount (address, cb = noop) {
-    const query = this._query
-    async.parallel({
-      balance: query.getBalance.bind(query, address),
-    }, cb)
+  async _updateAccount (address) {
+    // query balance
+    const balance = await this._query.getBalance(address)
+    const result = { address, balance }
+    // update accounts state
+    const { accounts } = this.store.getState()
+    // only populate if the entry is still present
+    if (!accounts[address]) return
+    accounts[address] = result
+    this.store.updateState({ accounts })
   }
 
 }
diff --git a/app/scripts/lib/events-proxy.js b/app/scripts/lib/events-proxy.js
deleted file mode 100644
index f83773ccc..000000000
--- a/app/scripts/lib/events-proxy.js
+++ /dev/null
@@ -1,42 +0,0 @@
-/**
- * Returns an EventEmitter that proxies events from the given event emitter
- * @param {any} eventEmitter
- * @param {object} listeners - The listeners to proxy to
- * @returns {any}
- */
-module.exports = function createEventEmitterProxy (eventEmitter, listeners) {
-  let target = eventEmitter
-  const eventHandlers = listeners || {}
-  const proxy = /** @type {any} */ (new Proxy({}, {
-    get: (_, name) => {
-      // intercept listeners
-      if (name === 'on') return addListener
-      if (name === 'setTarget') return setTarget
-      if (name === 'proxyEventHandlers') return eventHandlers
-      return (/** @type {any} */ (target))[name]
-    },
-    set: (_, name, value) => {
-      target[name] = value
-      return true
-    },
-  }))
-  function setTarget (/** @type {EventEmitter} */ eventEmitter) {
-    target = eventEmitter
-    // migrate listeners
-    Object.keys(eventHandlers).forEach((name) => {
-      /** @type {Array<Function>} */ (eventHandlers[name]).forEach((handler) => target.on(name, handler))
-    })
-  }
-  /**
-   * Attaches a function to be called whenever the specified event is emitted
-   * @param {string} name
-   * @param {Function} handler
-   */
-  function addListener (name, handler) {
-    if (!eventHandlers[name]) eventHandlers[name] = []
-    eventHandlers[name].push(handler)
-    target.on(name, handler)
-  }
-  if (listeners) proxy.setTarget(eventEmitter)
-  return proxy
-}
diff --git a/app/scripts/lib/message-manager.js b/app/scripts/lib/message-manager.js
index 901367f04..47925b94b 100644
--- a/app/scripts/lib/message-manager.js
+++ b/app/scripts/lib/message-manager.js
@@ -69,10 +69,39 @@ module.exports = class MessageManager extends EventEmitter {
    * new Message to this.messages, and to save the unapproved Messages from that list to this.memStore.
    *
    * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object possibly containing the origin
+   * @returns {promise} after signature has been
+   *
+   */
+  addUnapprovedMessageAsync (msgParams, req) {
+    return new Promise((resolve, reject) => {
+      const msgId = this.addUnapprovedMessage(msgParams, req)
+      // await finished
+      this.once(`${msgId}:finished`, (data) => {
+        switch (data.status) {
+          case 'signed':
+            return resolve(data.rawSig)
+          case 'rejected':
+            return reject(new Error('MetaMask Message Signature: User denied message signature.'))
+          default:
+            return reject(new Error(`MetaMask Message Signature: Unknown problem: ${JSON.stringify(msgParams)}`))
+        }
+      })
+    })
+  }
+
+  /**
+   * Creates a new Message with an 'unapproved' status using the passed msgParams. this.addMsg is called to add the
+   * new Message to this.messages, and to save the unapproved Messages from that list to this.memStore.
+   *
+   * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object where the origin may be specificied
    * @returns {number} The id of the newly created message.
    *
    */
-  addUnapprovedMessage (msgParams) {
+  addUnapprovedMessage (msgParams, req) {
+    // add origin from request
+    if (req) msgParams.origin = req.origin
     msgParams.data = normalizeMsgData(msgParams.data)
     // create txData obj with parameters and meta data
     var time = (new Date()).getTime()
diff --git a/app/scripts/lib/personal-message-manager.js b/app/scripts/lib/personal-message-manager.js
index e96ced1f2..fc2cccdf1 100644
--- a/app/scripts/lib/personal-message-manager.js
+++ b/app/scripts/lib/personal-message-manager.js
@@ -73,11 +73,43 @@ module.exports = class PersonalMessageManager extends EventEmitter {
    * this.memStore.
    *
    * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object possibly containing the origin
+   * @returns {promise} When the message has been signed or rejected
+   *
+   */
+  addUnapprovedMessageAsync (msgParams, req) {
+    return new Promise((resolve, reject) => {
+      if (!msgParams.from) {
+        reject(new Error('MetaMask Message Signature: from field is required.'))
+      }
+      const msgId = this.addUnapprovedMessage(msgParams, req)
+      this.once(`${msgId}:finished`, (data) => {
+        switch (data.status) {
+          case 'signed':
+            return resolve(data.rawSig)
+          case 'rejected':
+            return reject(new Error('MetaMask Message Signature: User denied message signature.'))
+          default:
+            return reject(new Error(`MetaMask Message Signature: Unknown problem: ${JSON.stringify(msgParams)}`))
+        }
+      })
+    })
+  }
+
+  /**
+   * Creates a new PersonalMessage with an 'unapproved' status using the passed msgParams. this.addMsg is called to add
+   * the new PersonalMessage to this.messages, and to save the unapproved PersonalMessages from that list to
+   * this.memStore.
+   *
+   * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object possibly containing the origin
    * @returns {number} The id of the newly created PersonalMessage.
    *
    */
-  addUnapprovedMessage (msgParams) {
+  addUnapprovedMessage (msgParams, req) {
     log.debug(`PersonalMessageManager addUnapprovedMessage: ${JSON.stringify(msgParams)}`)
+    // add origin from request
+    if (req) msgParams.origin = req.origin
     msgParams.data = this.normalizeMsgData(msgParams.data)
     // create txData obj with parameters and meta data
     var time = (new Date()).getTime()
@@ -257,4 +289,3 @@ module.exports = class PersonalMessageManager extends EventEmitter {
   }
 
 }
-
diff --git a/app/scripts/lib/setupRaven.js b/app/scripts/lib/setupRaven.js
index 3651524f1..e6e511640 100644
--- a/app/scripts/lib/setupRaven.js
+++ b/app/scripts/lib/setupRaven.js
@@ -70,11 +70,11 @@ function simplifyErrorMessages (report) {
 
 function rewriteErrorMessages (report, rewriteFn) {
   // rewrite top level message
-  if (report.message) report.message = rewriteFn(report.message)
+  if (typeof report.message === 'string') report.message = rewriteFn(report.message)
   // rewrite each exception message
   if (report.exception && report.exception.values) {
     report.exception.values.forEach(item => {
-      item.value = rewriteFn(item.value)
+      if (typeof item.value === 'string') item.value = rewriteFn(item.value)
     })
   }
 }
diff --git a/app/scripts/lib/typed-message-manager.js b/app/scripts/lib/typed-message-manager.js
index c58921610..e5e1c94b3 100644
--- a/app/scripts/lib/typed-message-manager.js
+++ b/app/scripts/lib/typed-message-manager.js
@@ -72,11 +72,40 @@ module.exports = class TypedMessageManager extends EventEmitter {
    * this.memStore. Before any of this is done, msgParams are validated
    *
    * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object possibly containing the origin
+   * @returns {promise} When the message has been signed or rejected
+   *
+   */
+  addUnapprovedMessageAsync (msgParams, req) {
+    return new Promise((resolve, reject) => {
+      const msgId = this.addUnapprovedMessage(msgParams, req)
+      this.once(`${msgId}:finished`, (data) => {
+        switch (data.status) {
+          case 'signed':
+            return resolve(data.rawSig)
+          case 'rejected':
+            return reject(new Error('MetaMask Message Signature: User denied message signature.'))
+          default:
+            return reject(new Error(`MetaMask Message Signature: Unknown problem: ${JSON.stringify(msgParams)}`))
+        }
+      })
+    })
+  }
+
+  /**
+   * Creates a new TypedMessage with an 'unapproved' status using the passed msgParams. this.addMsg is called to add
+   * the new TypedMessage to this.messages, and to save the unapproved TypedMessages from that list to
+   * this.memStore. Before any of this is done, msgParams are validated
+   *
+   * @param {Object} msgParams The params for the eth_sign call to be made after the message is approved.
+   * @param {Object} req (optional) The original request object possibly containing the origin
    * @returns {number} The id of the newly created TypedMessage.
    *
    */
-  addUnapprovedMessage (msgParams) {
+  addUnapprovedMessage (msgParams, req) {
     this.validateParams(msgParams)
+    // add origin from request
+    if (req) msgParams.origin = req.origin
 
     log.debug(`TypedMessageManager addUnapprovedMessage: ${JSON.stringify(msgParams)}`)
     // create txData obj with parameters and meta data
diff --git a/app/scripts/lib/util.js b/app/scripts/lib/util.js
index d7423f2ad..ea13b26be 100644
--- a/app/scripts/lib/util.js
+++ b/app/scripts/lib/util.js
@@ -127,7 +127,21 @@ function BnMultiplyByFraction (targetBN, numerator, denominator) {
   return targetBN.mul(numBN).div(denomBN)
 }
 
+function applyListeners (listeners, emitter) {
+  Object.keys(listeners).forEach((key) => {
+    emitter.on(key, listeners[key])
+  })
+}
+
+function removeListeners (listeners, emitter) {
+  Object.keys(listeners).forEach((key) => {
+    emitter.removeListener(key, listeners[key])
+  })
+}
+
 module.exports = {
+  removeListeners,
+  applyListeners,
   getPlatform,
   getStack,
   getEnvironmentType,