From 003d445a98164dac0c0529dfd69f5e3987d7d05c Mon Sep 17 00:00:00 2001
From: Whymarrh Whitby <whymarrh.whitby@gmail.com>
Date: Thu, 16 Aug 2018 12:59:39 -0230
Subject: [PATCH] Update unlock logic to not overwrite existing selected
 address

---
 app/scripts/metamask-controller.js | 42 +++++++++++++++++++++---------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 29838ad2d..4ee88186a 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -137,19 +137,7 @@ module.exports = class MetamaskController extends EventEmitter {
       encryptor: opts.encryptor || undefined,
     })
 
-    // If only one account exists, make sure it is selected.
-    this.keyringController.memStore.subscribe((state) => {
-      const addresses = state.keyrings.reduce((res, keyring) => {
-        return res.concat(keyring.accounts)
-      }, [])
-      if (addresses.length === 1) {
-        const address = addresses[0]
-        this.preferencesController.setSelectedAddress(address)
-      }
-      // ensure preferences + identities controller know about all addresses
-      this.preferencesController.addAddresses(addresses)
-      this.accountTracker.syncWithAddresses(addresses)
-    })
+    this.keyringController.memStore.subscribe((s) => this._onKeyringControllerUpdate(s))
 
     // detect tokens controller
     this.detectTokensController = new DetectTokensController({
@@ -1278,6 +1266,34 @@ module.exports = class MetamaskController extends EventEmitter {
     )
   }
 
+  /**
+   * Handle a KeyringController update
+   * @param {object} state the KC state
+   * @return {Promise<void>}
+   * @private
+   */
+  async _onKeyringControllerUpdate (state) {
+    const {isUnlocked, keyrings} = state
+    const addresses = keyrings.reduce((acc, {accounts}) => acc.concat(accounts), [])
+
+    if (!addresses.length) {
+      return
+    }
+
+    // Ensure preferences + identities controller know about all addresses
+    this.preferencesController.addAddresses(addresses)
+    this.accountTracker.syncWithAddresses(addresses)
+
+    const wasLocked = !isUnlocked
+    if (wasLocked) {
+      const oldSelectedAddress = this.preferencesController.getSelectedAddress()
+      if (!addresses.includes(oldSelectedAddress)) {
+        const address = addresses[0]
+        await this.preferencesController.setSelectedAddress(address)
+      }
+    }
+  }
+
   /**
    * A method for emitting the full MetaMask state to all registered listeners.
    * @private