refactor: improve preprocessing expression

Author: danpoletaev

packages/utilities/src/memory_calculator.ts

@@ -96,32 +96,49 @@ const roundToClosestPowerOf2 = (num: number): number | undefined => {
 };
 
 /**
- * Replaces `{{variable}}` placeholders in an expression string with `runOptions.variable`.
- *
- * This function also validates that the variable is one of the allowed 'runOptions' keys.
+ * Replaces `{{variable}}` placeholders in an expression string with the variable name.
+ * Enforces strict validation to only allow `input.*` paths or whitelisted `runOptions.*` keys.
  *
  * @example
  * // Returns "runOptions.memoryMbytes + 1024"
- * preprocessDefaultMemoryExpression("{{memoryMbytes}} + 1024");
+ * preprocessDefaultMemoryExpression("{{runOptions.memoryMbytes}} + 1024");
  *
- * @param defaultMemoryMbytes The raw string expression, e.g., "{{memoryMbytes}} * 2".
+ * @param defaultMemoryMbytes The raw string expression, e.g., "{{runOptions.memoryMbytes}} * 2".
  * @returns A safe, processed expression for evaluation, e.g., "runOptions.memoryMbytes * 2".
  */
 const preprocessDefaultMemoryExpression = (defaultMemoryMbytes: string): string => {
-    // This regex captures the variable name inside {{...}}
-    const variableRegex = /{{\s*([a-zA-Z0-9_]+)\s*}}/g;
+    const variableRegex = /{{\s*([a-zA-Z0-9_.]+)\s*}}/g;
 
     const processedExpression = defaultMemoryMbytes.replace(
         variableRegex,
         (_, variableName: string) => {
-            // Check if the captured variable name is in our allowlist
-            if (!ALLOWED_RUN_OPTION_KEYS.has(variableName as keyof ActorRunOptions)) {
+            // 1. Validate that the variable starts with either 'input.' or 'runOptions.'
+            if (!variableName.startsWith('runOptions.') && !variableName.startsWith('input.')) {
                 throw new Error(
-                    `Invalid variable '{{${variableName}}}' in expression.`,
+                    `Invalid variable '{{${variableName}}}' in expression. Variables must start with 'input.' or 'runOptions.'.`,
                 );
             }
 
-            return `runOptions.${variableName}`;
+            // 2. Check if the variable is accessing Input (e.g. {{input.someValue}})
+            // We do not validate the specific property name because input is dynamic.
+            if (variableName.startsWith('input.')) {
+                return variableName;
+            }
+
+            if (variableName.startsWith('runOptions.')) {
+                const key = variableName.slice('runOptions.'.length);
+                if (!ALLOWED_RUN_OPTION_KEYS.has(key as keyof ActorRunOptions)) {
+                    throw new Error(
+                        `Invalid variable '{{${variableName}}}' in expression. Only the following runOptions are allowed: ${Array.from(ALLOWED_RUN_OPTION_KEYS).map((k) => `runOptions.${k}`).join(', ')}.`,
+                    );
+                }
+                return variableName;
+            }
+
+            // 3. Throw error for unrecognized variables (e.g. {{process.env}})
+            throw new Error(
+                `Invalid variable '{{${variableName}}}' in expression.`,
+            );
         },
     );
 

test/memory_calculator.test.ts

@@ -56,19 +56,35 @@ describe('calculateDefaultMemoryFromExpression', () => {
     });
 
     describe('Preprocessing with {{variable}}', () => {
-        it('correctly replaces {{variable}} with valid runOptions.variable', () => {
+        it('should throw error if variable doesn\'t start with .runOptions or .input', () => {
             const context = { input: {}, runOptions: { memoryMbytes: 16 } };
-            const expr = '{{memoryMbytes}} * 1024';
+            const expr = '{{unexistingVariable}} * 1024';
+            expect(() => calculateDefaultMemoryFromExpression(expr, context))
+                .toThrow(`Invalid variable '{{unexistingVariable}}' in expression. Variables must start with 'input.' or 'runOptions.'.`);
+        });
+
+        it('correctly replaces {{runOptions.variable}} with valid runOptions.variable', () => {
+            const context = { input: {}, runOptions: { memoryMbytes: 16 } };
+            const expr = '{{runOptions.memoryMbytes}} * 1024';
             // 16 * 1024 = 16384, which is 2^14
             const result = calculateDefaultMemoryFromExpression(expr, context);
             expect(result).toBe(16384);
         });
 
-        it('should throw error for invalid variable in {{variable}} syntax', () => {
-            const context = { input: {}, runOptions: { memoryMbytes: 16 } };
-            const expr = '{{unexistingVariable}} * 1024';
+        it('correctly replaces {{input.variable}} with valid input.variable', () => {
+            const context = { input: { value: 16 }, runOptions: { } };
+            const expr = '{{input.value}} * 1024';
+            // 16 * 1024 = 16384, which is 2^14
+            const result = calculateDefaultMemoryFromExpression(expr, context);
+            expect(result).toBe(16384);
+        });
+
+        it('correctly throw error if runOptions variable is invalid', () => {
+            const context = { input: { value: 16 }, runOptions: { } };
+            const expr = '{{runOptions.customVariable}} * 1024';
+            // 16 * 1024 = 16384, which is 2^14
             expect(() => calculateDefaultMemoryFromExpression(expr, context))
-                .toThrow(`Invalid variable '{{unexistingVariable}}' in expression.`);
+                .toThrow(`Invalid variable '{{runOptions.customVariable}}' in expression. Only the following runOptions are allowed:`);
         });
     });