From b2bca575e1cda80ef9c2d1e097542d955632bb7a Mon Sep 17 00:00:00 2001
From: Michael Matloka <dev@twixes.com>
Date: Tue, 23 Mar 2021 13:49:21 +0100
Subject: [PATCH 1/2] Clean up extensions approach

---
 src/worker/vm/extensions/google.ts |  4 ++--
 src/worker/vm/extensions/index.ts  | 29 +++++++++++++++++++++++++++++
 src/worker/vm/vm.ts                | 15 ++++++---------
 3 files changed, 37 insertions(+), 11 deletions(-)
 create mode 100644 src/worker/vm/extensions/index.ts

diff --git a/src/worker/vm/extensions/google.ts b/src/worker/vm/extensions/google.ts
index c2faef4e4..9d5dd16d8 100644
--- a/src/worker/vm/extensions/google.ts
+++ b/src/worker/vm/extensions/google.ts
@@ -1,10 +1,10 @@
 import * as BigQuery from '@google-cloud/bigquery'
 
-type DummyCloud = {
+export interface DummyCloud {
     bigquery: typeof BigQuery
 }
 
-type DummyGoogle = {
+export interface DummyGoogle {
     cloud: DummyCloud
 }
 
diff --git a/src/worker/vm/extensions/index.ts b/src/worker/vm/extensions/index.ts
new file mode 100644
index 000000000..028a68c6e
--- /dev/null
+++ b/src/worker/vm/extensions/index.ts
@@ -0,0 +1,29 @@
+import { CacheExtension, ConsoleExtension, StorageExtension } from '@posthog/plugin-scaffold'
+import fetch, { RequestInfo, RequestInit, Response } from 'node-fetch'
+
+import { PluginConfig, PluginsServer } from '../../../types'
+import { createCache } from './cache'
+import { createConsole } from './console'
+import { createGoogle, DummyGoogle } from './google'
+import { createPosthog, DummyPostHog } from './posthog'
+import { createStorage } from './storage'
+
+export interface Extensions {
+    fetch: (url: RequestInfo, init?: RequestInit) => Promise<Response>
+    cache: CacheExtension
+    console: ConsoleExtension
+    google: DummyGoogle
+    posthog: DummyPostHog
+    storage: StorageExtension
+}
+
+export function createExtensions(server: PluginsServer, pluginConfig: PluginConfig): Extensions {
+    return {
+        fetch: async (url, init) => await fetch(url, init),
+        cache: createCache(server, pluginConfig.plugin_id, pluginConfig.team_id),
+        console: createConsole(),
+        google: createGoogle(),
+        posthog: createPosthog(server, pluginConfig),
+        storage: createStorage(server, pluginConfig),
+    }
+}
diff --git a/src/worker/vm/vm.ts b/src/worker/vm/vm.ts
index 25100f0ee..c51aada40 100644
--- a/src/worker/vm/vm.ts
+++ b/src/worker/vm/vm.ts
@@ -3,6 +3,7 @@ import fetch from 'node-fetch'
 import { VM } from 'vm2'
 
 import { PluginConfig, PluginConfigVMReponse, PluginsServer } from '../../types'
+import { createExtensions } from './extensions'
 import { createCache } from './extensions/cache'
 import { createConsole } from './extensions/console'
 import { createGoogle } from './extensions/google'
@@ -23,13 +24,10 @@ export async function createPluginConfigVM(
         sandbox: {},
     })
 
-    // Add PostHog utilities to virtual machine
-    vm.freeze(createConsole(), 'console')
-    vm.freeze(createPosthog(server, pluginConfig), 'posthog')
-
-    // Add non-PostHog utilities to virtual machine
-    vm.freeze(fetch, 'fetch')
-    vm.freeze(createGoogle(), 'google')
+    const extensions = createExtensions(server, pluginConfig)
+    for (const [key, extension] of Object.entries(extensions)) {
+        vm.freeze(extension, key) // Inject extensions as globals
+    }
 
     if (process.env.NODE_ENV === 'test') {
         vm.freeze(setTimeout, '__jestSetTimeout')
@@ -57,10 +55,9 @@ export async function createPluginConfigVM(
 
     vm.freeze(
         {
-            cache: createCache(server, pluginConfig.plugin_id, pluginConfig.team_id),
             config: pluginConfig.config,
             attachments: pluginConfig.attachments,
-            storage: createStorage(server, pluginConfig),
+            ...extensions, // Inject extensions into meta
         },
         '__pluginHostMeta'
     )

From 4dda1034730bd5f311a37b3c8b8ec998bb604b9d Mon Sep 17 00:00:00 2001
From: Michael Matloka <dev@twixes.com>
Date: Wed, 24 Mar 2021 15:37:08 +0100
Subject: [PATCH 2/2] Update tests

---
 tests/plugins.test.ts     | 4 ++++
 tests/postgres/vm.test.ts | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/plugins.test.ts b/tests/plugins.test.ts
index 11a4fb17b..bf1a39ff6 100644
--- a/tests/plugins.test.ts
+++ b/tests/plugins.test.ts
@@ -118,7 +118,11 @@ test('plugin meta has what it should have', async () => {
         'attachments',
         'cache',
         'config',
+        'console',
+        'fetch',
         'global',
+        'google',
+        'posthog',
         'storage',
     ])
     expect(returnedEvent!.properties!['attachments']).toEqual({
diff --git a/tests/postgres/vm.test.ts b/tests/postgres/vm.test.ts
index d5ca1e70a..c9692c072 100644
--- a/tests/postgres/vm.test.ts
+++ b/tests/postgres/vm.test.ts
@@ -600,7 +600,7 @@ test('fetch', async () => {
     }
 
     await vm.methods.processEvent(event)
-    expect(fetch).toHaveBeenCalledWith('https://google.com/results.json?query=fetched')
+    expect(fetch).toHaveBeenCalledWith('https://google.com/results.json?query=fetched', undefined)
 
     expect(event.properties).toEqual({ count: 2, query: 'bla', results: [true, true] })
 })