diff options
Diffstat (limited to 'packages/cli/src/services/prompt-processors/shellProcessor.test.ts')
| -rw-r--r-- | packages/cli/src/services/prompt-processors/shellProcessor.test.ts | 471 |
1 files changed, 438 insertions, 33 deletions
diff --git a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts index a2883923..fe4ae7a2 100644 --- a/packages/cli/src/services/prompt-processors/shellProcessor.test.ts +++ b/packages/cli/src/services/prompt-processors/shellProcessor.test.ts @@ -4,11 +4,32 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { vi, describe, it, expect, beforeEach } from 'vitest'; +import { describe, it, expect, beforeEach, vi } 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'; +import os from 'os'; +import { quote } from 'shell-quote'; + +// Helper function to determine the expected escaped string based on the current OS, +// mirroring the logic in the actual `escapeShellArg` implementation. This makes +// our tests robust and platform-agnostic. +function getExpectedEscapedArgForPlatform(arg: string): string { + if (os.platform() === 'win32') { + const comSpec = (process.env.ComSpec || 'cmd.exe').toLowerCase(); + const isPowerShell = + comSpec.endsWith('powershell.exe') || comSpec.endsWith('pwsh.exe'); + + if (isPowerShell) { + return `'${arg.replace(/'/g, "''")}'`; + } else { + return `"${arg.replace(/"/g, '""')}"`; + } + } else { + return quote([arg]); + } +} const mockCheckCommandPermissions = vi.hoisted(() => vi.fn()); const mockShellExecute = vi.hoisted(() => vi.fn()); @@ -24,6 +45,15 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }; }); +const SUCCESS_RESULT = { + stdout: 'default shell output', + stderr: '', + exitCode: 0, + error: null, + aborted: false, + signal: null, +}; + describe('ShellProcessor', () => { let context: CommandContext; let mockConfig: Partial<Config>; @@ -36,6 +66,11 @@ describe('ShellProcessor', () => { }; context = createMockCommandContext({ + invocation: { + raw: '/cmd default args', + name: 'cmd', + args: 'default args', + }, services: { config: mockConfig as Config, }, @@ -45,16 +80,29 @@ describe('ShellProcessor', () => { }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ - output: 'default shell output', - }), + result: Promise.resolve(SUCCESS_RESULT), }); + mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], }); }); + it('should throw an error if config is missing', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{ls}'; + const contextWithoutConfig = createMockCommandContext({ + services: { + config: null, + }, + }); + + await expect( + processor.process(prompt, contextWithoutConfig), + ).rejects.toThrow(/Security configuration not loaded/); + }); + 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.'; @@ -71,7 +119,7 @@ describe('ShellProcessor', () => { disallowedCommands: [], }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'On branch main' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'On branch main' }), }); const result = await processor.process(prompt, context); @@ -100,10 +148,13 @@ describe('ShellProcessor', () => { mockShellExecute .mockReturnValueOnce({ - result: Promise.resolve({ output: 'On branch main' }), + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'On branch main', + }), }) .mockReturnValueOnce({ - result: Promise.resolve({ output: '/usr/home' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '/usr/home' }), }); const result = await processor.process(prompt, context); @@ -226,8 +277,12 @@ describe('ShellProcessor', () => { }); mockShellExecute - .mockReturnValueOnce({ result: Promise.resolve({ output: 'output1' }) }) - .mockReturnValueOnce({ result: Promise.resolve({ output: 'output2' }) }); + .mockReturnValueOnce({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output1' }), + }) + .mockReturnValueOnce({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'output2' }), + }); const result = await processor.process(prompt, context); @@ -245,56 +300,406 @@ describe('ShellProcessor', () => { expect(result).toBe('Run output1 and output2'); }); - it('should trim whitespace from the command inside the injection', async () => { + it('should trim whitespace from the command inside the injection before interpolation', async () => { const processor = new ShellProcessor('test-command'); - const prompt = 'Files: !{ ls -l }'; + const prompt = 'Files: !{ ls {{args}} -l }'; + + const rawArgs = context.invocation!.args; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + + const expectedCommand = `ls ${expectedEscapedArgs} -l`; + mockCheckCommandPermissions.mockReturnValue({ allAllowed: true, disallowedCommands: [], }); mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'total 0' }), + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'total 0' }), }); await processor.process(prompt, context); expect(mockCheckCommandPermissions).toHaveBeenCalledWith( - 'ls -l', // Verifies that the command was trimmed + expectedCommand, expect.any(Object), context.session.sessionShellAllowlist, ); expect(mockShellExecute).toHaveBeenCalledWith( - 'ls -l', + expectedCommand, expect.any(String), expect.any(Function), expect.any(Object), ); }); - it('should handle an empty command inside the injection gracefully', async () => { + it('should handle an empty command inside the injection gracefully (skips execution)', async () => { const processor = new ShellProcessor('test-command'); const prompt = 'This is weird: !{}'; - mockCheckCommandPermissions.mockReturnValue({ - allAllowed: true, - disallowedCommands: [], + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).not.toHaveBeenCalled(); + expect(mockShellExecute).not.toHaveBeenCalled(); + + // It replaces !{} with an empty string. + expect(result).toBe('This is weird: '); + }); + + describe('Robust Parsing (Balanced Braces)', () => { + it('should correctly parse commands containing nested braces (e.g., awk)', async () => { + const processor = new ShellProcessor('test-command'); + const command = "awk '{print $1}' file.txt"; + const prompt = `Output: !{${command}}`; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'result' }), + }); + + const result = await processor.process(prompt, context); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + command, + expect.any(Object), + context.session.sessionShellAllowlist, + ); + expect(mockShellExecute).toHaveBeenCalledWith( + command, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('Output: result'); }); - mockShellExecute.mockReturnValue({ - result: Promise.resolve({ output: 'empty output' }), + + it('should handle deeply nested braces correctly', async () => { + const processor = new ShellProcessor('test-command'); + const command = "echo '{{a},{b}}'"; + const prompt = `!{${command}}`; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: '{{a},{b}}' }), + }); + + const result = await processor.process(prompt, context); + expect(mockShellExecute).toHaveBeenCalledWith( + command, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + expect(result).toBe('{{a},{b}}'); }); - const result = await processor.process(prompt, context); + it('should throw an error for unclosed shell injections', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'This prompt is broken: !{ls -l'; - 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'); + await expect(processor.process(prompt, context)).rejects.toThrow( + /Unclosed shell injection/, + ); + }); + + it('should throw an error for unclosed nested braces', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Broken: !{echo {a}'; + + await expect(processor.process(prompt, context)).rejects.toThrow( + /Unclosed shell injection/, + ); + }); + }); + + describe('Error Reporting', () => { + it('should append stderr information if the command produces it', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'some output', + stderr: 'some error', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe('some output\n--- STDERR ---\nsome error'); + }); + + it('should handle stderr-only output correctly', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: '', + stderr: 'error only', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe('--- STDERR ---\nerror only'); + }); + + it('should append exit code and command name if the command fails', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Run a failing command: !{exit 1}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'some error output', + stderr: '', + exitCode: 1, + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "Run a failing command: some error output\n[Shell command 'exit 1' exited with code 1]", + ); + }); + + it('should append signal info and command name if terminated by signal', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'output', + stderr: '', + exitCode: null, + signal: 'SIGTERM', + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "output\n[Shell command 'cmd' terminated by signal SIGTERM]", + ); + }); + + it('should append stderr and exit code information correctly', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{cmd}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'out', + stderr: 'err', + exitCode: 127, + }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe( + "out\n--- STDERR ---\nerr\n[Shell command 'cmd' exited with code 127]", + ); + }); + + it('should throw a detailed error if the shell fails to spawn', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{bad-command}'; + const spawnError = new Error('spawn EACCES'); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: '', + stderr: '', + exitCode: null, + error: spawnError, + aborted: false, + }), + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + "Failed to start shell command in 'test-command': spawn EACCES. Command: bad-command", + ); + }); + + it('should report abort status with command name if aborted', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{long-running-command}'; + const spawnError = new Error('Aborted'); + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ + ...SUCCESS_RESULT, + stdout: 'partial output', + stderr: '', + exitCode: null, + error: spawnError, + aborted: true, // Key difference + }), + }); + + const result = await processor.process(prompt, context); + expect(result).toBe( + "partial output\n[Shell command 'long-running-command' aborted]", + ); + }); + }); + + describe('Context-Aware Argument Interpolation ({{args}})', () => { + const rawArgs = 'user input'; + + beforeEach(() => { + // Update context for these tests to use specific arguments + context.invocation!.args = rawArgs; + }); + + it('should perform raw replacement if no shell injections are present (optimization path)', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'The user said: {{args}}'; + + const result = await processor.process(prompt, context); + + expect(result).toBe(`The user said: ${rawArgs}`); + expect(mockShellExecute).not.toHaveBeenCalled(); + }); + + it('should perform raw replacement outside !{} blocks', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Outside: {{args}}. Inside: !{echo "hello"}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'hello' }), + }); + + const result = await processor.process(prompt, context); + + expect(result).toBe(`Outside: ${rawArgs}. Inside: hello`); + }); + + it('should perform escaped replacement inside !{} blocks', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'Command: !{grep {{args}} file.txt}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'match found' }), + }); + + const result = await processor.process(prompt, context); + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedCommand = `grep ${expectedEscapedArgs} file.txt`; + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + + expect(result).toBe('Command: match found'); + }); + + it('should handle both raw (outside) and escaped (inside) injection simultaneously', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = 'User "({{args}})" requested search: !{search {{args}}}'; + mockShellExecute.mockReturnValue({ + result: Promise.resolve({ ...SUCCESS_RESULT, stdout: 'results' }), + }); + + const result = await processor.process(prompt, context); + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedCommand = `search ${expectedEscapedArgs}`; + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + + expect(result).toBe(`User "(${rawArgs})" requested search: results`); + }); + + it('should perform security checks on the final, resolved (escaped) command', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{rm {{args}}}'; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: [expectedResolvedCommand], + isHardDenial: false, + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + ConfirmationRequiredError, + ); + + expect(mockCheckCommandPermissions).toHaveBeenCalledWith( + expectedResolvedCommand, + expect.any(Object), + context.session.sessionShellAllowlist, + ); + }); + + it('should report the resolved command if a hard denial occurs', async () => { + const processor = new ShellProcessor('test-command'); + const prompt = '!{rm {{args}}}'; + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(rawArgs); + const expectedResolvedCommand = `rm ${expectedEscapedArgs}`; + mockCheckCommandPermissions.mockReturnValue({ + allAllowed: false, + disallowedCommands: [expectedResolvedCommand], + isHardDenial: true, + blockReason: 'It is forbidden.', + }); + + await expect(processor.process(prompt, context)).rejects.toThrow( + `Blocked command: "${expectedResolvedCommand}". Reason: It is forbidden.`, + ); + }); + }); + describe('Real-World Escaping Scenarios', () => { + it('should correctly handle multiline arguments', async () => { + const processor = new ShellProcessor('test-command'); + const multilineArgs = 'first line\nsecond line'; + context.invocation!.args = multilineArgs; + const prompt = 'Commit message: !{git commit -m {{args}}}'; + + const expectedEscapedArgs = + getExpectedEscapedArgForPlatform(multilineArgs); + const expectedCommand = `git commit -m ${expectedEscapedArgs}`; + + await processor.process(prompt, context); + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); + + it.each([ + { name: 'spaces', input: 'file with spaces.txt' }, + { name: 'double quotes', input: 'a "quoted" string' }, + { name: 'single quotes', input: "it's a string" }, + { name: 'command substitution (backticks)', input: '`reboot`' }, + { name: 'command substitution (dollar)', input: '$(reboot)' }, + { name: 'variable expansion', input: '$HOME' }, + { name: 'command chaining (semicolon)', input: 'a; reboot' }, + { name: 'command chaining (ampersand)', input: 'a && reboot' }, + ])('should safely escape args containing $name', async ({ input }) => { + const processor = new ShellProcessor('test-command'); + context.invocation!.args = input; + const prompt = '!{echo {{args}}}'; + + const expectedEscapedArgs = getExpectedEscapedArgForPlatform(input); + const expectedCommand = `echo ${expectedEscapedArgs}`; + + await processor.process(prompt, context); + + expect(mockShellExecute).toHaveBeenCalledWith( + expectedCommand, + expect.any(String), + expect.any(Function), + expect.any(Object), + ); + }); }); }); |
