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..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 @@ -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 () => { @@ -660,10 +673,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 +691,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(); }); @@ -749,9 +761,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 +779,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 +801,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 +819,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -922,9 +932,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 +950,7 @@ describe('user-storage/user-storage-controller - syncInternalAccountsWithUserSto await controller.syncInternalAccountsWithUserStorage(); mockAPI.mockEndpointGetUserStorage.done(); - mockAPI.mockEndpointUpsertUserStorage.done(); + mockAPI.mockEndpointBatchUpsertUserStorage.done(); expect( messengerMocks.mockAccountsUpdateAccountMetadata, @@ -1150,9 +1159,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 +1177,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 +1570,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( 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();