diff --git a/src/worker/vm/imports.ts b/src/worker/vm/imports.ts index 5b0dc41e0..f65406a5f 100644 --- a/src/worker/vm/imports.ts +++ b/src/worker/vm/imports.ts @@ -1,4 +1,4 @@ -import { BigQuery } from '@google-cloud/bigquery' +import * as bigquery from '@google-cloud/bigquery' import * as contrib from '@posthog/plugin-contrib' import * as scaffold from '@posthog/plugin-scaffold' import * as AWS from 'aws-sdk' @@ -19,7 +19,7 @@ export const imports = { 'generic-pool': genericPool, 'node-fetch': fetch, 'snowflake-sdk': snowflake, - '@google-cloud/bigquery': { BigQuery }, + '@google-cloud/bigquery': bigquery, '@posthog/plugin-scaffold': scaffold, '@posthog/plugin-contrib': contrib, 'aws-sdk': AWS, diff --git a/src/worker/vm/transforms/promise-timeout.ts b/src/worker/vm/transforms/promise-timeout.ts index 3eac20053..be03699a7 100644 --- a/src/worker/vm/transforms/promise-timeout.ts +++ b/src/worker/vm/transforms/promise-timeout.ts @@ -31,7 +31,12 @@ export const promiseTimeout: PluginGen = () => ({ types: t }) => ({ exit(path: any) { const { node } = path if (node && !node[REPLACED]) { - const newAwait = t.awaitExpression(t.callExpression(t.identifier('__asyncGuard'), [node.argument])) + const newAwait = t.awaitExpression( + t.callExpression(t.identifier('__asyncGuard'), [ + node.argument, + node.argument.callee || node.argument, + ]) + ) ;(newAwait as any)[REPLACED] = true path.replaceWith(newAwait) } diff --git a/src/worker/vm/vm.ts b/src/worker/vm/vm.ts index 6a9c4cde3..73d399b6a 100644 --- a/src/worker/vm/vm.ts +++ b/src/worker/vm/vm.ts @@ -14,6 +14,16 @@ import { imports } from './imports' import { transformCode } from './transforms' import { upgradeExportEvents } from './upgrades/export-events' +export class TimeoutError extends Error { + name = 'TimeoutError' + caller?: string = undefined + + constructor(message: string, caller?: string) { + super(message) + this.caller = caller + } +} + export async function createPluginConfigVM( hub: Hub, pluginConfig: PluginConfig, // NB! might have team_id = 0 @@ -46,7 +56,7 @@ export async function createPluginConfigVM( // Creating this outside the vm (so not in a babel plugin for example) // because `setTimeout` is not available inside the vm... and we don't want to // make it available for now, as it makes it easier to create malicious code - const asyncGuard = async (promise: () => Promise) => { + const asyncGuard = async (promise: Promise, name?: string) => { const timeout = hub.TASK_TIMEOUT return await Promise.race([ promise, @@ -55,7 +65,7 @@ export async function createPluginConfigVM( const message = `Script execution timed out after promise waited for ${timeout} second${ timeout === 1 ? '' : 's' }` - reject(new Error(message)) + reject(new TimeoutError(message, `${name}`)) }, timeout * 1000) ), ]) @@ -122,8 +132,8 @@ export async function createPluginConfigVM( // export helpers function __getExported (key) { return exportDestinations.find(a => a[key])?.[key] }; - function __asyncFunctionGuard (func) { - return func ? function __innerAsyncGuard${pluginConfigIdentifier}(...args) { return __asyncGuard(func(...args)) } : func + function __asyncFunctionGuard (func, name) { + return func ? function __innerAsyncGuard${pluginConfigIdentifier}(...args) { return __asyncGuard(func(...args), name) } : func }; // inject the meta object + shareable 'global' to the end of each exported function @@ -149,12 +159,12 @@ export async function createPluginConfigVM( // export various functions const __methods = { - setupPlugin: __asyncFunctionGuard(__bindMeta('setupPlugin')), - teardownPlugin: __asyncFunctionGuard(__bindMeta('teardownPlugin')), - exportEvents: __asyncFunctionGuard(__bindMeta('exportEvents')), - onEvent: __asyncFunctionGuard(__bindMeta('onEvent')), - onSnapshot: __asyncFunctionGuard(__bindMeta('onSnapshot')), - processEvent: __asyncFunctionGuard(__bindMeta('processEvent')), + setupPlugin: __asyncFunctionGuard(__bindMeta('setupPlugin'), 'setupPlugin'), + teardownPlugin: __asyncFunctionGuard(__bindMeta('teardownPlugin'), 'teardownPlugin'), + exportEvents: __asyncFunctionGuard(__bindMeta('exportEvents'), 'exportEvents'), + onEvent: __asyncFunctionGuard(__bindMeta('onEvent'), 'onEvent'), + onSnapshot: __asyncFunctionGuard(__bindMeta('onSnapshot'), 'onSnapshot'), + processEvent: __asyncFunctionGuard(__bindMeta('processEvent'), 'processEvent'), }; const __tasks = { diff --git a/tests/postgres/vm.timeout.test.ts b/tests/postgres/vm.timeout.test.ts index 4aa99e111..70aa3c08b 100644 --- a/tests/postgres/vm.timeout.test.ts +++ b/tests/postgres/vm.timeout.test.ts @@ -1,6 +1,7 @@ import { Hub } from '../../src/types' import { createHub } from '../../src/utils/db/hub' -import { createPluginConfigVM } from '../../src/worker/vm/vm' +import { delay } from '../../src/utils/utils' +import { createPluginConfigVM, TimeoutError } from '../../src/worker/vm/vm' import { pluginConfig39 } from '../helpers/plugins' import { resetTestDatabase } from '../helpers/sql' @@ -224,14 +225,18 @@ describe('vm timeout tests', () => { const vm = await createPluginConfigVM(hub, pluginConfig39, indexJs) const date = new Date() let errorMessage = undefined + let caller = undefined try { await vm.methods.processEvent!({ ...defaultEvent }) } catch (e) { + expect(e).toBeInstanceOf(TimeoutError) errorMessage = e.message + caller = e.caller } expect(new Date().valueOf() - date.valueOf()).toBeGreaterThanOrEqual(1000) expect(new Date().valueOf() - date.valueOf()).toBeLessThan(4000) expect(errorMessage!).toEqual('Script execution timed out after promise waited for 1 second') + expect(caller).toEqual('processEvent') }) test('small promises and overriding async guard', async () => { diff --git a/tests/transforms.test.ts b/tests/transforms.test.ts index 771256de2..a62ad0a8c 100644 --- a/tests/transforms.test.ts +++ b/tests/transforms.test.ts @@ -30,7 +30,47 @@ describe('transformCode', () => { "use strict"; async function x() { - await __asyncGuard(console.log()); + await __asyncGuard(console.log(), console.log); + } + `) + }) + + it('attaches caller information to awaits', () => { + const rawCode = code` + async function x() { + await anotherAsyncFunction('arg1', 'arg2') + } + ` + + const transformedCode = transformCode(rawCode, hub) + + expect(transformedCode).toStrictEqual(code` + "use strict"; + + async function x() { + await __asyncGuard(anotherAsyncFunction('arg1', 'arg2'), anotherAsyncFunction); + } + `) + }) + + it('attaches caller information to awaits for anonymous functions', () => { + const rawCode = code` + async function x() { + await (async () => {console.log()}) + } + ` + + const transformedCode = transformCode(rawCode, hub) + + expect(transformedCode).toStrictEqual(code` + "use strict"; + + async function x() { + await __asyncGuard(async () => { + console.log(); + }, async () => { + console.log(); + }); } `) })