feat: add unlock checks for notification and profile sync controllers (…

…#4569)

## Explanation

This adds wallet unlock checks to ensure we only can call these
controllers when the wallet is unlocked. This fixes potential issues in
extension where we call these controllers when the wallet is locked
(e.g. on browser startup)

## References

[Ticket](https://consensyssoftware.atlassian.net/browse/NOTIFY-865)

## Changelog

<!--
If you're making any consumer-facing changes, list those changes here as
if you were updating a changelog, using the template below as a guide.

(CATEGORY is one of BREAKING, ADDED, CHANGED, DEPRECATED, REMOVED, or
FIXED. For security-related issues, follow the Security Advisory
process.)

Please take care to name the exact pieces of the API you've added or
changed (e.g. types, interfaces, functions, or methods).

If there are any breaking changes, make sure to offer a solution for
consumers to follow once they upgrade to the changes.

Finally, if you're only making changes to development scripts or tests,
you may replace the template below with "None".
-->

### `@metamask/profile-sync-controller`

- **ADDED**: Actions and Events listening to the Keyring Controller
unlock requests
- **ADDED**: Unlock checks when the preinstalled snap is called.

### `@metamask/notification-services-controller`

- **ADDED**: Actions and Events listening to the Keyring Controller
unlock requests
- **ADDED**: Unlock checks when trying to setup push notifications

## 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 highlighted breaking changes using the "BREAKING" category
above as appropriate

packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts

@@ -40,6 +40,10 @@ import * as OnChainNotifications from './services/onchain-notifications';
 import type { UserStorage } from './types/user-storage/user-storage';
 import * as Utils from './utils/utils';
 
+// Mock type used for testing purposes
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type MockVar = any;
+
 const featureAnnouncementsEnv = {
   spaceId: ':space_id',
   accessToken: ':access_token',
@@ -664,6 +668,7 @@ function mockNotificationMessenger() {
     name: 'NotificationServicesController',
     allowedActions: [
       'KeyringController:getAccounts',
+      'KeyringController:getState',
       'AuthenticationController:getBearerToken',
       'AuthenticationController:isSignedIn',
       'NotificationServicesPushController:disablePushNotifications',
@@ -676,6 +681,8 @@ function mockNotificationMessenger() {
     ],
     allowedEvents: [
       'KeyringController:stateChange',
+      'KeyringController:lock',
+      'KeyringController:unlock',
       'NotificationServicesPushController:onNewNotifications',
     ],
   });
@@ -729,6 +736,10 @@ function mockNotificationMessenger() {
       return mockListAccounts();
     }
 
+    if (actionType === 'KeyringController:getState') {
+      return { isUnlocked: true } as MockVar;
+    }
+
     if (actionType === 'AuthenticationController:getBearerToken') {
       return mockGetBearerToken();
     }
@@ -774,12 +785,9 @@ function mockNotificationMessenger() {
       return mockPerformSetStorage(params[0], params[1]);
     }
 
-    const exhaustedMessengerMocks = (action: never) => {
-      return new Error(
-        `MOCK_FAIL - unsupported messenger call: ${action as string}`,
-      );
-    };
-    throw exhaustedMessengerMocks(actionType);
+    throw new Error(
+      `MOCK_FAIL - unsupported messenger call: ${actionType as string}`,
+    );
   });
 
   return {

packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts

@@ -9,6 +9,9 @@ import { toChecksumHexAddress } from '@metamask/controller-utils';
 import type {
   KeyringControllerGetAccountsAction,
   KeyringControllerStateChangeEvent,
+  KeyringControllerGetStateAction,
+  KeyringControllerLockEvent,
+  KeyringControllerUnlockEvent,
 } from '@metamask/keyring-controller';
 import type {
   AuthenticationController,
@@ -192,6 +195,7 @@ export type Actions =
 export type AllowedActions =
   // Keyring Controller Requests
   | KeyringControllerGetAccountsAction
+  | KeyringControllerGetStateAction
   // Auth Controller Requests
   | AuthenticationController.AuthenticationControllerGetBearerToken
   | AuthenticationController.AuthenticationControllerIsSignedIn
@@ -214,7 +218,11 @@ export type NotificationServicesControllerMessengerEvents =
 
 // Allowed Events
 export type AllowedEvents =
+  // Keyring Events
   | KeyringControllerStateChangeEvent
+  | KeyringControllerLockEvent
+  | KeyringControllerUnlockEvent
+  // Push Notification Events
   | NotificationServicesPushControllerOnNewNotification;
 
 // Type for the messenger of NotificationServicesController
@@ -244,6 +252,34 @@ export default class NotificationServicesController extends BaseController<
   // Temporary boolean as push notifications are not yet enabled on mobile
   #isPushIntegrated = true;
 
+  // Flag to check is notifications have been setup when the browser/extension is initialized.
+  // We want to re-initialize push notifications when the browser/extension is refreshed
+  // To ensure we subscribe to the most up-to-date notifications
+  #isPushNotificationsSetup = false;
+
+  #isUnlocked = false;
+
+  #keyringController = {
+    setupLockedStateSubscriptions: (onUnlock: () => Promise<void>) => {
+      const { isUnlocked } = this.messagingSystem.call(
+        'KeyringController:getState',
+      );
+      this.#isUnlocked = isUnlocked;
+
+      this.messagingSystem.subscribe('KeyringController:unlock', () => {
+        this.#isUnlocked = true;
+        // messaging system cannot await promises
+        // we don't need to wait for a result on this.
+        // eslint-disable-next-line @typescript-eslint/no-floating-promises
+        onUnlock();
+      });
+
+      this.messagingSystem.subscribe('KeyringController:lock', () => {
+        this.#isUnlocked = false;
+      });
+    },
+  };
+
   #auth = {
     getBearerToken: async () => {
       return await this.messagingSystem.call(
@@ -338,6 +374,12 @@ export default class NotificationServicesController extends BaseController<
       if (!this.state.isNotificationServicesEnabled) {
         return;
       }
+      if (this.#isPushNotificationsSetup) {
+        return;
+      }
+      if (!this.#isUnlocked) {
+        return;
+      }
 
       const storage = await this.#getUserStorage();
       if (!storage) {
@@ -346,6 +388,7 @@ export default class NotificationServicesController extends BaseController<
 
       const uuids = Utils.getAllUUIDs(storage);
       await this.#pushNotifications.enablePushNotifications(uuids);
+      this.#isPushNotificationsSetup = true;
     },
   };
 
@@ -463,10 +506,13 @@ export default class NotificationServicesController extends BaseController<
     });
 
     this.#isPushIntegrated = env.isPushIntegrated ?? true;
-
     this.#featureAnnouncementEnv = env.featureAnnouncements;
     this.#registerMessageHandlers();
     this.#clearLoadingStates();
+
+    this.#keyringController.setupLockedStateSubscriptions(
+      this.#pushNotifications.initializePushNotifications,
+    );
     // eslint-disable-next-line @typescript-eslint/no-floating-promises
     this.#accounts.initialize();
     // eslint-disable-next-line @typescript-eslint/no-floating-promises

