Skip to content

Commit 0f585db

Browse files
nirinchevnirinchev
authored
fix: wrap more tool responses in untrusted-user-data tags (#465)
1 parent 99a9b6a commit 0f585db

35 files changed

+305
-290
lines changed

‎package.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@
4848
"check:lint": "eslint .",
4949
"check:format": "prettier -c .",
5050
"check:types": "tsc --noEmit --project tsconfig.json",
51+
"fix": "npm run fix:lint && npm run reformat",
52+
"fix:lint": "eslint . --fix",
5153
"reformat": "prettier --write .",
5254
"generate": "./scripts/generate.sh",
5355
"test": "vitest --project unit-and-integration --coverage",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ export class CreateDBUserTool extends AtlasToolBase {
2121
.optional()
2222
.nullable()
2323
.describe(
24-
"Password for the new user. 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."
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."
2525
),
2626
roles: z
2727
.array(

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

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { formatUntrustedData } from "../../tool.js";
56

67
export class InspectAccessListTool extends AtlasToolBase {
78
public name = "atlas-inspect-access-list";
@@ -20,23 +21,25 @@ export class InspectAccessListTool extends AtlasToolBase {
2021
},
2122
});
2223

23-
if (!accessList?.results?.length) {
24-
throw new Error("No access list entries found.");
24+
const results = accessList.results ?? [];
25+
26+
if (!results.length) {
27+
return {
28+
content: [{ type: "text", text: "No access list entries found." }],
29+
};
2530
}
2631

2732
return {
28-
content: [
29-
{
30-
type: "text",
31-
text: `IP ADDRESS | CIDR | COMMENT
33+
content: formatUntrustedData(
34+
`Found ${results.length} access list entries`,
35+
`IP ADDRESS | CIDR | COMMENT
3236
------|------|------
33-
${(accessList.results || [])
37+
${results
3438
.map((entry) => {
3539
return `${entry.ipAddress} | ${entry.cidrBlock} | ${entry.comment}`;
3640
})
37-
.join("\n")}`,
38-
},
39-
],
41+
.join("\n")}`
42+
),
4043
};
4144
}
4245
}

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

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { formatUntrustedData } from "../../tool.js";
56
import type { Cluster } from "../../../common/atlas/cluster.js";
67
import { inspectCluster } from "../../../common/atlas/cluster.js";
78

@@ -22,14 +23,12 @@ export class InspectClusterTool extends AtlasToolBase {
2223

2324
private formatOutput(formattedCluster: Cluster): CallToolResult {
2425
return {
25-
content: [
26-
{
27-
type: "text",
28-
text: `Cluster Name | Cluster Type | Tier | State | MongoDB Version | Connection String
26+
content: formatUntrustedData(
27+
"Cluster details:",
28+
`Cluster Name | Cluster Type | Tier | State | MongoDB Version | Connection String
2929
----------------|----------------|----------------|----------------|----------------|----------------
30-
${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionString || "N/A"}`,
31-
},
32-
],
30+
${formattedCluster.name || "Unknown"} | ${formattedCluster.instanceType} | ${formattedCluster.instanceSize || "N/A"} | ${formattedCluster.state || "UNKNOWN"} | ${formattedCluster.mongoDBVersion || "N/A"} | ${formattedCluster.connectionString || "N/A"}`
31+
),
3332
};
3433
}
3534
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { formatUntrustedData } from "../../tool.js";
56

67
export class ListAlertsTool extends AtlasToolBase {
78
public name = "atlas-list-alerts";
@@ -39,7 +40,7 @@ export class ListAlertsTool extends AtlasToolBase {
3940
.join("\n");
4041

4142
return {
42-
content: [{ type: "text", text: output }],
43+
content: formatUntrustedData(`Found ${data.results.length} alerts in project ${projectId}`, output),
4344
};
4445
}
4546
}

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

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { formatUntrustedData } from "../../tool.js";
56
import type {
67
PaginatedClusterDescription20240805,
78
PaginatedOrgGroupView,
@@ -86,7 +87,9 @@ ${rows}`,
8687
): CallToolResult {
8788
// Check if both traditional clusters and flex clusters are absent
8889
if (!clusters?.results?.length && !flexClusters?.results?.length) {
89-
throw new Error("No clusters found.");
90+
return {
91+
content: [{ type: "text", text: "No clusters found." }],
92+
};
9093
}
9194
const formattedClusters = clusters?.results?.map((cluster) => formatCluster(cluster)) || [];
9295
const formattedFlexClusters = flexClusters?.results?.map((cluster) => formatFlexCluster(cluster)) || [];
@@ -96,18 +99,12 @@ ${rows}`,
9699
})
97100
.join("\n");
98101
return {
99-
content: [
100-
{
101-
type: "text",
102-
text: `Here are your MongoDB Atlas clusters in project "${project.name}" (${project.id}):`,
103-
},
104-
{
105-
type: "text",
106-
text: `Cluster Name | Cluster Type | Tier | State | MongoDB Version | Connection String
102+
content: formatUntrustedData(
103+
`Found ${rows.length} clusters in project "${project.name}" (${project.id}):`,
104+
`Cluster Name | Cluster Type | Tier | State | MongoDB Version | Connection String
107105
----------------|----------------|----------------|----------------|----------------|----------------
108-
${rows}`,
109-
},
110-
],
106+
${rows}`
107+
),
111108
};
112109
}
113110
}

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { z } from "zod";
22
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
33
import { AtlasToolBase } from "../atlasTool.js";
44
import type { ToolArgs, OperationType } from "../../tool.js";
5+
import { formatUntrustedData } from "../../tool.js";
56
import type { DatabaseUserRole, UserScope } from "../../../common/atlas/openapi.js";
67

78
export class ListDBUsersTool extends AtlasToolBase {
@@ -22,7 +23,9 @@ export class ListDBUsersTool extends AtlasToolBase {
2223
});
2324

2425
if (!data?.results?.length) {
25-
throw new Error("No database users found.");
26+
return {
27+
content: [{ type: "text", text: " No database users found" }],
28+
};
2629
}
2730

2831
const output =
@@ -35,7 +38,7 @@ export class ListDBUsersTool extends AtlasToolBase {
3538
})
3639
.join("\n");
3740
return {
38-
content: [{ type: "text", text: output }],
41+
content: formatUntrustedData(`Found ${data.results.length} database users in project ${projectId}`, output),
3942
};
4043
}
4144
}

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { AtlasToolBase } from "../atlasTool.js";
33
import type { OperationType } from "../../tool.js";
4+
import { formatUntrustedData } from "../../tool.js";
45

