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

2 files changed, 26 insertions, 19 deletions

diff --git a/packages/core/src/core/logger.test.ts b/packages/core/src/core/logger.test.ts
index c74b92cf..3f243b52 100644
--- a/packages/core/src/core/logger.test.ts
+++ b/packages/core/src/core/logger.test.ts
@@ -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'),
diff --git a/packages/core/src/core/logger.ts b/packages/core/src/core/logger.ts
index 2be9f1d4..9f4622e7 100644
--- a/packages/core/src/core/logger.ts
+++ b/packages/core/src/core/logger.ts
@@ -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 [];
     }
   }