Updated logic to be more secure and added more tests

This commit is contained in:
Dev Randalpura
2026-03-03 16:34:56 -05:00
parent dfc28b8f6f
commit aec5b0385d
3 changed files with 120 additions and 20 deletions

View File

@@ -5,11 +5,23 @@
*/
import { renderWithProviders } from '../../test-utils/render.js';
import { describe, it, expect, vi } from 'vitest';
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { SessionSummaryDisplay } from './SessionSummaryDisplay.js';
import * as SessionContext from '../contexts/SessionContext.js';
import type { SessionMetrics } from '../contexts/SessionContext.js';
import { ToolCallDecision } from '@google/gemini-cli-core';
import {
ToolCallDecision,
getShellConfiguration,
} from '@google/gemini-cli-core';
vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actual =
await importOriginal<typeof import('@google/gemini-cli-core')>();
return {
...actual,
getShellConfiguration: vi.fn(),
};
});
vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
const actual = await importOriginal<typeof SessionContext>();
@@ -19,12 +31,16 @@ vi.mock('../contexts/SessionContext.js', async (importOriginal) => {
};
});
const getShellConfigurationMock = vi.mocked(getShellConfiguration);
const useSessionStatsMock = vi.mocked(SessionContext.useSessionStats);
const renderWithMockedStats = async (metrics: SessionMetrics) => {
const renderWithMockedStats = async (
metrics: SessionMetrics,
sessionId = 'test-session',
) => {
useSessionStatsMock.mockReturnValue({
stats: {
sessionId: 'test-session',
sessionId,
sessionStartTime: new Date(),
metrics,
lastPromptTokenCount: 0,
@@ -46,8 +62,38 @@ const renderWithMockedStats = async (metrics: SessionMetrics) => {
};
describe('<SessionSummaryDisplay />', () => {
const emptyMetrics: SessionMetrics = {
models: {},
tools: {
totalCalls: 0,
totalSuccess: 0,
totalFail: 0,
totalDurationMs: 0,
totalDecisions: {
accept: 0,
reject: 0,
modify: 0,
[ToolCallDecision.AUTO_ACCEPT]: 0,
},
byName: {},
},
files: {
totalLinesAdded: 0,
totalLinesRemoved: 0,
},
};
beforeEach(() => {
getShellConfigurationMock.mockReturnValue({
executable: 'bash',
argsPrefix: ['-c'],
shell: 'bash',
});
});
it('renders the summary display with a title', async () => {
const metrics: SessionMetrics = {
...emptyMetrics,
models: {
'gemini-2.5-pro': {
api: { totalRequests: 10, totalErrors: 1, totalLatencyMs: 50234 },
@@ -63,19 +109,6 @@ describe('<SessionSummaryDisplay />', () => {
roles: {},
},
},
tools: {
totalCalls: 0,
totalSuccess: 0,
totalFail: 0,
totalDurationMs: 0,
totalDecisions: {
accept: 0,
reject: 0,
modify: 0,
[ToolCallDecision.AUTO_ACCEPT]: 0,
},
byName: {},
},
files: {
totalLinesAdded: 42,
totalLinesRemoved: 15,
@@ -89,4 +122,70 @@ describe('<SessionSummaryDisplay />', () => {
expect(output).toMatchSnapshot();
unmount();
});
describe('Session ID escaping', () => {
it('renders a standard UUID-formatted session ID in the footer (bash)', async () => {
const uuidSessionId = '1234-abcd-5678-efgh';
const { lastFrame, unmount } = await renderWithMockedStats(
emptyMetrics,
uuidSessionId,
);
const output = lastFrame();
// Standard UUID characters should not be escaped/quoted by default for bash.
expect(output).toContain('gemini --resume 1234-abcd-5678-efgh');
unmount();
});
it('sanitizes a malicious session ID in the footer (bash)', async () => {
const maliciousSessionId = "'; rm -rf / #";
const { lastFrame, unmount } = await renderWithMockedStats(
emptyMetrics,
maliciousSessionId,
);
const output = lastFrame();
// escapeShellArg (using shell-quote for bash) will wrap special characters in double quotes.
expect(output).toContain('gemini --resume "\'; rm -rf / #"');
unmount();
});
it('renders a standard UUID-formatted session ID in the footer (powershell)', async () => {
getShellConfigurationMock.mockReturnValue({
executable: 'powershell.exe',
argsPrefix: ['-NoProfile', '-Command'],
shell: 'powershell',
});
const uuidSessionId = '1234-abcd-5678-efgh';
const { lastFrame, unmount } = await renderWithMockedStats(
emptyMetrics,
uuidSessionId,
);
const output = lastFrame();
// PowerShell wraps strings in single quotes
expect(output).toContain("gemini --resume '1234-abcd-5678-efgh'");
unmount();
});
it('sanitizes a malicious session ID in the footer (powershell)', async () => {
getShellConfigurationMock.mockReturnValue({
executable: 'powershell.exe',
argsPrefix: ['-NoProfile', '-Command'],
shell: 'powershell',
});
const maliciousSessionId = "'; rm -rf / #";
const { lastFrame, unmount } = await renderWithMockedStats(
emptyMetrics,
maliciousSessionId,
);
const output = lastFrame();
// PowerShell wraps in single quotes and escapes internal single quotes by doubling them
expect(output).toContain("gemini --resume '''; rm -rf / #'");
unmount();
});
});
});

View File

@@ -7,6 +7,7 @@
import type React from 'react';
import { StatsDisplay } from './StatsDisplay.js';
import { useSessionStats } from '../contexts/SessionContext.js';
import { escapeShellArg, getShellConfiguration } from '@google/gemini-cli-core';
interface SessionSummaryDisplayProps {
duration: string;
@@ -16,7 +17,8 @@ export const SessionSummaryDisplay: React.FC<SessionSummaryDisplayProps> = ({
duration,
}) => {
const { stats } = useSessionStats();
const footer = `Resume this session with:\n gemini --resume ${stats.sessionId}`;
const { shell } = getShellConfiguration();
const footer = `To resume this session: gemini --resume ${escapeShellArg(stats.sessionId, shell)}`;
return (
<StatsDisplay

View File

@@ -24,8 +24,7 @@ exports[`<SessionSummaryDisplay /> > renders the summary display with a title 1`
│ │
│ Savings Highlight: 500 (50.0%) of input tokens were served from the cache, reducing costs. │
│ │
Resume this session with:
│ gemini --resume test-session │
To resume this session: gemini --resume test-session
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
"
`;