From 9cce6bec1502aaf86f7440758065be0f21a1b8c5 Mon Sep 17 00:00:00 2001
From: Ben White <ben@posthog.com>
Date: Thu, 20 Jun 2024 15:29:07 +0200
Subject: [PATCH] Fixes

---
 plugin-server/src/cdp/hog-executor.ts        | 16 ++++++++-
 plugin-server/tests/cdp/examples.ts          | 10 ++++--
 plugin-server/tests/cdp/hog-executor.test.ts | 37 +++++++++++++-------
 3 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/plugin-server/src/cdp/hog-executor.ts b/plugin-server/src/cdp/hog-executor.ts
index ac58af57b0ef8..f0a41a576edcc 100644
--- a/plugin-server/src/cdp/hog-executor.ts
+++ b/plugin-server/src/cdp/hog-executor.ts
@@ -16,7 +16,7 @@ import {
 import { convertToHogFunctionFilterGlobal } from './utils'
 
 const MAX_ASYNC_STEPS = 2
-
+const MAX_HOG_LOGS = 10
 const DEFAULT_TIMEOUT_MS = 100
 
 export const formatInput = (bytecode: any, globals: HogFunctionInvocation['globals']): any => {
@@ -240,6 +240,7 @@ export class HogExecutor {
             }
 
             try {
+                let hogLogs = 0
                 execRes = exec(state ?? hogFunction.bytecode, {
                     globals,
                     timeout: DEFAULT_TIMEOUT_MS, // TODO: Swap this to milliseconds when the package is updated
@@ -250,6 +251,19 @@ export class HogExecutor {
                     },
                     functions: {
                         print: (...args) => {
+                            hogLogs++
+                            if (hogLogs == MAX_HOG_LOGS) {
+                                addLog(
+                                    result,
+                                    'warn',
+                                    `Function exceeded maximum log entries. No more logs will be collected.`
+                                )
+                            }
+
+                            if (hogLogs >= MAX_HOG_LOGS) {
+                                return
+                            }
+
                             const message = args
                                 .map((arg) => (typeof arg !== 'string' ? JSON.stringify(arg) : arg))
                                 .join(', ')
diff --git a/plugin-server/tests/cdp/examples.ts b/plugin-server/tests/cdp/examples.ts
index 9684bf373477b..7303c0675c3bf 100644
--- a/plugin-server/tests/cdp/examples.ts
+++ b/plugin-server/tests/cdp/examples.ts
@@ -116,14 +116,20 @@ export const HOG_EXAMPLES: Record<string, Pick<HogFunctionType, 'hog' | 'bytecod
             35,
         ],
     },
-    long_function: {
-        hog: "fn fibonacci(number) {\n    if (number < 2) {\n        return number;\n    } else {\n        return fibonacci(number - 1) + fibonacci(number - 2);\n    }\n}\nprint(f'fib {fibonacci(30)}');",
+    malicious_function: {
+        hog: "fn fibonacci(number) {\n    print('I AM FIBONACCI')\n    if (number < 2) {\n        return number;\n    } else {\n        return fibonacci(number - 1) + fibonacci(number - 2);\n    }\n}\nprint(f'fib {fibonacci(30)}');",
         bytecode: [
             '_h',
             41,
             'fibonacci',
             1,
+            38,
             32,
+            'I AM FIBONACCI',
+            2,
+            'print',
+            1,
+            35,
             33,
             2,
             36,
diff --git a/plugin-server/tests/cdp/hog-executor.test.ts b/plugin-server/tests/cdp/hog-executor.test.ts
index cf74b321a81e2..49db36f109762 100644
--- a/plugin-server/tests/cdp/hog-executor.test.ts
+++ b/plugin-server/tests/cdp/hog-executor.test.ts
@@ -43,6 +43,7 @@ describe('Hog Executor', () => {
     const mockFunctionManager = {
         reloadAllHogFunctions: jest.fn(),
         getTeamHogFunctions: jest.fn(),
+        getTeamHogFunction: jest.fn(),
     }
 
     beforeEach(() => {
@@ -61,9 +62,8 @@ describe('Hog Executor', () => {
                 ...HOG_FILTERS_EXAMPLES.no_filters,
             })
 
-            mockFunctionManager.getTeamHogFunctions.mockReturnValue({
-                [hogFunction.id]: hogFunction,
-            })
+            mockFunctionManager.getTeamHogFunctions.mockReturnValue([hogFunction])
+            mockFunctionManager.getTeamHogFunction.mockReturnValue(hogFunction)
         })
 
         it('can parse incoming messages correctly', () => {
@@ -191,9 +191,7 @@ describe('Hog Executor', () => {
                 ...HOG_FILTERS_EXAMPLES.pageview_or_autocapture_filter,
             })
 
-            mockFunctionManager.getTeamHogFunctions.mockReturnValue({
-                [fn.id]: fn,
-            })
+            mockFunctionManager.getTeamHogFunctions.mockReturnValue([fn])
 
             const resultsShouldntMatch = executor.executeMatchingFunctions(createHogExecutionGlobals())
             expect(resultsShouldntMatch).toHaveLength(0)
@@ -220,9 +218,7 @@ describe('Hog Executor', () => {
                 ...HOG_FILTERS_EXAMPLES.no_filters,
             })
 
-            mockFunctionManager.getTeamHogFunctions.mockReturnValue({
-                [fn.id]: fn,
-            })
+            mockFunctionManager.getTeamHogFunctions.mockReturnValue([fn])
 
             // Simulate the recusive loop
             const results = executor.executeMatchingFunctions(createHogExecutionGlobals())
@@ -250,18 +246,33 @@ describe('Hog Executor', () => {
         })
         it('limits the execution time and exits appropriately', () => {
             const fn = createHogFunction({
-                ...HOG_EXAMPLES.long_function,
+                ...HOG_EXAMPLES.malicious_function,
                 ...HOG_INPUTS_EXAMPLES.simple_fetch,
                 ...HOG_FILTERS_EXAMPLES.no_filters,
             })
 
-            mockFunctionManager.getTeamHogFunctions.mockReturnValue({
-                [fn.id]: fn,
-            })
+            mockFunctionManager.getTeamHogFunctions.mockReturnValue([fn])
 
             const results = executor.executeMatchingFunctions(createHogExecutionGlobals())
             expect(results).toHaveLength(1)
             expect(results[0].error).toContain('Execution timed out after 0.1 seconds. Performed ')
+
+            expect(results[0].logs.map((log) => log.message)).toEqual([
+                'Executing function',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'I AM FIBONACCI',
+                'Function exceeded maximum log entries. No more logs will be collected.',
+                expect.stringContaining(
+                    'Error executing function: Error: Execution timed out after 0.1 seconds. Performed'
+                ),
+            ])
         })
     })
 })