fix: properly handle termination of Snaps while executing (#2304)

Fixes an issue where terminating a Snap that is currently executing
(e.g. when removing it) would cause unexpected errors to be thrown.

Fixes #2300

---------

Co-authored-by: Maarten Zuidhoorn <maarten@zuidhoorn.com>

packages/snaps-controllers/coverage.json

@@ -1,6 +1,6 @@
 {
   "branches": 91.55,
-  "functions": 96.62,
-  "lines": 97.88,
-  "statements": 97.55
+  "functions": 96.67,
+  "lines": 97.89,
+  "statements": 97.57
 }

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

@@ -1603,6 +1603,113 @@ describe('SnapController', () => {
     await service.terminateAllSnaps();
   });
 
+  // This isn't stable in CI unfortunately
+  it.skip('throws if the Snap is terminated while executing', async () => {
+    const { manifest, sourceCode, svgIcon } =
+      await getMockSnapFilesWithUpdatedChecksum({
+        sourceCode: `
+      module.exports.onRpcRequest = () => {
+        return new Promise((resolve) => {});
+      };
+    `,
+      });
+
+    const [snapController] = getSnapControllerWithEES(
+      getSnapControllerWithEESOptions({
+        detectSnapLocation: loopbackDetect({
+          manifest,
+          files: [sourceCode, svgIcon as VirtualFile],
+        }),
+      }),
+    );
+
+    await snapController.installSnaps(MOCK_ORIGIN, {
+      [MOCK_SNAP_ID]: {},
+    });
+
+    const snap = snapController.getExpect(MOCK_SNAP_ID);
+
+    expect(snapController.state.snaps[snap.id].status).toBe('running');
+
+    const promise = snapController.handleRequest({
+      snapId: snap.id,
+      origin: 'foo.com',
+      handler: HandlerType.OnRpcRequest,
+      request: {
+        jsonrpc: '2.0',
+        method: 'test',
+        params: {},
+        id: 1,
+      },
+    });
+
+    const results = await Promise.allSettled([
+      snapController.removeSnap(snap.id),
+      promise,
+    ]);
+
+    expect(results[0].status).toBe('fulfilled');
+    expect(results[1].status).toBe('rejected');
+    expect((results[1] as PromiseRejectedResult).reason.message).toBe(
+      `The snap "${snap.id}" has been terminated during execution.`,
+    );
+
+    snapController.destroy();
+  });
+
+  it('throws if unresponsive Snap is terminated while executing', async () => {
+    const { manifest, sourceCode, svgIcon } =
+      await getMockSnapFilesWithUpdatedChecksum({
+        sourceCode: `
+      module.exports.onRpcRequest = () => {
+        while(true) {}
+      };
+    `,
+      });
+
+    const [snapController] = getSnapControllerWithEES(
+      getSnapControllerWithEESOptions({
+        detectSnapLocation: loopbackDetect({
+          manifest,
+          files: [sourceCode, svgIcon as VirtualFile],
+        }),
+      }),
+    );
+
+    await snapController.installSnaps(MOCK_ORIGIN, {
+      [MOCK_SNAP_ID]: {},
+    });
+
+    const snap = snapController.getExpect(MOCK_SNAP_ID);
+
+    expect(snapController.state.snaps[snap.id].status).toBe('running');
+
+    const promise = snapController.handleRequest({
+      snapId: snap.id,
+      origin: 'foo.com',
+      handler: HandlerType.OnRpcRequest,
+      request: {
+        jsonrpc: '2.0',
+        method: 'test',
+        params: {},
+        id: 1,
+      },
+    });
+
+    const results = await Promise.allSettled([
+      snapController.removeSnap(snap.id),
+      promise,
+    ]);
+
+    expect(results[0].status).toBe('fulfilled');
+    expect(results[1].status).toBe('rejected');
+    expect((results[1] as PromiseRejectedResult).reason.message).toBe(
+      `${snap.id} failed to respond to the request in time.`,
+    );
+
+    snapController.destroy();
+  });
+
   it('does not kill snaps with open sessions', async () => {
     const sourceCode = `
       module.exports.onRpcRequest = () => 'foo bar';

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

@@ -1466,6 +1466,19 @@ export class SnapController extends BaseController<
    */
   async #terminateSnap(snapId: SnapId) {
     await this.messagingSystem.call('ExecutionService:terminateSnap', snapId);
+
+    // Hack to give up execution for a bit to let gracefully terminating Snaps return.
+    await new Promise((resolve) => setTimeout(resolve, 1));
+
+    const runtime = this.#getRuntimeExpect(snapId);
+    // Unresponsive requests may still be timed, time them out.
+    runtime.pendingInboundRequests
+      .filter((pendingRequest) => pendingRequest.timer.status !== 'finished')
+      .forEach((pendingRequest) => pendingRequest.timer.finish());
+
+    // Hack to give up execution for a bit to let timed out requests return.
+    await new Promise((resolve) => setTimeout(resolve, 1));
+
     this.messagingSystem.publish(
       'SnapController:snapTerminated',
       this.getTruncatedExpect(snapId),
@@ -3097,17 +3110,25 @@ export class SnapController extends BaseController<
 
         await this.#assertSnapRpcRequestResult(snapId, handlerType, result);
 
-        return this.#transformSnapRpcRequestResult(snapId, handlerType, result);
+        const transformedResult = await this.#transformSnapRpcRequestResult(
+          snapId,
+          handlerType,
+          result,
+        );
+
+        this.#recordSnapRpcRequestFinish(snapId, request.id);
+
+        return transformedResult;
       } catch (error) {
+        // We flag the RPC request as finished early since termination may affect pending requests
+        this.#recordSnapRpcRequestFinish(snapId, request.id);
         const [jsonRpcError, handled] = unwrapError(error);
 
         if (!handled) {
           await this.stopSnap(snapId, SnapStatusEvents.Crash);
         }
 
         throw jsonRpcError;
-      } finally {
-        this.#recordSnapRpcRequestFinish(snapId, request.id);
       }
     };
 

packages/snaps-controllers/src/snaps/Timer.test.ts

@@ -44,6 +44,17 @@ describe('Timer', () => {
     expect(timer.remaining).toBe(1000);
   });
 
+  it('can finish', () => {
+    const timer = new Timer(1000);
+    const callback = jest.fn();
+    timer.start(callback);
+    timer.finish();
+    expect(callback).toHaveBeenCalled();
+    expect(timer.status).toBe('finished');
+    jest.advanceTimersByTime(1000);
+    expect(callback).toHaveBeenCalledTimes(1);
+  });
+
   it('works with +Infinity', () => {
     const timer = new Timer(Infinity);
     expect(timer.status).toBe('stopped');

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

@@ -56,6 +56,19 @@ export class Timer {
     this.onFinish(false);
   }
 
+  /**
+   * Marks the timer as finished prematurely and triggers the callback.
+   *
+   * @throws {@link Error}. If it wasn't running or paused.
+   */
+  finish() {
+    assert(
+      this.status !== 'finished',
+      new Error('Tried to finish a finished Timer.'),
+    );
+    this.onFinish(true);
+  }
+
   /**
    * Pauses a currently running timer, allowing it to resume later.
    *