From aa1912c2e9fe65dfc943618f5a802dbdf9083044 Mon Sep 17 00:00:00 2001
From: Prithpal Sooriya <prithpal.sooriya@consensys.net>
Date: Thu, 18 Jul 2024 15:51:37 +0100
Subject: [PATCH] chore: profile sync controller - add snap caching (#4532)

## Explanation

making multiple calls to snaps can cause an internal 429 too many
requests as snap requests are queued. This adds simple in-mem caching to
prevent calling the snap with the same values.

## References

[NOTIFY-847](https://consensyssoftware.atlassian.net/browse/NOTIFY-847)

## Changelog

### `@metamask/profile-sync-controller`

- **ADDED**: Internal in-memory caching to pre-installed snap calls
- **<CATEGORY>**: Your change here

## 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
---
 .../AuthenticationController.ts               | 32 +++++++++++++++----
 .../user-storage/UserStorageController.ts     | 16 ++++++++--
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts b/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts
index e2285c667f..5eed4c09b6 100644
--- a/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts
+++ b/packages/profile-sync-controller/src/controllers/authentication/AuthenticationController.ts
@@ -311,28 +311,48 @@ export default class AuthenticationController extends BaseController<
     return THIRTY_MIN_MS > diffMs;
   }
 
+  #_snapPublicKeyCache: string | undefined;
+
   /**
    * Returns the auth snap public key.
    *
    * @returns The snap public key.
    */
-  #snapGetPublicKey(): Promise<string> {
-    return this.messagingSystem.call(
+  async #snapGetPublicKey(): Promise<string> {
+    if (this.#_snapPublicKeyCache) {
+      return this.#_snapPublicKeyCache;
+    }
+
+    const result = (await this.messagingSystem.call(
       'SnapController:handleRequest',
       createSnapPublicKeyRequest(),
-    ) as Promise<string>;
+    )) as string;
+
+    this.#_snapPublicKeyCache = result;
+
+    return result;
   }
 
+  #_snapSignMessageCache: Record<`metamask:${string}`, string> = {};
+
   /**
    * Signs a specific message using an underlying auth snap.
    *
    * @param message - A specific tagged message to sign.
    * @returns A Signature created by the snap.
    */
-  #snapSignMessage(message: `metamask:${string}`): Promise<string> {
-    return this.messagingSystem.call(
+  async #snapSignMessage(message: `metamask:${string}`): Promise<string> {
+    if (this.#_snapSignMessageCache[message]) {
+      return this.#_snapSignMessageCache[message];
+    }
+
+    const result = (await this.messagingSystem.call(
       'SnapController:handleRequest',
       createSnapSignMessageRequest(message),
-    ) as Promise<string>;
+    )) as string;
+
+    this.#_snapSignMessageCache[message] = result;
+
+    return result;
   }
 }
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 3aa9e3ce77..c4829d89d2 100644
--- a/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts
+++ b/packages/profile-sync-controller/src/controllers/user-storage/UserStorageController.ts
@@ -373,17 +373,27 @@ export default class UserStorageController extends BaseController<
     return storageKey;
   }
 
+  #_snapSignMessageCache: Record<`metamask:${string}`, string> = {};
+
   /**
    * Signs a specific message using an underlying auth snap.
    *
    * @param message - A specific tagged message to sign.
    * @returns A Signature created by the snap.
    */
-  #snapSignMessage(message: `metamask:${string}`): Promise<string> {
-    return this.messagingSystem.call(
+  async #snapSignMessage(message: `metamask:${string}`): Promise<string> {
+    if (this.#_snapSignMessageCache[message]) {
+      return this.#_snapSignMessageCache[message];
+    }
+
+    const result = (await this.messagingSystem.call(
       'SnapController:handleRequest',
       createSnapSignMessageRequest(message),
-    ) as Promise<string>;
+    )) as string;
+
+    this.#_snapSignMessageCache[message] = result;
+
+    return result;
   }
 
   #setIsProfileSyncingUpdateLoading(