feat: handle primary SRP switch cases

packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.test.ts

@@ -88,15 +88,14 @@ describe('authentication/authentication-controller - performSignIn() tests', ()
   it('leverages the _snapSignMessageCache', async () => {
     const metametrics = createMockAuthMetaMetrics();
     const mockEndpoints = arrangeAuthAPIs();
-    const { messenger, mockSnapGetPublicKey, mockSnapSignMessage } =
+    const { messenger, mockSnapSignMessage } =
       createMockAuthenticationMessenger();
 
     const controller = new AuthenticationController({ messenger, metametrics });
 
     await controller.performSignIn();
     controller.performSignOut();
     await controller.performSignIn();
-    expect(mockSnapGetPublicKey).toHaveBeenCalledTimes(1);
     expect(mockSnapSignMessage).toHaveBeenCalledTimes(1);
     mockEndpoints.mockNonceUrl.done();
     mockEndpoints.mockSrpLoginUrl.done();

packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts

@@ -219,19 +219,20 @@ export default class AuthenticationController extends BaseController<
       return null;
     }
 
-    return {
-      ...this.state.sessionData,
-      profile: {
-        ...this.state.sessionData.profile,
-        metaMetricsId: await this.#metametrics.getMetaMetricsId(),
-      },
-    };
+    return this.state.sessionData;
   }
 
   async #setLoginResponseToState(loginResponse: LoginResponse) {
+    const metaMetricsId = await this.#metametrics.getMetaMetricsId();
     this.update((state) => {
       state.isSignedIn = true;
-      state.sessionData = loginResponse;
+      state.sessionData = {
+        ...loginResponse,
+        profile: {
+          ...loginResponse.profile,
+          metaMetricsId,
+        },
+      };
     });
   }
 
@@ -280,27 +281,19 @@ export default class AuthenticationController extends BaseController<
     return this.state.isSignedIn;
   }
 
-  #_snapPublicKeyCache: string | undefined;
-
   /**
    * Returns the auth snap public key.
    *
    * @returns The snap public key.
    */
   async #snapGetPublicKey(): Promise<string> {
-    if (this.#_snapPublicKeyCache) {
-      return this.#_snapPublicKeyCache;
-    }
-
     this.#assertIsUnlocked('#snapGetPublicKey');
 
     const result = (await this.messagingSystem.call(
       'SnapController:handleRequest',
       createSnapPublicKeyRequest(),
     )) as string;
 
-    this.#_snapPublicKeyCache = result;
-
     return result;
   }
 

packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts

@@ -310,7 +310,7 @@ export default class UserStorageController extends BaseController<
 
   #isUnlocked = false;
 
-  #storageKeyCache: string | null = null;
+  #storageKeyCache: Record<`metamask:${string}`, string> = {};
 
   readonly #keyringController = {
     setupLockedStateSubscriptions: () => {
@@ -377,9 +377,10 @@ export default class UserStorageController extends BaseController<
       },
       {
         storage: {
-          getStorageKey: async () => this.#storageKeyCache,
-          setStorageKey: async (key) => {
-            this.#storageKeyCache = key;
+          getStorageKey: async (message) =>
+            this.#storageKeyCache[message] ?? null,
+          setStorageKey: async (message, key) => {
+            this.#storageKeyCache[message] = key;
           },
         },
       },
@@ -596,8 +597,13 @@ export default class UserStorageController extends BaseController<
     return await this.#userStorage.getStorageKey();
   }
 
+  /**
+   * Flushes the storage key cache.
+   * CAUTION: This is only public for testing purposes.
+   * It should not be used in production code.
+   */
   public flushStorageKeyCache(): void {
-    this.#storageKeyCache = null;
+    this.#storageKeyCache = {};
   }
 
   #_snapSignMessageCache: Record<`metamask:${string}`, string> = {};

packages/profile-sync-controller/src/sdk/user-storage.test.ts

@@ -534,6 +534,24 @@ describe('User Storage', () => {
     );
     expect(mockAuthSignMessage).toHaveBeenCalled(); // SignMessage called since generating new key
   });
+
+  it('uses existing storage key (in storage)', async () => {
+    const { auth } = arrangeAuth('SRP', MOCK_SRP);
+    const { userStorage, mockGetStorageKey } = arrangeUserStorage(auth);
+    mockGetStorageKey.mockResolvedValue(MOCK_STORAGE_KEY);
+
+    const mockAuthSignMessage = jest
+      .spyOn(auth, 'signMessage')
+      .mockResolvedValue(MOCK_STORAGE_KEY);
+
+    handleMockUserStoragePut();
+
+    await userStorage.setItem(
+      `${USER_STORAGE_FEATURE_NAMES.notifications}.notification_settings`,
+      'some fake data',
+    );
+    expect(mockAuthSignMessage).not.toHaveBeenCalled(); // SignMessage not called since key already exists
+  });
 });
 
 /**

packages/profile-sync-controller/src/sdk/user-storage.ts

@@ -22,8 +22,8 @@ export type UserStorageConfig = {
 };
 
 export type StorageOptions = {
-  getStorageKey: () => Promise<string | null>;
-  setStorageKey: (val: string) => Promise<void>;
+  getStorageKey: (message: `metamask:${string}`) => Promise<string | null>;
+  setStorageKey: (message: `metamask:${string}`, val: string) => Promise<void>;
 };
 
 export type UserStorageOptions = {
@@ -110,17 +110,20 @@ export class UserStorage {
   }
 
   async getStorageKey(): Promise<string> {
-    const storageKey = await this.options.storage?.getStorageKey();
+    const userProfile = await this.config.auth.getUserProfile();
+    const message = `metamask:${userProfile.profileId}` as const;
+
+    const storageKey = await this.options.storage?.getStorageKey(message);
     if (storageKey) {
       return storageKey;
     }
 
-    const userProfile = await this.config.auth.getUserProfile();
-    const storageKeySignature = await this.config.auth.signMessage(
-      `metamask:${userProfile.profileId}`,
-    );
+    const storageKeySignature = await this.config.auth.signMessage(message);
     const hashedStorageKeySignature = createSHA256Hash(storageKeySignature);
-    await this.options.storage?.setStorageKey(hashedStorageKeySignature);
+    await this.options.storage?.setStorageKey(
+      message,
+      hashedStorageKeySignature,
+    );
     return hashedStorageKeySignature;
   }