fix: Properly stop Snaps when clearing state (#3552)

Properly stop Snaps instead of only terminating their execution when
clearing state. This properly cleans up stuck requests that are not
terminated gracefully.

Author: FrederikBolding

packages/snaps-controllers/coverage.json

@@ -1,6 +1,6 @@
 {
   "branches": 95.22,
-  "functions": 98.72,
+  "functions": 98.49,
   "lines": 98.88,
   "statements": 98.71
 }

packages/snaps-controllers/src/snaps/SnapController.test.tsx

@@ -10211,10 +10211,18 @@ describe('SnapController', () => {
   });
 
   describe('clearState', () => {
-    it('clears the state and terminates running snaps', async () => {
+    it('clears the state, terminates running Snaps and cancels pending requests', async () => {
       const options = getSnapControllerWithEESOptions({
         state: {
-          snaps: getPersistedSnapsState(),
+          snaps: getPersistedSnapsState(
+            getPersistedSnapObject({
+              sourceCode: `
+      module.exports.onRpcRequest = () => {
+        while(true) {}
+      };
+    `,
+            }),
+          ),
           snapStates: {
             [MOCK_SNAP_ID]: JSON.stringify({ foo: 'bar' }),
           },
@@ -10231,19 +10239,39 @@ describe('SnapController', () => {
 
       expect(snapController.has(MOCK_SNAP_ID)).toBe(true);
 
+      const requestPromise = snapController.handleRequest({
+        snapId: MOCK_SNAP_ID,
+        origin: METAMASK_ORIGIN,
+        handler: HandlerType.OnRpcRequest,
+        request: {
+          method: 'foo',
+        },
+      });
+
+      await waitForStateChange(messenger);
+
+      expect(snapController.isRunning(MOCK_SNAP_ID)).toBe(true);
+
+      await new Promise((resolve) => setTimeout(resolve, 100));
+
       await snapController.clearState();
 
       expect(snapController.has(MOCK_SNAP_ID)).toBe(false);
 
       expect(callActionSpy).toHaveBeenCalledWith(
-        'ExecutionService:terminateAllSnaps',
+        'ExecutionService:terminateSnap',
+        MOCK_SNAP_ID,
       );
 
       expect(callActionSpy).toHaveBeenCalledWith(
         'PermissionController:revokeAllPermissions',
         MOCK_SNAP_ID,
       );
 
+      await expect(requestPromise).rejects.toThrow(
+        'npm:@metamask/example-snap was stopped and the request was cancelled. This is likely because the Snap crashed.',
+      );
+
       expect(snapController.state).toStrictEqual({
         snaps: {},
         snapStates: {},
@@ -10290,13 +10318,22 @@ describe('SnapController', () => {
       expect(snapController.has(MOCK_SNAP_ID)).toBe(true);
       expect(snapController.has(preinstalledSnapId)).toBe(true);
 
+      await snapController.startSnap(MOCK_SNAP_ID);
+      await snapController.startSnap(preinstalledSnapId);
+
       await snapController.clearState();
 
       expect(snapController.has(MOCK_SNAP_ID)).toBe(false);
       expect(snapController.has(preinstalledSnapId)).toBe(true);
 
       expect(callActionSpy).toHaveBeenCalledWith(
-        'ExecutionService:terminateAllSnaps',
+        'ExecutionService:terminateSnap',
+        MOCK_SNAP_ID,
+      );
+
+      expect(callActionSpy).toHaveBeenCalledWith(
+        'ExecutionService:terminateSnap',
+        preinstalledSnapId,
       );
 
       expect(callActionSpy).toHaveBeenCalledWith(

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

@@ -171,7 +171,6 @@ import type {
   ExecutionServiceEvents,
   HandleRpcRequestAction,
   SnapErrorJson,
-  TerminateAllSnapsAction,
   TerminateSnapAction,
 } from '../services';
 import type {
@@ -675,7 +674,6 @@ export type AllowedActions =
   | AddApprovalRequest
   | HandleRpcRequestAction
   | ExecuteSnapAction
-  | TerminateAllSnapsAction
   | TerminateSnapAction
   | UpdateCaveat
   | UpdateRequestState
@@ -2267,13 +2265,8 @@ export class SnapController extends BaseController<
    */
   async clearState() {
     const snapIds = Object.keys(this.state.snaps);
-    if (this.#closeAllConnections) {
-      snapIds.forEach((snapId) => {
-        this.#closeAllConnections?.(snapId);
-      });
-    }
 
-    await this.messagingSystem.call('ExecutionService:terminateAllSnaps');
+    await this.stopAllSnaps();
     snapIds.forEach((snapId) => this.#revokeAllSnapPermissions(snapId));
 
     this.update((state) => {