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' }] }, { role: 'model', parts: [{ text: 'Hi there' }] },
]; ];
it('should save a checkpoint to a tagged file when a tag is provided', async () => { it.each([
const tag = 'my-test-tag'; { 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); await logger.saveCheckpoint(conversation, tag);
const taggedFilePath = path.join( const taggedFilePath = path.join(
TEST_GEMINI_DIR, TEST_GEMINI_DIR,
`${CHECKPOINT_FILE_NAME.replace('.json', '')}-${tag}.json`, `checkpoint-${sanitizedTag}.json`,
); );
const fileContent = await fs.readFile(taggedFilePath, 'utf-8'); const fileContent = await fs.readFile(taggedFilePath, 'utf-8');
expect(JSON.parse(fileContent)).toEqual(conversation); 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 () => { it.each([
const tag = 'my-load-tag'; { 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 = [ const taggedConversation = [
...conversation, ...conversation,
{ role: 'user', parts: [{ text: 'Another message' }] }, { role: 'user', parts: [{ text: 'hello' }] },
]; ];
const taggedFilePath = path.join( const taggedFilePath = path.join(
TEST_GEMINI_DIR, TEST_GEMINI_DIR,
`${CHECKPOINT_FILE_NAME.replace('.json', '')}-${tag}.json`, `checkpoint-${sanitizedTag}.json`,
); );
await fs.writeFile( await fs.writeFile(
taggedFilePath, taggedFilePath,
@ -464,11 +472,16 @@ describe('Logger', () => {
}); });
it('should return an empty array if the file contains invalid JSON', async () => { 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 const consoleErrorSpy = vi
.spyOn(console, 'error') .spyOn(console, 'error')
.mockImplementation(() => {}); .mockImplementation(() => {});
const loadedCheckpoint = await logger.loadCheckpoint('missing'); const loadedCheckpoint = await logger.loadCheckpoint(tag);
expect(loadedCheckpoint).toEqual([]); expect(loadedCheckpoint).toEqual([]);
expect(consoleErrorSpy).toHaveBeenCalledWith( expect(consoleErrorSpy).toHaveBeenCalledWith(
expect.stringContaining('Failed to read or parse checkpoint file'), 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.'); throw new Error('Checkpoint file path not set.');
} }
// Sanitize tag to prevent directory traversal attacks // Sanitize tag to prevent directory traversal attacks
tag = tag.replace(/[^a-zA-Z0-9-_]/g, ''); let sanitizedTag = tag.replace(/[^a-zA-Z0-9-_]/g, '');
if (!tag) { if (!sanitizedTag) {
console.error('Sanitized tag is empty setting to "default".'); sanitizedTag = 'default';
tag = '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> { async saveCheckpoint(conversation: Content[], tag: string): Promise<void> {
@ -283,11 +282,6 @@ export class Logger {
return parsedContent as Content[]; return parsedContent as Content[];
} catch (error) { } catch (error) {
console.error(`Failed to read or parse checkpoint file ${path}:`, 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 []; return [];
} }
} }