feat: ensure extensions for file outputs (#1867)

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>

Author: 2 people

src/McpContext.ts

@@ -37,6 +37,7 @@ import type {
   Context,
   DevToolsData,
   ContextPage,
+  SupportedExtensions,
 } from './tools/ToolDefinition.js';
 import type {TraceResult} from './trace-processing/parse.js';
 import type {
@@ -50,7 +51,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 {
@@ -954,10 +955,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};

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,
             }),
           ),

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
         }

src/tools/ToolDefinition.ts

@@ -138,6 +138,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/*.
  */
@@ -172,7 +184,8 @@ export type Context = Readonly<{
   ): Promise<{filepath: string}>;
   saveFile(
     data: Uint8Array<ArrayBufferLike>,
-    filename: string,
+    clientProvidedFilePath: string,
+    extension: SupportedExtensions,
   ): Promise<{filename: string}>;
   waitForTextOnPage(
     text: string[],

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(

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(

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}.`,
       );

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;
 

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,

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;
+}