16834304026

Committed 08 Aug 2025 03:31PM UTC coverage: 81.151% (-0.2%) from 81.362%

Build # 16834304026

Build Type

Pull #430

github

Committed by

web-flow

Commit Message

Merge bf238a638 into 92687b86c

Pull Request Pull Request #430: fix: implement prompt poisoning mitigation

Run Details

661 of 857 branches covered (77.13%)

Branch coverage included in aggregate %.

28 of 28 new or added lines in 3 files covered. (100.0%)

8 existing lines in 1 file now uncovered.

3528 of 4305 relevant lines covered (81.95%)

58.12 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

73.31

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

import { z } from "zod";
import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
import { AtlasToolBase } from "../atlasTool.js";
import { 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 { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";

const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours

function sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
}

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"),
    };

    private queryConnection(
        projectId: string,
        clusterName: string
    ): "connected" | "disconnected" | "connecting" | "connected-to-other-cluster" | "unknown" {
        if (!this.session.connectedAtlasCluster) {
            if (this.session.isConnectedToMongoDB) {
                return "connected-to-other-cluster";
            }
            return "disconnected";
        }

        const currentConectionState = this.session.connectionManager.currentConnectionState;
        if (
            this.session.connectedAtlasCluster.projectId !== projectId ||
            this.session.connectedAtlasCluster.clusterName !== clusterName
        ) {
            return "connected-to-other-cluster";
        }

        switch (currentConectionState.tag) {
            case "connecting":
            case "disconnected": // we might still be calling Atlas APIs and not attempted yet to connect to MongoDB, but we are still "connecting"
                return "connecting";
            case "connected":
                return "connected";
            case "errored":
                this.session.logger.debug({
                    id: LogId.atlasConnectFailure,
                    context: "atlas-connect-cluster",
                    message: `error querying cluster: ${currentConectionState.errorReason}`,
                });
                return "unknown";
        }
    }

    private async prepareClusterConnection(
        projectId: string,
        clusterName: string
    ): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {
        const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);

        if (!cluster.connectionString) {
            throw new Error("Connection string not available");
        }

        const username = `mcpUser${Math.floor(Math.random() * 100000)}`;
        const password = await generateSecurePassword();

        const expiryDate = new Date(Date.now() + EXPIRY_MS);

        const readOnly =
            this.config.readOnly ||
            (this.config.disabledTools?.includes("create") &&
                this.config.disabledTools?.includes("update") &&
                this.config.disabledTools?.includes("delete") &&
                !this.config.disabledTools?.includes("read") &&
                !this.config.disabledTools?.includes("metadata"));

        const roleName = readOnly ? "readAnyDatabase" : "readWriteAnyDatabase";

        await this.session.apiClient.createDatabaseUser({
            params: {
                path: {
                    groupId: projectId,
                },
            },
            body: {
                databaseName: "admin",
                groupId: projectId,
                roles: [
                    {
                        roleName,
                        databaseName: "admin",
                    },
                ],
                scopes: [{ type: "CLUSTER", name: clusterName }],
                username,
                password,
                awsIAMType: "NONE",
                ldapAuthType: "NONE",
                oidcAuthType: "NONE",
                x509Type: "NONE",
                deleteAfterDate: expiryDate.toISOString(),
            },
        });

        const connectedAtlasCluster = {
            username,
            projectId,
            clusterName,
            expiryDate,
        };

        const cn = new URL(cluster.connectionString);
        cn.username = username;
        cn.password = password;
        cn.searchParams.set("authSource", "admin");

        return { connectionString: cn.toString(), atlas: connectedAtlasCluster };
    }

    private async connectToCluster(connectionString: string, atlas: AtlasClusterConnectionInfo): Promise<void> {
        let lastError: Error | undefined = undefined;

        this.session.logger.debug({
            id: LogId.atlasConnectAttempt,
            context: "atlas-connect-cluster",
            message: `attempting to connect to cluster: ${this.session.connectedAtlasCluster?.clusterName}`,
            noRedaction: true,
        });

        // try to connect for about 5 minutes
        for (let i = 0; i < 600; i++) {
            try {
                lastError = undefined;

                await this.session.connectToMongoDB({ connectionString, ...this.config.connectOptions, atlas });
                break;
            } catch (err: unknown) {
                const error = err instanceof Error ? err : new Error(String(err));

                lastError = error;

                this.session.logger.debug({
                    id: LogId.atlasConnectFailure,
                    context: "atlas-connect-cluster",
                    message: `error connecting to cluster: ${error.message}`,
                });

                await sleep(500); // wait for 500ms before retrying
            }

            if (
                !this.session.connectedAtlasCluster ||
                this.session.connectedAtlasCluster.projectId !== atlas.projectId ||
                this.session.connectedAtlasCluster.clusterName !== atlas.clusterName
            ) {
                throw new Error("Cluster connection aborted");
            }
        }

        if (lastError) {
            if (
                this.session.connectedAtlasCluster?.projectId === atlas.projectId &&
                this.session.connectedAtlasCluster?.clusterName === atlas.clusterName &&
                this.session.connectedAtlasCluster?.username
            ) {
                void this.session.apiClient
                    .deleteDatabaseUser({
                        params: {
                            path: {
                                groupId: this.session.connectedAtlasCluster.projectId,
                                username: this.session.connectedAtlasCluster.username,
                                databaseName: "admin",
                            },
                        },
                    })
                    .catch((err: unknown) => {
                        const error = err instanceof Error ? err : new Error(String(err));
                        this.session.logger.debug({
                            id: LogId.atlasConnectFailure,
                            context: "atlas-connect-cluster",
                            message: `error deleting database user: ${error.message}`,
                        });
                    });
            }
            throw lastError;
        }

        this.session.logger.debug({
            id: LogId.atlasConnectSucceeded,
            context: "atlas-connect-cluster",
            message: `connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}`,
            noRedaction: true,
        });
    }

    protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {
        await ensureCurrentIpInAccessList(this.session.apiClient, projectId);
        for (let i = 0; i < 60; i++) {
            const state = this.queryConnection(projectId, clusterName);
            switch (state) {
                case "connected": {
                    return {
                        content: [
                            {
                                type: "text",
                                text: `Connected to cluster "${clusterName}".`,
                            },
                        ],
                    };
                }
                case "connecting":
                case "unknown": {
                    break;
                }
                case "connected-to-other-cluster":
                case "disconnected":
                default: {
                    await this.session.disconnect();
                    const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);

                    // try to connect for about 5 minutes asynchronously
                    void this.connectToCluster(connectionString, atlas).catch((err: unknown) => {
                        const error = err instanceof Error ? err : new Error(String(err));
                        this.session.logger.error({
                            id: LogId.atlasConnectFailure,
                            context: "atlas-connect-cluster",
                            message: `error connecting to cluster: ${error.message}`,
                        });
                    });
                    break;
                }
            }

            await sleep(500);
        }

        return {
            content: [
                {
                    type: "text" as const,
                    text: `Attempting to connect to cluster "${clusterName}"...`,
                },
                {
                    type: "text" as const,
                    text: `Warning: Provisioning a user and connecting to the cluster may take more time, please check again in a few seconds.`,
                },
                {
                    type: "text" as const,
                    text: `Warning: Make sure your IP address was enabled in the allow list setting of the Atlas cluster.`,
                },
            ],
        };
    }
}

