refactor: use validate files on the tool level

Author: Lightning00Blade

src/ToolHandler.ts

@@ -220,6 +220,12 @@ export class ToolHandler {
 
       response.setRedactNetworkHeaders(this.serverArgs.redactNetworkHeaders);
       try {
+        if (this.tool.verifyFilesSchema) {
+          for (const key of this.tool.verifyFilesSchema) {
+            const filePath = params[key];
+            await context.validatePath(filePath as string);
+          }
+        }
         if (isPageScopedTool(this.tool)) {
           const pageId =
             typeof params.pageId === 'number' ? params.pageId : undefined;

src/tools/ToolDefinition.ts

@@ -32,9 +32,7 @@ import type {
   ToolDefinition as ThirdPartyDeveloperToolDefinition,
 } from './thirdPartyDeveloper.js';
 
-export interface BaseToolDefinition<
-  Schema extends zod.ZodRawShape = zod.ZodRawShape,
-> {
+export interface BaseToolDefinition<Schema extends zod.ZodRawShape = zod.ZodRawShape> {
   name: string;
   description: string;
   annotations: {
@@ -48,6 +46,7 @@ export interface BaseToolDefinition<
   };
   schema: Schema;
   blockedByDialog: boolean;
+  verifyFilesSchema?: Array<keyof Schema>;
 }
 
 export interface ToolDefinition<

src/tools/extensions.ts

@@ -22,9 +22,9 @@ export const installExtension = defineTool({
       .describe('Absolute path to the unpacked extension folder.'),
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['path'],
   handler: async (request, response, context) => {
     const {path} = request.params;
-    await context.validatePath(path);
     const id = await context.installExtension(path);
     response.appendResponseLine(`Extension installed. Id: ${id}`);
   },
@@ -80,7 +80,6 @@ export const reloadExtension = defineTool({
     if (!extension) {
       throw new Error(`Extension with ID ${id} not found.`);
     }
-    await context.validatePath(extension.path);
     await context.installExtension(extension.path);
     response.appendResponseLine('Extension reloaded.');
   },

src/tools/input.ts

@@ -461,9 +461,9 @@ export const uploadFile = definePageTool({
     includeSnapshot: includeSnapshotSchema,
   },
   blockedByDialog: true,
-  handler: async (request, response, context) => {
+  verifyFilesSchema: ['filePath'],
+  handler: async (request, response, _context) => {
     const {uid, filePath} = request.params;
-    await context.validatePath(filePath);
     const handle = (await request.page.getElementByUid(
       uid,
     )) as ElementHandle<HTMLInputElement>;

src/tools/lighthouse.ts

@@ -44,6 +44,7 @@ export const lighthouseAudit = definePageTool({
       .describe('Directory for reports. If omitted, uses temporary files.'),
   },
   blockedByDialog: true,
+  verifyFilesSchema: ['outputDirPath'],
   handler: async (request, response, context) => {
     const page = request.page;
     const categories = [
@@ -59,8 +60,6 @@ export const lighthouseAudit = definePageTool({
       outputDirPath,
     } = request.params;
 
-    await context.validatePath(outputDirPath);
-
     const flags: Flags = {
       onlyCategories: categories,
       output: formats,

src/tools/memory.ts

@@ -23,9 +23,9 @@ export const takeHeapSnapshot = definePageTool({
       .describe('A path to a .heapsnapshot file to save the heapsnapshot to.'),
   },
   blockedByDialog: true,
-  handler: async (request, response, context) => {
+  verifyFilesSchema: ['filePath'],
+  handler: async (request, response) => {
     const page = request.page;
-    await context.validatePath(request.params.filePath);
 
     await page.pptrPage.captureHeapSnapshot({
       path: ensureExtension(request.params.filePath, '.heapsnapshot'),
@@ -50,8 +50,8 @@ export const getHeapSnapshotSummary = defineTool({
     filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     const stats = await context.getHeapSnapshotStats(request.params.filePath);
     const staticData = await context.getHeapSnapshotStaticData(
       request.params.filePath,
@@ -82,8 +82,8 @@ export const getHeapSnapshotDetails = defineTool({
       .describe('The page size for pagination of aggregates.'),
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     const aggregates = await context.getHeapSnapshotAggregates(
       request.params.filePath,
     );
@@ -111,8 +111,8 @@ export const getHeapSnapshotClassNodes = defineTool({
     pageSize: zod.number().optional().describe('The page size for pagination.'),
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     const nodes = await context.getHeapSnapshotNodesById(
       request.params.filePath,
       request.params.id,
@@ -135,15 +135,14 @@ export const getHeapSnapshotRetainers = defineTool({
     conditions: ['experimentalMemory'],
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['filePath'],
   schema: {
     filePath: zod.string().describe('A path to a .heapsnapshot file to read.'),
     nodeId: zod.number().describe('The node ID to get retainers for.'),
     pageIdx: zod.number().optional().describe('The page index for pagination.'),
     pageSize: zod.number().optional().describe('The page size for pagination.'),
   },
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
-
     const retainers = await context.getHeapSnapshotRetainers(
       request.params.filePath,
       request.params.nodeId,

src/tools/network.ts

@@ -115,9 +115,8 @@ export const getNetworkRequest = definePageTool({
       ),
   },
   blockedByDialog: true,
+  verifyFilesSchema: ['requestFilePath', 'responseFilePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.requestFilePath);
-    await context.validatePath(request.params.responseFilePath);
     if (request.params.reqid) {
       response.attachNetworkRequest(request.params.reqid, {
         requestFilePath: request.params.requestFilePath,

src/tools/performance.ts

@@ -48,8 +48,8 @@ export const startTrace = definePageTool({
     filePath: filePathSchema,
   },
   blockedByDialog: true,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     if (context.isRunningPerformanceTrace()) {
       response.appendResponseLine(
         'Error: a performance trace is already running. Use performance_stop_trace to stop it. Only one trace can be running at any given time.',
@@ -127,8 +127,8 @@ export const stopTrace = definePageTool({
     filePath: filePathSchema,
   },
   blockedByDialog: true,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     if (!context.isRunningPerformanceTrace()) {
       return;
     }

src/tools/screencast.ts

@@ -39,8 +39,8 @@ export const startScreencast = definePageTool(args => ({
       ),
   },
   blockedByDialog: false,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     if (context.getScreenRecorder() !== null) {
       response.appendResponseLine(
         'Error: a screencast recording is already in progress. Use screencast_stop to stop it before starting a new one.',

src/tools/screenshot.ts

@@ -51,8 +51,8 @@ export const screenshot = definePageTool({
       ),
   },
   blockedByDialog: true,
+  verifyFilesSchema: ['filePath'],
   handler: async (request, response, context) => {
-    await context.validatePath(request.params.filePath);
     if (request.params.uid && request.params.fullPage) {
       throw new Error('Providing both "uid" and "fullPage" is not allowed.');
     }