diff --git a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts index 3a637e4662..87154aa8c0 100644 --- a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts +++ b/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 { diff --git a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts b/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts index 9775b5fcaf..3b432e8bd0 100644 --- a/packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts +++ b/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) => { + 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 diff --git a/packages/profile-sync-controller/package.json b/packages/profile-sync-controller/package.json index ee92038cfb..e81dc89ced 100644 --- a/packages/profile-sync-controller/package.json +++ b/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": { diff --git a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.test.ts b/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.test.ts index 66f650dd44..880a42f0f9 100644 --- a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.test.ts +++ b/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(); + 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, + }; } /** diff --git a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts b/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts index fa9730e0e4..351dbda53d 100644 --- a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts +++ b/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts @@ -3,6 +3,11 @@ import type { StateMetadata, } from '@metamask/base-controller'; import { BaseController } from '@metamask/base-controller'; +import type { + KeyringControllerGetStateAction, + KeyringControllerLockEvent, + KeyringControllerUnlockEvent, +} from '@metamask/keyring-controller'; import type { HandleSnapRequest } from '@metamask/snaps-controllers'; import { @@ -87,15 +92,21 @@ export type AuthenticationControllerGetSessionProfile = export type AuthenticationControllerIsSignedIn = ActionsObj['isSignedIn']; // Allowed Actions -export type AllowedActions = HandleSnapRequest; +export type AllowedActions = + | HandleSnapRequest + | KeyringControllerGetStateAction; + +export type AllowedEvents = + | KeyringControllerLockEvent + | KeyringControllerUnlockEvent; // Messenger export type AuthenticationControllerMessenger = RestrictedControllerMessenger< typeof controllerName, Actions | AllowedActions, - never, + AllowedEvents, AllowedActions['type'], - never + AllowedEvents['type'] >; /** @@ -109,6 +120,25 @@ export default class AuthenticationController extends BaseController< > { #metametrics: MetaMetricsAuth; + #isUnlocked = false; + + #keyringController = { + setupLockedStateSubscriptions: () => { + const { isUnlocked } = this.messagingSystem.call( + 'KeyringController:getState', + ); + this.#isUnlocked = isUnlocked; + + this.messagingSystem.subscribe('KeyringController:unlock', () => { + this.#isUnlocked = true; + }); + + this.messagingSystem.subscribe('KeyringController:lock', () => { + this.#isUnlocked = false; + }); + }, + }; + constructor({ messenger, state, @@ -135,6 +165,7 @@ export default class AuthenticationController extends BaseController< this.#metametrics = metametrics; + this.#keyringController.setupLockedStateSubscriptions(); this.#registerMessageHandlers(); } @@ -316,6 +347,12 @@ export default class AuthenticationController extends BaseController< return this.#_snapPublicKeyCache; } + if (!this.#isUnlocked) { + throw new Error( + '#snapGetPublicKey - unable to call snap, wallet is locked', + ); + } + const result = (await this.messagingSystem.call( 'SnapController:handleRequest', createSnapPublicKeyRequest(), @@ -339,6 +376,12 @@ export default class AuthenticationController extends BaseController< return this.#_snapSignMessageCache[message]; } + if (!this.#isUnlocked) { + throw new Error( + '#snapSignMessage - unable to call snap, wallet is locked', + ); + } + const result = (await this.messagingSystem.call( 'SnapController:handleRequest', createSnapSignMessageRequest(message), diff --git a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.test.ts b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.test.ts index 354e154b95..8aa8b1098d 100644 --- a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.test.ts +++ b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.test.ts @@ -18,6 +18,7 @@ import { } from './__fixtures__/mockStorage'; import type { AllowedActions, + AllowedEvents, NotificationServicesControllerDisableNotificationServices, NotificationServicesControllerSelectIsNotificationServicesEnabled, } from './UserStorageController'; @@ -305,10 +306,11 @@ describe('user-storage/user-storage-controller - enableProfileSyncing() tests', function mockUserStorageMessenger() { const messenger = new ControllerMessenger< AllowedActions, - never + AllowedEvents >().getRestricted({ name: 'UserStorageController', allowedActions: [ + 'KeyringController:getState', 'SnapController:handleRequest', 'AuthenticationController:getBearerToken', 'AuthenticationController:getSessionProfile', @@ -318,7 +320,7 @@ function mockUserStorageMessenger() { 'NotificationServicesController:disableNotificationServices', 'NotificationServicesController:selectIsNotificationServicesEnabled', ], - allowedEvents: [], + allowedEvents: ['KeyringController:lock', 'KeyringController:unlock'], }); const mockSnapGetPublicKey = jest.fn().mockResolvedValue('MOCK_PUBLIC_KEY'); @@ -415,6 +417,10 @@ function mockUserStorageMessenger() { return mockAuthPerformSignOut(); } + if (actionType === 'KeyringController:getState') { + return { isUnlocked: true }; + } + const exhaustedMessengerMocks = (action: never) => { throw new Error( `MOCK_FAIL - unsupported messenger call: ${action as string}`, diff --git a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts index 4ff0b6610f..f4d542ee03 100644 --- a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts +++ b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts @@ -3,6 +3,11 @@ import type { StateMetadata, } from '@metamask/base-controller'; import { BaseController } from '@metamask/base-controller'; +import type { + KeyringControllerGetStateAction, + KeyringControllerLockEvent, + KeyringControllerUnlockEvent, +} from '@metamask/keyring-controller'; import type { HandleSnapRequest } from '@metamask/snaps-controllers'; import { createSnapSignMessageRequest } from '../authentication/auth-snap-requests'; @@ -87,6 +92,8 @@ export type UserStorageControllerDisableProfileSyncing = // Allowed Actions export type AllowedActions = + // Keyring Requests + | KeyringControllerGetStateAction // Snap Requests | HandleSnapRequest // Auth Requests @@ -99,13 +106,17 @@ export type AllowedActions = | NotificationServicesControllerDisableNotificationServices | NotificationServicesControllerSelectIsNotificationServicesEnabled; +export type AllowedEvents = + | KeyringControllerLockEvent + | KeyringControllerUnlockEvent; + // Messenger export type UserStorageControllerMessenger = RestrictedControllerMessenger< typeof controllerName, Actions | AllowedActions, - never, + AllowedEvents, AllowedActions['type'], - never + AllowedEvents['type'] >; /** @@ -161,6 +172,25 @@ export default class UserStorageController extends BaseController< }, }; + #isUnlocked = false; + + #keyringController = { + setupLockedStateSubscriptions: () => { + const { isUnlocked } = this.messagingSystem.call( + 'KeyringController:getState', + ); + this.#isUnlocked = isUnlocked; + + this.messagingSystem.subscribe('KeyringController:unlock', () => { + this.#isUnlocked = true; + }); + + this.messagingSystem.subscribe('KeyringController:lock', () => { + this.#isUnlocked = false; + }); + }, + }; + getMetaMetricsState: () => boolean; constructor(params: { @@ -176,6 +206,7 @@ export default class UserStorageController extends BaseController< }); this.getMetaMetricsState = params.getMetaMetricsState; + this.#keyringController.setupLockedStateSubscriptions(); this.#registerMessageHandlers(); } @@ -260,7 +291,7 @@ export default class UserStorageController extends BaseController< const isMetaMetricsParticipation = this.getMetaMetricsState(); if (!isMetaMetricsParticipation) { - this.messagingSystem.call('AuthenticationController:performSignOut'); + await this.#auth.signOut(); } this.#setIsProfileSyncingUpdateLoading(false); @@ -389,6 +420,12 @@ export default class UserStorageController extends BaseController< return this.#_snapSignMessageCache[message]; } + if (!this.#isUnlocked) { + throw new Error( + '#snapSignMessage - unable to call snap, wallet is locked', + ); + } + const result = (await this.messagingSystem.call( 'SnapController:handleRequest', createSnapSignMessageRequest(message), diff --git a/packages/profile-sync-controller/tsconfig.build.json b/packages/profile-sync-controller/tsconfig.build.json index 8d4cf54b4c..d5dd6781e0 100644 --- a/packages/profile-sync-controller/tsconfig.build.json +++ b/packages/profile-sync-controller/tsconfig.build.json @@ -7,9 +7,8 @@ "skipLibCheck": true }, "references": [ - { - "path": "../base-controller/tsconfig.build.json" - } + { "path": "../base-controller/tsconfig.build.json" }, + { "path": "../keyring-controller/tsconfig.build.json" } ], "include": ["../../types", "./src"] } diff --git a/packages/profile-sync-controller/tsconfig.json b/packages/profile-sync-controller/tsconfig.json index 34354c4b09..092d51bd4d 100644 --- a/packages/profile-sync-controller/tsconfig.json +++ b/packages/profile-sync-controller/tsconfig.json @@ -3,6 +3,9 @@ "compilerOptions": { "baseUrl": "./" }, - "references": [{ "path": "../base-controller" }], + "references": [ + { "path": "../base-controller" }, + { "path": "../keyring-controller" } + ], "include": ["../../types", "./src"] } diff --git a/yarn.lock b/yarn.lock index 9966639ae7..4a0966e2cd 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3551,6 +3551,7 @@ __metadata: "@lavamoat/allow-scripts": "npm:^3.0.4" "@metamask/auto-changelog": "npm:^3.4.4" "@metamask/base-controller": "npm:^6.0.2" + "@metamask/keyring-controller": "npm:^17.1.2" "@metamask/snaps-controllers": "npm:^9.3.1" "@metamask/snaps-sdk": "npm:^6.1.1" "@metamask/snaps-utils": "npm:^7.8.1" @@ -3570,6 +3571,7 @@ __metadata: typedoc-plugin-missing-exports: "npm:^2.0.0" typescript: "npm:~5.0.4" peerDependencies: + "@metamask/keyring-controller": ^17.0.0 "@metamask/snaps-controllers": ^9.3.0 languageName: unknown linkType: soft