add logic to cleanup notification interfaces

MetaMask · Oct 25, 2024 · e005581 · e005581
1 parent 2ab88f0
commit e005581
Show file tree

Hide file tree

Showing 3 changed files with 854 additions and 23 deletions.
diff --git a/packages/snaps-controllers/package.json b/packages/snaps-controllers/package.json
@@ -83,6 +83,7 @@
     "@metamask/base-controller": "^6.0.2",
     "@metamask/json-rpc-engine": "^9.0.2",
     "@metamask/json-rpc-middleware-stream": "^8.0.2",
+    "@metamask/notification-services-controller": "^0.12.0",
     "@metamask/object-multiplex": "^2.0.0",
     "@metamask/permission-controller": "^11.0.0",
     "@metamask/phishing-controller": "^12.0.2",

diff --git a/packages/snaps-controllers/src/interface/SnapInterfaceController.ts b/packages/snaps-controllers/src/interface/SnapInterfaceController.ts
@@ -8,6 +8,11 @@ import type {
   ControllerStateChangeEvent,
 } from '@metamask/base-controller';
 import { BaseController } from '@metamask/base-controller';
+import type {
+  INotification,
+  NotificationListUpdatedEvent,
+} from '@metamask/notification-services-controller/notification-services';
+import { TRIGGER_TYPES } from '@metamask/notification-services-controller/notification-services';
 import type {
   MaybeUpdateState,
   TestOrigin,
@@ -95,7 +100,8 @@ export type SnapInterfaceControllerStateChangeEvent =
   >;
 
 export type SnapInterfaceControllerEvents =
-  SnapInterfaceControllerStateChangeEvent;
+  | SnapInterfaceControllerStateChangeEvent
+  | NotificationListUpdatedEvent;
 
 export type SnapInterfaceControllerMessenger = RestrictedControllerMessenger<
   typeof controllerName,
@@ -155,6 +161,11 @@ export class SnapInterfaceController extends BaseController<
       state: { interfaces: {}, ...state },
     });
 
+    this.messagingSystem.subscribe(
+      'NotificationServicesController:notificationsListUpdated',
+      /* eslint-disable @typescript-eslint/unbound-method */
+      this._onNotificationsListUpdated,
+    );
     this.#registerMessageHandlers();
   }
 
@@ -217,8 +228,6 @@ export class SnapInterfaceController extends BaseController<
     const componentState = constructState({}, element);
 
     this.update((draftState) => {
-      // @ts-expect-error - TS2589: Type instantiation is excessively deep and
-      // possibly infinite.
       draftState.interfaces[id] = {
         snapId,
         content: castDraft(element),
@@ -413,4 +422,55 @@ export class SnapInterfaceController extends BaseController<
       (id: string) => this.messagingSystem.call('SnapController:get', id),
     );
   }
+
+  _onNotificationsListUpdated(notificationsList: INotification[]) {
+    const snapNotificationsWithInterface = notificationsList.filter(
+      (notification) => {
+        return (
+          notification.type === TRIGGER_TYPES.SNAP &&
+          // @ts-expect-error detailedView can be undefined here, type needs to be updated in the core repo
+          notification.data?.detailedView
+        );
+      },
+    );
+
+    const interfaceIdSet = new Set(
+      snapNotificationsWithInterface.map(
+        // @ts-expect-error detailedView can be undefined here, type needs to be updated in the core repo
+        (notification) => notification.data.detailedView.interfaceId,
+      ),
+    );
+
+    const updatedState = Object.entries(this.state.interfaces).reduce<
+      Record<string, StoredInterface>
+    >((newState, [id, snapInterface]) => {
+      if (snapInterface.contentType === ContentType.Notification) {
+        if (interfaceIdSet.has(id)) {
+          newState[id] = snapInterface;
+        }
+      } else {
+        newState[id] = snapInterface;
+      }
+      return newState;
+    }, {});
+
+    this.update((state) => {
+      // @ts-expect-error - TS2589: Type instantiation is excessively deep and
+      // possibly infinite.
+      state.interfaces = updatedState;
+    });
+  }
+
+  /**
+   * Run controller teardown process and unsubscribe from notification service controller events.
+   */
+  destroy() {
+    super.destroy();
+
+    /* eslint-disable @typescript-eslint/unbound-method */
+    this.messagingSystem.unsubscribe(
+      'NotificationServicesController:notificationsListUpdated',
+      this._onNotificationsListUpdated,
+    );
+  }
 }