Skip to content
92 changes: 92 additions & 0 deletions packages/sdk/server-ai/__tests__/Judge.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -796,4 +796,96 @@ describe('Judge', () => {
);
});
});

describe('template injection prevention', () => {
/**
* Regression tests for template injection vulnerability.
*
* These tests verify that the judge's message interpolation uses simple string
* replacement instead of Mustache templating. Attacker-controlled values from
* pass 1 (e.g. Mustache delimiter-change tags) must be treated as inert literal
* text by pass 2.
*/

function makeJudge(content: string): Judge {
const config: LDAIJudgeConfig = {
key: 'test-judge',
enabled: true,
messages: [{ role: 'user', content }],
model: { name: 'gpt-4' },
provider: { name: 'openai' },
tracker: mockTracker,
evaluationMetricKey: 'metric',
};
return new Judge(config, mockTracker, mockProvider, mockLogger);
}

const injectionVariants = [
{ name: 'delimiter change brackets', payload: '{{=[ ]=}}' },
{ name: 'delimiter change angle', payload: '{{=<% %>=}}' },
{ name: 'partial', payload: '{{> evil}}' },
{ name: 'comment', payload: '{{! drop everything }}' },
{ name: 'triple stache', payload: '{{{raw}}}' },
{ name: 'section', payload: '{{#section}}inject{{/section}}' },
{ name: 'inverted section', payload: '{{^section}}inject{{/section}}' },
];

it.each(injectionVariants)(
'injection variant "$name" in message_history does not blind the judge',
({ payload }) => {
const afterPass1 = `Auditing ${payload}: {{message_history}}`;
const judge = makeJudge(afterPass1);

// eslint-disable-next-line no-underscore-dangle
const constructMessages = (judge as any)._constructEvaluationMessages.bind(judge);
const messages = constructMessages('ACTUAL HISTORY', 'some output');

expect(messages).toHaveLength(1);
expect(messages[0].content).toContain('ACTUAL HISTORY');
expect(messages[0].content).not.toContain('{{message_history}}');
},
);

it('injection via response is neutralized', () => {
const afterPass1 = 'History: {{message_history}}\nResponse: {{response_to_evaluate}}';
const judge = makeJudge(afterPass1);

// eslint-disable-next-line no-underscore-dangle
const constructMessages = (judge as any)._constructEvaluationMessages.bind(judge);
const maliciousResponse = '{{=[ ]=}} INJECTION ATTEMPT';
const messages = constructMessages('normal history', maliciousResponse);

expect(messages).toHaveLength(1);
expect(messages[0].content).toContain(maliciousResponse);
expect(messages[0].content).not.toContain('{{response_to_evaluate}}');
});

it('substitutes multiple occurrences of the same placeholder', () => {
const template = '{{message_history}} | {{message_history}}';
const judge = makeJudge(template);

// eslint-disable-next-line no-underscore-dangle
const constructMessages = (judge as any)._constructEvaluationMessages.bind(judge);
const messages = constructMessages('HISTORY', 'RESPONSE');

expect(messages).toHaveLength(1);
expect(messages[0].content).toBe('HISTORY | HISTORY');
});

it('preserves Mustache-like syntax inside history and response values', () => {
const template = 'History: {{message_history}}\nResponse: {{response_to_evaluate}}';
const judge = makeJudge(template);

// eslint-disable-next-line no-underscore-dangle
const constructMessages = (judge as any)._constructEvaluationMessages.bind(judge);
const historyWithMustache = 'How do I use {{user}} in Mustache?';
const responseWithMustache = 'Use {{user}} like this: {{#user}}Hello{{/user}}';

const messages = constructMessages(historyWithMustache, responseWithMustache);

expect(messages).toHaveLength(1);
expect(messages[0].content).toContain(historyWithMustache);
expect(messages[0].content).toContain(responseWithMustache);
});
});
});
10 changes: 6 additions & 4 deletions packages/sdk/server-ai/src/api/judge/Judge.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import Mustache from 'mustache';

import { LDLogger } from '@launchdarkly/js-server-sdk-common';

import { ChatResponse } from '../chat/types';
Expand Down Expand Up @@ -179,10 +177,14 @@ export class Judge {
}

/**
* Interpolates message content with variables using Mustache templating.
* Use string replacement to prevent context attributes like {{=[ ]=}}) from
Comment thread
joker23 marked this conversation as resolved.
Outdated
* influencing judge template parsing.
*/
private _interpolateMessage(content: string, variables: Record<string, string>): string {
return Mustache.render(content, variables, undefined, { escape: (item: any) => item });
return Object.entries(variables).reduce(
(result, [key, value]) => result.split(`{{${key}}}`).join(value),
content,
);
Comment thread
cursor[bot] marked this conversation as resolved.
}

/**
Expand Down
Loading