From ac22ada8d06ee714e0bb64a1e42b73172434daf6 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 2 Sep 2026 18:17:26 +0800 Subject: [PATCH] fix(mongodb): classify shell reads as read-only --- .../lib/__tests__/sql/sqlRisk.mongodb.spec.ts | 37 ++++++++++++++++ .../__tests__/productionSafety.spec.ts | 12 ++++++ apps/desktop/src/lib/sql/sqlRisk.ts | 43 +++++++++++++++++++ 3 files changed, 92 insertions(+) create mode 100644 apps/desktop/src/lib/__tests__/sql/sqlRisk.mongodb.spec.ts diff --git a/apps/desktop/src/lib/__tests__/sql/sqlRisk.mongodb.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlRisk.mongodb.spec.ts new file mode 100644 index 0000000000..7be9e26080 --- /dev/null +++ b/apps/desktop/src/lib/__tests__/sql/sqlRisk.mongodb.spec.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { classifySqlRisk } from "@/lib/sql/sqlRisk"; + +describe("MongoDB shell statement risk", () => { + it("treats supported read commands as reads", () => { + for (const sql of [ + 'db.getCollection("demo").find({}).skip(0).limit(100)', + 'db.getCollection("demo").findOne({"active": true})', + 'db.getCollection("demo").countDocuments({})', + 'db.getCollection("demo").aggregate([{$match: {"active": true}}])', + 'db.getCollection("demo").distinct("status")', + 'db.getCollection("demo").getIndexes()', + 'db.getCollection("demo").stats()', + "db.version()", + "show databases", + ]) { + expect(classifySqlRisk(sql, { dialect: "mongodb" }).risk, sql).toBe("read"); + } + }); + + it("keeps MongoDB writes and aggregate output stages unsafe", () => { + for (const sql of ['db.getCollection("demo").insertOne({"active": true})', 'db.getCollection("demo").updateOne({"_id": 1}, {$set: {"active": true}})', 'db.getCollection("demo").deleteOne({"_id": 1})', 'db.getCollection("demo").aggregate([{$merge: {into: "archive"}}])']) { + expect(classifySqlRisk(sql, { dialect: "mongodb" }).risk, sql).not.toBe("read"); + } + }); + + it("retains per-command assessments for shell batches", () => { + const assessment = classifySqlRisk('db.getCollection("demo").find({}); db.getCollection("demo").insertOne({"active": true})', { dialect: "mongodb" }); + expect(assessment.risk).toBe("write"); + expect(assessment.statements.map((statement) => statement.risk)).toEqual(["read", "write"]); + }); + + it("keeps unsupported Mongo shell commands conservative", () => { + expect(classifySqlRisk("db.runCommand({ping: 1})", { dialect: "mongodb" }).risk).toBe("unknown"); + expect(classifySqlRisk("db.customReadHelper()", { dialect: "mongodb" }).risk).toBe("unknown"); + }); +}); diff --git a/apps/desktop/src/lib/database/__tests__/productionSafety.spec.ts b/apps/desktop/src/lib/database/__tests__/productionSafety.spec.ts index b32b2e6990..6a4ec172e3 100644 --- a/apps/desktop/src/lib/database/__tests__/productionSafety.spec.ts +++ b/apps/desktop/src/lib/database/__tests__/productionSafety.spec.ts @@ -132,4 +132,16 @@ describe("production SQL safety", () => { it("does not require a production confirmation for reads", () => { expect(assessProductionSql("SELECT * FROM prod_app.orders", connection(), "staging")).toMatchObject({ active: false, isMutation: false }); }); + + it("does not classify Mongo shell reads as production mutations", () => { + const mongo = connection({ db_type: "mongodb", production_databases: ["app"] }); + expect(assessProductionSql('db.getCollection("demo").find({}).skip(0).limit(100)', mongo, "app")).toMatchObject({ active: true, isMutation: false }); + expect(assessProductionSql('db.getCollection("demo").find({}).skip(0).limit(100)', mongo, "scratch")).toMatchObject({ active: false, isMutation: false }); + }); + + it("keeps Mongo shell writes and aggregate output stages protected", () => { + const mongo = connection({ db_type: "mongodb", production_databases: ["app"] }); + expect(assessProductionSql('db.getCollection("demo").updateOne({"_id": 1}, {$set: {"active": true}})', mongo, "app")).toMatchObject({ active: true, isMutation: true }); + expect(assessProductionSql('db.getCollection("demo").aggregate([{$out: "archive"}])', mongo, "app")).toMatchObject({ active: true, isMutation: true }); + }); }); diff --git a/apps/desktop/src/lib/sql/sqlRisk.ts b/apps/desktop/src/lib/sql/sqlRisk.ts index d6575ed306..0055cb4261 100644 --- a/apps/desktop/src/lib/sql/sqlRisk.ts +++ b/apps/desktop/src/lib/sql/sqlRisk.ts @@ -1,4 +1,5 @@ import { classifyElasticsearchRequestRisk, classifyElasticsearchSourceRisk, type ElasticsearchRequestRisk } from "@/lib/elasticsearch/elasticsearchRequestRisk"; +import { mongoAggregateWriteStage, splitMongoCommandRanges, type MongoCommand } from "@/lib/mongo/mongoShellCommand"; import { isElasticsearchCompatibleDatabaseType, type DatabaseType } from "@/types/database"; export type SqlRiskLevel = "read" | "write" | "ddl" | "transaction" | "unknown"; @@ -45,6 +46,12 @@ export function splitSqlStatementsForSafety(sql: string): string[] { } export function classifySqlRisk(sql: string, options: SqlRiskOptions = {}): SqlRiskAssessment { + const mongoStatements = mongoShellStatements(sql, options.dialect); + if (mongoStatements) { + const highest = highestRiskStatement(mongoStatements); + return { ...highest, statements: mongoStatements }; + } + // REST requests carry a JSON body that must not be split on semicolons, and // every request in the text is classified so the highest risk wins. const searchEngineRisk = searchEngineAssessment(sql, options.dialect, classifyElasticsearchSourceRisk); @@ -57,6 +64,9 @@ export function classifySqlRisk(sql: string, options: SqlRiskOptions = {}): SqlR } export function classifySqlStatementRisk(sql: string, options: SqlRiskOptions = {}): SqlRiskStatementAssessment { + const mongoStatements = mongoShellStatements(sql, options.dialect); + if (mongoStatements) return highestRiskStatement(mongoStatements); + const dynamodbRisk = classifyDynamoDbStatementRisk(sql, options.dialect); if (dynamodbRisk) return dynamodbRisk; const searchEngineRisk = searchEngineAssessment(sql, options.dialect, classifyElasticsearchRequestRisk); @@ -82,6 +92,39 @@ function searchEngineAssessment(sql: string, dialect: DatabaseType | string | un return { risk: risk === "read" ? "read" : risk === "write" ? "write" : "ddl", firstKeyword: "rest" }; } +const MONGO_READ_KINDS = new Set(["find", "findOne", "countDocuments", "distinct", "getIndexes", "collectionStats", "version", "showDatabases", "use"]); +const MONGO_DDL_KINDS = new Set(["createIndex", "dropIndex", "dropIndexes", "dropCollection", "createUser"]); + +/** + * Mongo query tabs execute shell commands through the structured Mongo path, + * not the SQL tokenizer. Classify the same parsed command kinds here so + * production protection does not turn ordinary reads into write confirmations. + * Unknown or unparsed commands remain unsafe by falling through to the generic + * classifier. + */ +function mongoShellStatements(sql: string, dialect: DatabaseType | string | undefined): SqlRiskStatementAssessment[] | null { + if (dialect !== "mongodb") return null; + const commands = splitMongoCommandRanges(sql); + if (!commands.length) return null; + + return commands.map(({ command }) => ({ + risk: mongoCommandRisk(command), + firstKeyword: command.kind, + })); +} + +function highestRiskStatement(statements: SqlRiskStatementAssessment[]): SqlRiskStatementAssessment { + return statements.reduce((current, statement) => (RISK_ORDER[statement.risk] > RISK_ORDER[current.risk] ? statement : current), { risk: "read" }); +} + +function mongoCommandRisk(command: MongoCommand): SqlRiskLevel { + if (MONGO_READ_KINDS.has(command.kind)) return "read"; + if (command.kind === "aggregate") return mongoAggregateWriteStage(command.pipeline) ? "write" : "read"; + if (command.kind === "runCommand") return "unknown"; + if (MONGO_DDL_KINDS.has(command.kind)) return "ddl"; + return "write"; +} + function classifyDynamoDbStatementRisk(sql: string, dialect?: DatabaseType | string): SqlRiskStatementAssessment | null { if (dialect !== "dynamodb") return null; const header = sql