fix: use getState to get initial lock state

MetaMask · Sep 18, 2023 · db3bd03 · db3bd03
1 parent 6d86c42
commit db3bd03
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 30 deletions.
diff --git a/app/scripts/controllers/detect-tokens.js b/app/scripts/controllers/detect-tokens.js
@@ -28,7 +28,6 @@ export default class DetectTokensController {
    * @param config.interval
    * @param config.preferences
    * @param config.network
-   * @param config.keyringMemStore
    * @param config.tokenList
    * @param config.tokensController
    * @param config.assetsContractController
@@ -89,13 +88,8 @@ export default class DetectTokensController {
         this.restartTokenDetection({ chainId: this.chainId });
       }
     });
-    messenger.subscribe('KeyringController:unlock', () => {
-      this.isUnlocked = true;
-      this.restartTokenDetection();
-    });
-    messenger.subscribe('KeyringController:lock', () => {
-      this.isUnlocked = false;
-    });
+
+    this.#registerKeyringHandlers();
   }
 
   /**
@@ -260,4 +254,22 @@ export default class DetectTokensController {
     return this.isOpen && this.isUnlocked;
   }
   /* eslint-enable accessor-pairs */
+
+  /**
+   * Constructor helper to register listeners on the keyring
+   * locked state changes
+   */
+  #registerKeyringHandlers() {
+    const { isUnlocked } = this.messenger.call('KeyringController:getState');
+    this.isUnlocked = isUnlocked;
+
+    this.messenger.subscribe('KeyringController:unlock', () => {
+      this.isUnlocked = true;
+      this.restartTokenDetection();
+    });
+
+    this.messenger.subscribe('KeyringController:lock', () => {
+      this.isUnlocked = false;
+    });
+  }
 }
diff --git a/app/scripts/controllers/detect-tokens.test.js b/app/scripts/controllers/detect-tokens.test.js
@@ -15,28 +15,30 @@ import { toChecksumHexAddress } from '../../../shared/modules/hexstring-utils';
 import DetectTokensController from './detect-tokens';
 import PreferencesController from './preferences';
 
-function buildMessenger() {
-  return new ControllerMessenger().getRestricted({
-    name: 'DetectTokensController',
-    allowedEvents: [
-      'NetworkController:stateChange',
-      'KeyringController:lock',
-      'KeyringController:unlock',
-    ],
-  });
-}
-
 describe('DetectTokensController', function () {
   let sandbox,
     assetsContractController,
     network,
     preferences,
     provider,
     tokensController,
-    tokenListController;
+    tokenListController,
+    messenger;
 
   const noop = () => undefined;
 
+  const getRestrictedMessenger = () => {
+    return messenger.getRestricted({
+      name: 'DetectTokensController',
+      allowedActions: ['KeyringController:getState'],
+      allowedEvents: [
+        'NetworkController:stateChange',
+        'KeyringController:lock',
+        'KeyringController:unlock',
+      ],
+    });
+  };
+
   const networkControllerProviderConfig = {
     getAccounts: noop,
   };
@@ -200,6 +202,11 @@ describe('DetectTokensController', function () {
       .reply(200, { error: 'ChainId 3 is not supported' })
       .persist();
 
+    messenger = new ControllerMessenger();
+    messenger.registerActionHandler('KeyringController:getState', () => ({
+      isUnlocked: true,
+    }));
+
     const networkControllerMessenger = new ControllerMessenger();
     network = new NetworkController({
       messenger: networkControllerMessenger,
@@ -264,7 +271,11 @@ describe('DetectTokensController', function () {
 
   it('should poll on correct interval', async function () {
     const stub = sinon.stub(global, 'setInterval');
-    new DetectTokensController({ messenger: buildMessenger(), interval: 1337 }); // eslint-disable-line no-new
+    // eslint-disable-next-line no-new
+    new DetectTokensController({
+      messenger: getRestrictedMessenger(),
+      interval: 1337,
+    });
     assert.strictEqual(stub.getCall(0).args[1], 1337);
     stub.restore();
   });
@@ -273,7 +284,7 @@ describe('DetectTokensController', function () {
     const clock = sandbox.useFakeTimers();
     await network.setProviderType(NETWORK_TYPES.MAINNET);
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -309,7 +320,7 @@ describe('DetectTokensController', function () {
     });
     await tokenListController.start();
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -332,7 +343,7 @@ describe('DetectTokensController', function () {
     sandbox.useFakeTimers();
     await network.setProviderType(NETWORK_TYPES.MAINNET);
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -383,7 +394,7 @@ describe('DetectTokensController', function () {
     sandbox.useFakeTimers();
     await network.setProviderType(NETWORK_TYPES.MAINNET);
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -441,7 +452,7 @@ describe('DetectTokensController', function () {
   it('should trigger detect new tokens when change address', async function () {
     sandbox.useFakeTimers();
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -459,9 +470,8 @@ describe('DetectTokensController', function () {
 
   it('should trigger detect new tokens when submit password', async function () {
     sandbox.useFakeTimers();
-    const messenger = buildMessenger();
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -481,7 +491,7 @@ describe('DetectTokensController', function () {
     const clock = sandbox.useFakeTimers();
     await network.setProviderType(NETWORK_TYPES.MAINNET);
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokenList: tokenListController,
@@ -502,7 +512,7 @@ describe('DetectTokensController', function () {
     const clock = sandbox.useFakeTimers();
     await network.setProviderType(NETWORK_TYPES.MAINNET);
     const controller = new DetectTokensController({
-      messenger: buildMessenger(),
+      messenger: getRestrictedMessenger(),
       preferences,
       network,
       tokensController,

diff --git a/app/scripts/metamask-controller.js b/app/scripts/metamask-controller.js
@@ -1142,6 +1142,7 @@ export default class MetamaskController extends EventEmitter {
     const detectTokensControllerMessenger =
       this.controllerMessenger.getRestricted({
         name: 'DetectTokensController',
+        allowedActions: ['KeyringController:getState'],
         allowedEvents: [
           'NetworkController:stateChange',
           'KeyringController:lock',