feat(NOTIFY-1140): add batch PUT endpoint for account syncing (#4724)

## Explanation This PR replaces the singular PUTs by batch PUTs for account syncing. ## References [NOTIFY-1140](https://consensyssoftware.atlassian.net/jira/software/projects/NOTIFY/boards/616?assignee=712020%3A5843b7e2-a7fe-4c45-9fbd-e1f2b2eb58c2&selectedIssue=NOTIFY-1140) ## Changelog ### `@metamask/profile-sync-controller` - **CHANGED**: use batch endpoint instead of singular requests to update user storage for account syncing ## 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 - [x] I've prepared draft pull requests for clients and consumer packages to resolve any breaking changes [NOTIFY-1140]: https://consensyssoftware.atlassian.net/browse/NOTIFY-1140?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
MetaMask · Sep 23, 2024 · c525459 · c525459
1 parent ea44c16
commit c525459
Show file tree

Hide file tree

Showing 5 changed files with 119 additions and 56 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
@@ -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<GetUserStorageAllFeatureEntriesResponse> {
   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<string>(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
@@ -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
@@ -202,6 +202,10 @@ export async function batchUpsertUserStorage(
   data: [UserStoragePathWithKeyOnly, string][],
   opts: UserStorageBatchUpsertOptions,
 ): Promise<void> {
+  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
@@ -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
@@ -136,6 +136,10 @@ export class UserStorage {
     data: [UserStoragePathWithKeyOnly, string][],
   ): Promise<void> {
     try {
+      if (!data.length) {
+        return;
+      }
+
       const headers = await this.#getAuthorizationHeader();
       const storageKey = await this.getStorageKey();