Updated Ecma sandbox connector to validate main function existance in code

Author: zubairirfan96

packages/core/src/helpers/ECMASandbox.helper.ts

@@ -30,7 +30,6 @@ async function setupIsolate() {
         const context = await isolate.createContext();
         const jail = context.global;
         await jail.set('global', jail.derefInto());
-
         // Define a SafeBuffer object
         const ___internal = {
             b64decode: (str) => Buffer.from(str, 'base64').toString('utf8'),
@@ -62,45 +61,66 @@ export async function runJs(code: string) {
 
         if (!code.endsWith(';')) code += ';';
 
-        let scriptCode = '';
         const { isolate, context, jail } = await setupIsolate();
         const remoteUrls = await extractFetchUrls(code);
         for (const url of remoteUrls) {
             const remoteCode = await fetchCodeFromCDN(url);
-            context.eval(`${remoteCode}`);
+            await context.eval(`${remoteCode}`);
         }
-        const randomId = Math.random().toString(36).substring(2, 15);
-        const resId = `res${randomId}`;
-        scriptCode = `
-      var ${resId}; 
-      ${code};
-
-      ${resId} = JSON.stringify(_output); 
-      ${resId};
-    `;
-        const script: any = await isolate.compileScript(scriptCode).catch((err) => {
+
+        const executionCode = `
+    (async () => {
+        ${code}
+        globalThis.__finalResult = result;
+    })();
+        `;
+
+        // Execute the original code
+        const executeScript = await isolate.compileScript(executionCode).catch((err) => {
             console.error(err);
             return { error: 'Compile Error - ' + err.message };
         });
-        if (script?.error) {
-            throw new Error(script.error);
+        if ('error' in executeScript) {
+            throw new Error(executeScript.error);
         }
 
-        const rawResult = await script.run(context).catch((err) => {
+        await executeScript.run(context).catch((err) => {
             console.error(err);
-            return { error: 'Run Error - ' + err.message };
+            throw new Error('Run Error - ' + err.message);
         });
+
+        // Try to get the result from the global variable first, then fallback to 'result'
+        let rawResult = await context.eval('globalThis.__finalResult').catch((err) => {
+            console.error('Failed to get __finalResult:', err);
+            return null;
+        });
+
         if (rawResult?.error) {
             throw new Error(rawResult.error);
         }
-
-        // Transfer the result out of the isolate and parse it
-        //const serializedResult = rawResult.copySync();
-        const Output = JSON.parse(rawResult);
-
-        return Output;
+        return { Output: rawResult };
     } catch (error) {
         console.error(error);
         throw new Error(error.message);
     }
 }
+
+function getParametersString(parameters: string[], inputs: Record<string, any>) {
+    let params = [];
+    for (const parameter of parameters) {
+        if (typeof inputs[parameter] === 'string') {
+            params.push(`'${inputs[parameter]}'`);
+        } else {
+            params.push(`${inputs[parameter]};`);
+        }
+    }
+    return params.join(',');
+}
+
+export function generateExecutableCode(code: string, parameters: string[], inputs: Record<string, any>) {
+    const executableCode = `
+    ${code}
+        const result = await main(${getParametersString(parameters, inputs)});
+    `
+    return executableCode;
+}

packages/core/src/subsystems/ComputeManager/Code.service/connectors/ECMASandbox.class.ts

@@ -4,7 +4,8 @@ import { CodeConfig, CodePreparationResult, CodeConnector, CodeInput, CodeDeploy
 import { AccessRequest } from '@sre/Security/AccessControl/AccessRequest.class';
 import { Logger } from '@sre/helpers/Log.helper';
 import axios from 'axios';
-import { runJs } from '@sre/helpers/ECMASandbox.helper';
+import { generateExecutableCode, runJs } from '@sre/helpers/ECMASandbox.helper';
+import { validateAsyncMainFunction } from '@sre/helpers/AWSLambdaCode.helper';
 
 const console = Logger('ECMASandbox');
 export class ECMASandbox extends CodeConnector {
@@ -34,10 +35,20 @@ export class ECMASandbox extends CodeConnector {
 
     public async execute(acRequest: AccessRequest, codeUID: string, inputs: Record<string, any>, config: CodeConfig): Promise<CodeExecutionResult> {
         try {
+            const { isValid, error, parameters } = validateAsyncMainFunction(inputs.code);
+            if (!isValid) {
+                return {
+                    output: undefined,
+                    executionTime: 0,
+                    success: false,
+                    errors: [error],
+                }
+            }
+            const executableCode = generateExecutableCode(inputs.code, parameters, inputs.inputs);
             if (!this.sandboxUrl) {
                 // run js code in isolated vm
                 console.debug('Running code in isolated vm');
-                const result = await runJs(inputs.code);
+                const result = await runJs(executableCode);
                 console.debug(`Code result: ${result}`);
                 return {
                     output: result?.Output,
@@ -47,7 +58,7 @@ export class ECMASandbox extends CodeConnector {
                 };
             } else {
                 console.debug('Running code in remote sandbox');
-                const result: any = await axios.post(this.sandboxUrl, { code: inputs.code }).catch((error) => ({ error }));
+                const result: any = await axios.post(this.sandboxUrl, { code: executableCode }).catch((error) => ({ error }));
                 if (result.error) {
 
                     const error = result.error?.response?.data || result.error?.message || result.error.toString() || 'Unknown error';

packages/core/tests/unit/core/ecma-sandbox.test.ts

@@ -7,7 +7,7 @@ setupSRE({
     Code: {
         Connector: 'ECMASandbox',
         Settings: {
-            sandboxUrl: 'http://localhost:6100/run-js',
+            sandboxUrl: 'http://localhost:6100/run-js/v2',
         },
     },
     Log: {
@@ -23,16 +23,36 @@ describe('ECMASandbox Tests', () => {
                 id: 'test-user',
                 role: TAccessRole.User,
             };
-            
+
             const codeConnector = ConnectorService.getCodeConnector('ECMASandbox');
             const result = await codeConnector.agent(mockCandidate.id).execute(Date.now().toString(), {
-                code: 'let _output=undefined;\nconsole.log("Hello, world!");\n_output=1;',
-            }, {});
+                code: `async function main(prompt) { return prompt + ' ' + 'Hello World'; }`,
+                inputs: {
+                    prompt: 'Say'
+                }
+            });
 
             const output = result.output;
-
-            expect(output).toBe(1);
+            expect(output).toBe('Say Hello World');
         },
     );
+    it(
+        'Try to run a simple code without main function',
+        async () => {
+            const mockCandidate: IAccessCandidate = {
+                id: 'test-user',
+                role: TAccessRole.User,
+            };
 
+            const codeConnector = ConnectorService.getCodeConnector('ECMASandbox');
+            const result = await codeConnector.agent(mockCandidate.id).execute(Date.now().toString(), {
+                code: `async function testFunction(prompt) { return prompt + ' ' + 'Hello World'; }`,
+                inputs: {
+                    prompt: 'Say'
+                }
+            });
+            const error = result.errors;
+            expect(error).toContain('No main function found at root level');
+        },
+    );
 });