fix: add argument validation - MCP-188

src/tools/args.ts

@@ -0,0 +1,70 @@
+import { z, type ZodString } from "zod";
+
+const NO_UNICODE_REGEX = /^[\x20-\x7E]*$/;
+export const NO_UNICODE_ERROR = "String cannot contain special characters or Unicode symbols";
+
+const ALLOWED_USERNAME_CHARACTERS_REGEX = /^[a-zA-Z0-9._-]+$/;
+export const ALLOWED_USERNAME_CHARACTERS_ERROR =
+    "Username can only contain letters, numbers, dots, hyphens, and underscores";
+
+const ALLOWED_REGION_CHARACTERS_REGEX = /^[a-zA-Z0-9_-]+$/;
+export const ALLOWED_REGION_CHARACTERS_ERROR = "Region can only contain letters, numbers, hyphens, and underscores";
+
+const ALLOWED_CLUSTER_NAME_CHARACTERS_REGEX = /^[a-zA-Z0-9_-]+$/;
+export const ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR =
+    "Cluster names can only contain ASCII letters, numbers, and hyphens.";
+
+const ALLOWED_PROJECT_NAME_CHARACTERS_REGEX = /^[a-zA-Z0-9\s()@&+:._',-]+$/;
+export const ALLOWED_PROJECT_NAME_CHARACTERS_ERROR =
+    "Project names can't be longer than 64 characters and can only contain letters, numbers, spaces, and the following symbols: ( ) @ & + : . _ - ' ,";
+export const CommonArgs = {
+    string: (): ZodString => z.string().regex(NO_UNICODE_REGEX, NO_UNICODE_ERROR),
+
+    objectId: (fieldName: string): z.ZodString =>
+        z
+            .string()
+            .min(1, `${fieldName} is required`)
+            .length(24, `${fieldName} must be exactly 24 characters`)
+            .regex(/^[0-9a-fA-F]+$/, `${fieldName} must contain only hexadecimal characters`),
+};
+
+export const AtlasArgs = {
+    projectId: (): z.ZodString => CommonArgs.objectId("projectId"),
+
+    organizationId: (): z.ZodString => CommonArgs.objectId("organizationId"),
+
+    clusterName: (): z.ZodString =>
+        z
+            .string()
+            .min(1, "Cluster name is required")
+            .max(64, "Cluster name must be 64 characters or less")
+            .regex(ALLOWED_CLUSTER_NAME_CHARACTERS_REGEX, ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR),
+
+    projectName: (): z.ZodString =>
+        z
+            .string()
+            .min(1, "Project name is required")
+            .max(64, "Project name must be 64 characters or less")
+            .regex(ALLOWED_PROJECT_NAME_CHARACTERS_REGEX, ALLOWED_PROJECT_NAME_CHARACTERS_ERROR),
+
+    username: (): z.ZodString =>
+        z
+            .string()
+            .min(1, "Username is required")
+            .max(100, "Username must be 100 characters or less")
+            .regex(ALLOWED_USERNAME_CHARACTERS_REGEX, ALLOWED_USERNAME_CHARACTERS_ERROR),
+
+    ipAddress: (): z.ZodString => z.string().ip({ version: "v4" }),
+
+    cidrBlock: (): z.ZodString => z.string().cidr(),
+
+    region: (): z.ZodString =>
+        z
+            .string()
+            .min(1, "Region is required")
+            .max(50, "Region must be 50 characters or less")
+            .regex(ALLOWED_REGION_CHARACTERS_REGEX, ALLOWED_REGION_CHARACTERS_ERROR),
+
+    password: (): z.ZodString =>
+        z.string().min(1, "Password is required").max(100, "Password must be 100 characters or less"),
+};

src/tools/atlas/atlasTool.ts

@@ -1,6 +1,5 @@
-import type { ToolCategory, TelemetryToolMetadata, ToolArgs } from "../tool.js";
-import { ToolBase } from "../tool.js";
 import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { ToolBase, type ToolArgs, type ToolCategory, type TelemetryToolMetadata } from "../tool.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { LogId } from "../../common/logger.js";
 import { z } from "zod";

src/tools/atlas/connect/connectCluster.ts

@@ -1,13 +1,13 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type OperationType, type ToolArgs } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
 import { generateSecurePassword } from "../../../helpers/generatePassword.js";
 import { LogId } from "../../../common/logger.js";
 import { inspectCluster } from "../../../common/atlas/cluster.js";
 import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
 import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
 import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
+import { AtlasArgs } from "../../args.js";
 
 const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
 const addedIpAccessListMessage =
