From efd34343c8077acaf57a44847e344ca404943af6 Mon Sep 17 00:00:00 2001
From: Michele Esposito <34438276+mikesposito@users.noreply.github.com>
Date: Sat, 19 Aug 2023 00:07:38 +0200
Subject: [PATCH] Use KeyringController messenger events (#20341)

* refactor: use keyring-controller events and unlock methods

* test: add tests for setLocked

* rollback: call trezor dispose directly

* Update app/scripts/metamask-controller.js

Co-authored-by: Mark Stacey <markjstacey@gmail.com>

---------

Co-authored-by: Mark Stacey <markjstacey@gmail.com>
---
 .../metamask-controller.actions.test.js       |  8 ++++
 app/scripts/metamask-controller.js            | 39 ++++++++++---------
 app/scripts/metamask-controller.test.js       | 16 +++++++-
 3 files changed, 43 insertions(+), 20 deletions(-)

diff --git a/app/scripts/metamask-controller.actions.test.js b/app/scripts/metamask-controller.actions.test.js
index 8c015c9c5..a0dba444b 100644
--- a/app/scripts/metamask-controller.actions.test.js
+++ b/app/scripts/metamask-controller.actions.test.js
@@ -220,6 +220,14 @@ describe('MetaMaskController', function () {
     });
   });
 
+  describe('#setLocked', function () {
+    it('should lock the wallet', async function () {
+      const { isUnlocked, keyrings } = await metamaskController.setLocked();
+      assert(!isUnlocked);
+      assert.deepEqual(keyrings, []);
+    });
+  });
+
   describe('#addToken', function () {
     const address = '0x514910771af9ca656af840dff83e8264ecf986ca';
     const symbol = 'LINK';
diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
index 6bc0d9531..daf2a74e1 100644
--- a/app/scripts/metamask-controller.js
+++ b/app/scripts/metamask-controller.js
@@ -877,15 +877,21 @@ export default class MetamaskController extends EventEmitter {
       ),
     });
 
-    this.keyringController =
-      this.coreKeyringController.getEthKeyringController();
-
-    this.keyringController.memStore.subscribe((state) =>
-      this._onKeyringControllerUpdate(state),
+    this.controllerMessenger.subscribe('KeyringController:unlock', () =>
+      this._onUnlock(),
+    );
+    this.controllerMessenger.subscribe('KeyringController:lock', () =>
+      this._onLock(),
+    );
+    this.controllerMessenger.subscribe(
+      'KeyringController:stateChange',
+      (state) => {
+        this._onKeyringControllerUpdate(state);
+      },
     );
 
-    this.keyringController.on('unlock', () => this._onUnlock());
-    this.keyringController.on('lock', () => this._onLock());
+    this.keyringController =
+      this.coreKeyringController.getEthKeyringController();
 
     const getIdentities = () =>
       this.preferencesController.store.getState().identities;
@@ -3060,10 +3066,9 @@ export default class MetamaskController extends EventEmitter {
    * is up to date with known accounts once the vault is decrypted.
    *
    * @param {string} password - The user's password
-   * @returns {Promise<object>} The keyringController update.
    */
   async submitPassword(password) {
-    await this.keyringController.submitPassword(password);
+    await this.coreKeyringController.submitPassword(password);
 
     ///: BEGIN:ONLY_INCLUDE_IN(build-mmi)
     this.mmiController.onSubmitPassword();
@@ -3083,8 +3088,6 @@ export default class MetamaskController extends EventEmitter {
       this.preferencesController.getLedgerTransportPreference();
 
     this.setLedgerTransportPreference(transportPreference);
-
-    return this.keyringController.fullUpdate();
   }
 
   async _loginUser() {
@@ -3123,7 +3126,7 @@ export default class MetamaskController extends EventEmitter {
       const { loginToken, loginSalt } =
         await this.extension.storage.session.get(['loginToken', 'loginSalt']);
       if (loginToken && loginSalt) {
-        const { vault } = this.keyringController.store.getState();
+        const { vault } = this.coreKeyringController.state;
 
         const jsonVault = JSON.parse(vault);
 
@@ -3135,7 +3138,10 @@ export default class MetamaskController extends EventEmitter {
           return;
         }
 
-        await this.keyringController.submitEncryptionKey(loginToken, loginSalt);
+        await this.coreKeyringController.submitEncryptionKey(
+          loginToken,
+          loginSalt,
+        );
       }
     } catch (e) {
       // If somehow this login token doesn't work properly,
@@ -4709,16 +4715,11 @@ export default class MetamaskController extends EventEmitter {
       trezorKeyring.dispose();
     }
 
-    const [ledgerKeyring] = this.coreKeyringController.getKeyringsByType(
-      KeyringType.ledger,
-    );
-    ledgerKeyring?.destroy?.();
-
     if (isManifestV3) {
       this.clearLoginArtifacts();
     }
 
-    return this.keyringController.setLocked();
+    return this.coreKeyringController.setLocked();
   }
 
   removePermissionsFor = (subjects) => {
diff --git a/app/scripts/metamask-controller.test.js b/app/scripts/metamask-controller.test.js
index 20f496aba..ec7f40643 100644
--- a/app/scripts/metamask-controller.test.js
+++ b/app/scripts/metamask-controller.test.js
@@ -370,7 +370,7 @@ describe('MetaMaskController', function () {
           metamaskController.preferencesController.store.getState().identities,
         );
         const addresses =
-          await metamaskController.keyringController.getAccounts();
+          await metamaskController.coreKeyringController.getAccounts();
 
         identities.forEach((identity) => {
           assert.ok(
@@ -388,6 +388,20 @@ describe('MetaMaskController', function () {
       });
     });
 
+    describe('setLocked', function () {
+      it('should lock KeyringController', async function () {
+        sandbox.spy(metamaskController.coreKeyringController, 'setLocked');
+
+        await metamaskController.setLocked();
+
+        assert(metamaskController.coreKeyringController.setLocked.called);
+        assert.equal(
+          metamaskController.coreKeyringController.state.isUnlocked,
+          false,
+        );
+      });
+    });
+
     describe('#createNewVaultAndKeychain', function () {
       it('can only create new vault on keyringController once', async function () {
         const selectStub = sandbox.stub(