1 files changed, 56 insertions, 54 deletions

diff --git a/app/scripts/controllers/network.js b/app/scripts/controllers/network.js
index 0a3e5e26b..0f9db4d53 100644
--- a/app/scripts/controllers/network.js
+++ b/app/scripts/controllers/network.js
@@ -1,71 +1,40 @@
+const assert = require('assert')
 const EventEmitter = require('events')
-const MetaMaskProvider = require('web3-provider-engine/zero.js')
+const createMetamaskProvider = require('web3-provider-engine/zero.js')
 const ObservableStore = require('obs-store')
 const ComposedStore = require('obs-store/lib/composed')
 const extend = require('xtend')
 const EthQuery = require('eth-query')
+const createEventEmitterProxy = require('../lib/events-proxy.js')
 const RPC_ADDRESS_LIST = require('../config.js').network
 const DEFAULT_RPC = RPC_ADDRESS_LIST['rinkeby']
 
 module.exports = class NetworkController extends EventEmitter {
+
   constructor (config) {
     super()
-    this.networkStore = new ObservableStore('loading')
     config.provider.rpcTarget = this.getRpcAddressForType(config.provider.type, config.provider)
+    this.networkStore = new ObservableStore('loading')
     this.providerStore = new ObservableStore(config.provider)
     this.store = new ComposedStore({ provider: this.providerStore, network: this.networkStore })
-    this._providerListeners = {}
+    this._proxy = createEventEmitterProxy()
 
     this.on('networkDidChange', this.lookupNetwork)
-    this.providerStore.subscribe((state) => this.switchNetwork({rpcUrl: state.rpcTarget}))
-  }
-
-  get provider () {
-    return this._proxy
-  }
-
-  set provider (provider) {
-    this._provider = provider
   }
 
-  initializeProvider (opts, providerContructor = MetaMaskProvider) {
-    this.providerInit = opts
-    this._provider = providerContructor(opts)
-    this._proxy = new Proxy(this._provider, {
-      get: (obj, name) => {
-        if (name === 'on') return this._on.bind(this)
-        return this._provider[name]
-      },
-      set: (obj, name, value) => {
-        this._provider[name] = value
-        return value
-      },
-    })
-    this.provider.on('block', this._logBlock.bind(this))
-    this.provider.on('error', this.verifyNetwork.bind(this))
-    this.ethQuery = new EthQuery(this.provider)
+  initializeProvider (_providerParams) {
+    this._baseProviderParams = _providerParams
+    const rpcUrl = this.getCurrentRpcAddress()
+    this._configureStandardProvider({ rpcUrl })
+    this._proxy.on('block', this._logBlock.bind(this))
+    this._proxy.on('error', this.verifyNetwork.bind(this))
+    this.ethQuery = new EthQuery(this._proxy)
     this.lookupNetwork()
-    return this.provider
-  }
-
-  switchNetwork (providerInit) {
-    this.setNetworkState('loading')
-    const newInit = extend(this.providerInit, providerInit)
-    this.providerInit = newInit
-
-    this._provider.removeAllListeners()
-    this._provider.stop()
-    this.provider = MetaMaskProvider(newInit)
-    // apply the listners created by other controllers
-    Object.keys(this._providerListeners).forEach((key) => {
-      this._providerListeners[key].forEach((handler) => this._provider.addListener(key, handler))
-    })
-    this.emit('networkDidChange')
+    return this._proxy
   }
 
-
   verifyNetwork () {
-  // Check network when restoring connectivity:
+    // Check network when restoring connectivity:
     if (this.isNetworkLoading()) this.lookupNetwork()
   }
 
@@ -94,6 +63,7 @@ module.exports = class NetworkController extends EventEmitter {
       type: 'rpc',
       rpcTarget: rpcUrl,
     })
+    this._switchNetwork({ rpcUrl })
   }
 
   getCurrentRpcAddress () {
@@ -102,10 +72,14 @@ module.exports = class NetworkController extends EventEmitter {
     return this.getRpcAddressForType(provider.type)
   }
 
-  setProviderType (type) {
+  async setProviderType (type) {
+    assert(type !== 'rpc', `NetworkController.setProviderType - cannot connect by type "rpc"`)
+    // skip if type already matches
     if (type === this.getProviderConfig().type) return
     const rpcTarget = this.getRpcAddressForType(type)
-    this.providerStore.updateState({type, rpcTarget})
+    assert(rpcTarget, `NetworkController - unknown rpc address for type "${type}"`)
+    this.providerStore.updateState({ type, rpcTarget })
+    this._switchNetwork({ rpcUrl: rpcTarget })
   }
 
   getProviderConfig () {
@@ -117,14 +91,42 @@ module.exports = class NetworkController extends EventEmitter {
     return provider && provider.rpcTarget ? provider.rpcTarget : DEFAULT_RPC
   }
 
+  //
+  // Private
+  //
+
+  _switchNetwork (providerParams) {
+    this.setNetworkState('loading')
+    this._configureStandardProvider(providerParams)
+    this.emit('networkDidChange')
+  }
+
+  _configureStandardProvider(_providerParams) {
+    const providerParams = extend(this._baseProviderParams, _providerParams)
+    const provider = createMetamaskProvider(providerParams)
+    this._setProvider(provider)
+  }
+
+  _setProvider (provider) {
+    // collect old block tracker events
+    const oldProvider = this._provider
+    let blockTrackerHandlers
+    if (oldProvider) {
+      // capture old block handlers
+      blockTrackerHandlers = oldProvider._blockTracker.proxyEventHandlers
+      // tear down
+      oldProvider.removeAllListeners()
+      oldProvider.stop()
+    }
+    // override block tracler
+    provider._blockTracker = createEventEmitterProxy(provider._blockTracker, blockTrackerHandlers)
+    // set as new provider
+    this._provider = provider
+    this._proxy.setTarget(provider)
+  }
+
   _logBlock (block) {
     log.info(`BLOCK CHANGED: #${block.number.toString('hex')} 0x${block.hash.toString('hex')}`)
     this.verifyNetwork()
   }
-
-  _on (event, handler) {
-    if (!this._providerListeners[event]) this._providerListeners[event] = []
-    this._providerListeners[event].push(handler)
-    this._provider.on(event, handler)
-  }
 }