Feat(NOTIFY-852): Enforce user storage schema usage (#4543)

## Explanation This PR enforces the`UserStorageSchema` usage in every function that gets or sets user storage. This will help keep user storage entries consistent, and improve DX when using user storage. ## References [NOTIFY-852](https://consensyssoftware.atlassian.net/browse/NOTIFY-852) ## Changelog ### `@metamask/profile-sync-controller` - **CHANGED**: User storage schema - **ADDED**: User storage path compile-time and runtime validation ### `@metamask/notification-services-controller` - **CHANGED**: Update argument to follow on `@metamask/profile-sync-controller` changes when calling `UserStorageController:performGetStorage` and `UserStorageController:performSetStorage` ## 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 [NOTIFY-852]: https://consensyssoftware.atlassian.net/browse/NOTIFY-852?atlOrigin=eyJpIjoiNWRkNTljNzYxNjVmNDY3MDlhMDU5Y2ZhYzA5YTRkZjUiLCJwIjoiZ2l0aHViLWNvbS1KU1cifQ
MetaMask · Jul 22, 2024 · 51f7584 · 51f7584
1 parent cc09e4e
commit 51f7584
Show file tree

Hide file tree

Showing 11 changed files with 191 additions and 102 deletions.
diff --git a/...-services-controller/src/NotificationServicesController/NotificationServicesController.ts b/...-services-controller/src/NotificationServicesController/NotificationServicesController.ts
@@ -267,13 +267,13 @@ export default class NotificationServicesController extends BaseController<
     getNotificationStorage: async () => {
       return await this.messagingSystem.call(
         'UserStorageController:performGetStorage',
-        'notificationSettings',
+        'notifications.notificationSettings',
       );
     },
     setNotificationStorage: async (state: string) => {
       return await this.messagingSystem.call(
         'UserStorageController:performSetStorage',
-        'notificationSettings',
+        'notifications.notificationSettings',
         state,
       );
     },

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
@@ -59,7 +59,9 @@ describe('user-storage/user-storage-controller - performGetStorage() tests', ()
       getMetaMetricsState: () => true,
     });
 
-    const result = await controller.performGetStorage('notificationSettings');
+    const result = await controller.performGetStorage(
+      'notifications.notificationSettings',
+    );
     mockAPI.done();
     expect(result).toBe(MOCK_STORAGE_DATA);
   });
@@ -76,7 +78,7 @@ describe('user-storage/user-storage-controller - performGetStorage() tests', ()
     });
 
     await expect(
-      controller.performGetStorage('notificationSettings'),
+      controller.performGetStorage('notifications.notificationSettings'),
     ).rejects.toThrow(expect.any(Error));
   });
 
@@ -111,7 +113,7 @@ describe('user-storage/user-storage-controller - performGetStorage() tests', ()
       });
 
       await expect(
-        controller.performGetStorage('notificationSettings'),
+        controller.performGetStorage('notifications.notificationSettings'),
       ).rejects.toThrow(expect.any(Error));
     },
   );
@@ -132,7 +134,10 @@ describe('user-storage/user-storage-controller - performSetStorage() tests', ()
       getMetaMetricsState: () => true,
     });
 
-    await controller.performSetStorage('notificationSettings', 'new data');
+    await controller.performSetStorage(
+      'notifications.notificationSettings',
+      'new data',
+    );
     expect(mockAPI.isDone()).toBe(true);
   });
 
@@ -148,7 +153,10 @@ describe('user-storage/user-storage-controller - performSetStorage() tests', ()
     });
 
     await expect(
-      controller.performSetStorage('notificationSettings', 'new data'),
+      controller.performSetStorage(
+        'notifications.notificationSettings',
+        'new data',
+      ),
     ).rejects.toThrow(expect.any(Error));
   });
 
@@ -183,7 +191,10 @@ describe('user-storage/user-storage-controller - performSetStorage() tests', ()
       });
 
       await expect(
-        controller.performSetStorage('notificationSettings', 'new data'),
+        controller.performSetStorage(
+          'notifications.notificationSettings',
+          'new data',
+        ),
       ).rejects.toThrow(expect.any(Error));
     },
   );
