feat: ensure extensions for file outputs (#1867)

OrKoN · Lightning00Blade · web-flow · commit e7a0d5097785 · 2026-04-16T06:08:22.000Z
This PR ensures the extensions for the file outputs of different types minimizing the chance of misuse. The input filePath, thus, might be modified but it should not be an issue for clients as the final output path is returned to the clients in the response. Closes #1864 --------- Co-authored-by: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>
diff --git a/src/McpContext.ts b/src/McpContext.ts
@@ -33,7 +33,11 @@ import {Locator} from './third_party/index.js';
 import {PredefinedNetworkConditions} from './third_party/index.js';
 import {listPages} from './tools/pages.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
-import type {Context, DevToolsData} from './tools/ToolDefinition.js';
+import type {
+  Context,
+  DevToolsData,
+  SupportedExtensions,
+} from './tools/ToolDefinition.js';
 import type {TraceResult} from './trace-processing/parse.js';
 import type {
   EmulationSettings,
@@ -46,7 +50,7 @@ import {
   ExtensionRegistry,
   type InstalledExtension,
 } from './utils/ExtensionRegistry.js';
-import {saveTemporaryFile} from './utils/files.js';
+import {ensureExtension, saveTemporaryFile} from './utils/files.js';
 import {getNetworkMultiplierFromString} from './WaitForHelper.js';
 
 interface McpContextOptions {
@@ -801,10 +805,14 @@ export class McpContext implements Context {
   }
   async saveFile(
     data: Uint8Array<ArrayBufferLike>,
-    filename: string,
+    clientProvidedFilePath: string,
+    extension: SupportedExtensions,
   ): Promise<{filename: string}> {
     try {
-      const filePath = path.resolve(filename);
+      const filePath = ensureExtension(
+        path.resolve(clientProvidedFilePath),
+        extension,
+      );
       await fs.mkdir(path.dirname(filePath), {recursive: true});
       await fs.writeFile(filePath, data);
       return {filename: filePath};
diff --git a/src/McpResponse.ts b/src/McpResponse.ts
@@ -403,11 +403,12 @@ export class McpResponse implements Response {
       if (textSnapshot) {
         const formatter = new SnapshotFormatter(textSnapshot);
         if (this.#snapshotParams.filePath) {
-          await context.saveFile(
+          const result = await context.saveFile(
             new TextEncoder().encode(formatter.toString()),
             this.#snapshotParams.filePath,
+            '.txt',
           );
-          snapshot = this.#snapshotParams.filePath;
+          snapshot = result.filename;
         } else {
           snapshot = formatter;
         }
@@ -429,7 +430,8 @@ export class McpResponse implements Response {
         fetchData: true,
         requestFilePath: this.#attachedNetworkRequestOptions?.requestFilePath,
         responseFilePath: this.#attachedNetworkRequestOptions?.responseFilePath,
-        saveFile: (data, filename) => context.saveFile(data, filename),
+        saveFile: (data, filename, extension) =>
+          context.saveFile(data, filename, extension),
         redactNetworkHeaders: this.#redactNetworkHeaders,
       });
       detailedNetworkRequest = formatter;
@@ -573,7 +575,8 @@ export class McpResponse implements Response {
                 context.getNetworkRequestStableId(request) ===
                 this.#networkRequestsOptions?.networkRequestIdInDevToolsUI,
               fetchData: false,
-              saveFile: (data, filename) => context.saveFile(data, filename),
+              saveFile: (data, filename, extension) =>
+                context.saveFile(data, filename, extension),
               redactNetworkHeaders: this.#redactNetworkHeaders,
             }),
           ),
diff --git a/src/formatters/NetworkFormatter.ts b/src/formatters/NetworkFormatter.ts
@@ -24,6 +24,7 @@ export interface NetworkFormatterOptions {
   saveFile?: (
     data: Uint8Array<ArrayBufferLike>,
     filename: string,
+    extension: '.network-request' | '.network-response',
   ) => Promise<{filename: string}>;
   redactNetworkHeaders: boolean;
 }
@@ -88,11 +89,12 @@ export class NetworkFormatter {
           throw new Error('saveFile is not provided');
         }
         if (data) {
-          await this.#options.saveFile(
+          const result = await this.#options.saveFile(
             Buffer.from(data),
             this.#options.requestFilePath,
+            '.network-request',
           );
-          this.#requestBodyFilePath = this.#options.requestFilePath;
+          this.#requestBodyFilePath = result.filename;
         } else {
           this.#requestBody = requestBodyNotAvailableMessage;
         }
@@ -119,8 +121,12 @@ export class NetworkFormatter {
           if (!this.#options.saveFile) {
             throw new Error('saveFile is not provided');
           }
-          await this.#options.saveFile(buffer, this.#options.responseFilePath);
-          this.#responseBodyFilePath = this.#options.responseFilePath;
+          const result = await this.#options.saveFile(
+            buffer,
+            this.#options.responseFilePath,
+            '.network-response',
+          );
+          this.#responseBodyFilePath = result.filename;
         } catch {
           // Flatten error handling for buffer() failure and save failure
         }
diff --git a/src/tools/ToolDefinition.ts b/src/tools/ToolDefinition.ts
@@ -136,6 +136,18 @@ export interface Response {
   setListInPageTools(): void;
 }
 
+export type SupportedExtensions =
+  | '.png'
+  | '.jpeg'
+  | '.webp'
+  | '.json'
+  | '.network-response'
+  | '.network-request'
+  | '.html'
+  | '.txt'
+  | '.csv'
+  | '.json.gz';
+
 /**
  * Only add methods required by tools/*.
  */
@@ -170,7 +182,8 @@ export type Context = Readonly<{
   ): Promise<{filepath: string}>;
   saveFile(
     data: Uint8Array<ArrayBufferLike>,
-    filename: string,
+    clientProvidedFilePath: string,
+    extension: SupportedExtensions,
   ): Promise<{filename: string}>;
   waitForTextOnPage(
     text: string[],
diff --git a/src/tools/lighthouse.ts b/src/tools/lighthouse.ts
@@ -107,8 +107,12 @@ export const lighthouseAudit = definePageTool({
       const report = generateReport(lhr, format);
       const data = encoder.encode(report);
       if (outputDirPath) {
-        const reportPath = path.join(outputDirPath, `report.${format}`);
-        const {filename} = await context.saveFile(data, reportPath);
+        const reportPath = path.join(outputDirPath, `report`);
+        const {filename} = await context.saveFile(
+          data,
+          reportPath,
+          `.${format}`,
+        );
         reportPaths.push(filename);
       } else {
         const {filepath} = await context.saveTemporaryFile(
diff --git a/src/tools/memory.ts b/src/tools/memory.ts
@@ -5,6 +5,7 @@
  */
 
 import {zod} from '../third_party/index.js';
+import {ensureExtension} from '../utils/files.js';
 
 import {ToolCategory} from './categories.js';
 import {definePageTool} from './ToolDefinition.js';
@@ -25,7 +26,7 @@ export const takeMemorySnapshot = definePageTool({
     const page = request.page;
 
     await page.pptrPage.captureHeapSnapshot({
-      path: request.params.filePath,
+      path: ensureExtension(request.params.filePath, '.heapsnapshot'),
     });
 
     response.appendResponseLine(
diff --git a/src/tools/performance.ts b/src/tools/performance.ts
@@ -197,7 +197,11 @@ async function stopTracingAndAppendOutput(
           });
         });
       }
-      const file = await context.saveFile(dataToWrite, filePath);
+      const file = await context.saveFile(
+        dataToWrite,
+        filePath,
+        filePath.endsWith('.gz') ? '.json.gz' : '.json',
+      );
       response.appendResponseLine(
         `The raw trace data was saved to ${file.filename}.`,
       );
diff --git a/src/tools/screencast.ts b/src/tools/screencast.ts
@@ -10,6 +10,7 @@ import path from 'node:path';
 
 import {zod} from '../third_party/index.js';
 import type {ScreenRecorder} from '../third_party/index.js';
+import {ensureExtension} from '../utils/files.js';
 
 import {ToolCategory} from './categories.js';
 import {definePageTool} from './ToolDefinition.js';
@@ -46,7 +47,7 @@ export const startScreencast = definePageTool({
     }
 
     const filePath = request.params.path ?? (await generateTempFilePath());
-    const resolvedPath = path.resolve(filePath);
+    const resolvedPath = ensureExtension(path.resolve(filePath), '.mp4');
 
     const page = request.page;
 
diff --git a/src/tools/screenshot.ts b/src/tools/screenshot.ts
@@ -87,8 +87,12 @@ export const screenshot = definePageTool({
     }
 
     if (request.params.filePath) {
-      const file = await context.saveFile(screenshot, request.params.filePath);
-      response.appendResponseLine(`Saved screenshot to ${file.filename}.`);
+      const result = await context.saveFile(
+        screenshot,
+        request.params.filePath,
+        `.${format}`,
+      );
+      response.appendResponseLine(`Saved screenshot to ${result.filename}.`);
     } else if (screenshot.length >= 2_000_000) {
       const {filepath} = await context.saveTemporaryFile(
         screenshot,
diff --git a/src/utils/files.ts b/src/utils/files.ts
@@ -24,3 +24,11 @@ export async function saveTemporaryFile(
     throw new Error('Could not save a file', {cause: err});
   }
 }
+
+export function ensureExtension(
+  filepath: string,
+  extension: `.${string}`,
+): string {
+  const ext = path.extname(filepath);
+  return filepath.slice(0, filepath.length - ext.length) + extension;
+}
diff --git a/tests/McpResponse.test.ts b/tests/McpResponse.test.ts
@@ -156,7 +156,7 @@ describe('McpResponse', () => {
   });
 
   it('saves snapshot to file and returns structured content', async t => {
-    const filePath = join(tmpdir(), 'test-screenshot.png');
+    const filePath = join(tmpdir(), 'test-snapshot.txt');
     try {
       await withMcpContext(async (response, context) => {
         const page = context.getSelectedPptrPage();
diff --git a/tests/utils/files.test.ts b/tests/utils/files.test.ts
@@ -0,0 +1,43 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'node:assert';
+import {describe, it} from 'node:test';
+
+import {ensureExtension} from '../../src/utils/files.js';
+
+describe('ensureExtension', () => {
+  it('should add an extension to a filename without one', () => {
+    assert.strictEqual(ensureExtension('filename', '.txt'), 'filename.txt');
+  });
+
+  it('should replace an existing extension', () => {
+    assert.strictEqual(ensureExtension('filename.jpg', '.txt'), 'filename.txt');
+  });
+
+  it('should handle extension without a leading dot', () => {
+    assert.strictEqual(ensureExtension('filename', '.txt'), 'filename.txt');
+  });
+
+  it('should not add a second dot if already present', () => {
+    assert.strictEqual(ensureExtension('filename.txt', '.txt'), 'filename.txt');
+  });
+
+  it('should handle paths with directories', () => {
+    assert.strictEqual(
+      ensureExtension('/path/to/file.jpg', '.png'),
+      '/path/to/file.png',
+    );
+  });
+
+  it('should handle hidden files (starting with dot)', () => {
+    assert.strictEqual(ensureExtension('.bashrc', '.txt'), '.bashrc.txt');
+  });
+
+  it('should handle complex extensions (like .tar.gz) - path.extname only gets the last one', () => {
+    assert.strictEqual(ensureExtension('file.tar.gz', '.zip'), 'file.tar.zip');
+  });
+});

Original file line number	Diff line number	Diff line change
`@@ -24,3 +24,11 @@ export async function saveTemporaryFile(`
`24`	`24`	`throw new Error('Could not save a file', {cause: err});`
`25`	`25`	`}`
`26`	`26`	`}`
	`27`	`+`
	`28`	`+export function ensureExtension(`
	`29`	`+ filepath: string,`
	`30`	+ extension: `.${string}`,
	`31`	`+): string {`
	`32`	`+ const ext = path.extname(filepath);`
	`33`	`+ return filepath.slice(0, filepath.length - ext.length) + extension;`
	`34`	`+}`