1	import { z } from "zod";	2✔
2	import { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
3	import { AtlasToolBase } from "../atlasTool.js";	2✔
4	import { ToolArgs, OperationType } from "../../tool.js";
5	import { generateSecurePassword } from "../../../helpers/generatePassword.js";	2✔
6	import { LogId } from "../../../common/logger.js";	2✔
7	import { inspectCluster } from "../../../common/atlas/cluster.js";	2✔
8	import { ensureCurrentIpInAccessList } from "../../../common/atlas/accessListUtils.js";	2✔
9	import { AtlasClusterConnectionInfo } from "../../../common/connectionManager.js";
10
11	const EXPIRY_MS = 1000 * 60 * 60 * 12; // 12 hours	2✔
12
13	function sleep(ms: number): Promise<void> {	9✔
14	return new Promise((resolve) => setTimeout(resolve, ms));	9✔
15	}	9✔
16
17	export class ConnectClusterTool extends AtlasToolBase {	2✔
18	public name = "atlas-connect-cluster";	68✔
19	protected description = "Connect to MongoDB Atlas cluster";	68✔
20	public operationType: OperationType = "connect";	68✔
21	protected argsShape = {	68✔
22	projectId: z.string().describe("Atlas project ID"),	68✔
23	clusterName: z.string().describe("Atlas cluster name"),	68✔
24	};	68✔
25
26	private queryConnection(	2✔
27	projectId: string,	3✔
28	clusterName: string	3✔
29	): "connected" \| "disconnected" \| "connecting" \| "connected-to-other-cluster" \| "unknown" {	3✔
30	if (!this.session.connectedAtlasCluster) {	3✔
31	if (this.session.isConnectedToMongoDB) {	2!
32	return "connected-to-other-cluster";	×
33	}	×
34	return "disconnected";	2✔
35	}	2✔
36
37	const currentConectionState = this.session.connectionManager.currentConnectionState;	1✔
38	if (	1✔
39	this.session.connectedAtlasCluster.projectId !== projectId \|\|	1✔
40	this.session.connectedAtlasCluster.clusterName !== clusterName	1✔
41	) {	3!
42	return "connected-to-other-cluster";	×
43	}	✔
44
45	switch (currentConectionState.tag) {	1✔
46	case "connecting":	3!
47	case "disconnected": // we might still be calling Atlas APIs and not attempted yet to connect to MongoDB, but we are still "connecting"	3!
48	return "connecting";	×
49	case "connected":	3✔
50	return "connected";	1✔
51	case "errored":	3!
UNCOV 52	this.session.logger.debug({	×
UNCOV 53	id: LogId.atlasConnectFailure,	×
UNCOV 54	context: "atlas-connect-cluster",	×
UNCOV 55	message: `error querying cluster: ${currentConectionState.errorReason}`,	×
UNCOV 56	});	×
UNCOV 57	return "unknown";	×
58	}	3✔
59	}	3✔
60
61	private async prepareClusterConnection(	2✔
62	projectId: string,	2✔
63	clusterName: string	2✔
64	): Promise<{ connectionString: string; atlas: AtlasClusterConnectionInfo }> {	2✔
65	const cluster = await inspectCluster(this.session.apiClient, projectId, clusterName);	2✔
66
67	if (!cluster.connectionString) {	2!
68	throw new Error("Connection string not available");	×
69	}	×
70
71	const username = `mcpUser${Math.floor(Math.random() * 100000)}`;	2✔
72	const password = await generateSecurePassword();	2✔
73
74	const expiryDate = new Date(Date.now() + EXPIRY_MS);	2✔
75
76	const readOnly =	2✔
77	this.config.readOnly \|\|	2✔
78	(this.config.disabledTools?.includes("create") &&	2!
79	this.config.disabledTools?.includes("update") &&	×
80	this.config.disabledTools?.includes("delete") &&	×
81	!this.config.disabledTools?.includes("read") &&	×
82	!this.config.disabledTools?.includes("metadata"));	×
83
84	const roleName = readOnly ? "readAnyDatabase" : "readWriteAnyDatabase";	2!
85
86	await this.session.apiClient.createDatabaseUser({	2✔
87	params: {	2✔
88	path: {	2✔
89	groupId: projectId,	2✔
90	},	2✔
91	},	2✔
92	body: {	2✔
93	databaseName: "admin",	2✔
94	groupId: projectId,	2✔
95	roles: [	2✔
96	{	2✔
97	roleName,	2✔
98	databaseName: "admin",	2✔
99	},	2✔
100	],	2✔
101	scopes: [{ type: "CLUSTER", name: clusterName }],	2✔
102	username,	2✔
103	password,	2✔
104	awsIAMType: "NONE",	2✔
105	ldapAuthType: "NONE",	2✔
106	oidcAuthType: "NONE",	2✔
107	x509Type: "NONE",	2✔
108	deleteAfterDate: expiryDate.toISOString(),	2✔
109	},	2✔
110	});	2✔
111
112	const connectedAtlasCluster = {	2✔
113	username,	2✔
114	projectId,	2✔
115	clusterName,	2✔
116	expiryDate,	2✔
117	};	2✔
118
119	const cn = new URL(cluster.connectionString);	2✔
120	cn.username = username;	2✔
121	cn.password = password;	2✔
122	cn.searchParams.set("authSource", "admin");	2✔
123
124	return { connectionString: cn.toString(), atlas: connectedAtlasCluster };	2✔
125	}	2✔
126
127	private async connectToCluster(connectionString: string, atlas: AtlasClusterConnectionInfo): Promise<void> {	2✔
128	let lastError: Error \| undefined = undefined;	2✔
129
130	this.session.logger.debug({	2✔
131	id: LogId.atlasConnectAttempt,	2✔
132	context: "atlas-connect-cluster",	2✔
133	message: `attempting to connect to cluster: ${this.session.connectedAtlasCluster?.clusterName}`,	2!
134	noRedaction: true,	2✔
135	});	2✔
136
137	// try to connect for about 5 minutes
138	for (let i = 0; i < 600; i++) {	2✔
139	try {	9✔
140	lastError = undefined;	9✔
141
142	await this.session.connectToMongoDB({ connectionString, ...this.config.connectOptions, atlas });	9✔
143	break;	2✔
144	} catch (err: unknown) {	9✔
145	const error = err instanceof Error ? err : new Error(String(err));	7!
146
147	lastError = error;	7✔
148
149	this.session.logger.debug({	7✔
150	id: LogId.atlasConnectFailure,	7✔
151	context: "atlas-connect-cluster",	7✔
152	message: `error connecting to cluster: ${error.message}`,	7✔
153	});	7✔
154
155	await sleep(500); // wait for 500ms before retrying	7✔
156	}	7✔
157
158	if (	7✔
159	!this.session.connectedAtlasCluster \|\|	7✔
160	this.session.connectedAtlasCluster.projectId !== atlas.projectId \|\|	7✔
161	this.session.connectedAtlasCluster.clusterName !== atlas.clusterName	7✔
162	) {	9!
163	throw new Error("Cluster connection aborted");	×
164	}	×
165	}	9✔
166
167	if (lastError) {	2!
168	if (	×
169	this.session.connectedAtlasCluster?.projectId === atlas.projectId &&	×
170	this.session.connectedAtlasCluster?.clusterName === atlas.clusterName &&	×
171	this.session.connectedAtlasCluster?.username	×
172	) {	×
173	void this.session.apiClient	×
174	.deleteDatabaseUser({	×
175	params: {	×
176	path: {	×
177	groupId: this.session.connectedAtlasCluster.projectId,	×
178	username: this.session.connectedAtlasCluster.username,	×
179	databaseName: "admin",	×
180	},	×
181	},	×
182	})	×
183	.catch((err: unknown) => {	×
184	const error = err instanceof Error ? err : new Error(String(err));	×
185	this.session.logger.debug({	×
186	id: LogId.atlasConnectFailure,	×
187	context: "atlas-connect-cluster",	×
188	message: `error deleting database user: ${error.message}`,	×
189	});	×
190	});	×
191	}	×
192	throw lastError;	×
193	}	×
194
195	this.session.logger.debug({	2✔
196	id: LogId.atlasConnectSucceeded,	2✔
197	context: "atlas-connect-cluster",	2✔
198	message: `connected to cluster: ${this.session.connectedAtlasCluster?.clusterName}`,	2✔
199	noRedaction: true,	2✔
200	});	2✔
201	}	2✔
202
203	protected async execute({ projectId, clusterName }: ToolArgs<typeof this.argsShape>): Promise<CallToolResult> {	2✔
204	await ensureCurrentIpInAccessList(this.session.apiClient, projectId);	1✔
205	for (let i = 0; i < 60; i++) {	1✔
206	const state = this.queryConnection(projectId, clusterName);	3✔
207	switch (state) {	3✔
208	case "connected": {	3✔
209	return {	1✔
210	content: [	1✔
211	{	1✔
212	type: "text",	1✔
213	text: `Connected to cluster "${clusterName}".`,	1✔
214	},	1✔
215	],	1✔
216	};	1✔
217	}	1✔
218	case "connecting":	3!
219	case "unknown": {	3!
UNCOV 220	break;	×
UNCOV 221	}	×
222	case "connected-to-other-cluster":	3!
223	case "disconnected":	3✔
224	default: {	3✔
225	await this.session.disconnect();	2✔
226	const { connectionString, atlas } = await this.prepareClusterConnection(projectId, clusterName);	2✔
227
228	// try to connect for about 5 minutes asynchronously
229	void this.connectToCluster(connectionString, atlas).catch((err: unknown) => {	2✔
230	const error = err instanceof Error ? err : new Error(String(err));	×
231	this.session.logger.error({	×
232	id: LogId.atlasConnectFailure,	×
233	context: "atlas-connect-cluster",	×
234	message: `error connecting to cluster: ${error.message}`,	×
235	});	×
236	});	2✔
237	break;	2✔
238	}	2✔
239	}	3✔
240
241	await sleep(500);	2✔
242	}	2✔
243
244	return {	2✔
245	content: [	2✔
246	{	2✔
247	type: "text" as const,	2✔
248	text: `Attempting to connect to cluster "${clusterName}"...`,	2✔
249	},	2✔
250	{	2✔
251	type: "text" as const,	2✔
252	text: `Warning: Provisioning a user and connecting to the cluster may take more time, please check again in a few seconds.`,	2✔
253	},	2✔
254	{	2✔
255	type: "text" as const,	2✔
256	text: `Warning: Make sure your IP address was enabled in the allow list setting of the Atlas cluster.`,	2✔
257	},	2✔
258	],	2✔
259	};	2✔
260	}	1✔
261	}	2✔

mongodb-js / mongodb-mcp-server / 16834304026

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous