Preload encrypted state of preinstalled Snaps

Author: Mrtenz

packages/snaps-controllers/src/snaps/SnapController.ts

@@ -462,6 +462,11 @@ export type GetSnapFile = {
   handler: SnapController['getSnapFile'];
 };
 
+export type PreloadSnap = {
+  type: `${typeof controllerName}:preloadSnap`;
+  handler: SnapController['preloadSnap'];
+};
+
 export type SnapControllerGetStateAction = ControllerGetStateAction<
   typeof controllerName,
   SnapControllerState
@@ -489,7 +494,8 @@ export type SnapControllerActions =
   | RevokeDynamicPermissions
   | GetSnapFile
   | SnapControllerGetStateAction
-  | StopAllSnaps;
+  | StopAllSnaps
+  | PreloadSnap;
 
 // Controller Messenger Events
 
@@ -602,6 +608,11 @@ export type SnapControllerStateChangeEvent = ControllerStateChangeEvent<
   SnapControllerState
 >;
 
+type KeyringControllerUnlock = {
+  type: 'KeyringController:unlock';
+  payload: [];
+};
+
 type KeyringControllerLock = {
   type: 'KeyringController:lock';
   payload: [];
@@ -652,6 +663,7 @@ export type AllowedEvents =
   | ExecutionServiceEvents
   | SnapInstalled
   | SnapUpdated
+  | KeyringControllerUnlock
   | KeyringControllerLock;
 
 type SnapControllerMessenger = RestrictedMessenger<
@@ -1010,6 +1022,11 @@ export class SnapController extends BaseController<
       },
     );
 
+    this.messagingSystem.subscribe(
+      'KeyringController:unlock',
+      this.#handleUnlock.bind(this),
+    );
+
     this.messagingSystem.subscribe(
       'KeyringController:lock',
       this.#handleLock.bind(this),
@@ -1204,6 +1221,11 @@ export class SnapController extends BaseController<
       `${controllerName}:stopAllSnaps`,
       async (...args) => this.stopAllSnaps(...args),
     );
+
+    this.messagingSystem.registerActionHandler(
+      `${controllerName}:preloadSnap`,
+      async (...args) => this.preloadSnap(...args),
+    );
   }
 
   #handlePreinstalledSnaps(preinstalledSnaps: PreinstalledSnap[]) {
@@ -1569,6 +1591,24 @@ export class SnapController extends BaseController<
     });
   }
 
+  /**
+   * Preload the state of the given Snap. This is a no-op if the Snap is already
+   * loaded.
+   *
+   * @param snapId - The id of the Snap to preload.
+   */
+  async preloadSnap(snapId: SnapId): Promise<void> {
+    this.#assertCanUsePlatform();
+    const runtime = this.#getRuntimeExpect(snapId);
+
+    if (runtime.state !== undefined) {
+      return;
+    }
+
+    // Calling `getSnapState` caches the state in the Snap runtime data.
+    await this.getSnapState(snapId, true);
+  }
+
   /**
    * Starts the given snap. Throws an error if no such snap exists
    * or if it is already running.
@@ -4245,6 +4285,40 @@ export class SnapController extends BaseController<
     });
   }
 
+  /**
+   * Handle the `KeyringController:unlock` event.
+   *
+   * Currently this preloads all preinstalled Snaps.
+   */
+  #handleUnlock() {
+    if (this.#preinstalledSnaps) {
+      const promises = this.#preinstalledSnaps.map(async ({ snapId }) =>
+        this.preloadSnap(snapId)
+          .then(() => ({ snapId, status: 'fulfilled' }))
+          .catch((error) => ({ snapId, status: 'rejected', reason: error })),
+      );
+
+      Promise.all(promises)
+        .then((results) => {
+          const failedSnaps = results.filter(
+            (result) => result.status === 'rejected',
+          );
+
+          if (failedSnaps.length > 0) {
+            logWarning(
+              `Failed to preload ${failedSnaps.length} preinstalled Snap(s):`,
+              failedSnaps,
+            );
+          }
+        })
+        .catch((error) => {
+          // This should never happen since we catch all errors in the promises
+          // above, but just in case, we log it.
+          logError('An unknown error occurred while preloading Snaps.', error);
+        });
+    }
+  }
+
   /**
    * Handle the `KeyringController:lock` event.
    *