refactor: Refactor initialisation of Snaps controllers

app/scripts/controller-init/controller-list.ts

@@ -11,6 +11,18 @@ import SmartTransactionsController from '@metamask/smart-transactions-controller
 import { TransactionController } from '@metamask/transaction-controller';
 import { TransactionUpdateController } from '@metamask-institutional/transaction-update';
 import { AccountsController } from '@metamask/accounts-controller';
+import {
+  CronjobController,
+  ExecutionService,
+  JsonSnapsRegistry,
+  SnapController,
+  SnapInsightsController,
+  SnapInterfaceController,
+} from '@metamask/snaps-controllers';
+import {
+  RateLimitController,
+  RateLimitedApiMap,
+} from '@metamask/rate-limit-controller';
 import OnboardingController from '../controllers/onboarding';
 import { PreferencesController } from '../controllers/preferences-controller';
 import SwapsController from '../controllers/swaps';
@@ -19,7 +31,10 @@ import SwapsController from '../controllers/swaps';
  * Union of all controllers supporting or required by modular initialization.
  */
 export type Controller =
+  | CronjobController
+  | ExecutionService
   | GasFeeController
+  | JsonSnapsRegistry
   | KeyringController
   | NetworkController
   | OnboardingController
@@ -29,7 +44,11 @@ export type Controller =
     >
   | PPOMController
   | PreferencesController
+  | RateLimitController<RateLimitedApiMap>
   | SmartTransactionsController
