From b65492ae293cf5f01a755b3eb1ecf48b6c3d7a9d Mon Sep 17 00:00:00 2001
From: Juanfran <juanfran.ag@gmail.com>
Date: Fri, 2 Aug 2024 14:27:12 +0200
Subject: [PATCH] fix(plugins-runtime): prevent plugin execution after close

---
 .../src/lib/load-plugin.spec.ts               | 72 +++++++++----------
 libs/plugins-runtime/src/lib/load-plugin.ts   |  9 ++-
 2 files changed, 44 insertions(+), 37 deletions(-)

diff --git a/libs/plugins-runtime/src/lib/load-plugin.spec.ts b/libs/plugins-runtime/src/lib/load-plugin.spec.ts
index 19309c4..4fd340a 100644
--- a/libs/plugins-runtime/src/lib/load-plugin.spec.ts
+++ b/libs/plugins-runtime/src/lib/load-plugin.spec.ts
@@ -1,10 +1,9 @@
-import { describe, it, vi, expect, beforeEach, afterEach, Mock } from 'vitest';
+import { describe, it, vi, expect, beforeEach, afterEach } from 'vitest';
 import { loadPlugin, setContextBuilder } from './load-plugin.js';
 import { loadManifestCode } from './parse-manifest.js';
 import { createApi, themeChange } from './api/index.js';
 import type { PenpotContext, PenpotTheme } from '@penpot/plugin-types';
 import type { Manifest } from './models/manifest.model.js';
-import { ses } from './ses.js';
 
 vi.mock('./parse-manifest.js', () => ({
   loadManifestCode: vi.fn(),
@@ -15,17 +14,25 @@ vi.mock('./api/index.js', () => ({
   themeChange: vi.fn(),
 }));
 
+const evaluateMock = vi.fn();
+
 vi.mock('./ses.js', () => ({
   ses: {
     hardenIntrinsics: vi.fn().mockReturnValue(null),
-    createCompartment: vi.fn().mockReturnValue({
-      evaluate: vi.fn(),
+    createCompartment: vi.fn().mockImplementation((publicApi) => {
+      return {
+        evaluate: evaluateMock,
+        globalThis: publicApi,
+      };
     }),
     harden: vi.fn().mockImplementation((obj) => obj),
   },
 }));
 
 describe('loadPlugin', () => {
+  vi.spyOn(globalThis, 'clearTimeout');
+  vi.spyOn(globalThis.console, 'error').mockImplementation(() => {});
+
   let mockContext: PenpotContext;
   let manifest: Manifest = {
     pluginId: 'test-plugin',
@@ -54,8 +61,8 @@ describe('loadPlugin', () => {
       closePlugin: vi.fn(),
     } as unknown as ReturnType<typeof createApi>;
 
-    (createApi as Mock).mockReturnValue(mockApi);
-    (loadManifestCode as Mock).mockResolvedValue(
+    vi.mocked(createApi).mockReturnValue(mockApi);
+    vi.mocked(loadManifestCode).mockResolvedValue(
       'console.log("Plugin loaded");'
     );
     setContextBuilder(() => mockContext);
@@ -96,20 +103,6 @@ describe('loadPlugin', () => {
     expect(mockApi.closePlugin).toHaveBeenCalledTimes(1);
   });
 
-  it('should clear timeouts on plugin close', async () => {
-    await loadPlugin(manifest);
-
-    const clearTimeoutSpy = vi.spyOn(globalThis, 'clearTimeout');
-    const setTimeoutSpy = vi.spyOn(globalThis, 'setTimeout');
-
-    const timeoutCallback = vi.fn();
-    const timeoutId = setTimeout(timeoutCallback, 1000);
-    clearTimeout(timeoutId);
-
-    expect(clearTimeoutSpy).toHaveBeenCalledWith(timeoutId);
-    expect(setTimeoutSpy).toHaveBeenCalled();
-  });
-
   it('should remove finish event listener on plugin finish', async () => {
     await loadPlugin(manifest);
 
@@ -123,16 +116,6 @@ describe('loadPlugin', () => {
   });
 
   it('shoud clean setTimeout when plugin is closed', async () => {
-    vi.spyOn(globalThis, 'clearTimeout');
-
-    let closedCallback = () => {};
-
-    (createApi as Mock).mockImplementation((context, manifest, closed) => {
-      closedCallback = closed;
-
-      return mockApi;
-    });
-
     const plugin = await loadPlugin(manifest);
 
     if (!plugin) {
@@ -144,6 +127,7 @@ describe('loadPlugin', () => {
 
     expect(plugin.timeouts.size).toBe(2);
 
+    const closedCallback = vi.mocked(createApi).mock.calls[0][2];
     closedCallback();
 
     expect(plugin.timeouts.size).toBe(0);
@@ -151,16 +135,32 @@ describe('loadPlugin', () => {
   });
 
   it('should close plugin on evaluation error', async () => {
-    ses.createCompartment = vi.fn().mockImplementation(() => {
-      return {
-        evaluate: vi.fn().mockImplementation(() => {
-          throw new Error('Error in plugin');
-        }),
-      };
+    evaluateMock.mockImplementationOnce(() => {
+      throw new Error('Evaluation error');
     });
 
     await loadPlugin(manifest);
 
     expect(mockApi.closePlugin).toHaveBeenCalled();
+    expect(console.error).toHaveBeenCalled();
+  });
+
+  it('should prevent using the api after closing the plugin', async () => {
+    const plugin = await loadPlugin(manifest);
+
+    if (!plugin) {
+      throw new Error('Plugin not loaded');
+    }
+
+    expect(
+      Object.keys(plugin.compartment.globalThis).filter((it) => !!it).length
+    ).toBeGreaterThan(0);
+
+    const closedCallback = vi.mocked(createApi).mock.calls[0][2];
+    closedCallback();
+
+    expect(
+      Object.keys(plugin.compartment.globalThis).filter((it) => !!it).length
+    ).toBe(0);
   });
 });
diff --git a/libs/plugins-runtime/src/lib/load-plugin.ts b/libs/plugins-runtime/src/lib/load-plugin.ts
index 3ce9470..7a77383 100644
--- a/libs/plugins-runtime/src/lib/load-plugin.ts
+++ b/libs/plugins-runtime/src/lib/load-plugin.ts
@@ -44,8 +44,15 @@ export const loadPlugin = async function (manifest: Manifest) {
     }
 
     const pluginApi = createApi(context, manifest, () => {
+      createdApis = createdApis.filter((api) => api !== pluginApi);
+
       timeouts.forEach(clearTimeout);
       timeouts.clear();
+
+      // Remove all public API from globalThis
+      Object.keys(publicPluginApi).forEach((key) => {
+        delete c.globalThis[key];
+      });
     });
 
     createdApis.push(pluginApi);
@@ -53,7 +60,7 @@ export const loadPlugin = async function (manifest: Manifest) {
     const timeouts = new Set<ReturnType<typeof setTimeout>>();
 
     const publicPluginApi = {
-      penpot: ses.harden(pluginApi),
+      penpot: ses.harden(pluginApi) as typeof pluginApi,
       fetch: ses.harden((...args: Parameters<typeof fetch>) => {
         const requestArgs: RequestInit = {
           ...args[1],