From cf2e5694ca7aabf1185f3d239bdfa2f7199d13a8 Mon Sep 17 00:00:00 2001 From: Mathieu Artu Date: Mon, 23 Sep 2024 11:13:41 +0200 Subject: [PATCH 1/2] feat(NOTIFY-1140): add batch PUT endpoint for account syncing --- .../UserStorageController.test.ts | 137 ++++++++++++++---- .../user-storage/UserStorageController.ts | 48 +++--- 2 files changed, 130 insertions(+), 55 deletions(-) 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 5996653c22..5444e8cf5e 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 @@ -2,7 +2,7 @@ import { ControllerMessenger } from '@metamask/base-controller'; import type { InternalAccount } from '@metamask/keyring-api'; import type nock from 'nock'; -import encryption from '../../shared/encryption'; +import encryption, { createSHA256Hash } from '../../shared/encryption'; import type { AuthenticationControllerGetBearerToken, AuthenticationControllerGetSessionProfile, @@ -14,6 +14,7 @@ import { MOCK_USER_STORAGE_ACCOUNTS, } from './__fixtures__/mockAccounts'; import { + mockEndpointBatchUpsertUserStorage, mockEndpointGetUserStorage, mockEndpointGetUserStorageAllFeatureEntries, mockEndpointUpsertUserStorage, @@ -504,12 +505,27 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', mockUserStorageAccountsResponse, ), - mockEndpointUpsertUserStorageAccount1: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ALL[0].address}`, - ), - mockEndpointUpsertUserStorageAccount2: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ALL[1].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage( + 'accounts', + undefined, + async (_uri, requestBody) => { + const decryptedBody = await decryptBatchUpsertBody( + requestBody, + MOCK_STORAGE_KEY, + ); + + const expectedBody = createExpectedAccountSyncBatchUpsertBody( + MOCK_INTERNAL_ACCOUNTS.ALL.slice(0, 2).map((account) => [ + account.address, + account as InternalAccount, + ]), + MOCK_STORAGE_KEY, + ); + + expect(decryptedBody).toStrictEqual(expectedBody); + }, + ), }, }; }; @@ -525,12 +541,9 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorageAccount1.done(); - mockAPI.mockEndpointUpsertUserStorageAccount2.done(); expect(mockAPI.mockEndpointGetUserStorage.isDone()).toBe(true); - expect(mockAPI.mockEndpointUpsertUserStorageAccount1.isDone()).toBe(true); - expect(mockAPI.mockEndpointUpsertUserStorageAccount2.isDone()).toBe(true); + expect(mockAPI.mockEndpointBatchUpsertUserStorage.isDone()).toBe(true); }); it('creates internal accounts if user storage has more accounts', async () => { @@ -556,6 +569,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -572,8 +587,10 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect(mockAPI.mockEndpointGetUserStorage.isDone()).toBe(true); + expect(mockAPI.mockEndpointBatchUpsertUserStorage.isDone()).toBe(true); expect(messengerMocks.mockKeyringAddNewAccount).toHaveBeenCalledTimes( MOCK_USER_STORAGE_ACCOUNTS.SAME_AS_INTERNAL_ALL.length - @@ -604,6 +621,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -660,10 +679,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - - mockEndpointUpsertUserStorageAccount2: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ALL[1].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -680,9 +697,10 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorageAccount2.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect(mockAPI.mockEndpointGetUserStorage.isDone()).toBe(true); + expect(mockAPI.mockEndpointBatchUpsertUserStorage.isDone()).toBe(true); expect(messengerMocks.mockKeyringAddNewAccount).not.toHaveBeenCalled(); }); @@ -712,6 +730,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -749,9 +769,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointUpsertUserStorage: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ONE_CUSTOM_NAME_WITHOUT_LAST_UPDATED[0].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -768,7 +787,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -790,9 +809,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointUpsertUserStorage: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ONE_CUSTOM_NAME_WITH_LAST_UPDATED[0].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -809,7 +827,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -842,6 +860,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -885,6 +905,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -922,9 +944,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointUpsertUserStorage: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ONE_CUSTOM_NAME_WITH_LAST_UPDATED[0].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -941,7 +962,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -963,6 +984,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1016,6 +1039,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1059,6 +1084,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1105,6 +1132,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', mockGetEntriesResponse, ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1150,9 +1179,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointUpsertUserStorage: mockEndpointUpsertUserStorage( - `accounts.${MOCK_INTERNAL_ACCOUNTS.ONE_CUSTOM_NAME_WITH_LAST_UPDATED_MOST_RECENT[0].address}`, - ), + mockEndpointBatchUpsertUserStorage: + mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1169,7 +1197,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -1562,3 +1590,48 @@ async function createMockUserStorageEntries( ): Promise { return await Promise.all(data.map((d) => createMockUserStorageEntry(d))); } + +/** + * Test Utility - decrypts a realistic batch upsert payload + * @param requestBody - nock body + * @param storageKey - storage key + * @returns decrypted body + */ +async function decryptBatchUpsertBody( + requestBody: nock.Body, + storageKey: string, +) { + if (typeof requestBody === 'string') { + return requestBody; + } + return await Promise.all( + Object.entries(requestBody.data).map( + async ([entryKey, entryValue]) => { + return [ + entryKey, + await encryption.decryptString(entryValue, storageKey), + ]; + }, + ), + ); +} + +/** + * Test Utility - creates a realistic expected batch upsert payload + * @param data - data supposed to be upserted + * @param storageKey - storage key + * @returns expected body + */ +function createExpectedAccountSyncBatchUpsertBody( + data: [string, InternalAccount][], + storageKey: string, +) { + return data.map(([entryKey, entryValue]) => [ + createSHA256Hash(String(entryKey) + storageKey), + JSON.stringify( + AccountsUserStorageModule.mapInternalAccountToUserStorageAccount( + entryValue, + ), + ), + ]); +} 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 19e8824995..35a2c79459 100644 --- a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts +++ b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts @@ -342,17 +342,15 @@ export default class UserStorageController extends BaseController< return; } - const userStorageAccountsList = internalAccountsList.map( - mapInternalAccountToUserStorageAccount, - ); - - await Promise.all( - userStorageAccountsList.map(async (userStorageAccount) => { - await this.performSetStorage( - `accounts.${userStorageAccount.a}`, - JSON.stringify(userStorageAccount), - ); - }), + const internalAccountsListFormattedForUserStorage = + internalAccountsList.map(mapInternalAccountToUserStorageAccount); + + await this.performBatchSetStorage( + 'accounts', + internalAccountsListFormattedForUserStorage.map((account) => [ + account.a, + JSON.stringify(account), + ]), ); }, }; @@ -774,6 +772,9 @@ export default class UserStorageController extends BaseController< return; } + // Prepare an array of internal accounts to be saved to the user storage + const internalAccountsToBeSavedToUserStorage: InternalAccount[] = []; + // Compare internal accounts list with user storage accounts list // First step: compare lengths let internalAccountsList = await this.#accounts.getInternalAccountsList(); @@ -812,9 +813,7 @@ export default class UserStorageController extends BaseController< ); if (!userStorageAccount) { - await this.#accounts.saveInternalAccountToUserStorage( - internalAccount, - ); + internalAccountsToBeSavedToUserStorage.push(internalAccount); continue; } @@ -844,9 +843,7 @@ export default class UserStorageController extends BaseController< // Internal account has custom name but user storage account has default name if (isUserStorageAccountNameDefault) { - await this.#accounts.saveInternalAccountToUserStorage( - internalAccount, - ); + internalAccountsToBeSavedToUserStorage.push(internalAccount); continue; } @@ -861,9 +858,7 @@ export default class UserStorageController extends BaseController< userStorageAccount.nlu; if (isInternalAccountNameNewer) { - await this.#accounts.saveInternalAccountToUserStorage( - internalAccount, - ); + internalAccountsToBeSavedToUserStorage.push(internalAccount); continue; } } @@ -881,12 +876,19 @@ export default class UserStorageController extends BaseController< continue; } else if (internalAccount.metadata.nameLastUpdatedAt !== undefined) { - await this.#accounts.saveInternalAccountToUserStorage( - internalAccount, - ); + internalAccountsToBeSavedToUserStorage.push(internalAccount); continue; } } + + // Save the internal accounts list to the user storage + await this.performBatchSetStorage( + 'accounts', + internalAccountsToBeSavedToUserStorage.map((account) => [ + account.address, + JSON.stringify(mapInternalAccountToUserStorageAccount(account)), + ]), + ); } catch (e) { const errorMessage = e instanceof Error ? e.message : JSON.stringify(e); throw new Error( From 9735ac79f8ae2e045ad6501e94f0b4c12d30e9cc Mon Sep 17 00:00:00 2001 From: Mathieu Artu Date: Mon, 23 Sep 2024 11:58:36 +0200 Subject: [PATCH 2/2] fix: add empty check for batch upserts --- .../UserStorageController.test.ts | 20 ------------------- .../src/controllers/user-storage/services.ts | 4 ++++ .../src/sdk/user-storage.test.ts | 2 +- .../src/sdk/user-storage.ts | 4 ++++ 4 files changed, 9 insertions(+), 21 deletions(-) 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 5444e8cf5e..4dfdb2ccf4 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 @@ -569,8 +569,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -587,10 +585,8 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect(mockAPI.mockEndpointGetUserStorage.isDone()).toBe(true); - expect(mockAPI.mockEndpointBatchUpsertUserStorage.isDone()).toBe(true); expect(messengerMocks.mockKeyringAddNewAccount).toHaveBeenCalledTimes( MOCK_USER_STORAGE_ACCOUNTS.SAME_AS_INTERNAL_ALL.length - @@ -621,8 +617,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -730,8 +724,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -860,8 +852,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -905,8 +895,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -984,8 +972,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1039,8 +1025,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1084,8 +1068,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', await mockUserStorageAccountsResponse(), ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; @@ -1132,8 +1114,6 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto 'accounts', mockGetEntriesResponse, ), - mockEndpointBatchUpsertUserStorage: - mockEndpointBatchUpsertUserStorage('accounts'), }, }; }; diff --git a/packages/profile-sync-controller/src/controllers/user-storage/services.ts b/packages/profile-sync-controller/src/controllers/user-storage/services.ts index f3abe1bcb7..2db2c933c8 100644 --- a/packages/profile-sync-controller/src/controllers/user-storage/services.ts +++ b/packages/profile-sync-controller/src/controllers/user-storage/services.ts @@ -202,6 +202,10 @@ export async function batchUpsertUserStorage( data: [UserStoragePathWithKeyOnly, string][], opts: UserStorageBatchUpsertOptions, ): Promise { + if (!data.length) { + return; + } + const { bearerToken, path, storageKey, nativeScryptCrypto } = opts; const encryptedData = await Promise.all( diff --git a/packages/profile-sync-controller/src/sdk/user-storage.test.ts b/packages/profile-sync-controller/src/sdk/user-storage.test.ts index ea240e4a58..39777a226e 100644 --- a/packages/profile-sync-controller/src/sdk/user-storage.test.ts +++ b/packages/profile-sync-controller/src/sdk/user-storage.test.ts @@ -163,7 +163,7 @@ describe('User Storage', () => { }); await expect( - userStorage.batchSetItems('notifications', []), + userStorage.batchSetItems('notifications', [['key', 'value']]), ).rejects.toThrow(UserStorageError); }); diff --git a/packages/profile-sync-controller/src/sdk/user-storage.ts b/packages/profile-sync-controller/src/sdk/user-storage.ts index cc1d8185dc..838842be1e 100644 --- a/packages/profile-sync-controller/src/sdk/user-storage.ts +++ b/packages/profile-sync-controller/src/sdk/user-storage.ts @@ -136,6 +136,10 @@ export class UserStorage { data: [UserStoragePathWithKeyOnly, string][], ): Promise { try { + if (!data.length) { + return; + } + const headers = await this.#getAuthorizationHeader(); const storageKey = await this.getStorageKey();