feat: Add tests for checkpoint tag sanitization (#4882)
This commit is contained in:
parent
b08679c906
commit
1c1aa047ff
|
@ -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'),
|
||||||
|
|
|
@ -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 [];
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue