Add StorageService and handle preinstalled snap setup in init

Author: GuillaumeRx

packages/snaps-controllers/package.json

@@ -94,6 +94,7 @@
     "@metamask/snaps-rpc-methods": "workspace:^",
     "@metamask/snaps-sdk": "workspace:^",
     "@metamask/snaps-utils": "workspace:^",
+    "@metamask/storage-service": "^0.0.1",
     "@metamask/utils": "^11.8.1",
     "@xstate/fsm": "^2.0.0",
     "async-mutex": "^0.5.0",

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

@@ -111,6 +111,12 @@ import {
   OnAssetsConversionResponseStruct,
   OnAssetsMarketDataResponseStruct,
 } from '@metamask/snaps-utils';
+import type {
+  StorageServiceGetItemAction,
+  StorageServiceSetItemAction,
+  StorageServiceRemoveItemAction,
+  StorageServiceClearAction,
+} from '@metamask/storage-service';
 import type {
   Json,
   NonEmptyArray,
@@ -669,7 +675,11 @@ export type AllowedActions =
   | Update
   | ResolveVersion
   | CreateInterface
-  | GetInterface;
+  | GetInterface
+  | StorageServiceSetItemAction
+  | StorageServiceGetItemAction
+  | StorageServiceRemoveItemAction
+  | StorageServiceClearAction;
 
 export type AllowedEvents =
   | ExecutionServiceEvents
@@ -941,6 +951,8 @@ export class SnapController extends BaseController<
 
   readonly #ensureOnboardingComplete: () => Promise<void>;
 
+  readonly #controllerSetup = createDeferredPromise();
+
   constructor({
     closeAllConnections,
     messenger,
@@ -983,7 +995,6 @@ export class SnapController extends BaseController<
             return Object.values(snaps).reduce<Record<SnapId, Partial<Snap>>>(
               (acc, snap) => {
                 const snapCopy: Partial<Snap> = { ...snap };
-                delete snapCopy.sourceCode;
                 delete snapCopy.auxiliaryFiles;
                 acc[snap.id] = snapCopy;
                 return acc;
@@ -1107,10 +1118,6 @@ export class SnapController extends BaseController<
       this.#setupRuntime(snap.id),
     );
 
-    if (this.#preinstalledSnaps) {
-      this.#handlePreinstalledSnaps(this.#preinstalledSnaps);
-    }
-
     this.#trackSnapExport = throttleTracking(
       (snapId: SnapId, handler: string, success: boolean, origin: string) => {
         const snapMetadata = this.messenger.call(
@@ -1207,8 +1214,9 @@ export class SnapController extends BaseController<
    * actions.
    */
   #registerMessageHandlers(): void {
-    this.messenger.registerActionHandler(`${controllerName}:init`, (...args) =>
-      this.init(...args),
+    this.messenger.registerActionHandler(
+      `${controllerName}:init`,
+      async (...args) => this.init(...args),
     );
 
     this.messenger.registerActionHandler(
@@ -1323,14 +1331,20 @@ export class SnapController extends BaseController<
   /**
    * Initialise the SnapController.
    *
-   * Currently this method calls the `onStart` lifecycle hook for all
+   * Currently this method sets up the preinstalled snaps and calls the `onStart` lifecycle hook for all
    * runnable Snaps.
    */
-  init() {
+  async init() {
+    if (this.#preinstalledSnaps) {
+      await this.#handlePreinstalledSnaps(this.#preinstalledSnaps);
+    }
+
     this.#callLifecycleHooks(METAMASK_ORIGIN, HandlerType.OnStart);
+
+    this.#controllerSetup.resolve();
   }
 
-  #handlePreinstalledSnaps(preinstalledSnaps: PreinstalledSnap[]) {
+  async #handlePreinstalledSnaps(preinstalledSnaps: PreinstalledSnap[]) {
     for (const {
       snapId,
       manifest,
@@ -1402,7 +1416,7 @@ export class SnapController extends BaseController<
       };
 
       // Add snap to the SnapController state
-      this.#set({
+      await this.#set({
         id: snapId,
         origin: METAMASK_ORIGIN,
         files: filesObject,
@@ -1721,6 +1735,9 @@ export class SnapController extends BaseController<
     // Ensure the user has onboarded before allowing access to Snaps.
     await this.#ensureOnboardingComplete();
 
+    // Ensure the controller has finished setting up.
+    await this.#controllerSetup.promise;
+
     const flags = this.#getFeatureFlags();
     assert(
       flags.disableSnaps !== true,
@@ -1810,9 +1827,11 @@ export class SnapController extends BaseController<
       throw new Error(`Snap "${snapId}" is disabled.`);
     }
 
+    const sourceCode = await this.#getSourceCode(snapId);
+
     await this.#startSnap({
       snapId,
-      sourceCode: snap.sourceCode,
+      sourceCode,
     });
   }
 
@@ -2395,9 +2414,11 @@ export class SnapController extends BaseController<
     this.#snapsRuntimeData.clear();
     this.#rollbackSnapshots.clear();
 
+    await this.#clearStorageService();
+
     // We want to remove all snaps & permissions, except for preinstalled snaps
     if (this.#preinstalledSnaps) {
-      this.#handlePreinstalledSnaps(this.#preinstalledSnaps);
+      await this.#handlePreinstalledSnaps(this.#preinstalledSnaps);
     }
   }
 
@@ -2448,6 +2469,8 @@ export class SnapController extends BaseController<
           delete state.unencryptedSnapStates[snapId];
         });
 
+        await this.#removeSourceCode(snapId);
+
         // If the snap has been fully installed before, also emit snapUninstalled.
         if (snap.status !== SnapStatus.Installing) {
           this.messenger.publish(`SnapController:snapUninstalled`, truncated);
@@ -3096,7 +3119,7 @@ export class SnapController extends BaseController<
 
       this.#transition(snapId, SnapStatusEvents.Update);
 
-      this.#set({
+      await this.#set({
         origin,
         id: snapId,
         files: newSnap,
@@ -3351,7 +3374,7 @@ export class SnapController extends BaseController<
    * @param args - The add snap args.
    * @returns The resulting snap object.
    */
-  #set(args: SetSnapArgs): PersistedSnap {
+  async #set(args: SetSnapArgs): Promise<PersistedSnap> {
     const {
       id: snapId,
       origin,
@@ -3424,7 +3447,6 @@ export class SnapController extends BaseController<
       initialPermissions: manifest.result.initialPermissions,
       manifest: manifest.result,
       status: this.#statusMachine.config.initial as StatusStates['value'],
-      sourceCode,
       version,
       versionHistory,
       auxiliaryFiles,
@@ -3434,13 +3456,16 @@ export class SnapController extends BaseController<
     // If the snap was blocked, it isn't any longer
     delete snap.blockInformation;
 
+    await this.#setSourceCode(snapId, sourceCode);
+
     // store the snap back in state
     const { inversePatches } = this.update((state: any) => {
       state.snaps[snapId] = snap;
     });
 
     // checking for isUpdate here as this function is also used in
     // the install flow, we do not care to create snapshots for installs
+    // @TODO (guillaumerx): Find a way to add the sourceCode to the rollback snapshot
     if (isUpdate) {
       const rollbackSnapshot = this.#getRollbackSnapshot(snapId);
       if (rollbackSnapshot !== undefined) {
@@ -4642,4 +4667,53 @@ export class SnapController extends BaseController<
       runtime.state = undefined;
     }
   }
+
+  /**
+   * Retrieve the source code for a snap from storage.
+   *
+   * @param snapId - The snap ID.
+   * @returns The source code for the snap.
+   */
+  async #getSourceCode(snapId: SnapId) {
+    const { result } = await this.messenger.call(
+      'StorageService:getItem',
+      this.name,
+      snapId,
+    );
+
+    assert(result, `Source code for snap "${snapId}" not found.`);
+
+    return result as string;
+  }
+
+  /**
+   * Store the source code for a snap in storage.
+   *
+   * @param snapId - The snap ID.
+   * @param sourceCode - The source code for the snap.
+   */
+  async #setSourceCode(snapId: SnapId, sourceCode: string) {
+    await this.messenger.call(
+      'StorageService:setItem',
+      this.name,
+      snapId,
+      sourceCode,
+    );
+  }
+
+  /**
+   * Remove the source code for a snap from storage.
+   *
+   * @param snapId - The snap ID.
+   */
+  async #removeSourceCode(snapId: SnapId) {
+    await this.messenger.call('StorageService:removeItem', this.name, snapId);
+  }
+
+  /**
+   * Clear all snap source code from storage.
+   */
+  async #clearStorageService() {
+    await this.messenger.call('StorageService:clear', this.name);
+  }
 }

packages/snaps-simulation/src/methods/hooks/get-snap.ts

@@ -23,7 +23,6 @@ export function getGetSnapImplementation(preinstalled: boolean = true) {
       status: SnapStatus.Running,
       versionHistory: [],
       initialPermissions: {},
-      sourceCode: '',
       manifest: {
         version: '0.1.0' as SemVerVersion,
         proposedName: 'Test Snap',

packages/snaps-utils/src/snaps.ts

@@ -91,7 +91,9 @@ export type SnapAuxiliaryFile = {
  */
 export type SnapAuxilaryFile = SnapAuxiliaryFile;
 
-export type PersistedSnap = Snap;
+export type PersistedSnap = Snap & {
+  sourceCode: string;
+};
 
 /**
  * A Snap as it exists in {@link SnapController} state.
@@ -107,11 +109,6 @@ export type Snap = TruncatedSnap & {
    */
   initialPermissions: SnapPermissions;
 
-  /**
-   * The source code of the Snap.
-   */
-  sourceCode: string;
-
   /**
    * The Snap's manifest file.
    */

yarn.lock

@@ -4269,6 +4269,7 @@ __metadata:
     "@metamask/snaps-rpc-methods": "workspace:^"
     "@metamask/snaps-sdk": "workspace:^"
     "@metamask/snaps-utils": "workspace:^"
+    "@metamask/storage-service": "npm:^0.0.1"
     "@metamask/utils": "npm:^11.8.1"
     "@noble/hashes": "npm:^1.7.1"
     "@swc/core": "npm:1.11.31"
@@ -4681,6 +4682,16 @@ __metadata:
   languageName: unknown
   linkType: soft
 
+"@metamask/storage-service@npm:^0.0.1":
+  version: 0.0.1
+  resolution: "@metamask/storage-service@npm:0.0.1"
+  dependencies:
+    "@metamask/messenger": "npm:^0.3.0"
+    "@metamask/utils": "npm:^11.8.1"
+  checksum: 10/ba2443e4bab7a4ef64bf3e0a10403ef94d50225e5ae5931a6fd32a21bfa8e276916ab6dc377d2db30c15c50acaeaee74a3c0b418a563311da9f98b9172fba300
+  languageName: node
+  linkType: hard
+
 "@metamask/superstruct@npm:^3.1.0, @metamask/superstruct@npm:^3.2.1":
   version: 3.2.1
   resolution: "@metamask/superstruct@npm:3.2.1"