fix: validate extension-enforced output paths (#2269)

Fixes output path validation so tools cannot validate one path and then
write to a different canonical target after extension enforcement.

Changes:
- Resolve dangling symlinks to their target path during
canonicalization.
- Validate the final extension-enforced output path before writing.
- Apply the same final-path validation to heap snapshots and
screencasts.
- Add regression coverage for dangling symlinks that point outside
configured roots.

Validation:
- npm run format
- npm run check-format
- npm run test tests/utils/files.test.ts
- npm run test tests/roots.test.ts
- npm run test tests/tools/memory.test.ts tests/tools/screencast.test.ts

Note: I also ran the full npm test suite locally. The targeted tests
above passed, but the full suite hit local WSL daemon/e2e startup
timeouts while waiting for daemon.pid / server_start, which appear
unrelated to this path-validation change.

---------

Co-authored-by: huynhtrungcsc <huynhtrungcsc@users.noreply.github.com>

Author: huynhtrungcsc

src/McpContext.ts

@@ -54,11 +54,7 @@ import type {
   GeolocationOptions,
   ExtensionServiceWorker,
 } from './types.js';
-import {
-  ensureExtension,
-  getTempFilePath,
-  resolveCanonicalPath,
-} from './utils/files.js';
+import {getTempFilePath, resolveCanonicalPath} from './utils/files.js';
 import {getNetworkMultiplierFromString} from './WaitForHelper.js';
 
 interface McpContextOptions {
@@ -258,6 +254,20 @@ export class McpContext implements Context {
     }
   }
 
+  async ensureExtension<Extension extends `.${string}`>(
+    filePath: string,
+    extension: Extension,
+  ): Promise<`${string}${Extension}`> {
+    const resolvedPath = path.resolve(filePath);
+    const currentExtension = path.extname(resolvedPath);
+    const outputPath: `${string}${Extension}` = `${resolvedPath.slice(
+      0,
+      resolvedPath.length - currentExtension.length,
+    )}${extension}`;
+    await this.validatePath(outputPath);
+    return outputPath;
+  }
+
   resolveCdpRequestId(page: McpPage, cdpRequestId: string): number | undefined {
     if (!cdpRequestId) {
       this.logger?.('no network request');
@@ -800,12 +810,11 @@ export class McpContext implements Context {
     clientProvidedFilePath: string,
     extension: SupportedExtensions,
   ): Promise<{filename: string}> {
-    await this.validatePath(clientProvidedFilePath);
+    const filePath = await this.ensureExtension(
+      clientProvidedFilePath,
+      extension,
+    );
     try {
-      const filePath = ensureExtension(
-        path.resolve(clientProvidedFilePath),
-        extension,
-      );
       await fs.mkdir(path.dirname(filePath), {recursive: true});
       await fs.writeFile(filePath, data);
       return {filename: filePath};

src/tools/ToolDefinition.ts

@@ -187,6 +187,10 @@ export type SupportedExtensions =
  */
 export type Context = Readonly<{
   validatePath(filePath?: string): Promise<void>;
+  ensureExtension<Extension extends `.${string}`>(
+    filePath: string,
+    extension: Extension,
+  ): Promise<`${string}${Extension}`>;
   isRunningPerformanceTrace(): boolean;
   setIsRunningPerformanceTrace(x: boolean): void;
   isCruxEnabled(): boolean;

src/tools/memory.ts

@@ -5,7 +5,6 @@
  */
 
 import {zod} from '../third_party/index.js';
-import {ensureExtension} from '../utils/files.js';
 
 import {ToolCategory} from './categories.js';
 import {definePageTool, defineTool} from './ToolDefinition.js';
@@ -24,16 +23,18 @@ export const takeHeapSnapshot = definePageTool({
   },
   blockedByDialog: true,
   verifyFilesSchema: ['filePath'],
-  handler: async (request, response) => {
+  handler: async (request, response, context) => {
     const page = request.page;
+    const snapshotPath = await context.ensureExtension(
+      request.params.filePath,
+      '.heapsnapshot',
+    );
 
     await page.pptrPage.captureHeapSnapshot({
-      path: ensureExtension(request.params.filePath, '.heapsnapshot'),
+      path: snapshotPath,
     });
 
-    response.appendResponseLine(
-      `Heap snapshot saved to ${request.params.filePath}`,
-    );
+    response.appendResponseLine(`Heap snapshot saved to ${snapshotPath}`);
   },
 });
 

src/tools/screencast.ts

@@ -10,7 +10,6 @@ import path from 'node:path';
 
 import {zod} from '../third_party/index.js';
 import type {ScreenRecorder, VideoFormat} from '../third_party/index.js';
-import {ensureExtension} from '../utils/files.js';
 
 import {ToolCategory} from './categories.js';
 import {definePageTool} from './ToolDefinition.js';
@@ -20,7 +19,9 @@ async function generateTempFilePath(): Promise<string> {
   return path.join(dir, `screencast.mp4`);
 }
 
-const supportedExtensions: Array<`.${string}`> = ['.webm', '.mp4'];
+type SupportedVideoExtension = '.webm' | '.mp4';
+
+const supportedExtensions: SupportedVideoExtension[] = ['.webm', '.mp4'];
 
 export const startScreencast = definePageTool(args => ({
   name: 'screencast_start',
@@ -68,14 +69,15 @@ export const startScreencast = definePageTool(args => ({
           `Supported formats: ${supportedExtensions.join(', ')} (case-insensitive).`,
       );
     }
-    const enforcedExtension: `.${string}` = matchedExtension ?? '.mp4';
+    const enforcedExtension: SupportedVideoExtension =
+      matchedExtension ?? '.mp4';
     const format: VideoFormat = (matchedExtension?.substring(1) ??
       'mp4') as VideoFormat;
 
-    const resolvedPath = ensureExtension(
-      path.resolve(filePath),
+    const resolvedPath = await context.ensureExtension(
+      filePath,
       enforcedExtension,
-    ) as `${string}.webm`;
+    );
 
     const page = request.page;
 

src/utils/files.ts

@@ -15,14 +15,6 @@ export async function getTempFilePath(filename: string) {
   return filepath;
 }
 
-export function ensureExtension(
-  filepath: string,
-  extension: `.${string}`,
-): string {
-  const ext = path.extname(filepath);
-  return filepath.slice(0, filepath.length - ext.length) + extension;
-}
-
 export async function resolveCanonicalPath(filePath: string): Promise<string> {
   const absolutePath = path.resolve(filePath);
   try {

tests/roots.test.ts

@@ -54,4 +54,86 @@ describe('McpContext Roots', () => {
       }
     });
   });
+
+  it('should enforce extensions and validate the output path', async () => {
+    await withMcpContext(async (_response, context) => {
+      const workspacePath = await fs.mkdtemp(
+        path.join(os.tmpdir(), 'workspace-root-'),
+      );
+      try {
+        context.setRoots([
+          {uri: pathToFileURL(workspacePath).href, name: 'workspace'},
+        ]);
+
+        const testCases: Array<{
+          filePath: string;
+          extension: '.json' | '.txt' | '.png' | '.zip';
+          expected: string;
+        }> = [
+          {
+            filePath: 'result',
+            extension: '.json',
+            expected: 'result.json',
+          },
+          {
+            filePath: 'result.jpg',
+            extension: '.txt',
+            expected: 'result.txt',
+          },
+          {
+            filePath: 'nested/result.jpg',
+            extension: '.png',
+            expected: 'nested/result.png',
+          },
+          {
+            filePath: '.bashrc',
+            extension: '.txt',
+            expected: '.bashrc.txt',
+          },
+          {
+            filePath: 'file.tar.gz',
+            extension: '.zip',
+            expected: 'file.tar.zip',
+          },
+        ];
+
+        for (const testCase of testCases) {
+          const resolvedPath = await context.ensureExtension(
+            path.join(workspacePath, testCase.filePath),
+            testCase.extension,
+          );
+
+          assert.strictEqual(
+            resolvedPath,
+            path.join(workspacePath, testCase.expected),
+          );
+        }
+      } finally {
+        await fs.rm(workspacePath, {recursive: true, force: true});
+      }
+    });
+  });
+
+  it('should deny extension-enforced paths outside roots', async () => {
+    await withMcpContext(async (_response, context) => {
+      const workspacePath = await fs.mkdtemp(
+        path.join(os.tmpdir(), 'workspace-root-'),
+      );
+      try {
+        context.setRoots([
+          {uri: pathToFileURL(workspacePath).href, name: 'workspace'},
+        ]);
+
+        await assert.rejects(
+          context.ensureExtension(
+            path.join(os.homedir(), 'outside-root-result'),
+            '.json',
+          ),
+          /Access denied/,
+        );
+      } finally {
+        await fs.rm(workspacePath, {recursive: true, force: true});
+      }
+    });
+  });
 });

tests/utils/files.test.ts

@@ -10,40 +10,7 @@ import os from 'node:os';
 import path from 'node:path';
 import {describe, it} from 'node:test';
 
-import {ensureExtension, resolveCanonicalPath} 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');
-  });
-});
+import {resolveCanonicalPath} from '../../src/utils/files.js';
 
 describe('resolveCanonicalPath', () => {
   it('should resolve an existing standard file path', async () => {