Skip to content

Commit 17b544b

Browse files
blvablva
authored
fix: add argument validation - MCP-188 (#542)
1 parent ce02189 commit 17b544b

14 files changed

+573
-73
lines changed

‎src/tools/args.ts‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import { z, type ZodString } from "zod";
2+
3+
const NO_UNICODE_REGEX = /^[\x20-\x7E]*$/;
4+
export const NO_UNICODE_ERROR = "String cannot contain special characters or Unicode symbols";
5+
6+
const ALLOWED_USERNAME_CHARACTERS_REGEX = /^[a-zA-Z0-9._-]+$/;
7+
export const ALLOWED_USERNAME_CHARACTERS_ERROR =
8+
"Username can only contain letters, numbers, dots, hyphens, and underscores";
9+
10+
const ALLOWED_REGION_CHARACTERS_REGEX = /^[a-zA-Z0-9_-]+$/;
11+
export const ALLOWED_REGION_CHARACTERS_ERROR = "Region can only contain letters, numbers, hyphens, and underscores";
12+
13+
const ALLOWED_CLUSTER_NAME_CHARACTERS_REGEX = /^[a-zA-Z0-9_-]+$/;
14+
export const ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR =
15+
"Cluster names can only contain ASCII letters, numbers, and hyphens.";
16+
17+
const ALLOWED_PROJECT_NAME_CHARACTERS_REGEX = /^[a-zA-Z0-9\s()@&+:._',-]+$/;
18+
export const ALLOWED_PROJECT_NAME_CHARACTERS_ERROR =
19+
"Project names can't be longer than 64 characters and can only contain letters, numbers, spaces, and the following symbols: ( ) @ & + : . _ - ' ,";
20+
export const CommonArgs = {
21+
string: (): ZodString => z.string().regex(NO_UNICODE_REGEX, NO_UNICODE_ERROR),
22+
23+
objectId: (fieldName: string): z.ZodString =>
24+
z
25+
.string()
26+
.min(1, `${fieldName} is required`)
27+
.length(24, `${fieldName} must be exactly 24 characters`)
28+
.regex(/^[0-9a-fA-F]+$/, `${fieldName} must contain only hexadecimal characters`),
29+
};
30+
31+
export const AtlasArgs = {
32+
projectId: (): z.ZodString => CommonArgs.objectId("projectId"),
33+
34+
organizationId: (): z.ZodString => CommonArgs.objectId("organizationId"),
35+
36+
clusterName: (): z.ZodString =>
37+
z
38+
.string()
39+
.min(1, "Cluster name is required")
40+
.max(64, "Cluster name must be 64 characters or less")
41+
.regex(ALLOWED_CLUSTER_NAME_CHARACTERS_REGEX, ALLOWED_CLUSTER_NAME_CHARACTERS_ERROR),
42+
43+
projectName: (): z.ZodString =>
44+
z
45+
.string()
46+
.min(1, "Project name is required")
47+
.max(64, "Project name must be 64 characters or less")
48+
.regex(ALLOWED_PROJECT_NAME_CHARACTERS_REGEX, ALLOWED_PROJECT_NAME_CHARACTERS_ERROR),
49+
50+
username: (): z.ZodString =>
51+
z
52+
.string()
53+
.min(1, "Username is required")
54+
.max(100, "Username must be 100 characters or less")
55+
.regex(ALLOWED_USERNAME_CHARACTERS_REGEX, ALLOWED_USERNAME_CHARACTERS_ERROR),
56+
57+
ipAddress: (): z.ZodString => z.string().ip({ version: "v4" }),
58+
59+
cidrBlock: (): z.ZodString => z.string().cidr(),
60+
61+
region: (): z.ZodString =>
62+
z
63+
.string()
64+
.min(1, "Region is required")
65+
.max(50, "Region must be 50 characters or less")
66+
.regex(ALLOWED_REGION_CHARACTERS_REGEX, ALLOWED_REGION_CHARACTERS_ERROR),
67+
68+
password: (): z.ZodString =>
69+
z.string().min(1, "Password is required").max(100, "Password must be 100 characters or less"),
70+
};

‎src/tools/atlas/atlasTool.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
1-
import type { ToolCategory, TelemetryToolMetadata, ToolArgs } from "../tool.js";
2-
import { ToolBase } from "../tool.js";
31
import type { ToolCallback } from "@modelcontextprotocol/sdk/server/mcp.js";
2+
import { ToolBase, type ToolArgs, type ToolCategory, type TelemetryToolMetadata } from "../tool.js";
43
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
54
import { LogId } from "../../common/logger.js";
65
import { z } from "zod";

‎src/tools/atlas/connect/connectCluster.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type OperationType, type ToolArgs } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
54
import { generateSecurePassword } from "../../../helpers/generatePassword.js";
65
import { LogId } from "../../../common/logger.js";
76
import { inspectCluster } from "../../../common/atlas/cluster.js";
87
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
98
import type { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
109
import { getDefaultRoleFromConfig } from "../../../common/atlas/roles.js";
10+
import { AtlasArgs } from "../../args.js";
1111

1212
const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours
1313
const addedIpAccessListMessage =
@@ -20,13 +20,17 @@ function sleep(ms: number): Promise<void> {
2020
return new Promise((resolve) => setTimeout(resolve, ms));
2121
}
2222

23+
export const ConnectClusterArgs = {
24+
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
25+
clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
26+
};
27+
2328
export class ConnectClusterTool extends AtlasToolBase {
2429
public name = "atlas-connect-cluster";
2530
protected description = "Connect to MongoDB Atlas cluster";
2631
public operationType: OperationType = "connect";
2732
protected argsShape = {
28-
projectId: z.string().describe("Atlas project ID"),
29-
clusterName: z.string().describe("Atlas cluster name"),
33+
...ConnectClusterArgs,
3034
};
3135

3236
private queryConnection(

‎src/tools/atlas/create/createAccessList.ts‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,27 @@
11
import { z } from "zod";
2+
import { type OperationType, type ToolArgs } from "../../tool.js";
23
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
34
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
55
import { makeCurrentIpAccessListEntry, DEFAULT_ACCESS_LIST_COMMENT } from "../../../common/atlas/accessListUtils.js";
6+
import { AtlasArgs, CommonArgs } from "../../args.js";
7+
8+
export const CreateAccessListArgs = {
9+
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
10+
ipAddresses: z.array(AtlasArgs.ipAddress()).describe("IP addresses to allow access from").optional(),
11+
cidrBlocks: z.array(AtlasArgs.cidrBlock()).describe("CIDR blocks to allow access from").optional(),
12+
currentIpAddress: z.boolean().describe("Add the current IP address").default(false),
13+
comment: CommonArgs.string()
14+
.describe("Comment for the access list entries")
15+
.default(DEFAULT_ACCESS_LIST_COMMENT)
16+
.optional(),
17+
};
618

719
export class CreateAccessListTool extends AtlasToolBase {
820
public name = "atlas-create-access-list";
921
protected description = "Allow Ip/CIDR ranges to access your MongoDB Atlas clusters.";
1022
public operationType: OperationType = "create";
1123
protected argsShape = {
12-
projectId: z.string().describe("Atlas project ID"),
13-
ipAddresses: z
14-
.array(z.string().ip({ version: "v4" }))
15-
.describe("IP addresses to allow access from")
16-
.optional(),
17-
cidrBlocks: z.array(z.string().cidr()).describe("CIDR blocks to allow access from").optional(),
18-
currentIpAddress: z.boolean().describe("Add the current IP address").default(false),
19-
comment: z
20-
.string()
21-
.describe("Comment for the access list entries")
22-
.default(DEFAULT_ACCESS_LIST_COMMENT)
23-
.optional(),
24+
...CreateAccessListArgs,
2425
};
2526

2627
protected async execute({

‎src/tools/atlas/create/createDBUser.ts‎

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,45 @@
11
import { z } from "zod";
2+
import type { ToolArgs, OperationType } from "../../tool.js";
23
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
34
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
55
import type { CloudDatabaseUser, DatabaseUserRole } from "../../../common/atlas/openapi.js";
66
import { generateSecurePassword } from "../../../helpers/generatePassword.js";
77
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
8+
import { AtlasArgs, CommonArgs } from "../../args.js";
9+
10+
export const CreateDBUserArgs = {
11+
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
12+
username: AtlasArgs.username().describe("Username for the new user"),
13+
// Models will generate overly simplistic passwords like SecurePassword123 or
14+
// AtlasPassword123, which are easily guessable and exploitable. We're instructing
15+
// the model not to try and generate anything and instead leave the field unset.
16+
password: AtlasArgs.password()
17+
.optional()
18+
.nullable()
19+
.describe(
20+
"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."
21+
),
22+
roles: z
23+
.array(
24+
z.object({
25+
roleName: CommonArgs.string().describe("Role name"),
26+
databaseName: CommonArgs.string().describe("Database name").default("admin"),
27+
collectionName: CommonArgs.string().describe("Collection name").optional(),
28+
})
29+
)
30+
.describe("Roles for the new user"),
31+
clusters: z
32+
.array(AtlasArgs.clusterName())
33+
.describe("Clusters to assign the user to, leave empty for access to all clusters")
34+
.optional(),
35+
};
836

937
export class CreateDBUserTool extends AtlasToolBase {
1038
public name = "atlas-create-db-user";
1139
protected description = "Create an MongoDB Atlas database user";
1240
public operationType: OperationType = "create";
1341
protected argsShape = {
14-
projectId: z.string().describe("Atlas project ID"),
15-
username: z.string().describe("Username for the new user"),
16-
// Models will generate overly simplistic passwords like SecurePassword123 or
17-
// AtlasPassword123, which are easily guessable and exploitable. We're instructing
18-
// the model not to try and generate anything and instead leave the field unset.
19-
password: z
20-
.string()
21-
.optional()
22-
.nullable()
23-
.describe(
24-
"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."
25-
),
26-
roles: z
27-
.array(
28-
z.object({
29-
roleName: z.string().describe("Role name"),
30-
databaseName: z.string().describe("Database name").default("admin"),
31-
collectionName: z.string().describe("Collection name").optional(),
32-
})
33-
)
34-
.describe("Roles for the new user"),
35-
clusters: z
36-
.array(z.string())
37-
.describe("Clusters to assign the user to, leave empty for access to all clusters")
38-
.optional(),
42+
...CreateDBUserArgs,
3943
};
4044

4145
protected async execute({

‎src/tools/atlas/create/createFreeCluster.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,18 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type ToolArgs, type OperationType } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
54
import type { ClusterDescription20240805 } from "../../../common/atlas/openapi.js";
65
import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";
6+
import { AtlasArgs } from "../../args.js";
77

88
export class CreateFreeClusterTool extends AtlasToolBase {
99
public name = "atlas-create-free-cluster";
1010
protected description = "Create a free MongoDB Atlas cluster";
1111
public operationType: OperationType = "create";
1212
protected argsShape = {
13-
projectId: z.string().describe("Atlas project ID to create the cluster in"),
14-
name: z.string().describe("Name of the cluster"),
15-
region: z.string().describe("Region of the cluster").default("US_EAST_1"),
13+
projectId: AtlasArgs.projectId().describe("Atlas project ID to create the cluster in"),
14+
name: AtlasArgs.clusterName().describe("Name of the cluster"),
15+
region: AtlasArgs.region().describe("Region of the cluster").default("US_EAST_1"),
1616
};
1717

1818
protected async execute({ projectId, name, region }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

‎src/tools/atlas/create/createProject.ts‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type OperationType, type ToolArgs } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
54
import type { Group } from "../../../common/atlas/openapi.js";
5+
import { AtlasArgs } from "../../args.js";
6+
7+
export const CreateProjectArgs = {
8+
projectName: AtlasArgs.projectName().optional().describe("Name for the new project"),
9+
organizationId: AtlasArgs.organizationId().optional().describe("Organization ID for the new project"),
10+
};
611

712
export class CreateProjectTool extends AtlasToolBase {
813
public name = "atlas-create-project";
914
protected description = "Create a MongoDB Atlas project";
1015
public operationType: OperationType = "create";
1116
protected argsShape = {
12-
projectName: z.string().optional().describe("Name for the new project"),
13-
organizationId: z.string().optional().describe("Organization ID for the new project"),
17+
...CreateProjectArgs,
1418
};
1519

1620
protected async execute({ projectName, organizationId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

‎src/tools/atlas/read/inspectAccessList.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
5-
import { formatUntrustedData } from "../../tool.js";
4+
import { AtlasArgs } from "../../args.js";
5+
6+
export const InspectAccessListArgs = {
7+
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
8+
};
69

710
export class InspectAccessListTool extends AtlasToolBase {
811
public name = "atlas-inspect-access-list";
912
protected description = "Inspect Ip/CIDR ranges with access to your MongoDB Atlas clusters.";
1013
public operationType: OperationType = "read";
1114
protected argsShape = {
12-
projectId: z.string().describe("Atlas project ID"),
15+
...InspectAccessListArgs,
1316
};
1417

1518
protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

‎src/tools/atlas/read/inspectCluster.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,21 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
5-
import { formatUntrustedData } from "../../tool.js";
64
import type { Cluster } from "../../../common/atlas/cluster.js";
75
import { inspectCluster } from "../../../common/atlas/cluster.js";
6+
import { AtlasArgs } from "../../args.js";
7+
8+
export const InspectClusterArgs = {
9+
projectId: AtlasArgs.projectId().describe("Atlas project ID"),
10+
clusterName: AtlasArgs.clusterName().describe("Atlas cluster name"),
11+
};
812

913
export class InspectClusterTool extends AtlasToolBase {
1014
public name = "atlas-inspect-cluster";
1115
protected description = "Inspect MongoDB Atlas cluster";
1216
public operationType: OperationType = "read";
1317
protected argsShape = {
14-
projectId: z.string().describe("Atlas project ID"),
15-
clusterName: z.string().describe("Atlas cluster name"),
18+
...InspectClusterArgs,
1619
};
1720

1821
protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

‎src/tools/atlas/read/listAlerts.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,18 @@
1-
import { z } from "zod";
21
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
2+
import { type OperationType, type ToolArgs, formatUntrustedData } from "../../tool.js";
33
import { AtlasToolBase } from "../atlasTool.js";
4-
import type { ToolArgs, OperationType } from "../../tool.js";
5-
import { formatUntrustedData } from "../../tool.js";
4+
import { AtlasArgs } from "../../args.js";
5+
6+
export const ListAlertsArgs = {
7+
projectId: AtlasArgs.projectId().describe("Atlas project ID to list alerts for"),
8+
};
69

710
export class ListAlertsTool extends AtlasToolBase {
811
public name = "atlas-list-alerts";
912
protected description = "List MongoDB Atlas alerts";
1013
public operationType: OperationType = "read";
1114
protected argsShape = {
12-
projectId: z.string().describe("Atlas project ID to list alerts for"),
15+
...ListAlertsArgs,
1316
};
1417

1518
protected async execute({ projectId }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {

0 commit comments

Comments
 (0)