@@ -197,7 +208,10 @@ describe('user-storage/user-storage-controller - performSetStorage() tests', ()
       getMetaMetricsState: () => true,
     });
     await expect(
-      controller.performSetStorage('notificationSettings', 'new data'),
+      controller.performSetStorage(
+        'notifications.notificationSettings',
+        'new data',
+      ),
     ).rejects.toThrow(expect.any(Error));
   });
 });

diff --git a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts
@@ -14,7 +14,7 @@ import type {
   AuthenticationControllerPerformSignOut,
 } from '../authentication/AuthenticationController';
 import { createSHA256Hash } from './encryption';
-import type { UserStorageEntryKeys } from './schema';
+import type { UserStoragePath } from './schema';
 import { getUserStorage, upsertUserStorage } from './services';
 
 // TODO: fix external dependencies
@@ -281,17 +281,19 @@ export default class UserStorageController extends BaseController<
    * Allows retrieval of stored data. Data stored is string formatted.
    * Developers can extend the entry path and entry name through the `schema.ts` file.
    *
-   * @param entryKey - entry key that matches schema
+   * @param path - string in the form of `${feature}.${key}` that matches schema
    * @returns the decrypted string contents found from user storage (or null if not found)
    */
   public async performGetStorage(
-    entryKey: UserStorageEntryKeys,
+    path: UserStoragePath,
   ): Promise<string | null> {
     this.#assertProfileSyncingEnabled();
+
     const { bearerToken, storageKey } =
       await this.#getStorageKeyAndBearerToken();
+
     const result = await getUserStorage({
-      entryKey,
+      path,
       bearerToken,
       storageKey,
     });
@@ -303,20 +305,21 @@ export default class UserStorageController extends BaseController<
    * Allows storage of user data. Data stored must be string formatted.
    * Developers can extend the entry path and entry name through the `schema.ts` file.
    *
-   * @param entryKey - entry key that matches schema
+   * @param path - string in the form of `${feature}.${key}` that matches schema
    * @param value - The string data you want to store.
    * @returns nothing. NOTE that an error is thrown if fails to store data.
    */
   public async performSetStorage(
-    entryKey: UserStorageEntryKeys,
+    path: UserStoragePath,
     value: string,
   ): Promise<void> {
     this.#assertProfileSyncingEnabled();
+
     const { bearerToken, storageKey } =
       await this.#getStorageKeyAndBearerToken();
 
     await upsertUserStorage(value, {
-      entryKey,
+      path,
       bearerToken,
       storageKey,
     });

diff --git a/packages/profile-sync-controller/src/controllers/user-storage/__fixtures__/mockResponses.ts b/packages/profile-sync-controller/src/controllers/user-storage/__fixtures__/mockResponses.ts
@@ -10,7 +10,7 @@ type MockResponse = {
 };
 
 export const MOCK_USER_STORAGE_NOTIFICATIONS_ENDPOINT = `${USER_STORAGE_ENDPOINT}${createEntryPath(
-  'notificationSettings',
+  'notifications.notificationSettings',
   MOCK_STORAGE_KEY,
 )}`;
 

diff --git a/packages/profile-sync-controller/src/controllers/user-storage/schema.test.ts b/packages/profile-sync-controller/src/controllers/user-storage/schema.test.ts
@@ -0,0 +1,45 @@
+import { getFeatureAndKeyFromPath, USER_STORAGE_SCHEMA } from './schema';
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type ErroneousUserStoragePath = any;
+
+describe('user-storage/schema.ts', () => {
+  describe('getFeatureAndKeyFromPath', () => {
+    it('should throw error if the feature.key format is incorrect', () => {
+      const path = 'feature/key';
+      expect(() =>
+        getFeatureAndKeyFromPath(path as ErroneousUserStoragePath),
+      ).toThrow(
+        "user-storage - path is not in the correct format. Correct format: 'feature.key'",
+      );
+    });
+
+    it('should throw error if feature is invalid', () => {
+      const path = 'invalid.feature';
+      expect(() =>
+        getFeatureAndKeyFromPath(path as ErroneousUserStoragePath),
+      ).toThrow('user-storage - invalid feature provided: invalid');
+    });
+
+    it('should throw error if key is invalid', () => {
+      const feature = 'notifications';
+      const path = `${feature}.invalid`;
+      const validKeys = USER_STORAGE_SCHEMA[feature].join(', ');
+
+      expect(() =>
+        getFeatureAndKeyFromPath(path as ErroneousUserStoragePath),
+      ).toThrow(
+        `user-storage - invalid key provided for this feature: invalid. Valid keys: ${validKeys}`,
+      );
+    });
+
+    it('should return feature and key from path', () => {
+      const path = 'notifications.notificationSettings';
+      const result = getFeatureAndKeyFromPath(path);
+      expect(result).toStrictEqual({
+        feature: 'notifications',
+        key: 'notificationSettings',
+      });
+    });
+  });
+});
diff --git a/packages/profile-sync-controller/src/controllers/user-storage/schema.ts b/packages/profile-sync-controller/src/controllers/user-storage/schema.ts
@@ -1,38 +1,76 @@
 import { createSHA256Hash } from './encryption';
 
-type UserStorageEntry = { path: string; entryName: string };
-
 /**
- * The User Storage Endpoint requires a path and an entry name.
- * Developers can provide additional paths by extending this variable below
+ * The User Storage Endpoint requires a feature name and a namespace key.
+ * Developers can provide additional features and keys by extending these types below.
  */
-export const USER_STORAGE_ENTRIES = {
-  notificationSettings: {
-    path: 'notifications',
-    entryName: 'notificationSettings',
-  },
-} satisfies Record<string, UserStorageEntry>;
 
-export type UserStorageEntryKeys = keyof typeof USER_STORAGE_ENTRIES;
+export const USER_STORAGE_SCHEMA = {
+  notifications: ['notificationSettings'],
+} as const;
+
+type UserStorageSchema = typeof USER_STORAGE_SCHEMA;
+type UserStorageFeatures = keyof UserStorageSchema;
+type UserStorageFeatureKeys<Feature extends UserStorageFeatures> =
+  UserStorageSchema[Feature][number];
+
+type UserStorageFeatureAndKey = {
+  feature: UserStorageFeatures;
+  key: UserStorageFeatureKeys<UserStorageFeatures>;
+};
+
+export type UserStoragePath = {
+  [K in keyof UserStorageSchema]: `${K}.${UserStorageSchema[K][number]}`;
+}[keyof UserStorageSchema];
+
+export const getFeatureAndKeyFromPath = (
+  path: UserStoragePath,
+): UserStorageFeatureAndKey => {
+  const pathRegex = /^\w+\.\w+$/u;
+
+  if (!pathRegex.test(path)) {
+    throw new Error(
+      `user-storage - path is not in the correct format. Correct format: 'feature.key'`,
+    );
+  }
+
+  const [feature, key] = path.split('.') as [
+    UserStorageFeatures,
+    UserStorageFeatureKeys<UserStorageFeatures>,
+  ];
+
+  if (!(feature in USER_STORAGE_SCHEMA)) {
+    throw new Error(`user-storage - invalid feature provided: ${feature}`);
+  }
+
+  const validFeature = USER_STORAGE_SCHEMA[feature] as readonly string[];
+
+  if (!validFeature.includes(key)) {
+    const validKeys = USER_STORAGE_SCHEMA[feature].join(', ');
+
+    throw new Error(
+      `user-storage - invalid key provided for this feature: ${key}. Valid keys: ${validKeys}`,
+    );
+  }
+
+  return { feature, key };
+};
 
 /**
  * Constructs a unique entry path for a user.
  * This can be done due to the uniqueness of the storage key (no users will share the same storage key).
  * The users entry is a unique hash that cannot be reversed.
  *
- * @param entryKey - storage schema entry key
+ * @param path - string in the form of `${feature}.${key}` that matches schema
  * @param storageKey - users storage key
  * @returns path to store entry
  */
 export function createEntryPath(
-  entryKey: UserStorageEntryKeys,
+  path: UserStoragePath,
   storageKey: string,
 ): string {
-  const entry = USER_STORAGE_ENTRIES[entryKey];
-  if (!entry) {
-    throw new Error(`user-storage - invalid entry provided: ${entryKey}`);
-  }
+  const { feature, key } = getFeatureAndKeyFromPath(path);
+  const hashedKey = createSHA256Hash(key + storageKey);
 
-  const hashedKey = createSHA256Hash(entry.entryName + storageKey);
-  return `/${entry.path}/${hashedKey}`;
+  return `/${feature}/${hashedKey}`;
 }
diff --git a/packages/profile-sync-controller/src/controllers/user-storage/services.test.ts b/packages/profile-sync-controller/src/controllers/user-storage/services.test.ts
@@ -14,7 +14,7 @@ describe('user-storage/services.ts - getUserStorage() tests', () => {
   const actCallGetUserStorage = () => {
     return getUserStorage({
       bearerToken: 'MOCK_BEARER_TOKEN',
-      entryKey: 'notificationSettings',
+      path: 'notifications.notificationSettings',
       storageKey: MOCK_STORAGE_KEY,
     });
   };
@@ -63,7 +63,7 @@ describe('user-storage/services.ts - upsertUserStorage() tests', () => {
   const actCallUpsertUserStorage = () => {
     return upsertUserStorage(MOCK_ENCRYPTED_STORAGE_DATA, {
       bearerToken: 'MOCK_BEARER_TOKEN',
-      entryKey: 'notificationSettings',
+      path: 'notifications.notificationSettings',
       storageKey: MOCK_STORAGE_KEY,
     });
   };

diff --git a/packages/profile-sync-controller/src/controllers/user-storage/services.ts b/packages/profile-sync-controller/src/controllers/user-storage/services.ts
@@ -2,7 +2,7 @@ import log from 'loglevel';
 
 import { Env, getEnvUrls } from '../../sdk';
 import encryption from './encryption';
-import type { UserStorageEntryKeys } from './schema';
+import type { UserStoragePath } from './schema';
 import { createEntryPath } from './schema';
 
 const ENV_URLS = getEnvUrls(Env.PRD);
@@ -21,8 +21,8 @@ export type GetUserStorageResponse = {
 };
 
 export type UserStorageOptions = {
+  path: UserStoragePath;
   bearerToken: string;
-  entryKey: UserStorageEntryKeys;
   storageKey: string;
 };
 
@@ -36,13 +36,15 @@ export async function getUserStorage(
   opts: UserStorageOptions,
 ): Promise<string | null> {
   try {
-    const path = createEntryPath(opts.entryKey, opts.storageKey);
-    const url = new URL(`${USER_STORAGE_ENDPOINT}${path}`);
+    const { bearerToken, path, storageKey } = opts;
+
+    const encryptedPath = createEntryPath(path, storageKey);
+    const url = new URL(`${USER_STORAGE_ENDPOINT}${encryptedPath}`);
 
     const userStorageResponse = await fetch(url.toString(), {
       headers: {
         'Content-Type': 'application/json',
-        Authorization: `Bearer ${opts.bearerToken}`,
+        Authorization: `Bearer ${bearerToken}`,
       },
     });
 
@@ -85,15 +87,17 @@ export async function upsertUserStorage(
   data: string,
   opts: UserStorageOptions,
 ): Promise<void> {
+  const { bearerToken, path, storageKey } = opts;
+
   const encryptedData = encryption.encryptString(data, opts.storageKey);
-  const path = createEntryPath(opts.entryKey, opts.storageKey);
-  const url = new URL(`${USER_STORAGE_ENDPOINT}${path}`);
+  const encryptedPath = createEntryPath(path, storageKey);
+  const url = new URL(`${USER_STORAGE_ENDPOINT}${encryptedPath}`);
 
   const res = await fetch(url.toString(), {
     method: 'PUT',
     headers: {
       'Content-Type': 'application/json',
-      Authorization: `Bearer ${opts.bearerToken}`,
+      Authorization: `Bearer ${bearerToken}`,
     },
     body: JSON.stringify({ data: encryptedData }),
   });

diff --git a/packages/profile-sync-controller/src/sdk/__fixtures__/mock-userstorage.ts b/packages/profile-sync-controller/src/sdk/__fixtures__/mock-userstorage.ts
@@ -10,7 +10,10 @@ type MockReply = {
 };
 
 // Example mock notifications storage entry (wildcard)
-const MOCK_STORAGE_URL = STORAGE_URL(Env.DEV, 'notifications', '');
+const MOCK_STORAGE_URL = STORAGE_URL(
+  Env.DEV,
+  'notifications/notificationSettings',
+);
 
 export const MOCK_STORAGE_KEY = 'MOCK_STORAGE_KEY';
 // TODO: Either fix this lint violation or explain why it's necessary to ignore.