From 908c37b400d0cb7a99d75f3a0266a527ff80633f Mon Sep 17 00:00:00 2001
From: Hassan Malik <41640681+hmalik88@users.noreply.github.com>
Date: Fri, 30 Aug 2024 06:30:51 -0400
Subject: [PATCH] Add `stopAllSnaps` function to `SnapController` (#2674)

---
 packages/snaps-controllers/coverage.json      |  8 ++---
 .../src/snaps/SnapController.test.tsx         | 35 +++++++++++++++++++
 .../src/snaps/SnapController.ts               | 34 +++++++++++++++++-
 3 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/packages/snaps-controllers/coverage.json b/packages/snaps-controllers/coverage.json
index aa297f891c..34b7a7350f 100644
--- a/packages/snaps-controllers/coverage.json
+++ b/packages/snaps-controllers/coverage.json
@@ -1,6 +1,6 @@
 {
-  "branches": 92.59,
-  "functions": 96.91,
-  "lines": 98.01,
-  "statements": 97.71
+  "branches": 92.6,
+  "functions": 96.95,
+  "lines": 98.02,
+  "statements": 97.72
 }
diff --git a/packages/snaps-controllers/src/snaps/SnapController.test.tsx b/packages/snaps-controllers/src/snaps/SnapController.test.tsx
index cd755ab439..2a0b0381a8 100644
--- a/packages/snaps-controllers/src/snaps/SnapController.test.tsx
+++ b/packages/snaps-controllers/src/snaps/SnapController.test.tsx
@@ -9525,4 +9525,39 @@ describe('SnapController', () => {
       snapController.destroy();
     });
   });
+
+  describe('SnapController:stopAllSnaps', () => {
+    it('calls SnapController.stopAllSnaps()', async () => {
+      const messenger = getSnapControllerMessenger();
+      const mockSnap = getMockSnapData({
+        id: MOCK_SNAP_ID,
+        origin: MOCK_ORIGIN,
+      });
+      const mockSnap2 = getMockSnapData({
+        id: `${MOCK_SNAP_ID}2` as SnapId,
+        origin: MOCK_ORIGIN,
+      });
+
+      const snapController = getSnapController(
+        getSnapControllerOptions({
+          messenger,
+          state: {
+            snaps: getPersistedSnapsState(
+              mockSnap.stateObject,
+              mockSnap2.stateObject,
+            ),
+          },
+        }),
+      );
+
+      await snapController.startSnap(mockSnap.id);
+      await snapController.startSnap(mockSnap2.id);
+
+      await messenger.call('SnapController:stopAllSnaps');
+      expect(snapController.state.snaps[mockSnap.id].status).toBe('stopped');
+      expect(snapController.state.snaps[mockSnap2.id].status).toBe('stopped');
+
+      snapController.destroy();
+    });
+  });
 });
diff --git a/packages/snaps-controllers/src/snaps/SnapController.ts b/packages/snaps-controllers/src/snaps/SnapController.ts
index b4c03b1942..2597155230 100644
--- a/packages/snaps-controllers/src/snaps/SnapController.ts
+++ b/packages/snaps-controllers/src/snaps/SnapController.ts
@@ -368,6 +368,11 @@ export type GetAllSnaps = {
   handler: SnapController['getAllSnaps'];
 };
 
+export type StopAllSnaps = {
+  type: `${typeof controllerName}:stopAllSnaps`;
+  handler: SnapController['stopAllSnaps'];
+};
+
 export type IncrementActiveReferences = {
   type: `${typeof controllerName}:incrementActiveReferences`;
   handler: SnapController['incrementActiveReferences'];
@@ -428,7 +433,8 @@ export type SnapControllerActions =
   | DisconnectOrigin
   | RevokeDynamicPermissions
   | GetSnapFile
-  | SnapControllerGetStateAction;
+  | SnapControllerGetStateAction
+  | StopAllSnaps;
 
 // Controller Messenger Events
 
@@ -1105,6 +1111,11 @@ export class SnapController extends BaseController<
       `${controllerName}:getFile`,
       async (...args) => this.getSnapFile(...args),
     );
+
+    this.messagingSystem.registerActionHandler(
+      `${controllerName}:stopAllSnaps`,
+      async (...args) => this.stopAllSnaps(...args),
+    );
   }
 
   #handlePreinstalledSnaps(preinstalledSnaps: PreinstalledSnap[]) {
@@ -1551,6 +1562,27 @@ export class SnapController extends BaseController<
     }
   }
 
+  /**
+   * Stops all running snaps, removes all hooks, closes all connections, and
+   * terminates their workers.
+   *
+   * @param statusEvent - The Snap status event that caused the snap to be
+   * stopped.
+   */
+  public async stopAllSnaps(
+    statusEvent:
+      | SnapStatusEvents.Stop
+      | SnapStatusEvents.Crash = SnapStatusEvents.Stop,
+  ): Promise<void> {
+    const snaps = Object.values(this.state.snaps).filter((snap) =>
+      this.isRunning(snap.id),
+    );
+    const promises = snaps.map(async (snap) =>
+      this.stopSnap(snap.id, statusEvent),
+    );
+    await Promise.allSettled(promises);
+  }
+
   /**
    * Terminates the specified snap and emits the `snapTerminated` event.
    *