packages/profile-sync-controller/package.json

@@ -53,6 +53,7 @@
   "devDependencies": {
     "@lavamoat/allow-scripts": "^3.0.4",
     "@metamask/auto-changelog": "^3.4.4",
+    "@metamask/keyring-controller": "^17.1.2",
     "@metamask/snaps-controllers": "^9.3.1",
     "@types/jest": "^27.4.1",
     "deepmerge": "^4.2.2",
@@ -66,6 +67,7 @@
     "typescript": "~5.0.4"
   },
   "peerDependencies": {
+    "@metamask/keyring-controller": "^17.0.0",
     "@metamask/snaps-controllers": "^9.3.0"
   },
   "engines": {

packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.test.ts

@@ -12,6 +12,7 @@ import {
 import type {
   Actions,
   AllowedActions,
+  AllowedEvents,
   AuthenticationControllerState,
 } from './AuthenticationController';
 import AuthenticationController from './AuthenticationController';
@@ -32,7 +33,7 @@ describe('authentication/authentication-controller - constructor() tests', () =>
   it('should initialize with default state', () => {
     const metametrics = createMockAuthMetaMetrics();
     const controller = new AuthenticationController({
-      messenger: createAuthenticationMessenger(),
+      messenger: createMockAuthenticationMessenger().messenger,
       metametrics,
     });
 
@@ -43,7 +44,7 @@ describe('authentication/authentication-controller - constructor() tests', () =>
   it('should initialize with override state', () => {
     const metametrics = createMockAuthMetaMetrics();
     const controller = new AuthenticationController({
-      messenger: createAuthenticationMessenger(),
+      messenger: createMockAuthenticationMessenger().messenger,
       state: mockSignedInState(),
       metametrics,
     });
@@ -90,6 +91,20 @@ describe('authentication/authentication-controller - performSignIn() tests', ()
     await testAndAssertFailingEndpoints('token');
   });
 
+  // When the wallet is locked, we are unable to call the snap
+  it('should error when wallet is locked', async () => {
+    const { messenger, mockKeyringControllerGetState } =
+      createMockAuthenticationMessenger();
+    const metametrics = createMockAuthMetaMetrics();
+
+    // Mock wallet is locked
+    mockKeyringControllerGetState.mockReturnValue({ isUnlocked: false });
+
+    const controller = new AuthenticationController({ messenger, metametrics });
+
+    await expect(controller.performSignIn()).rejects.toThrow(expect.any(Error));
+  });
+
   /**
    * Jest Test & Assert Utility - for testing and asserting endpoint failures
    *
@@ -208,6 +223,38 @@ describe('authentication/authentication-controller - getBearerToken() tests', ()
     expect(result).toBeDefined();
     expect(result).toBe(MOCK_ACCESS_TOKEN);
   });
+
+  // If the state is invalid, we need to re-login.
+  // But as wallet is locked, we will not be able to call the snap
+  it('should throw error if wallet is locked', async () => {
+    const metametrics = createMockAuthMetaMetrics();
+    const { messenger, mockKeyringControllerGetState } =
+      createMockAuthenticationMessenger();
+    mockAuthenticationFlowEndpoints();
+
+    // Invalid/old state
+    const originalState = mockSignedInState();
+    if (originalState.sessionData) {
+      originalState.sessionData.accessToken = 'ACCESS_TOKEN_1';
+
+      const d = new Date();
+      d.setMinutes(d.getMinutes() - 31); // expires at 30 mins
+      originalState.sessionData.expiresIn = d.toString();
+    }
+
+    // Mock wallet is locked
+    mockKeyringControllerGetState.mockReturnValue({ isUnlocked: false });
+
+    const controller = new AuthenticationController({
+      messenger,
+      state: originalState,
+      metametrics,
+    });
+
+    await expect(controller.getBearerToken()).rejects.toThrow(
+      expect.any(Error),
+    );
+  });
 });
 
 describe('authentication/authentication-controller - getSessionProfile() tests', () => {
@@ -264,6 +311,38 @@ describe('authentication/authentication-controller - getSessionProfile() tests',
     expect(result.identifierId).toBe(MOCK_LOGIN_RESPONSE.profile.identifier_id);
     expect(result.profileId).toBe(MOCK_LOGIN_RESPONSE.profile.profile_id);
   });
+
+  // If the state is invalid, we need to re-login.
+  // But as wallet is locked, we will not be able to call the snap
+  it('should throw error if wallet is locked', async () => {
+    const metametrics = createMockAuthMetaMetrics();
+    const { messenger, mockKeyringControllerGetState } =
+      createMockAuthenticationMessenger();
+    mockAuthenticationFlowEndpoints();
+
+    // Invalid/old state
+    const originalState = mockSignedInState();
+    if (originalState.sessionData) {
+      originalState.sessionData.profile.identifierId = 'ID_1';
+
+      const d = new Date();
+      d.setMinutes(d.getMinutes() - 31); // expires at 30 mins
+      originalState.sessionData.expiresIn = d.toString();
+    }
+
+    // Mock wallet is locked
+    mockKeyringControllerGetState.mockReturnValue({ isUnlocked: false });
+
+    const controller = new AuthenticationController({
+      messenger,
+      state: originalState,
+      metametrics,
+    });
+
+    await expect(controller.getSessionProfile()).rejects.toThrow(
+      expect.any(Error),
+    );
+  });
 });
 
 /**
@@ -272,11 +351,17 @@ describe('authentication/authentication-controller - getSessionProfile() tests',
  * @returns Auth Messenger
  */
 function createAuthenticationMessenger() {
-  const messenger = new ControllerMessenger<Actions | AllowedActions, never>();
+  const messenger = new ControllerMessenger<
+    Actions | AllowedActions,
+    AllowedEvents
+  >();
   return messenger.getRestricted({
     name: 'AuthenticationController',
-    allowedActions: [`SnapController:handleRequest`],
-    allowedEvents: [],
+    allowedActions: [
+      'KeyringController:getState',
+      'SnapController:handleRequest',
+    ],
+    allowedEvents: ['KeyringController:lock', 'KeyringController:unlock'],
   });
 }
 
@@ -293,6 +378,10 @@ function createMockAuthenticationMessenger() {
     .fn()
     .mockResolvedValue('MOCK_SIGNED_MESSAGE');
 
+  const mockKeyringControllerGetState = jest
+    .fn()
+    .mockReturnValue({ isUnlocked: true });
+
   mockCall.mockImplementation((...args) => {
     const [actionType, params] = args;
     if (actionType === 'SnapController:handleRequest') {
@@ -311,12 +400,21 @@ function createMockAuthenticationMessenger() {
       );
     }
 
+    if (actionType === 'KeyringController:getState') {
+      return mockKeyringControllerGetState();
+    }
+
     throw new Error(
       `MOCK_FAIL - unsupported messenger call: ${actionType as string}`,
     );
   });
 
-  return { messenger, mockSnapGetPublicKey, mockSnapSignMessage };
+  return {
+    messenger,
+    mockSnapGetPublicKey,
+    mockSnapSignMessage,
+    mockKeyringControllerGetState,
+  };
 }
 
 /**