diff options
Diffstat (limited to 'packages/cli/src/services')
5 files changed, 680 insertions, 11 deletions
diff --git a/packages/cli/src/services/FileCommandLoader.test.ts b/packages/cli/src/services/FileCommandLoader.test.ts index 0e3d781b..e3cbceb2 100644 --- a/packages/cli/src/services/FileCommandLoader.test.ts +++ b/packages/cli/src/services/FileCommandLoader.test.ts @@ -11,12 +11,68 @@ import { getUserCommandsDir, } from '@google/gemini-cli-core'; import mock from 'mock-fs'; -import { assert } from 'vitest'; +import { assert, vi } from 'vitest'; import { createMockCommandContext } from '../test-utils/mockCommandContext.js'; +import { + SHELL_INJECTION_TRIGGER, + SHORTHAND_ARGS_PLACEHOLDER, +} from './prompt-processors/types.js'; +import { + ConfirmationRequiredError, + ShellProcessor, +} from './prompt-processors/shellProcessor.js'; +import { ShorthandArgumentProcessor } from './prompt-processors/argumentProcessor.js'; + +const mockShellProcess = vi.hoisted(() => vi.fn()); +vi.mock('./prompt-processors/shellProcessor.js', () => ({ + ShellProcessor: vi.fn().mockImplementation(() => ({ + process: mockShellProcess, + })), + ConfirmationRequiredError: class extends Error { + constructor( + message: string, + public commandsToConfirm: string[], + ) { + super(message); + this.name = 'ConfirmationRequiredError'; + } + }, +})); + +vi.mock('./prompt-processors/argumentProcessor.js', async (importOriginal) => { + const original = + await importOriginal< + typeof import('./prompt-processors/argumentProcessor.js') + >(); + return { + ShorthandArgumentProcessor: vi + .fn() + .mockImplementation(() => new original.ShorthandArgumentProcessor()), + DefaultArgumentProcessor: vi + .fn() + .mockImplementation(() => new original.DefaultArgumentProcessor()), + }; +}); +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = + await importOriginal<typeof import('@google/gemini-cli-core')>(); + return { + ...original, + isCommandAllowed: vi.fn(), + ShellExecutionService: { + execute: vi.fn(), + }, + }; +}); describe('FileCommandLoader', () => { const signal: AbortSignal = new AbortController().signal; + beforeEach(() => { + vi.clearAllMocks(); + mockShellProcess.mockImplementation((prompt) => Promise.resolve(prompt)); + }); + afterEach(() => { mock.restore(); }); @@ -371,4 +427,180 @@ describe('FileCommandLoader', () => { } }); }); + + describe('Shell Processor Integration', () => { + it('instantiates ShellProcessor if the trigger is present', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shell.toml': `prompt = "Run this: ${SHELL_INJECTION_TRIGGER}echo hello}"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).toHaveBeenCalledWith('shell'); + }); + + it('does not instantiate ShellProcessor if trigger is missing', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'regular.toml': `prompt = "Just a regular prompt"`, + }, + }); + + const loader = new FileCommandLoader(null as unknown as Config); + await loader.loadCommands(signal); + + expect(ShellProcessor).not.toHaveBeenCalled(); + }); + + it('returns a "submit_prompt" action if shell processing succeeds', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shell.toml': `prompt = "Run !{echo 'hello'}"`, + }, + }); + mockShellProcess.mockResolvedValue('Run hello'); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'shell'); + expect(command).toBeDefined(); + + const result = await command!.action!( + createMockCommandContext({ + invocation: { raw: '/shell', name: 'shell', args: '' }, + }), + '', + ); + + expect(result?.type).toBe('submit_prompt'); + if (result?.type === 'submit_prompt') { + expect(result.content).toBe('Run hello'); + } + }); + + it('returns a "confirm_shell_commands" action if shell processing requires it', async () => { + const userCommandsDir = getUserCommandsDir(); + const rawInvocation = '/shell rm -rf /'; + mock({ + [userCommandsDir]: { + 'shell.toml': `prompt = "Run !{rm -rf /}"`, + }, + }); + + // Mock the processor to throw the specific error + const error = new ConfirmationRequiredError('Confirmation needed', [ + 'rm -rf /', + ]); + mockShellProcess.mockRejectedValue(error); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'shell'); + expect(command).toBeDefined(); + + const result = await command!.action!( + createMockCommandContext({ + invocation: { raw: rawInvocation, name: 'shell', args: 'rm -rf /' }, + }), + 'rm -rf /', + ); + + expect(result?.type).toBe('confirm_shell_commands'); + if (result?.type === 'confirm_shell_commands') { + expect(result.commandsToConfirm).toEqual(['rm -rf /']); + expect(result.originalInvocation.raw).toBe(rawInvocation); + } + }); + + it('re-throws other errors from the processor', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'shell.toml': `prompt = "Run !{something}"`, + }, + }); + + const genericError = new Error('Something else went wrong'); + mockShellProcess.mockRejectedValue(genericError); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'shell'); + expect(command).toBeDefined(); + + await expect( + command!.action!( + createMockCommandContext({ + invocation: { raw: '/shell', name: 'shell', args: '' }, + }), + '', + ), + ).rejects.toThrow('Something else went wrong'); + }); + + it('assembles the processor pipeline in the correct order (Shell -> Argument)', async () => { + const userCommandsDir = getUserCommandsDir(); + mock({ + [userCommandsDir]: { + 'pipeline.toml': ` + prompt = "Shell says: ${SHELL_INJECTION_TRIGGER}echo foo} and user says: ${SHORTHAND_ARGS_PLACEHOLDER}" + `, + }, + }); + + // Mock the process methods to track call order + const argProcessMock = vi + .fn() + .mockImplementation((p) => `${p}-arg-processed`); + + // Redefine the mock for this specific test + mockShellProcess.mockImplementation((p) => + Promise.resolve(`${p}-shell-processed`), + ); + + vi.mocked(ShorthandArgumentProcessor).mockImplementation( + () => + ({ + process: argProcessMock, + }) as unknown as ShorthandArgumentProcessor, + ); + + const loader = new FileCommandLoader(null as unknown as Config); + const commands = await loader.loadCommands(signal); + const command = commands.find((c) => c.name === 'pipeline'); + expect(command).toBeDefined(); + + await command!.action!( + createMockCommandContext({ + invocation: { + raw: '/pipeline bar', + name: 'pipeline', + args: 'bar', + }, + }), + 'bar', + ); + + // Verify that the shell processor was called before the argument processor + expect(mockShellProcess.mock.invocationCallOrder[0]).toBeLessThan( + argProcessMock.mock.invocationCallOrder[0], + ); + + // Also verify the flow of the prompt through the processors + expect(mockShellProcess).toHaveBeenCalledWith( + expect.any(String), + expect.any(Object), + ); + expect(argProcessMock).toHaveBeenCalledWith( + expect.stringContaining('-shell-processed'), // It receives the output of the shell processor + expect.any(Object), + ); + }); + }); }); diff --git a/packages/cli/src/services/FileCommandLoader.ts b/packages/cli/src/services/FileCommandLoader.ts index 23d5af19..c96acead 100644 --- a/packages/cli/src/services/FileCommandLoader.ts +++ b/packages/cli/src/services/FileCommandLoader.ts @@ -19,7 +19,7 @@ import { CommandContext, CommandKind, SlashCommand, - SubmitPromptActionReturn, + SlashCommandActionReturn, } from '../ui/commands/types.js'; import { DefaultArgumentProcessor, @@ -28,7 +28,12 @@ import { import { IPromptProcessor, SHORTHAND_ARGS_PLACEHOLDER, + SHELL_INJECTION_TRIGGER, } from './prompt-processors/types.js'; +import { + ConfirmationRequiredError, + ShellProcessor, +} from './prompt-processors/shellProcessor.js'; /** * Defines the Zod schema for a command definition file. This serves as the @@ -172,6 +177,11 @@ export class FileCommandLoader implements ICommandLoader { const processors: IPromptProcessor[] = []; + // Add the Shell Processor if needed. + if (validDef.prompt.includes(SHELL_INJECTION_TRIGGER)) { + processors.push(new ShellProcessor(commandName)); + } + // The presence of '{{args}}' is the switch that determines the behavior. if (validDef.prompt.includes(SHORTHAND_ARGS_PLACEHOLDER)) { processors.push(new ShorthandArgumentProcessor()); @@ -188,7 +198,7 @@ export class FileCommandLoader implements ICommandLoader { action: async ( context: CommandContext, _args: string, - ): Promise<SubmitPromptActionReturn> => { + ): Promise<SlashCommandActionReturn> => { if (!context.invocation) { console.error( `[FileCommandLoader] Critical error: Command '${commandName}' was executed without invocation context.`, @@ -199,15 +209,31 @@ export class FileCommandLoader implements ICommandLoader { }; } - let processedPrompt = validDef.prompt; - for (const processor of processors) { - processedPrompt = await processor.process(processedPrompt, context); - } + try { + let processedPrompt = validDef.prompt; + for (const processor of processors) { + processedPrompt = await processor.process(processedPrompt, context); + } - return { - type: 'submit_prompt', - content: processedPrompt, - }; + return { + type: 'submit_prompt', + content: processedPrompt, + }; + } catch (e) { + // Check if it's our specific error type + if (e instanceof ConfirmationRequiredError) { + // Halt and request confirmation from the UI layer. + return { + type: 'confirm_shell_commands', + commandsToConfirm: e.commandsToConfirm, + originalInvocation: { + raw: context.invocation.raw, + }, + }; + } + // Re-throw other errors to be handled by the global error handler. + throw e; + } }, }; } diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts new file mode 100644 index 00000000..a2883923 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts @@ -0,0 +1,300 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { ConfirmationRequiredError, ShellProcessor } from './shellProcessor.js'; +import { createMockCommandContext } from '../../test-utils/mockCommandContext.js'; +import { CommandContext } from '../../ui/commands/types.js'; +import { Config } from '@google/gemini-cli-core'; + +const mockCheckCommandPermissions = vi.hoisted(() => vi.fn()); +const mockShellExecute = vi.hoisted(() => vi.fn()); + +vi.mock('@google/gemini-cli-core', async (importOriginal) => { + const original = await importOriginal<object>(); + return { + ...original, + checkCommandPermissions: mockCheckCommandPermissions, + ShellExecutionService: { + execute: mockShellExecute, + }, + }; +}); + +describe('ShellProcessor', () => { + let context: CommandContext; + let mockConfig: Partial<Config>; + + beforeEach(() => { + vi.clearAllMocks(); + + mockConfig = { + getTargetDir: vi.fn().mockReturnValue('/test/dir'), + }; + + context = createMockCommandContext({ + services: { + config: mockConfig as Config, + }, + session: { + sessionShellAllowlist: new Set(), + }, + }); + + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + output: 'default shell output', + }), + }); + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + }); + + it('should not change the prompt if no shell injections are present', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This is a simple prompt with no injections.'; + const result = await processor.process(prompt, context); + expect(result).toBe(prompt); + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should process a single valid shell injection if allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'The current status is: !{git status}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'On branch main' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'git status', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + 'git status', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('The current status is: On branch main'); + }); + + it('should process multiple valid shell injections if all are allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{git status} in !{pwd}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + + mockShellExecute + .mockReturnValueOnce({ + result: Promise.resolve({ output: 'On branch main' }), + }) + .mockReturnValueOnce({ + result: Promise.resolve({ output: '/usr/home' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledTimes(2); + expect(mockShellExecute).toHaveBeenCalledTimes(2); + expect(result).toBe('On branch main in /usr/home'); + }); + + it('should throw ConfirmationRequiredError if a command is not allowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Do something dangerous: !{rm -rf /}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: ['rm -rf /'], + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + }); + + it('should throw ConfirmationRequiredError with the correct command', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Do something dangerous: !{rm -rf /}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: ['rm -rf /'], + }); + + try { + await processor.process(prompt, context); + // Fail if it doesn't throw + expect(true).toBe(false); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['rm -rf /']); + } + } + + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should throw ConfirmationRequiredError with multiple commands if multiple are disallowed', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd1} and !{cmd2}'; + mockCheckCommandPermissions.mockImplementation((cmd) => { + if (cmd === 'cmd1') { + return { allAllowed: false, disallowedCommands: ['cmd1'] }; + } + if (cmd === 'cmd2') { + return { allAllowed: false, disallowedCommands: ['cmd2'] }; + } + return { allAllowed: true, disallowedCommands: [] }; + }); + + try { + await processor.process(prompt, context); + // Fail if it doesn't throw + expect(true).toBe(false); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['cmd1', 'cmd2']); + } + } + }); + + it('should not execute any commands if at least one requires confirmation', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'First: !{echo "hello"}, Second: !{rm -rf /}'; + + mockCheckCommandPermissions.mockImplementation((cmd) => { + if (cmd.includes('rm')) { + return { allAllowed: false, disallowedCommands: [cmd] }; + } + return { allAllowed: true, disallowedCommands: [] }; + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + + // Ensure no commands were executed because the pipeline was halted. + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should only request confirmation for disallowed commands in a mixed prompt', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Allowed: !{ls -l}, Disallowed: !{rm -rf /}'; + + mockCheckCommandPermissions.mockImplementation((cmd) => ({ + allAllowed: !cmd.includes('rm'), + disallowedCommands: cmd.includes('rm') ? [cmd] : [], + })); + + try { + await processor.process(prompt, context); + expect.fail('Should have thrown ConfirmationRequiredError'); + } catch (e) { + expect(e).toBeInstanceOf(ConfirmationRequiredError); + if (e instanceof ConfirmationRequiredError) { + expect(e.commandsToConfirm).toEqual(['rm -rf /']); + } + } + }); + + it('should execute all commands if they are on the session allowlist', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Run !{cmd1} and !{cmd2}'; + + // Add commands to the session allowlist + context.session.sessionShellAllowlist = new Set(['cmd1', 'cmd2']); + + // checkCommandPermissions should now pass for these + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + + mockShellExecute + .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) }) + .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'cmd1', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'cmd2', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledTimes(2); + expect(result).toBe('Run output1 and output2'); + }); + + it('should trim whitespace from the command inside the injection', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Files: !{ ls -l }'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'total 0' }), + }); + + await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + 'ls -l', // Verifies that the command was trimmed + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + 'ls -l', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); + + it('should handle an empty command inside the injection gracefully', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This is weird: !{}'; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: true, + disallowedCommands: [], + }); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ output: 'empty output' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + '', + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + '', + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('This is weird: empty output'); + }); +}); diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.ts b/packages/cli/src/services/prompt-processors/shellProcessor.ts new file mode 100644 index 00000000..bf811d66 --- /dev/null +++ b/packages/cli/src/services/prompt-processors/shellProcessor.ts @@ -0,0 +1,106 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + checkCommandPermissions, + ShellExecutionService, +} from '@google/gemini-cli-core'; + +import { CommandContext } from '../../ui/commands/types.js'; +import { IPromptProcessor } from './types.js'; + +export class ConfirmationRequiredError extends Error { + constructor( + message: string, + public commandsToConfirm: string[], + ) { + super(message); + this.name = 'ConfirmationRequiredError'; + } +} + +/** + * Finds all instances of shell command injections (`!{...}`) in a prompt, + * executes them, and replaces the injection site with the command's output. + * + * This processor ensures that only allowlisted commands are executed. If a + * disallowed command is found, it halts execution and reports an error. + */ +export class ShellProcessor implements IPromptProcessor { + /** + * A regular expression to find all instances of `!{...}`. The inner + * capture group extracts the command itself. + */ + private static readonly SHELL_INJECTION_REGEX = /!\{([^}]*)\}/g; + + /** + * @param commandName The name of the custom command being executed, used + * for logging and error messages. + */ + constructor(private readonly commandName: string) {} + + async process(prompt: string, context: CommandContext): Promise<string> { + const { config, sessionShellAllowlist } = { + ...context.services, + ...context.session, + }; + const commandsToExecute: Array<{ fullMatch: string; command: string }> = []; + const commandsToConfirm = new Set<string>(); + + const matches = [...prompt.matchAll(ShellProcessor.SHELL_INJECTION_REGEX)]; + if (matches.length === 0) { + return prompt; // No shell commands, nothing to do. + } + + // Discover all commands and check permissions. + for (const match of matches) { + const command = match[1].trim(); + const { allAllowed, disallowedCommands, blockReason, isHardDenial } = + checkCommandPermissions(command, config!, sessionShellAllowlist); + + if (!allAllowed) { + // If it's a hard denial, this is a non-recoverable security error. + if (isHardDenial) { + throw new Error( + `${this.commandName} cannot be run. ${blockReason || 'A shell command in this custom command is explicitly blocked in your config settings.'}`, + ); + } + + // Add each soft denial disallowed command to the set for confirmation. + disallowedCommands.forEach((uc) => commandsToConfirm.add(uc)); + } + commandsToExecute.push({ fullMatch: match[0], command }); + } + + // If any commands require confirmation, throw a special error to halt the + // pipeline and trigger the UI flow. + if (commandsToConfirm.size > 0) { + throw new ConfirmationRequiredError( + 'Shell command confirmation required', + Array.from(commandsToConfirm), + ); + } + + // Execute all commands (only runs if no confirmation was needed). + let processedPrompt = prompt; + for (const { fullMatch, command } of commandsToExecute) { + const { result } = ShellExecutionService.execute( + command, + config!.getTargetDir(), + () => {}, // No streaming needed. + new AbortController().signal, // For now, we don't support cancellation from here. + ); + + const executionResult = await result; + processedPrompt = processedPrompt.replace( + fullMatch, + executionResult.output, + ); + } + + return processedPrompt; + } +} diff --git a/packages/cli/src/services/prompt-processors/types.ts b/packages/cli/src/services/prompt-processors/types.ts index 2ca61062..2653d2b7 100644 --- a/packages/cli/src/services/prompt-processors/types.ts +++ b/packages/cli/src/services/prompt-processors/types.ts @@ -35,3 +35,8 @@ export interface IPromptProcessor { * The placeholder string for shorthand argument injection in custom commands. */ export const SHORTHAND_ARGS_PLACEHOLDER = '{{args}}'; + +/** + * The trigger string for shell command injection in custom commands. + */ +export const SHELL_INJECTION_TRIGGER = '!{'; |