[plugin-server] Allow whitelisted imports in babel (PostHog/plugin-se…

…rver#284)

* Allow whitelisted imports in babel

* add test for fetching via `import`

* Reword whitelisted → provided

* also support "require" and handle defaults nodejs style

* support requires

* fix test

Co-authored-by: Michael Matloka <[email protected]>

src/worker/vm/transforms/common.ts

@@ -3,4 +3,4 @@ import * as types from '@babel/types'
 
 import { PluginsServer } from '../../../types'
 
-export type PluginGen = (server: PluginsServer) => (param: { types: typeof types }) => PluginObj
+export type PluginGen = (server: PluginsServer, ...args: any[]) => (param: { types: typeof types }) => PluginObj

src/worker/vm/transforms/index.ts

@@ -3,10 +3,11 @@ import { transform } from '@babel/standalone'
 import { PluginsServer } from '../../../types'
 import { loopTimeout } from './loop-timeout'
 import { promiseTimeout } from './promise-timeout'
+import { replaceImports } from './replace-imports'
 
 const memoize: Record<string, string> = {}
 
-export function transformCode(rawCode: string, server: PluginsServer): string {
+export function transformCode(rawCode: string, server: PluginsServer, imports?: Record<string, any>): string {
     if (process.env.NODE_ENV === 'test' && memoize[rawCode]) {
         // Memoizing in tests for speed, not in production though due to reliability concerns
         return memoize[rawCode]
@@ -19,7 +20,7 @@ export function transformCode(rawCode: string, server: PluginsServer): string {
         configFile: false,
         filename: 'index.ts',
         presets: ['typescript', ['env', { targets: { node: process.versions.node } }]],
-        plugins: [loopTimeout(server), promiseTimeout(server)],
+        plugins: [replaceImports(server, imports), loopTimeout(server), promiseTimeout(server)],
     })
     if (!code) {
         throw new Error(`Babel transform gone wrong! Could not process the following code:\n${rawCode}`)

src/worker/vm/transforms/replace-imports.ts

@@ -0,0 +1,73 @@
+import { PluginsServer } from '../../../types'
+import { PluginGen } from './common'
+
+export const replaceImports: PluginGen = (server: PluginsServer, imports: Record<string, any> = {}) => ({
+    types: t,
+}) => ({
+    visitor: {
+        ImportDeclaration: {
+            exit(path: any) {
+                const { node } = path
+                const importSource = node.source.value
+                const importedVars = new Map<string, string>()
+
+                if (typeof imports[importSource] === 'undefined') {
+                    throw new Error(
+                        `Cannot import '${importSource}'! This package is not provided by PostHog in plugins.`
+                    )
+                }
+
+                for (const specifier of node.specifiers) {
+                    if (t.isImportSpecifier(specifier)) {
+                        if (t.isStringLiteral(specifier.imported)) {
+                            importedVars.set(specifier.local.name, specifier.imported.value)
+                        } else {
+                            importedVars.set(specifier.local.name, specifier.imported.name)
+                        }
+                    } else if (t.isImportDefaultSpecifier(specifier)) {
+                        importedVars.set(specifier.local.name, 'default')
+                    } else if (t.isImportNamespaceSpecifier(specifier)) {
+                        importedVars.set(specifier.local.name, 'default')
+                    }
+                }
+
+                path.replaceWith(
+                    t.variableDeclaration(
+                        'const',
+                        Array.from(importedVars.entries()).map(([varName, sourceName]) => {
+                            const importExpression = t.memberExpression(
+                                t.identifier('__pluginHostImports'),
+                                t.stringLiteral(importSource),
+                                true
+                            )
+                            return t.variableDeclarator(
+                                t.identifier(varName),
+                                sourceName === 'default'
+                                    ? importExpression
+                                    : t.memberExpression(importExpression, t.stringLiteral(sourceName), true)
+                            )
+                        })
+                    )
+                )
+            },
+        },
+        CallExpression: {
+            exit(path: any) {
+                const { node } = path
+                if (t.isIdentifier(node.callee) && node.callee.name === 'require' && node.arguments.length === 1) {
+                    const importSource = node.arguments[0].value
+
+                    if (typeof imports[importSource] === 'undefined') {
+                        throw new Error(
+                            `Cannot import '${importSource}'! This package is not provided by PostHog in plugins.`
+                        )
+                    }
+
+                    path.replaceWith(
+                        t.memberExpression(t.identifier('__pluginHostImports'), t.stringLiteral(importSource), true)
+                    )
+                }
+            },
+        },
+    },
+})

src/worker/vm/vm.ts

@@ -1,3 +1,4 @@
+import { BigQuery } from '@google-cloud/bigquery'
 import { randomBytes } from 'crypto'
 import fetch from 'node-fetch'
 import { VM } from 'vm2'
@@ -16,7 +17,11 @@ export async function createPluginConfigVM(
     pluginConfig: PluginConfig, // NB! might have team_id = 0
     indexJs: string
 ): Promise<PluginConfigVMReponse> {
-    const transformedCode = transformCode(indexJs, server)
+    const imports = {
+        'node-fetch': fetch,
+        '@google-cloud/bigquery': { BigQuery },
+    }
+    const transformedCode = transformCode(indexJs, server, imports)
 
     // Create virtual machine
     const vm = new VM({
@@ -32,6 +37,8 @@ export async function createPluginConfigVM(
     vm.freeze(fetch, 'fetch')
     vm.freeze(createGoogle(), 'google')
 
+    vm.freeze(imports, '__pluginHostImports')
+
     if (process.env.NODE_ENV === 'test') {
         vm.freeze(setTimeout, '__jestSetTimeout')
     }

tests/postgres/vm.test.ts

@@ -605,6 +605,49 @@ test('fetch', async () => {
     expect(event.properties).toEqual({ count: 2, query: 'bla', results: [true, true] })
 })
 
+test('fetch via import', async () => {
+    const indexJs = `
+        import importedFetch from 'node-fetch'
+        async function processEvent (event, meta) {
+            const response = await importedFetch('https://google.com/results.json?query=' + event.event)
+            event.properties = await response.json()
+            return event
+        }
+    `
+    await resetTestDatabase(indexJs)
+    const vm = await createPluginConfigVM(mockServer, pluginConfig39, indexJs)
+    const event: PluginEvent = {
+        ...defaultEvent,
+        event: 'fetched',
+    }
+
+    await vm.methods.processEvent(event)
+    expect(fetch).toHaveBeenCalledWith('https://google.com/results.json?query=fetched')
+
+    expect(event.properties).toEqual({ count: 2, query: 'bla', results: [true, true] })
+})
+
+test('fetch via require', async () => {
+    const indexJs = `
+        async function processEvent (event, meta) {
+            const response = await require('node-fetch')('https://google.com/results.json?query=' + event.event)
+            event.properties = await response.json()
+            return event
+        }
+    `
+    await resetTestDatabase(indexJs)
+    const vm = await createPluginConfigVM(mockServer, pluginConfig39, indexJs)
+    const event: PluginEvent = {
+        ...defaultEvent,
+        event: 'fetched',
+    }
+
+    await vm.methods.processEvent(event)
+    expect(fetch).toHaveBeenCalledWith('https://google.com/results.json?query=fetched')
+
+    expect(event.properties).toEqual({ count: 2, query: 'bla', results: [true, true] })
+})
+
 test('attachments', async () => {
     const indexJs = `
         async function processEvent (event, meta) {

tests/transforms.test.ts

@@ -149,4 +149,71 @@ describe('transformCode', () => {
             }));
         `)
     })
+
+    it('replaces imports', () => {
+        const rawCode = code`
+            import { bla, bla2, bla3 as bla4 } from 'node-fetch'
+            import fetch1 from 'node-fetch'
+            import * as fetch2 from 'node-fetch'
+            console.log(bla, bla2, bla4, fetch1, fetch2);
+        `
+
+        const transformedCode = transformCode(rawCode, server, { 'node-fetch': { bla: () => true } })
+
+        expect(transformedCode).toStrictEqual(code`
+            "use strict";
+
+            const bla = __pluginHostImports["node-fetch"]["bla"],
+                  bla2 = __pluginHostImports["node-fetch"]["bla2"],
+                  bla4 = __pluginHostImports["node-fetch"]["bla3"];
+            const fetch1 = __pluginHostImports["node-fetch"];
+            const fetch2 = __pluginHostImports["node-fetch"];
+            console.log(bla, bla2, bla4, fetch1, fetch2);
+        `)
+    })
+
+    it('only replaces provided imports', () => {
+        const rawCode = code`
+            import { kea } from 'kea'
+            console.log(kea)
+        `
+
+        expect(() => {
+            transformCode(rawCode, server, { 'node-fetch': { default: () => true } })
+        }).toThrow("/index.ts: Cannot import 'kea'! This package is not provided by PostHog in plugins.")
+    })
+
+    it('replaces requires', () => {
+        const rawCode = code`
+            const fetch = require('node-fetch')
+            const { BigQuery } = require('@google-cloud/bigquery')
+            console.log(fetch, BigQuery);
+        `
+
+        const transformedCode = transformCode(rawCode, server, {
+            'node-fetch': { bla: () => true },
+            '@google-cloud/bigquery': { BigQuery: () => true },
+        })
+
+        expect(transformedCode).toStrictEqual(code`
+            "use strict";
+
+            const fetch = __pluginHostImports["node-fetch"];
+            const {
+              BigQuery
+            } = __pluginHostImports["@google-cloud/bigquery"];
+            console.log(fetch, BigQuery);
+        `)
+    })
+
+    it('only replaces provided requires', () => {
+        const rawCode = code`
+            const { kea } = require('kea')
+            console.log(kea)
+        `
+
+        expect(() => {
+            transformCode(rawCode, server, { 'node-fetch': { default: () => true } })
+        }).toThrow("/index.ts: Cannot import 'kea'! This package is not provided by PostHog in plugins.")
+    })
 })