add template security validation

Author: chad-syntax

__tests__/lib/template-utils.test.ts

@@ -1,4 +1,9 @@
-import { extractTemplateVariables, validateVariables } from '../../src/utils/template-utils';
+import {
+  extractTemplateVariables,
+  validateVariables,
+  validateTemplate,
+  compilePrompt,
+} from '../../src/utils/template-utils';
 import { Database } from '@/app/__generated__/supabase.types';
 
 type PromptVariable = Database['public']['Tables']['prompt_variables']['Row'];
@@ -825,3 +830,139 @@ describe('validateVariables', () => {
     expect(variablesWithDefaults).toEqual({ opt1: 'a', opt2: '1' });
   });
 });
+
+describe('validateTemplate', () => {
+  it('should return isValid: true for a valid template', () => {
+    const template = 'Hello {{ name }}!';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(true);
+    expect(result.error).toBeUndefined();
+  });
+
+  it("should return isValid: true for a template using 'global'", () => {
+    const template = 'Version: {{ global.version }}';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(true);
+    expect(result.error).toBeUndefined();
+  });
+
+  it('should return isValid: false for Nunjucks syntax errors', () => {
+    const template = 'Hello {{ name !'; // Invalid Nunjucks syntax
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(false);
+    expect(result.error).toMatch(/expected variable end/i);
+  });
+
+  it('should return isValid: false for disallowed identifiers', () => {
+    const template = 'Data: {{ process.env.SECRET }}';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(false);
+    expect(result.error).toBe("Security: Disallowed identifier 'process' found.");
+  });
+
+  it('should return isValid: false for another disallowed identifier (eval)', () => {
+    const template = '{{ eval("alert(1)") }}';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(false);
+    expect(result.error).toBe("Security: Disallowed identifier 'eval' found.");
+  });
+
+  it('should return isValid: false for disallowed property lookups', () => {
+    const template = 'Access: {{ myObj.__proto__ }}';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(false);
+    expect(result.error).toBe("Security: Disallowed property lookup '__proto__' found.");
+  });
+
+  it('should return isValid: false for disallowed property lookups via string literal', () => {
+    const template = "Access: {{ myObj['constructor'] }}";
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(false);
+    expect(result.error).toBe("Security: Disallowed property lookup 'constructor' found.");
+  });
+
+  it('should allow legitimate property lookups that are not disallowed', () => {
+    const template = 'Value: {{ myObj.legitProperty }}';
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(true);
+    expect(result.error).toBeUndefined();
+  });
+
+  it('should handle complex but valid templates correctly', () => {
+    const template = `
+      {% if user.isAdmin %}
+        Admin: {{ user.name }}
+        {% for item in user.items %}
+          {{ item.id }} - {{ item.value }}
+        {% endfor %}
+      {% else %}
+        User: {{ user.name }}
+      {% endif %}
+    `;
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(true);
+    expect(result.error).toBeUndefined();
+  });
+});
+
+describe('compilePrompt', () => {
+  const baseMockVariables = {
+    name: 'Tester',
+    myObj: { legitProperty: 'hello' },
+    global: { version: '1.0' },
+  };
+
+  it('should compile a valid prompt', () => {
+    const template = 'Hello {{ name }}!';
+    expect(() => compilePrompt(template, baseMockVariables)).not.toThrow();
+    expect(compilePrompt(template, baseMockVariables)).toBe('Hello Tester!');
+  });
+
+  it("should compile a prompt using 'global' as it was removed from blacklist by user", () => {
+    const template = 'Version: {{ global.version }}';
+    expect(() => compilePrompt(template, baseMockVariables)).not.toThrow();
+    // Ensure the output matches the global.version from baseMockVariables
+    expect(compilePrompt(template, baseMockVariables)).toBe('Version: 1.0');
+  });
+
+  it('should throw error for disallowed identifiers', () => {
+    const template = 'Data: {{ process.env.SECRET }}';
+    expect(() => compilePrompt(template, baseMockVariables)).toThrow(
+      "Security: Disallowed identifier 'process' found.",
+    );
+  });
+
+  it('should throw error for disallowed property lookups', () => {
+    const template = 'Access: {{ myObj.__proto__ }}';
+    expect(() => compilePrompt(template, baseMockVariables)).toThrow(
+      "Security: Disallowed property lookup '__proto__' found.",
+    );
+  });
+
+  it('should throw error for Nunjucks syntax errors which cause parse fail in ensureSecureAst', () => {
+    const template = 'Hello {{ name !'; // Invalid Nunjucks syntax
+    expect(() => compilePrompt(template, baseMockVariables)).toThrow(
+      'Template parsing failed: expected variable end',
+    );
+  });
+
+  it("should allow mixed-case identifiers like 'Process' if not explicitly blacklisted (current check is case-sensitive)", () => {
+    const template = 'Data: {{ Process.env.SECRET }}'; // 'Process' is not in DISALLOWED_IDENTIFIERS
+    const result = validateTemplate(template);
+    expect(result.isValid).toBe(true);
+    expect(result.error).toBeUndefined();
+
+    const specificMockVariables = {
+      ...baseMockVariables,
+      Process: { env: { SECRET: 'oops' } },
+    };
+    expect(() => compilePrompt(template, specificMockVariables)).not.toThrow();
+    expect(compilePrompt(template, specificMockVariables)).toBe('Data: oops');
+  });
+
+  it('should allow deeply nested valid lookups', () => {
+    const template = 'Deep: {{ a.b.c.d }}';
+    const vars = { ...baseMockVariables, a: { b: { c: { d: 'value' } } } };
+    expect(compilePrompt(template, vars)).toBe('Deep: value');
+  });
+});