56
export class ListOrganizationsTool extends AtlasToolBase {
67
public name = "atlas-list-orgs";
@@ -12,10 +13,12 @@ export class ListOrganizationsTool extends AtlasToolBase {
1213
const data = await this.session.apiClient.listOrganizations();
1314

1415
if (!data?.results?.length) {
15-
throw new Error("No projects found in your MongoDB Atlas account.");
16+
return {
17+
content: [{ type: "text", text: "No organizations found in your MongoDB Atlas account." }],
18+
};
1619
}
1720

18-
// Format projects as a table
21+
// Format organizations as a table
1922
const output =
2023
`Organization Name | Organization ID
2124
----------------| ----------------
@@ -26,7 +29,10 @@ export class ListOrganizationsTool extends AtlasToolBase {
2629
})
2730
.join("\n");
2831
return {
29-
content: [{ type: "text", text: output }],
32+
content: formatUntrustedData(
33+
`Found ${data.results.length} organizations in your MongoDB Atlas account.`,
34+
output
35+
),
3036
};
3137
}
3238
}

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { AtlasToolBase } from "../atlasTool.js";
33
import type { OperationType } from "../../tool.js";
4+
import { formatUntrustedData } from "../../tool.js";
45
import { z } from "zod";
56
import type { ToolArgs } from "../../tool.js";
67

@@ -16,7 +17,9 @@ export class ListProjectsTool extends AtlasToolBase {
1617
const orgData = await this.session.apiClient.listOrganizations();
1718

1819
if (!orgData?.results?.length) {
19-
throw new Error("No organizations found in your MongoDB Atlas account.");
20+
return {
21+
content: [{ type: "text", text: "No organizations found in your MongoDB Atlas account." }],
22+
};
2023
}
2124

2225
const orgs: Record<string, string> = orgData.results
@@ -35,7 +38,9 @@ export class ListProjectsTool extends AtlasToolBase {
3538
: await this.session.apiClient.listProjects();
3639

3740
if (!data?.results?.length) {
38-
throw new Error("No projects found in your MongoDB Atlas account.");
41+
return {
42+
content: [{ type: "text", text: `No projects found in organization ${orgId}.` }],
43+
};
3944
}
4045

4146
// Format projects as a table
@@ -50,7 +55,7 @@ export class ListProjectsTool extends AtlasToolBase {
5055
----------------| ----------------| ----------------| ----------------| ----------------
5156
${rows}`;
5257
return {
53-
content: [{ type: "text", text: formattedProjects }],
58+
content: formatUntrustedData(`Found ${rows.length} projects`, formattedProjects),
5459
};
5560
}
5661
}

‎src/tools/mongodb/metadata/collectionSchema.ts

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
22
import { DbOperationArgs, MongoDBToolBase } from "../mongodbTool.js";
33
import type { ToolArgs, OperationType } from "../../tool.js";
4+
import { formatUntrustedData } from "../../tool.js";
45
import { getSimplifiedSchema } from "mongodb-schema";
56

67
export class CollectionSchemaTool extends MongoDBToolBase {
@@ -28,16 +29,10 @@ export class CollectionSchemaTool extends MongoDBToolBase {
2829
}
2930

3031
return {
31-
content: [
32-
{
33-
text: `Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
34-
type: "text",
35-
},
36-
{
37-
text: JSON.stringify(schema),
38-
type: "text",
39-
},
40-
],
32+
content: formatUntrustedData(
33+
`Found ${fieldsCount} fields in the schema for "${database}.${collection}"`,
34+
JSON.stringify(schema)
35+
),
4136
};
4237
}
4338
}

0 commit comments

Comments
 (0)