From 1adf53ce51eb35d8bb207fe315d44c7e2428cc22 Mon Sep 17 00:00:00 2001
From: Frederik Bolding <frederik.bolding@gmail.com>
Date: Fri, 18 Oct 2024 12:49:35 +0200
Subject: [PATCH] Short-circuit when cached encryption key already exists

---
 packages/snaps-controllers/coverage.json        |  2 +-
 .../src/snaps/SnapController.ts                 | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/packages/snaps-controllers/coverage.json b/packages/snaps-controllers/coverage.json
index cb15fa2386..58563e2d34 100644
--- a/packages/snaps-controllers/coverage.json
+++ b/packages/snaps-controllers/coverage.json
@@ -1,5 +1,5 @@
 {
-  "branches": 92.63,
+  "branches": 92.69,
   "functions": 96.65,
   "lines": 97.97,
   "statements": 97.67
diff --git a/packages/snaps-controllers/src/snaps/SnapController.ts b/packages/snaps-controllers/src/snaps/SnapController.ts
index aa34994e7a..1f850673d7 100644
--- a/packages/snaps-controllers/src/snaps/SnapController.ts
+++ b/packages/snaps-controllers/src/snaps/SnapController.ts
@@ -1734,6 +1734,17 @@ export class SnapController extends BaseController<
     return { key: encryptionKey, salt };
   }
 
+  /**
+   * Check if a given Snap has a cached encryption key stored in the runtime.
+   *
+   * @param snapId - The Snap ID.
+   * @returns True if the Snap has a cached encryption key, otherwise false.
+   */
+  #hasCachedEncryptionKey(snapId: SnapId) {
+    const runtime = this.#getRuntimeExpect(snapId);
+    return runtime.encryptionKey !== null && runtime.encryptionSalt !== null;
+  }
+
   /**
    * Decrypt the encrypted state for a given Snap.
    *
@@ -1748,7 +1759,11 @@ export class SnapController extends BaseController<
       // This lets us skip JSON validation.
       const parsed = JSON.parse(state) as EncryptionResult;
       const { salt, keyMetadata } = parsed;
-      const useCache = this.#encryptor.isVaultUpdated(state);
+
+      // We only cache encryption keys if they are already cached or if the encryption key is using the latest key derivation params.
+      const useCache =
+        this.#hasCachedEncryptionKey(snapId) ||
+        this.#encryptor.isVaultUpdated(state);
       const { key } = await this.#getSnapEncryptionKey({
         snapId,
         salt,