src/lib/PromptsService.ts

@@ -1,9 +1,7 @@
-import { Database, Json } from '@/app/__generated__/supabase.types';
+import { Database } from '@/app/__generated__/supabase.types';
 import type { SupabaseClient } from '@supabase/supabase-js';
 import { AgentsmithSupabaseService } from './AgentsmithSupabaseService';
 import { AgentsmithServicesDirectory } from './AgentsmithServices';
-// @ts-ignore needs to be browser version so nextjs can import it
-import nunjucks from 'nunjucks/browser/nunjucks';
 import {
   OpenrouterRequestBody,
   CompletionConfig,
@@ -16,14 +14,7 @@ import {
 import { routes } from '@/utils/routes';
 import { ORGANIZATION_KEYS, SEMVER_PATTERN } from '@/app/constants';
 import { compareSemanticVersions, incrementVersion } from '@/utils/versioning';
-import {
-  compilePrompt,
-  extractTemplateVariables,
-  findMissingGlobalContext,
-  processUniqueValues,
-} from '@/utils/template-utils';
-
-type PromptVariable = Database['public']['Tables']['prompt_variables']['Row'];
+import { compilePrompt } from '@/utils/template-utils';
 
 type PromptVariableBase = Omit<
   Database['public']['Tables']['prompt_variables']['Row'],

src/utils/template-utils.ts

@@ -231,21 +231,108 @@ export const extractTemplateVariables = (
   }
 };
 
+const DISALLOWED_IDENTIFIERS = [
+  'process',
+  'require',
+  'eval',
+  'Function',
+  'this',
+  'window',
+  'document',
+  'self',
+  'globalThis',
+  'constructor',
+  'prototype',
+  '__proto__',
+];
+
+const DISALLOWED_PROPERTY_LOOKUPS = ['__proto__', 'constructor', 'prototype'];
+
+const ensureSecureAstRecursive = (node: any): void => {
+  if (!node) return;
+
+  if (Array.isArray(node)) {
+    node.forEach(ensureSecureAstRecursive);
+    return;
+  }
+
+  switch (node.typename) {
+    case 'Symbol':
+      if (DISALLOWED_IDENTIFIERS.includes(node.value)) {
+        throw new Error(`Security: Disallowed identifier '${node.value}' found.`);
+      }
+      break;
+    case 'LookupVal':
+      if (
+        node.val &&
+        node.val.typename === 'Literal' &&
+        typeof node.val.value === 'string' &&
+        DISALLOWED_PROPERTY_LOOKUPS.includes(node.val.value)
+      ) {
+        throw new Error(`Security: Disallowed property lookup '${node.val.value}' found.`);
+      }
+      ensureSecureAstRecursive(node.target);
+      ensureSecureAstRecursive(node.val);
+      break;
+    default:
+      if (typeof node === 'object') {
+        for (const key in node) {
+          if (
+            Object.prototype.hasOwnProperty.call(node, key) &&
+            !key.startsWith('_') &&
+            key !== 'parent'
+          ) {
+            ensureSecureAstRecursive(node[key]);
+          }
+        }
+      }
+      break;
+  }
+};
+
+const ensureSecureAst = (content: string): void => {
+  try {
+    // @ts-ignore - parser exists but is not in type definitions
+    const ast = nunjucks.parser.parse(content);
+    ensureSecureAstRecursive(ast);
+  } catch (error) {
+    if (error instanceof Error && error.message.startsWith('Security:')) {
+      throw error; // Re-throw specific security errors
+    }
+    // For other parsing errors, wrap them if necessary or let them bubble up as Nunjucks errors.
+    // For now, re-throwing to ensure `validateTemplate` can catch them as generic parsing issues if not security related.
+    throw new Error(
+      `Template parsing failed: ${error instanceof Error ? error.message : String(error)}`,
+    );
+  }
+};
+
 /**
  * Validates if a Nunjucks template is syntactically correct
  * Note: Currently not used in the application, but kept for potential future use
  */
 export const validateTemplate = (content: string): { isValid: boolean; error?: string } => {
   try {
+    // First, check basic Nunjucks syntax
     transform(nunjucks.parser.parse(content));
-
-    return { isValid: true };
   } catch (error) {
     return {
       isValid: false,
       error: error instanceof Error ? error.message : 'Invalid template syntax',
     };
   }
+
+  try {
+    // Then, perform our security AST check
+    ensureSecureAst(content);
+  } catch (error) {
+    return {
+      isValid: false,
+      error: error instanceof Error ? error.message : 'Template violates security policy',
+    };
+  }
+
+  return { isValid: true };
 };
 
 type FindMissingGlobalContextOptions = {
@@ -348,8 +435,8 @@ export const compilePrompt = (
   promptContent: string,
   variables: Record<string, any> & { global: Record<string, any> },
 ): string => {
+  ensureSecureAst(promptContent); // Perform security check first
   const processedVariables = processUniqueValues(variables);
-
   nunjucks.configure({ autoescape: false });
   return nunjucks.renderString(promptContent, processedVariables);
 };