From fc660e94cefe200a63cbd6bc6ef1da797d035a2b Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 29 Jul 2024 16:31:26 +0100 Subject: [PATCH] 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 ### `@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 --- .../NotificationServicesController.test.ts | 20 +++- .../NotificationServicesController.ts | 48 +++++++- packages/profile-sync-controller/package.json | 2 + .../AuthenticationController.test.ts | 110 +++++++++++++++++- .../AuthenticationController.ts | 49 +++++++- .../UserStorageController.test.ts | 10 +- .../user-storage/UserStorageController.ts | 43 ++++++- .../tsconfig.build.json | 5 +- .../profile-sync-controller/tsconfig.json | 5 +- yarn.lock | 2 + 10 files changed, 269 insertions(+), 25 deletions(-) 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