diff --git a/extensions/copilot/src/extension/tools/common/test/toolService.spec.ts b/extensions/copilot/src/extension/tools/common/test/toolService.spec.ts index 844fb0b74cc..ad59dee1b60 100644 --- a/extensions/copilot/src/extension/tools/common/test/toolService.spec.ts +++ b/extensions/copilot/src/extension/tools/common/test/toolService.spec.ts @@ -230,5 +230,183 @@ describe('Tool Service', () => { error: expect.stringContaining('ERROR: Your input to the tool was invalid') }); }); + + test('should reconstruct flattened path keys', () => { + const askQuestionsTool: vscode.LanguageModelToolInformation = { + name: 'askQuestionsTool', + description: 'A tool that expects an array of nested question objects', + inputSchema: { + type: 'object', + properties: { + questions: { + type: 'array', + items: { + type: 'object', + properties: { + header: { type: 'string' }, + question: { type: 'string' }, + allowFreeformInput: { type: 'boolean' }, + options: { + type: 'array', + items: { + type: 'object', + properties: { + label: { type: 'string' }, + description: { type: 'string' }, + recommended: { type: 'boolean' } + }, + required: ['label'] + } + } + }, + required: ['header', 'question'] + } + } + }, + required: ['questions'] + }, + tags: [], + source: undefined + }; + + (toolsService.tools as vscode.LanguageModelToolInformation[]).push(askQuestionsTool); + + // Gemini-style flattened path keys instead of a nested object/array. + const flattenedInput = JSON.stringify({ + 'questions[0].allowFreeformInput': true, + 'questions[0].header': 'repro_question_1', + 'questions[0].options[0].description': 'First option description', + 'questions[0].options[0].label': 'Option A', + 'questions[0].options[0].recommended': true, + 'questions[0].options[1].description': 'Second option description', + 'questions[0].options[1].label': 'Option B', + 'questions[0].question': 'Which option do you prefer?', + 'questions[1].allowFreeformInput': false, + 'questions[1].header': 'repro_question_2', + 'questions[1].options[0].label': 'Yes', + 'questions[1].options[1].label': 'No', + 'questions[1].question': 'Do you want to continue?' + }); + + const result = toolsService.validateToolInput('askQuestionsTool', flattenedInput); + expect(result).toEqual({ + inputObj: { + questions: [ + { + allowFreeformInput: true, + header: 'repro_question_1', + question: 'Which option do you prefer?', + options: [ + { description: 'First option description', label: 'Option A', recommended: true }, + { description: 'Second option description', label: 'Option B' } + ] + }, + { + allowFreeformInput: false, + header: 'repro_question_2', + question: 'Do you want to continue?', + options: [ + { label: 'Yes' }, + { label: 'No' } + ] + } + ] + } + }); + }); + + test('should not pollute prototype when reconstructing flattened keys', () => { + const pollutionTool: vscode.LanguageModelToolInformation = { + name: 'pollutionTool', + description: 'A tool whose flattened input contains unsafe property names', + inputSchema: { + type: 'object', + properties: { + data: { + type: 'object', + properties: { value: { type: 'string' } } + } + }, + required: ['data'] + }, + tags: [], + source: undefined + }; + + (toolsService.tools as vscode.LanguageModelToolInformation[]).push(pollutionTool); + + const malicious = JSON.stringify({ + '__proto__.polluted': 'yes', + 'data.value': 'ok' + }); + + const result = toolsService.validateToolInput('pollutionTool', malicious); + + // The unsafe key makes reconstruction bail out, so validation fails + // rather than mutating Object.prototype. + expect(result).toMatchObject({ + error: expect.stringContaining('ERROR: Your input to the tool was invalid') + }); + expect(({} as Record).polluted).toBeUndefined(); + }); + + test('should bail out on conflicting flattened keys', () => { + const conflictTool: vscode.LanguageModelToolInformation = { + name: 'conflictTool', + description: 'A tool whose flattened input has conflicting paths', + inputSchema: { + type: 'object', + properties: { + a: { type: 'object' } + }, + required: ['a'] + }, + tags: [], + source: undefined + }; + + (toolsService.tools as vscode.LanguageModelToolInformation[]).push(conflictTool); + + // `a` is both a primitive and a parent of `a.b` — unresolvable. + const conflicting = JSON.stringify({ + 'a': 'primitive', + 'a.b': 'nested' + }); + + const result = toolsService.validateToolInput('conflictTool', conflicting); + expect(result).toMatchObject({ + error: expect.stringContaining('ERROR: Your input to the tool was invalid') + }); + }); + + test('should reject out-of-range array indices in flattened keys', () => { + const indexTool: vscode.LanguageModelToolInformation = { + name: 'indexTool', + description: 'A tool whose flattened input has an enormous array index', + inputSchema: { + type: 'object', + properties: { + items: { + type: 'array', + items: { type: 'string' } + } + }, + required: ['items'] + }, + tags: [], + source: undefined + }; + + (toolsService.tools as vscode.LanguageModelToolInformation[]).push(indexTool); + + // A huge index would create a massive sparse array; reconstruction + // must bail rather than produce one. + const huge = JSON.stringify({ 'items[999999999999]': 'value' }); + + const result = toolsService.validateToolInput('indexTool', huge); + expect(result).toMatchObject({ + error: expect.stringContaining('ERROR: Your input to the tool was invalid') + }); + }); }); }); diff --git a/extensions/copilot/src/extension/tools/common/toolsService.ts b/extensions/copilot/src/extension/tools/common/toolsService.ts index 4e66b020e33..2eb611f9b5b 100644 --- a/extensions/copilot/src/extension/tools/common/toolsService.ts +++ b/extensions/copilot/src/extension/tools/common/toolsService.ts @@ -130,6 +130,109 @@ function getObjectPropertyByPath(obj: any, jsonPointerPath: string): { parent: a return null; } +/** + * Property names that must never be used as path segments when reconstructing + * objects from untrusted tool input, to avoid prototype pollution. + */ +const UNSAFE_PROPERTY_NAMES = new Set(['__proto__', 'constructor', 'prototype']); + +/** + * Upper bound for array indices accepted when reconstructing flattened tool + * input. Caps the reconstructed array length to avoid huge sparse arrays from + * untrusted input (e.g. `items[999999999999]`) that would make subsequent Ajv + * validation pathologically slow. + */ +const MAX_FLATTENED_ARRAY_INDEX = 1000; + +/** + * Parses a flattened path key (e.g. `questions[0].options[1].label`) into an + * ordered list of segments (`['questions', 0, 'options', 1, 'label']`). Object + * properties are returned as strings and array indices as numbers. Returns + * `undefined` if the key is not a well-formed, contiguous path expression, if + * it contains an unsafe property name (e.g. `__proto__`), or if an array index + * exceeds {@link MAX_FLATTENED_ARRAY_INDEX}. + */ +function parseFlattenedPath(key: string): (string | number)[] | undefined { + const segments: (string | number)[] = []; + const re = /\.?([^.[\]]+)|\[(\d+)\]/g; + let lastIndex = 0; + let match: RegExpExecArray | null; + while ((match = re.exec(key)) !== null) { + // Bail if there is an unexpected character between tokens (e.g. `a..b`). + if (match.index !== lastIndex) { + return undefined; + } + if (match[2] !== undefined) { + const index = Number(match[2]); + // Reject out-of-range indices to avoid huge sparse arrays. + if (!Number.isSafeInteger(index) || index > MAX_FLATTENED_ARRAY_INDEX) { + return undefined; + } + segments.push(index); + } else { + // Reject prototype-pollution keys from untrusted tool input. + if (UNSAFE_PROPERTY_NAMES.has(match[1])) { + return undefined; + } + segments.push(match[1]); + } + lastIndex = re.lastIndex; + } + if (lastIndex !== key.length || segments.length === 0) { + return undefined; + } + return segments; +} + +/** + * Reconstructs a nested object/array structure from an object whose keys are + * flattened path expressions. Some models (notably Gemini) serialize nested + * tool-call arguments as flat keys like `questions[0].header` instead of a + * proper nested object. Returns `undefined` when none of the keys use path + * notation (so normal inputs are left untouched), when a key is malformed, or + * when keys conflict (e.g. both `a` and `a.b`). + */ +function tryUnflattenObject(obj: Record): Record | undefined { + const keys = Object.keys(obj); + if (!keys.some(key => /\.|\[\d+\]/.test(key))) { + return undefined; + } + + // Use null-prototype containers so untrusted keys cannot reach Object.prototype. + const result: Record = Object.create(null); + for (const key of keys) { + const path = parseFlattenedPath(key); + if (!path) { + return undefined; + } + + let current: any = result; + for (let i = 0; i < path.length - 1; i++) { + const segment = path[i]; + const nextSegment = path[i + 1]; + const child = current[segment]; + if (child === undefined) { + current[segment] = typeof nextSegment === 'number' ? [] : Object.create(null); + } else if (typeof child !== 'object' || child === null) { + // Conflicting keys (e.g. both `a` and `a.b`) would require + // overwriting a primitive with a container; bail out instead. + return undefined; + } + current = current[segment]; + } + + const leaf = path[path.length - 1]; + if (typeof current[leaf] === 'object' && current[leaf] !== null) { + // A container already exists at this leaf (e.g. both `a` and `a.b` + // where `a` is assigned last); refuse to clobber it. + return undefined; + } + current[leaf] = obj[key]; + } + + return result; +} + function ajvValidateForTool(toolName: string, fn: ValidateFunction, inputObj: unknown): IToolValidationResult { // Empty output can be valid when the schema only has optional properties if (fn(inputObj ?? {})) { @@ -168,6 +271,16 @@ function ajvValidateForTool(toolName: string, fn: ValidateFunction, inputObj: un } } + // Recovery: some models (notably Gemini) serialize nested arguments as + // flattened path keys like `questions[0].header` instead of nested + // objects/arrays. Reconstruct the nested structure and re-validate. + if (typeof inputObj === 'object' && inputObj !== null && !Array.isArray(inputObj)) { + const unflattened = tryUnflattenObject(inputObj as Record); + if (unflattened) { + return ajvValidateForTool(toolName, fn, unflattened); + } + } + const errors = fn.errors!.map(e => e.message || `${e.instancePath} is invalid}`); return { error: `ERROR: Your input to the tool was invalid (${errors.join(', ')})` }; }