+  | SnapController
+  | SnapInterfaceController
+  | SnapInsightsController
   | TransactionController
   | (TransactionUpdateController & {
       name: 'TransactionUpdateController';
@@ -41,7 +60,9 @@ export type Controller =
  * e.g. `{ transactions: [] }`.
  */
 export type ControllerFlatState = AccountsController['state'] &
+  CronjobController['state'] &
   GasFeeController['state'] &
+  JsonSnapsRegistry['state'] &
   KeyringController['state'] &
   NetworkController['state'] &
   OnboardingController['state'] &
@@ -51,6 +72,10 @@ export type ControllerFlatState = AccountsController['state'] &
   >['state'] &
   PPOMController['state'] &
   PreferencesController['state'] &
+  RateLimitController<RateLimitedApiMap>['state'] &
   SmartTransactionsController['state'] &
+  SnapController['state'] &
+  SnapInsightsController['state'] &
+  SnapInterfaceController['state'] &
   TransactionController['state'] &
   SwapsController['state'];

app/scripts/controller-init/messengers/index.ts

@@ -1,3 +1,14 @@
+import {
+  getCronjobControllerMessenger,
+  getExecutionServiceMessenger,
+  getRateLimitControllerInitMessenger,
+  getRateLimitControllerMessenger,
+  getSnapControllerInitMessenger,
+  getSnapControllerMessenger,
+  getSnapInsightsControllerMessenger,
+  getSnapInterfaceControllerMessenger,
+  getSnapsRegistryMessenger,
+} from '../snaps';
 import {
   getPPOMControllerMessenger,
   getPPOMControllerInitMessenger,
@@ -8,6 +19,29 @@ import {
 } from './transaction-controller-messenger';
 
 export const CONTROLLER_MESSENGERS = {
+  CronjobController: {
+    getMessenger: getCronjobControllerMessenger,
+  },
+  ExecutionService: {
+    getMessenger: getExecutionServiceMessenger,
+  },
+  RateLimitController: {
+    getMessenger: getRateLimitControllerMessenger,
+    getInitMessenger: getRateLimitControllerInitMessenger,
+  },
+  SnapsRegistry: {
+    getMessenger: getSnapsRegistryMessenger,
+  },
+  SnapController: {
+    getMessenger: getSnapControllerMessenger,
+    getInitMessenger: getSnapControllerInitMessenger,
+  },
+  SnapInsightsController: {
+    getMessenger: getSnapInsightsControllerMessenger,
+  },
+  SnapInterfaceController: {
+    getMessenger: getSnapInterfaceControllerMessenger,
+  },
   PPOMController: {
     getMessenger: getPPOMControllerMessenger,
     getInitMessenger: getPPOMControllerInitMessenger,

app/scripts/controller-init/snaps/cronjob-controller-init.test.ts

@@ -0,0 +1,34 @@
+import { CronjobController, GetAllSnaps } from '@metamask/snaps-controllers';
+import { Messenger } from '@metamask/base-controller';
+import { ControllerInitRequest } from '../types';
+import { buildControllerInitRequestMock } from '../test/utils';
+import { CronjobControllerInit } from './cronjob-controller-init';
+import {
+  CronjobControllerMessenger,
+  getCronjobControllerMessenger,
+} from './cronjob-controller-messenger';
+
+function getInitRequestMock(): jest.Mocked<
+  ControllerInitRequest<CronjobControllerMessenger>
+> {
+  const baseMessenger = new Messenger<GetAllSnaps, never>();
+
+  baseMessenger.registerActionHandler(
+    'SnapController:getAll',
+    jest.fn().mockReturnValue([]),
+  );
+
+  const requestMock = {
+    ...buildControllerInitRequestMock(),
+    controllerMessenger: getCronjobControllerMessenger(baseMessenger),
+  };
+
+  return requestMock;
+}
+
+describe('CronjobControllerInit', () => {
+  it('initializes the controller', () => {
+    const { controller } = CronjobControllerInit(getInitRequestMock());
+    expect(controller).toBeInstanceOf(CronjobController);
+  });
+});

app/scripts/controller-init/snaps/cronjob-controller-init.ts

@@ -0,0 +1,28 @@
+import { CronjobController } from '@metamask/snaps-controllers';
+import { ControllerInitFunction } from '../types';
+import { CronjobControllerMessenger } from './cronjob-controller-messenger';
+
+/**
+ * Initialize the cronjob controller.
+ *
+ * @param request - The request object.
+ * @param request.controllerMessenger - The messenger to use for the controller.
+ * @param request.persistedState - The persisted state of the extension.
+ * @returns The initialized controller.
+ */
+export const CronjobControllerInit: ControllerInitFunction<
+  CronjobController,
+  CronjobControllerMessenger
+> = ({ controllerMessenger, persistedState }) => {
+  const controller = new CronjobController({
+    // @ts-expect-error: `persistedState.CronjobController` is not compatible
+    // with the expected type.
+    // TODO: Look into the type mismatch.
+    state: persistedState.CronjobController,
+    messenger: controllerMessenger,
+  });
+
+  return {
+    controller,
+  };
+};

app/scripts/controller-init/snaps/cronjob-controller-messenger.test.ts

@@ -0,0 +1,12 @@
+import { Messenger, RestrictedMessenger } from '@metamask/base-controller';
+import { getCronjobControllerMessenger } from './cronjob-controller-messenger';
+
+describe('getCronjobControllerMessenger', () => {
+  it('returns a restricted controller messenger', () => {
+    const controllerMessenger = new Messenger<never, never>();
+    const cronjobControllerMessenger =
+      getCronjobControllerMessenger(controllerMessenger);
+
+    expect(cronjobControllerMessenger).toBeInstanceOf(RestrictedMessenger);
+  });
+});

app/scripts/controller-init/snaps/cronjob-controller-messenger.ts

@@ -0,0 +1,51 @@
+import { Messenger } from '@metamask/base-controller';
+import {
+  SnapInstalled,
+  SnapUpdated,
+  SnapDisabled,
+  SnapEnabled,
+  SnapUninstalled,
+  HandleSnapRequest,
+  GetAllSnaps,
+} from '@metamask/snaps-controllers';
+import { GetPermissions } from '@metamask/permission-controller';
+
+type Actions = GetPermissions | HandleSnapRequest | GetAllSnaps;
+
+type Events =
+  | SnapInstalled
+  | SnapUpdated
+  | SnapUninstalled
+  | SnapEnabled
+  | SnapDisabled;
+
+export type CronjobControllerMessenger = ReturnType<
+  typeof getCronjobControllerMessenger
+>;
+
+/**
+ * Get a restricted messenger for the cronjob controller. This is scoped to the
+ * actions and events that the cronjob controller is allowed to handle.
+ *
+ * @param messenger - The controller messenger to restrict.
+ * @returns The restricted controller messenger.
+ */
+export function getCronjobControllerMessenger(
+  messenger: Messenger<Actions, Events>,
+) {
+  return messenger.getRestricted({
+    name: 'CronjobController',
+    allowedEvents: [
+      'SnapController:snapInstalled',
+      'SnapController:snapUpdated',
+      'SnapController:snapUninstalled',
+      'SnapController:snapEnabled',
+      'SnapController:snapDisabled',
+    ],
+    allowedActions: [
+      `PermissionController:getPermissions`,
+      'SnapController:handleRequest',
+      'SnapController:getAll',
+    ],
+  });
+}

app/scripts/controller-init/snaps/execution-service-init.test.ts

@@ -0,0 +1,65 @@
+import {
+  IframeExecutionService,
+  OffscreenExecutionService,
+} from '@metamask/snaps-controllers';
+import { Messenger } from '@metamask/base-controller';
+import { ControllerInitRequest } from '../types';
+import { buildControllerInitRequestMock } from '../test/utils';
+import {
+  ExecutionServiceMessenger,
+  getExecutionServiceMessenger,
+} from './execution-service-messenger';
+import { ExecutionServiceInit } from './execution-service-init';
+
+jest.mock('@metamask/snaps-controllers');
+jest.mock('../../../../shared/modules/mv3.utils', () => ({
+  isManifestV3: true,
+}));
+
+function getInitRequestMock(): jest.Mocked<
+  ControllerInitRequest<ExecutionServiceMessenger>
+> {
+  const baseMessenger = new Messenger<never, never>();
+
+  const requestMock = {
+    ...buildControllerInitRequestMock(),
+    controllerMessenger: getExecutionServiceMessenger(baseMessenger),
+  };
+
+  return requestMock;
+}
+
+describe('ExecutionServiceInit', () => {
+  it('initializes the iframe execution service if `chrome.offscreen` is not available', () => {
+    const { controller } = ExecutionServiceInit(getInitRequestMock());
+    expect(controller).toBeInstanceOf(IframeExecutionService);
+  });
+
+  it('initializes the offscreen execution service if `chrome.offscreen` is available', () => {
+    Object.defineProperty(global, 'chrome', {
+      value: {
+        offscreen: {},
+      },
+    });
+
+    const { controller } = ExecutionServiceInit(getInitRequestMock());
+    expect(controller).toBeInstanceOf(OffscreenExecutionService);
+  });
+
+  it('passes the proper arguments to the service', () => {
+    Object.defineProperty(global, 'chrome', {
+      value: {
+        offscreen: {},
+      },
+    });
+
+    ExecutionServiceInit(getInitRequestMock());
+
+    const controllerMock = jest.mocked(OffscreenExecutionService);
+    expect(controllerMock).toHaveBeenCalledWith({
+      messenger: expect.any(Object),
+      offscreenPromise: expect.any(Promise),
+      setupSnapProvider: expect.any(Function),
+    });
+  });
+});

app/scripts/controller-init/snaps/execution-service-init.ts

@@ -0,0 +1,77 @@
+import {
+  AbstractExecutionService,
+  IframeExecutionService,
+  OffscreenExecutionService,
+  ExecutionServiceArgs,
+} from '@metamask/snaps-controllers';
+import { assert } from '@metamask/utils';
+import { SubjectType } from '@metamask/permission-controller';
+import { Duplex } from 'readable-stream';
+import { ControllerInitFunction } from '../types';
+import { isManifestV3 } from '../../../../shared/modules/mv3.utils';
+import { ExecutionServiceMessenger } from './execution-service-messenger';
+
+/**
+ * Initialize the Snaps execution service.
+ *
+ * @param request - The request object.
+ * @param request.controllerMessenger - The messenger to use for the service.
+ * @param request.offscreenPromise - The promise that resolves when the
+ * offscreen document is ready.
+ * @param request.setupUntrustedCommunicationEip1193 - The setup function for
+ * EIP-1193 communication.
+ * @returns The initialized controller.
+ */
+export const ExecutionServiceInit: ControllerInitFunction<
+  AbstractExecutionService<unknown>,
+  ExecutionServiceMessenger
+> = ({
+  controllerMessenger,
+  offscreenPromise,
+  setupUntrustedCommunicationEip1193,
+}) => {
+  const useOffscreenDocument =
+    isManifestV3 &&
+    typeof chrome !== 'undefined' &&
+    typeof chrome.offscreen !== 'undefined';
+
+  /**
+   * Set up the EIP-1193 provider for the given Snap.
+   *
+   * @param snapId - The ID of the Snap.
+   * @param connectionStream - The stream to connect to the Snap.
+   */
+  function setupSnapProvider(snapId: string, connectionStream: Duplex) {
+    return setupUntrustedCommunicationEip1193({
+      connectionStream,
+      sender: { snapId },
+      subjectType: SubjectType.Snap,
+    });
+  }
+
+  const args: ExecutionServiceArgs = {
+    messenger: controllerMessenger,
+    setupSnapProvider,
+  };
+
+  if (useOffscreenDocument) {
+    return {
+      persistedStateKey: null,
+      controller: new OffscreenExecutionService({
+        ...args,
+        offscreenPromise,
+      }),
+    };
+  }
+
+  const iframeUrl = process.env.IFRAME_EXECUTION_ENVIRONMENT_URL;
+  assert(iframeUrl, 'Missing iframe URL.');
+
+  return {
+    persistedStateKey: null,
+    controller: new IframeExecutionService({
+      ...args,
+      iframeUrl: new URL(iframeUrl),
+    }),
+  };
+};

app/scripts/controller-init/snaps/execution-service-messenger.test.ts

@@ -0,0 +1,11 @@
+import { Messenger, RestrictedMessenger } from '@metamask/base-controller';
+import { getExecutionServiceMessenger } from './execution-service-messenger';
+
+describe('getExecutionServiceMessenger', () => {
+  it('returns a restricted messenger', () => {
+    const messenger = new Messenger<never, never>();
+    const executionServiceMessenger = getExecutionServiceMessenger(messenger);
+
+    expect(executionServiceMessenger).toBeInstanceOf(RestrictedMessenger);
+  });
+});