feat: apply mutex on controller lock (#6292)

## Explanation

This PR adds `lock (mutex)` mechanism to the
seedless-onboarding-controller `setLocked` method to achieve the atomic
state updates and prevent unexpected issues (state out of sync) while
locking the controller and another operation (such as `changePassword`,
`addNewSecretData` etc) is in progress on the other hand.

As a result, we can't simply lock the seedless-onboarding-controller,
syncing with keyring's. Syncing with keyring's lock will cause more
issues in the controller communications as seedless-onboarding
`setLocked`, now, has some side effects and it won't get locked
immediately in some scenarios (when another seedless operation in
progress).
Hence, `setLocked` becomes independent of the Keyring's and this PR also
removes the `KeyringController:lock` and `KeyringController:unlock`
events from the seedless-onboarding messenger.

**BREAKING:**
- Removed `Keyring:lock` and `Keyring:unlock` events from the controller
allowed events.
- 

## References

<!--
Are there any issues that this pull request is tied to?
Are there other links that reviewers should consult to understand these
changes better?
Are there client or consumer pull requests to adopt any breaking
changes?

For example:

* Fixes #12345
* Related to #67890
-->

## Checklist

- [x] I've updated the test suite for new or updated code as appropriate
- [x] I've updated documentation (JSDoc, Markdown, etc.) for new or
updated code as appropriate
- [x] I've communicated my changes to consumers by [updating changelogs
for packages I've
changed](https://github.com/MetaMask/core/tree/main/docs/contributing.md#updating-changelogs),
highlighting breaking changes as necessary
- [x] I've prepared draft pull requests for clients and consumer
packages to resolve any breaking changes

Author: lwin-kyaw

packages/seedless-onboarding-controller/CHANGELOG.md

@@ -15,10 +15,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- **BREAKING:** Updated ControllerMessenger `AllowedEvents`. ([#6292](https://github.com/MetaMask/core/pull/6292))
+- Update `setLocked()` method with `mutex` and it becomes `async` method. ([#6292](https://github.com/MetaMask/core/pull/6292))
 - Bump `@metamask/base-controller` from `^8.1.0` to `^8.3.0` ([#6355](https://github.com/MetaMask/core/pull/6355), [#6465](https://github.com/MetaMask/core/pull/6465))
 
 ### Removed
 
+- **BREAKING:** Removed `Keyring:lock` and `Keyring:unlock` events from the controller allowed events. ([#6292](https://github.com/MetaMask/core/pull/6292))
 - Removed `revokeRefreshToken` method ([#6275](https://github.com/MetaMask/core/pull/6275))
 
 ## [3.0.0]

packages/seedless-onboarding-controller/src/SeedlessOnboardingController.test.ts

@@ -28,6 +28,7 @@ import {
 import { gcm } from '@noble/ciphers/aes';
 import { utf8ToBytes } from '@noble/ciphers/utils';
 import { managedNonce } from '@noble/ciphers/webcrypto';
+import { Mutex } from 'async-mutex';
 import type { webcrypto } from 'node:crypto';
 
 import {
@@ -2665,29 +2666,6 @@ describe('SeedlessOnboardingController', () => {
       );
     });
 
-    it('should throw an error if the old password is incorrect', async () => {
-      await withController(
-        {
-          state: getMockInitialControllerState({
-            vault: MOCK_VAULT,
-            withMockAuthenticatedUser: true,
-          }),
-        },
-        async ({ controller, encryptor, baseMessenger }) => {
-          // unlock the controller
-          baseMessenger.publish('KeyringController:unlock');
-          await new Promise((resolve) => setTimeout(resolve, 100));
-
-          jest
-            .spyOn(encryptor, 'decrypt')
-            .mockRejectedValueOnce(new Error('Incorrect password'));
-          await expect(
-            controller.changePassword(NEW_MOCK_PASSWORD, 'INCORRECT_PASSWORD'),
-          ).rejects.toThrow('Incorrect password');
-        },
-      );
-    });
-
     it('should throw an error if failed to change password', async () => {
       await withController(
         {
@@ -2777,40 +2755,6 @@ describe('SeedlessOnboardingController', () => {
       );
     });
 
-    it('should throw error when authentication info is missing for assertPasswordInSync', async () => {
-      await withController(
-        {
-          state: {
-            // Create a state with vault but missing auth info
-            vault: JSON.stringify({ mockVault: 'data' }),
-            authPubKey: MOCK_AUTH_PUB_KEY,
-            socialBackupsMetadata: [],
-            // Intentionally missing nodeAuthTokens, authConnectionId, userId
-          },
-        },
-        async ({ controller, baseMessenger, encryptor }) => {
-          // Mock the encryptor to pass verifyVaultPassword
-          jest
-            .spyOn(encryptor, 'decrypt')
-            .mockResolvedValueOnce('mock decrypted data');
-
-          // unlock the controller
-          baseMessenger.publish('KeyringController:unlock');
-          await new Promise((resolve) => setTimeout(resolve, 100));
-
-          await expect(
-            controller.changePassword(NEW_MOCK_PASSWORD, MOCK_PASSWORD),
-          ).rejects.toThrow(
-            SeedlessOnboardingControllerErrorMessage.MissingAuthUserInfo,
-          );
-
-          expect(controller.state.isSeedlessOnboardingUserAuthenticated).toBe(
-            false,
-          );
-        },
-      );
-    });
-
     it('should call recoverEncKey when keyIndex is missing', async () => {
       await withController(
         {
@@ -2991,42 +2935,17 @@ describe('SeedlessOnboardingController', () => {
     const MOCK_PASSWORD = 'mock-password';
 
     it('should lock the controller', async () => {
-      await withController(
-        {
-          state: getMockInitialControllerState({
-            withMockAuthenticatedUser: true,
-          }),
-        },
-        async ({ controller, toprfClient }) => {
-          await mockCreateToprfKeyAndBackupSeedPhrase(
-            toprfClient,
-            controller,
-            MOCK_PASSWORD,
-            MOCK_SEED_PHRASE,
-            MOCK_KEYRING_ID,
-          );
-
-          controller.setLocked();
+      const mutexAcquireSpy = jest
+        .spyOn(Mutex.prototype, 'acquire')
+        .mockResolvedValueOnce(jest.fn());
 
-          await expect(
-            controller.addNewSecretData(MOCK_SEED_PHRASE, SecretType.Mnemonic, {
-              keyringId: MOCK_KEYRING_ID,
-            }),
-          ).rejects.toThrow(
-            SeedlessOnboardingControllerErrorMessage.ControllerLocked,
-          );
-        },
-      );
-    });
-
-    it('should lock the controller when the keyring is locked', async () => {
       await withController(
         {
           state: getMockInitialControllerState({
             withMockAuthenticatedUser: true,
           }),
         },
-        async ({ controller, baseMessenger, toprfClient }) => {
+        async ({ controller, toprfClient }) => {
           await mockCreateToprfKeyAndBackupSeedPhrase(
             toprfClient,
             controller,
@@ -3035,53 +2954,18 @@ describe('SeedlessOnboardingController', () => {
             MOCK_KEYRING_ID,
           );
 
-          baseMessenger.publish('KeyringController:lock');
+          await controller.setLocked();
 
-          await expect(
-            controller.addNewSecretData(MOCK_SEED_PHRASE, SecretType.Mnemonic, {
-              keyringId: MOCK_KEYRING_ID,
-            }),
-          ).rejects.toThrow(
-            SeedlessOnboardingControllerErrorMessage.ControllerLocked,
-          );
-        },
-      );
-    });
+          // verify that the mutex acquire was called
+          expect(mutexAcquireSpy).toHaveBeenCalled();
 
-    it('should unlock the controller when the keyring is unlocked', async () => {
-      await withController(
-        {
-          state: getMockInitialControllerState({
-            withMockAuthenticatedUser: true,
-          }),
-        },
-        async ({ controller, baseMessenger }) => {
           await expect(
             controller.addNewSecretData(MOCK_SEED_PHRASE, SecretType.Mnemonic, {
               keyringId: MOCK_KEYRING_ID,
             }),
           ).rejects.toThrow(
             SeedlessOnboardingControllerErrorMessage.ControllerLocked,
           );
-
-          baseMessenger.publish('KeyringController:unlock');
-
-          await new Promise((resolve) => setTimeout(resolve, 100));
-
-          controller.updateBackupMetadataState({
-            keyringId: MOCK_KEYRING_ID,
-            data: MOCK_SEED_PHRASE,
-            type: SecretType.Mnemonic,
-          });
-
-          const MOCK_SEED_PHRASE_HASH = keccak256AndHexify(MOCK_SEED_PHRASE);
-          expect(controller.state.socialBackupsMetadata).toStrictEqual([
-            {
-              type: SecretType.Mnemonic,
-              keyringId: MOCK_KEYRING_ID,
-              hash: MOCK_SEED_PHRASE_HASH,
-            },
-          ]);
         },
       );
     });
@@ -3306,7 +3190,7 @@ describe('SeedlessOnboardingController', () => {
             pwEncKey: recoveredPwEncKey,
           });
 
-          controller.setLocked();
+          await controller.setLocked();
 
           await controller.submitGlobalPassword({
             globalPassword: GLOBAL_PASSWORD,
@@ -3474,7 +3358,7 @@ describe('SeedlessOnboardingController', () => {
             pwEncKey: recoveredPwEncKey,
           });
 
-          controller.setLocked();
+          await controller.setLocked();
 
           await controller.submitGlobalPassword({
             globalPassword: GLOBAL_PASSWORD,
@@ -3791,7 +3675,7 @@ describe('SeedlessOnboardingController', () => {
           // We still need verifyPassword to work conceptually, even if unlock is bypassed
           // verifyPasswordSpy.mockResolvedValueOnce(); // Don't mock, let the real one run inside syncLatestGlobalPassword
 
-          controller.setLocked();
+          await controller.setLocked();
 
           // Mock recoverEncKey for the global password
           const encKey = mockToprfEncryptor.deriveEncKey(GLOBAL_PASSWORD);
@@ -4023,7 +3907,7 @@ describe('SeedlessOnboardingController', () => {
         },
         async ({ controller, toprfClient }) => {
           // Ensure the controller is locked
-          controller.setLocked();
+          await controller.setLocked();
 
           // Mock fetchAuthPubKey to return a valid response
           jest.spyOn(toprfClient, 'fetchAuthPubKey').mockResolvedValue({

packages/seedless-onboarding-controller/src/SeedlessOnboardingController.ts

@@ -260,15 +260,6 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     this.#refreshJWTToken = refreshJWTToken;
     this.#revokeRefreshToken = revokeRefreshToken;
     this.#renewRefreshToken = renewRefreshToken;
-
-    // setup subscriptions to the keyring lock event
-    // when the keyring is locked (wallet is locked), the controller will be cleared of its credentials
-    this.messagingSystem.subscribe('KeyringController:lock', () => {
-      this.setLocked();
-    });
-    this.messagingSystem.subscribe('KeyringController:unlock', () => {
-      this.#setUnlocked();
-    });
   }
 
   async fetchMetadataAccessCreds(): Promise<{
@@ -697,17 +688,21 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
    * Set the controller to locked state, and deallocate the secrets (vault encryption key and salt).
    *
    * When the controller is locked, the user will not be able to perform any operations on the controller/vault.
+   *
+   * @returns A promise that resolves to the success of the operation.
    */
-  setLocked() {
-    this.update((state) => {
-      delete state.vaultEncryptionKey;
-      delete state.vaultEncryptionSalt;
-      delete state.revokeToken;
-      delete state.accessToken;
-    });
+  async setLocked() {
+    return await this.#withControllerLock(async () => {
+      this.update((state) => {
+        delete state.vaultEncryptionKey;
+        delete state.vaultEncryptionSalt;
+        delete state.revokeToken;
+        delete state.accessToken;
+      });
 
-    this.#cachedDecryptedVaultData = undefined;
-    this.#isUnlocked = false;
+      this.#cachedDecryptedVaultData = undefined;
+      this.#isUnlocked = false;
+    });
   }
 
   /**
@@ -1693,17 +1688,13 @@ export class SeedlessOnboardingController<EncryptionKey> extends BaseController<
     authPubKey: SEC1EncodedPublicKey;
     latestKeyIndex: number;
   }> {
+    this.#assertIsAuthenticatedUser(this.state);
     const {
       nodeAuthTokens,
       authConnectionId,
       groupedAuthConnectionId,
       userId,
     } = this.state;
-    if (!nodeAuthTokens || !authConnectionId || !userId) {
-      throw new Error(
-        SeedlessOnboardingControllerErrorMessage.MissingAuthUserInfo,
-      );
-    }
 
     const { authPubKey, keyIndex: latestKeyIndex } = await this.toprfClient
       .fetchAuthPubKey({

packages/seedless-onboarding-controller/src/types.ts

@@ -1,11 +1,7 @@
 import type { RestrictedMessenger } from '@metamask/base-controller';
 import type { ControllerGetStateAction } from '@metamask/base-controller';
 import type { ControllerStateChangeEvent } from '@metamask/base-controller';
-import type {
-  ExportableKeyEncryptor,
-  KeyringControllerLockEvent,
-  KeyringControllerUnlockEvent,
-} from '@metamask/keyring-controller';
+import type { ExportableKeyEncryptor } from '@metamask/keyring-controller';
 import type { KeyPair, NodeAuthTokens } from '@metamask/toprf-secure-backup';
 import type { MutexInterface } from 'async-mutex';
 
@@ -206,7 +202,7 @@ export type SeedlessOnboardingControllerStateChangeEvent =
 export type SeedlessOnboardingControllerEvents =
   SeedlessOnboardingControllerStateChangeEvent;
 
-type AllowedEvents = KeyringControllerLockEvent | KeyringControllerUnlockEvent;
+type AllowedEvents = never;
 
 // Messenger
 export type SeedlessOnboardingControllerMessenger = RestrictedMessenger<

packages/seedless-onboarding-controller/tests/__fixtures__/mockMessenger.ts

@@ -19,7 +19,7 @@ export function createCustomSeedlessOnboardingMessenger() {
   const messenger = baseMessenger.getRestricted({
     name: 'SeedlessOnboardingController',
     allowedActions: [],
-    allowedEvents: ['KeyringController:lock', 'KeyringController:unlock'],
+    allowedEvents: [],
   });
 
   return {