@@ -20,13 +20,17 @@ function sleep(ms: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+export const ConnectClusterArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID"),
+    clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
+};
+
 export class ConnectClusterTool extends AtlasToolBase {
     public name = "atlas-connect-cluster";
     protected description = "Connect to MongoDB Atlas cluster";
     public operationType: OperationType = "connect";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID"),
-        clusterName: z.string().describe("Atlas cluster name"),
+        ...ConnectClusterArgs,
     };
 
     private queryConnection(

src/tools/atlas/create/createAccessList.ts

@@ -1,26 +1,27 @@
 import { z } from "zod";
+import { type OperationType, type ToolArgs } from "../../tool.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
 import { makeCurrentIpAccessListEntry, DEFAULT_ACCESS_LIST_COMMENT } from "../../../common/atlas/accessListUtils.js";
+import { AtlasArgs, CommonArgs } from "../../args.js";
+
+export const CreateAccessListArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID"),
+    ipAddresses: z.array(AtlasArgs.ipAddress()).describe("IP addresses to allow access from").optional(),
+    cidrBlocks: z.array(AtlasArgs.cidrBlock()).describe("CIDR blocks to allow access from").optional(),
+    currentIpAddress: z.boolean().describe("Add the current IP address").default(false),
+    comment: CommonArgs.string()
+        .describe("Comment for the access list entries")
+        .default(DEFAULT_ACCESS_LIST_COMMENT)
+        .optional(),
+};
 
 export class CreateAccessListTool extends AtlasToolBase {
     public name = "atlas-create-access-list";
     protected description = "Allow Ip/CIDR ranges to access your MongoDB Atlas clusters.";
     public operationType: OperationType = "create";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID"),
-        ipAddresses: z
-            .array(z.string().ip({ version: "v4" }))
-            .describe("IP addresses to allow access from")
-            .optional(),
-        cidrBlocks: z.array(z.string().cidr()).describe("CIDR blocks to allow access from").optional(),
-        currentIpAddress: z.boolean().describe("Add the current IP address").default(false),
-        comment: z
-            .string()
-            .describe("Comment for the access list entries")
-            .default(DEFAULT_ACCESS_LIST_COMMENT)
-            .optional(),
+        ...CreateAccessListArgs,
     };
 
     protected async execute({

src/tools/atlas/create/createDBUser.ts

@@ -1,41 +1,45 @@
 import { z } from "zod";
+import type { ToolArgs, OperationType } from "../../tool.js";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
 import type { CloudDatabaseUser, DatabaseUserRole } from "../../../common/atlas/openapi.js";
 import { generateSecurePassword } from "../../../helpers/generatePassword.js";
 import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
+import { AtlasArgs, CommonArgs } from "../../args.js";
+
+export const CreateDBUserArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID"),
+    username: AtlasArgs.username().describe("Username for the new user"),
+    // Models will generate overly simplistic passwords like SecurePassword123 or
+    // AtlasPassword123, which are easily guessable and exploitable. We're instructing
+    // the model not to try and generate anything and instead leave the field unset.
+    password: AtlasArgs.password()
+        .optional()
+        .nullable()
+        .describe(
+            "Password for the new user. IMPORTANT: If the user hasn't supplied an explicit password, leave it unset and under no circumstances try to generate a random one. A secure password will be generated by the MCP server if necessary."
+        ),
+    roles: z
+        .array(
+            z.object({
+                roleName: CommonArgs.string().describe("Role name"),
+                databaseName: CommonArgs.string().describe("Database name").default("admin"),
+                collectionName: CommonArgs.string().describe("Collection name").optional(),
+            })
+        )
+        .describe("Roles for the new user"),
+    clusters: z
+        .array(AtlasArgs.clusterName())
+        .describe("Clusters to assign the user to, leave empty for access to all clusters")
+        .optional(),
+};
 
 export class CreateDBUserTool extends AtlasToolBase {
     public name = "atlas-create-db-user";
     protected description = "Create an MongoDB Atlas database user";
     public operationType: OperationType = "create";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID"),
-        username: z.string().describe("Username for the new user"),
-        // Models will generate overly simplistic passwords like SecurePassword123 or
-        // AtlasPassword123, which are easily guessable and exploitable. We're instructing
-        // the model not to try and generate anything and instead leave the field unset.
-        password: z
-            .string()
-            .optional()
-            .nullable()
-            .describe(
-                "Password for the new user. IMPORTANT: If the user hasn't supplied an explicit password, leave it unset and under no circumstances try to generate a random one. A secure password will be generated by the MCP server if necessary."
-            ),
-        roles: z
-            .array(
-                z.object({
-                    roleName: z.string().describe("Role name"),
-                    databaseName: z.string().describe("Database name").default("admin"),
-                    collectionName: z.string().describe("Collection name").optional(),
-                })
-            )
-            .describe("Roles for the new user"),
-        clusters: z
-            .array(z.string())
-            .describe("Clusters to assign the user to, leave empty for access to all clusters")
-            .optional(),
+        ...CreateDBUserArgs,
     };
 
     protected async execute({

src/tools/atlas/create/createFreeCluster.ts

@@ -1,18 +1,18 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type ToolArgs, type OperationType } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
 import type { ClusterDescription20240805 } from "../../../common/atlas/openapi.js";
 import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
+import { AtlasArgs } from "../../args.js";
 
 export class CreateFreeClusterTool extends AtlasToolBase {
     public name = "atlas-create-free-cluster";
     protected description = "Create a free MongoDB Atlas cluster";
     public operationType: OperationType = "create";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID to create the cluster in"),
-        name: z.string().describe("Name of the cluster"),
-        region: z.string().describe("Region of the cluster").default("US_EAST_1"),
+        projectId: AtlasArgs.projectId().describe("Atlas project ID to create the cluster in"),
+        name: AtlasArgs.clusterName().describe("Name of the cluster"),
+        region: AtlasArgs.region().describe("Region of the cluster").default("US_EAST_1"),
     };
 
     protected async execute({ projectId, name, region }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/create/createProject.ts

@@ -1,16 +1,20 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type OperationType, type ToolArgs } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
 import type { Group } from "../../../common/atlas/openapi.js";
+import { AtlasArgs } from "../../args.js";
+
+export const CreateProjectArgs = {
+    projectName: AtlasArgs.projectName().optional().describe("Name for the new project"),
+    organizationId: AtlasArgs.organizationId().optional().describe("Organization ID for the new project"),
+};
 
 export class CreateProjectTool extends AtlasToolBase {
     public name = "atlas-create-project";
     protected description = "Create a MongoDB Atlas project";
     public operationType: OperationType = "create";
     protected argsShape = {
-        projectName: z.string().optional().describe("Name for the new project"),
-        organizationId: z.string().optional().describe("Organization ID for the new project"),
+        ...CreateProjectArgs,
     };
 
     protected async execute({ projectName, organizationId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/read/inspectAccessList.ts

@@ -1,15 +1,18 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
-import { formatUntrustedData } from "../../tool.js";
+import { AtlasArgs } from "../../args.js";
+
+export const InspectAccessListArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID"),
+};
 
 export class InspectAccessListTool extends AtlasToolBase {
     public name = "atlas-inspect-access-list";
     protected description = "Inspect Ip/CIDR ranges with access to your MongoDB Atlas clusters.";
     public operationType: OperationType = "read";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID"),
+        ...InspectAccessListArgs,
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/read/inspectCluster.ts

@@ -1,18 +1,21 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
-import { formatUntrustedData } from "../../tool.js";
 import type { Cluster } from "../../../common/atlas/cluster.js";
 import { inspectCluster } from "../../../common/atlas/cluster.js";
+import { AtlasArgs } from "../../args.js";
+
+export const InspectClusterArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID"),
+    clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
+};
 
 export class InspectClusterTool extends AtlasToolBase {
     public name = "atlas-inspect-cluster";
     protected description = "Inspect MongoDB Atlas cluster";
     public operationType: OperationType = "read";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID"),
-        clusterName: z.string().describe("Atlas cluster name"),
+        ...InspectClusterArgs,
     };
 
     protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/read/listAlerts.ts

@@ -1,15 +1,18 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
+import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
 import { AtlasToolBase } from "../atlasTool.js";
-import type { ToolArgs, OperationType } from "../../tool.js";
-import { formatUntrustedData } from "../../tool.js";
+import { AtlasArgs } from "../../args.js";
+
+export const ListAlertsArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID to list alerts for"),
+};
 
 export class ListAlertsTool extends AtlasToolBase {
     public name = "atlas-list-alerts";
     protected description = "List MongoDB Atlas alerts";
     public operationType: OperationType = "read";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID to list alerts for"),
+        ...ListAlertsArgs,
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/read/listClusters.ts

@@ -1,4 +1,3 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasToolBase } from "../atlasTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
@@ -10,13 +9,18 @@ import type {
     PaginatedFlexClusters20241113,
 } from "../../../common/atlas/openapi.js";
 import { formatCluster, formatFlexCluster } from "../../../common/atlas/cluster.js";
+import { AtlasArgs } from "../../args.js";
+
+export const ListClustersArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID to filter clusters").optional(),
+};
 
 export class ListClustersTool extends AtlasToolBase {
     public name = "atlas-list-clusters";
     protected description = "List MongoDB Atlas clusters";
     public operationType: OperationType = "read";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID to filter clusters").optional(),
+        ...ListClustersArgs,
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

src/tools/atlas/read/listDBUsers.ts

@@ -1,16 +1,20 @@
-import { z } from "zod";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
 import { AtlasToolBase } from "../atlasTool.js";
 import type { ToolArgs, OperationType } from "../../tool.js";
 import { formatUntrustedData } from "../../tool.js";
 import type { DatabaseUserRole, UserScope } from "../../../common/atlas/openapi.js";
+import { AtlasArgs } from "../../args.js";
+
+export const ListDBUsersArgs = {
+    projectId: AtlasArgs.projectId().describe("Atlas project ID to filter DB users"),
+};
 
 export class ListDBUsersTool extends AtlasToolBase {
     public name = "atlas-list-db-users";
     protected description = "List MongoDB Atlas database users";
     public operationType: OperationType = "read";
     protected argsShape = {
-        projectId: z.string().describe("Atlas project ID to filter DB users"),
+        ...ListDBUsersArgs,
     };
 
     protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {