fix: add untrusted data wrapper to the export resource MCP-197 (#550)

Author: nirinchev

src/common/exportsManager.ts

@@ -27,6 +27,7 @@ interface CommonExportData {
 interface ReadyExport extends CommonExportData {
     exportStatus: "ready";
     exportCreatedAt: number;
+    docsTransformed: number;
 }
 
 interface InProgressExport extends CommonExportData {
@@ -124,7 +125,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
         }
     }
 
-    public async readExport(exportName: string): Promise<string> {
+    public async readExport(exportName: string): Promise<{ content: string; docsTransformed: number }> {
         try {
             this.assertIsNotShuttingDown();
             exportName = decodeAndNormalize(exportName);
@@ -137,9 +138,12 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
                 throw new Error("Requested export is still being generated. Try again later.");
             }
 
-            const { exportPath } = exportHandle;
+            const { exportPath, docsTransformed } = exportHandle;
 
-            return fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal });
+            return {
+                content: await fs.readFile(exportPath, { encoding: "utf8", signal: this.shutdownController.signal }),
+                docsTransformed,
+            };
         } catch (error) {
             this.logger.error({
                 id: LogId.exportReadError,
@@ -202,17 +206,15 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
     }): Promise<void> {
         try {
             let pipeSuccessful = false;
+            let docsTransformed = 0;
             try {
                 await fs.mkdir(this.exportsDirectoryPath, { recursive: true });
                 const outputStream = createWriteStream(inProgressExport.exportPath);
-                await pipeline(
-                    [
-                        input.stream(),
-                        this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat)),
-                        outputStream,
-                    ],
-                    { signal: this.shutdownController.signal }
-                );
+                const ejsonTransform = this.docToEJSONStream(this.getEJSONOptionsForFormat(jsonExportFormat));
+                await pipeline([input.stream(), ejsonTransform, outputStream], {
+                    signal: this.shutdownController.signal,
+                });
+                docsTransformed = ejsonTransform.docsTransformed;
                 pipeSuccessful = true;
             } catch (error) {
                 // If the pipeline errors out then we might end up with
@@ -231,6 +233,7 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
                         ...inProgressExport,
                         exportCreatedAt: Date.now(),
                         exportStatus: "ready",
+                        docsTransformed,
                     };
                     this.emit("export-available", inProgressExport.exportURI);
                 }
@@ -256,33 +259,39 @@ export class ExportsManager extends EventEmitter<ExportsManagerEvents> {
         }
     }
 
-    private docToEJSONStream(ejsonOptions: EJSONOptions | undefined): Transform {
+    private docToEJSONStream(ejsonOptions: EJSONOptions | undefined): Transform & { docsTransformed: number } {
         let docsTransformed = 0;
-        return new Transform({
-            objectMode: true,
-            transform(chunk: unknown, encoding, callback): void {
-                try {
-                    const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions);
+        const result = Object.assign(
+            new Transform({
+                objectMode: true,
+                transform(chunk: unknown, encoding, callback): void {
+                    try {
+                        const doc = EJSON.stringify(chunk, undefined, undefined, ejsonOptions);
+                        if (docsTransformed === 0) {
+                            this.push("[" + doc);
+                        } else {
+                            this.push(",\n" + doc);
+                        }
+                        docsTransformed++;
+                        callback();
+                    } catch (err) {
+                        callback(err as Error);
+                    }
+                },
+                flush(callback): void {
                     if (docsTransformed === 0) {
-                        this.push("[" + doc);
+                        this.push("[]");
                     } else {
-                        this.push(",\n" + doc);
+                        this.push("]");
                     }
-                    docsTransformed++;
+                    result.docsTransformed = docsTransformed;
                     callback();
-                } catch (err) {
-                    callback(err as Error);
-                }
-            },
-            flush(callback): void {
-                if (docsTransformed === 0) {
-                    this.push("[]");
-                } else {
-                    this.push("]");
-                }
-                callback();
-            },
-        });
+                },
+            }),
+            { docsTransformed }
+        );
+
+        return result;
     }
 
     private async cleanupExpiredExports(): Promise<void> {

src/resources/common/exportedData.ts

@@ -7,6 +7,7 @@ import { ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { Server } from "../../server.js";
 import { LogId } from "../../common/logger.js";
 import type { Session } from "../../common/session.js";
+import { formatUntrustedData } from "../../tools/tool.js";
 
 export class ExportedData {
     private readonly name = "exported-data";
@@ -95,13 +96,17 @@ export class ExportedData {
                 throw new Error("Cannot retrieve exported data, exportName not provided.");
             }
 
-            const content = await this.session.exportsManager.readExport(exportName);
+            const { content, docsTransformed } = await this.session.exportsManager.readExport(exportName);
+
+            const text = formatUntrustedData(`The exported data contains ${docsTransformed} documents.`, content)
+                .map((t) => t.text)
+                .join("\n");
 
             return {
                 contents: [
                     {
                         uri: url.href,
-                        text: content,
+                        text,
                         mimeType: "application/json",
                     },
                 ],

tests/integration/resources/exportedData.test.ts

@@ -1,9 +1,9 @@
 import path from "path";
 import fs from "fs/promises";
-import { Long } from "bson";
+import { EJSON, Long, ObjectId } from "bson";
 import { describe, expect, it, beforeEach, afterAll } from "vitest";
 import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";
-import { defaultTestConfig, resourceChangedNotification, timeout } from "../helpers.js";
+import { defaultTestConfig, getDataFromUntrustedContent, resourceChangedNotification, timeout } from "../helpers.js";
 import { describeWithMongoDB } from "../tools/mongodb/mongodbHelpers.js";
 import { contentWithResourceURILink } from "../tools/mongodb/read/export.test.js";
 import type { UserConfig } from "../../../src/lib.js";
@@ -18,15 +18,17 @@ const userConfig: UserConfig = {
 describeWithMongoDB(
     "exported-data resource",
     (integration) => {
+        let docs: { _id: ObjectId; name: string; longNumber?: Long; bigInt?: Long }[];
+        let collection: string;
+
         beforeEach(async () => {
             const mongoClient = integration.mongoClient();
-            await mongoClient
-                .db("db")
-                .collection("coll")
-                .insertMany([
-                    { name: "foo", longNumber: new Long(1234) },
-                    { name: "bar", bigInt: new Long(123412341234) },
-                ]);
+            collection = new ObjectId().toString();
+            docs = [
+                { name: "foo", longNumber: new Long(1234), _id: new ObjectId() },
+                { name: "bar", bigInt: new Long(123412341234), _id: new ObjectId() },
+            ];
+            await mongoClient.db("db").collection(collection).insertMany(docs);
         });
 
         afterAll(async () => {
@@ -67,7 +69,7 @@ describeWithMongoDB(
                     name: "export",
                     arguments: {
                         database: "db",
-                        collection: "coll",
+                        collection,
                         exportTitle: "Export for db.coll",
                         exportTarget: [{ name: "find", arguments: {} }],
                     },
@@ -106,7 +108,7 @@ describeWithMongoDB(
                     name: "export",
                     arguments: {
                         database: "db",
-                        collection: "coll",
+                        collection,
                         exportTitle: "Export for db.coll",
                         exportTarget: [{ name: "find", arguments: {} }],
                     },
@@ -125,7 +127,16 @@ describeWithMongoDB(
                 });
                 expect(response.isError).toBeFalsy();
                 expect(response.contents[0]?.mimeType).toEqual("application/json");
-                expect(response.contents[0]?.text).toContain("foo");
+
+                expect(response.contents[0]?.text).toContain(`The exported data contains ${docs.length} documents.`);
+                expect(response.contents[0]?.text).toContain("<untrusted-user-data");
+                const exportContent = getDataFromUntrustedContent((response.contents[0]?.text as string) || "");
+                const exportedDocs = EJSON.parse(exportContent) as { name: string; _id: ObjectId }[];
+                const expectedDocs = docs as unknown as { name: string; _id: ObjectId }[];
+                expect(exportedDocs[0]?.name).toEqual(expectedDocs[0]?.name);
+                expect(exportedDocs[0]?._id).toEqual(expectedDocs[0]?._id);
+                expect(exportedDocs[1]?.name).toEqual(expectedDocs[1]?.name);
+                expect(exportedDocs[1]?._id).toEqual(expectedDocs[1]?._id);
             });
 
             it("should be able to autocomplete the resource", async () => {
@@ -134,7 +145,7 @@ describeWithMongoDB(
                     name: "export",
                     arguments: {
                         database: "big",
-                        collection: "coll",
+                        collection,
                         exportTitle: "Export for big.coll",
                         exportTarget: [{ name: "find", arguments: {} }],
                     },

tests/unit/common/exportsManager.test.ts

@@ -235,7 +235,9 @@ describe("ExportsManager unit test", () => {
                 jsonExportFormat: "relaxed",
             });
             await exportAvailableNotifier;
-            expect(await manager.readExport(exportName)).toEqual("[]");
+            const { content, docsTransformed } = await manager.readExport(exportName);
+            expect(content).toEqual("[]");
+            expect(docsTransformed).toEqual(0);
         });
 
         it("should handle encoded name", async () => {
@@ -249,7 +251,9 @@ describe("ExportsManager unit test", () => {
                 jsonExportFormat: "relaxed",
             });
             await exportAvailableNotifier;
-            expect(await manager.readExport(encodeURIComponent(exportName))).toEqual("[]");
+            const { content, docsTransformed } = await manager.readExport(encodeURIComponent(exportName));
+            expect(content).toEqual("[]");
+            expect(docsTransformed).toEqual(0);
         });
     });
 
@@ -332,7 +336,7 @@ describe("ExportsManager unit test", () => {
                 expect(emitSpy).toHaveBeenCalledWith("export-available", exportURI);
 
                 // Exports relaxed json
-                const jsonData = JSON.parse(await manager.readExport(exportName)) as unknown[];
+                const jsonData = JSON.parse((await manager.readExport(exportName)).content) as unknown[];
                 expect(jsonData).toEqual([]);
             });
         });
@@ -366,7 +370,7 @@ describe("ExportsManager unit test", () => {
                 expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`);
 
                 // Exports relaxed json
-                const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[];
+                const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[];
                 expect(jsonData).toContainEqual(expect.objectContaining({ name: "foo", longNumber: 12 }));
                 expect(jsonData).toContainEqual(expect.objectContaining({ name: "bar", longNumber: 123456 }));
             });
@@ -401,7 +405,7 @@ describe("ExportsManager unit test", () => {
                 expect(emitSpy).toHaveBeenCalledWith("export-available", `exported-data://${expectedExportName}`);
 
                 // Exports relaxed json
-                const jsonData = JSON.parse(await manager.readExport(expectedExportName)) as unknown[];
+                const jsonData = JSON.parse((await manager.readExport(expectedExportName)).content) as unknown[];
                 expect(jsonData).toContainEqual(
                     expect.objectContaining({ name: "foo", longNumber: { $numberLong: "12" } })
                 );