feat: Add tests for checkpoint tag sanitization (#4882)

This commit is contained in:
Seth Troisi 2025-07-28 13:43:39 -07:00 committed by GitHub
parent b08679c906
commit 1c1aa047ff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 26 additions and 19 deletions

View File

@ -393,12 +393,16 @@ describe('Logger', () => {
{ role: 'model', parts: [{ text: 'Hi there' }] },
];
it('should save a checkpoint to a tagged file when a tag is provided', async () => {
const tag = 'my-test-tag';
it.each([
{ tag: 'test-tag', sanitizedTag: 'test-tag' },
{ tag: 'invalid/?*!', sanitizedTag: 'invalid' },
{ tag: '/?*!', sanitizedTag: 'default' },
{ tag: '../../secret', sanitizedTag: 'secret' },
])('should save a checkpoint', async ({ tag, sanitizedTag }) => {
await logger.saveCheckpoint(conversation, tag);
const taggedFilePath = path.join(
TEST_GEMINI_DIR,
`${CHECKPOINT_FILE_NAME.replace('.json', '')}-${tag}.json`,
`checkpoint-${sanitizedTag}.json`,
);
const fileContent = await fs.readFile(taggedFilePath, 'utf-8');
expect(JSON.parse(fileContent)).toEqual(conversation);
@ -433,15 +437,19 @@ describe('Logger', () => {
);
});
it('should load from a tagged checkpoint file when a tag is provided', async () => {
const tag = 'my-load-tag';
it.each([
{ tag: 'load-tag', sanitizedTag: 'load-tag' },
{ tag: 'inv/load?*!', sanitizedTag: 'invload' },
{ tag: '/?*!', sanitizedTag: 'default' },
{ tag: '../../secret', sanitizedTag: 'secret' },
])('should load from a checkpoint', async ({ tag, sanitizedTag }) => {
const taggedConversation = [
...conversation,
{ role: 'user', parts: [{ text: 'Another message' }] },
{ role: 'user', parts: [{ text: 'hello' }] },
];
const taggedFilePath = path.join(
TEST_GEMINI_DIR,
`${CHECKPOINT_FILE_NAME.replace('.json', '')}-${tag}.json`,
`checkpoint-${sanitizedTag}.json`,
);
await fs.writeFile(
taggedFilePath,
@ -464,11 +472,16 @@ describe('Logger', () => {
});
it('should return an empty array if the file contains invalid JSON', async () => {
await fs.writeFile(TEST_CHECKPOINT_FILE_PATH, 'invalid json');
const tag = 'invalid-json-tag';
const taggedFilePath = path.join(
TEST_GEMINI_DIR,
`checkpoint-${tag}.json`,
);
await fs.writeFile(taggedFilePath, 'invalid json');
const consoleErrorSpy = vi
.spyOn(console, 'error')
.mockImplementation(() => {});
const loadedCheckpoint = await logger.loadCheckpoint('missing');
const loadedCheckpoint = await logger.loadCheckpoint(tag);
expect(loadedCheckpoint).toEqual([]);
expect(consoleErrorSpy).toHaveBeenCalledWith(
expect.stringContaining('Failed to read or parse checkpoint file'),

View File

@ -239,12 +239,11 @@ export class Logger {
throw new Error('Checkpoint file path not set.');
}
// Sanitize tag to prevent directory traversal attacks
tag = tag.replace(/[^a-zA-Z0-9-_]/g, '');
if (!tag) {
console.error('Sanitized tag is empty setting to "default".');
tag = 'default';
let sanitizedTag = tag.replace(/[^a-zA-Z0-9-_]/g, '');
if (!sanitizedTag) {
sanitizedTag = 'default';
}
return path.join(this.geminiDir, `checkpoint-${tag}.json`);
return path.join(this.geminiDir, `checkpoint-${sanitizedTag}.json`);
}
async saveCheckpoint(conversation: Content[], tag: string): Promise<void> {
@ -283,11 +282,6 @@ export class Logger {
return parsedContent as Content[];
} catch (error) {
console.error(`Failed to read or parse checkpoint file ${path}:`, error);
const nodeError = error as NodeJS.ErrnoException;
if (nodeError.code === 'ENOENT') {
// File doesn't exist, which is fine. Return empty array.
return [];
}
return [];
}
}