From 3e707d7777644db5daee8db5cbee70b872db4aa8 Mon Sep 17 00:00:00 2001 From: Loc Nguyen <113083148+thailoc-dev@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:18:22 +0700 Subject: [PATCH] fix(mongodb): support zero-arg and numeric shell constructors --- crates/dbx-core/src/mongo_shell.rs | 147 ++++++++++++++++--- packages/app-tests/mongoShellCommand.test.ts | 44 ++++++ packages/mongo-shell/src/json.ts | 114 ++++++++++++-- 3 files changed, 275 insertions(+), 30 deletions(-) diff --git a/crates/dbx-core/src/mongo_shell.rs b/crates/dbx-core/src/mongo_shell.rs index 3eee085e0e..95fc2e19f5 100644 --- a/crates/dbx-core/src/mongo_shell.rs +++ b/crates/dbx-core/src/mongo_shell.rs @@ -1,3 +1,5 @@ +use chrono::{SecondsFormat, Utc}; +use mongodb::bson::oid::ObjectId; use serde::Serialize; use serde_json::Value; @@ -1045,6 +1047,11 @@ fn is_shell_regex_value_prefix(previous_significant: Option) -> bool { previous_significant.is_none_or(|character| matches!(character, ':' | '[' | ',' | '(')) } +/// Shell value constructors rewritten to extended JSON before json5 parsing. +/// `Date` is only recognised after `new`, matching the shell where a bare `Date()` +/// returns a string rather than a date. +const SHELL_CONSTRUCTORS: [&str; 6] = ["ObjectId", "ISODate", "Date", "NumberLong", "NumberInt", "NumberDecimal"]; + fn transform_shell_constructors(input: &str) -> Result { let mut output = String::with_capacity(input.len()); let mut index = 0; @@ -1069,29 +1076,86 @@ fn transform_shell_constructors(input: &str) -> Result { output.push_str(&input[start..index]); continue; } - let constructor = if rest.starts_with("ObjectId(") { - Some("ObjectId(") - } else if rest.starts_with("ISODate(") { - Some("ISODate(") - } else { - None - }; - let Some(constructor) = constructor else { - output.push(ch); - index += ch.len_utf8(); + if let Some((json, end)) = shell_constructor_call(input, index)? { + output.push_str(&json); + index = end; continue; - }; - let open = index + constructor.len() - 1; - let close = matching_paren(input, open).ok_or("Unclosed MongoDB value constructor.")?; - let inner = input[open + 1..close].trim(); - let value = parse_string_arg(inner)?; - let key = if constructor.starts_with("ObjectId") { "$oid" } else { "$date" }; - output.push_str(&format!("{{\"{key}\":{}}}", serde_json::to_string(&value).unwrap())); - index = close + 1; + } + output.push(ch); + index += ch.len_utf8(); } Ok(output) } +/// Rewrite one `Name(...)` / `new Name(...)` call at `index`, or None when it is not a known constructor. +fn shell_constructor_call(input: &str, index: usize) -> Result, String> { + let rest = &input[index..]; + let (name_start, is_new) = match rest.strip_prefix("new") { + Some(after) if after.starts_with(|ch: char| ch.is_whitespace()) => { + (index + 3 + after.len() - after.trim_start().len(), true) + } + _ => (index, false), + }; + let name_len = input[name_start..] + .find(|ch: char| !ch.is_alphanumeric() && ch != '_' && ch != '$') + .unwrap_or(input.len() - name_start); + let name = &input[name_start..name_start + name_len]; + if !SHELL_CONSTRUCTORS.contains(&name) { + return Ok(None); + } + if name == "Date" && !is_new { + return Ok(None); + } + let after_name = &input[name_start + name_len..]; + let paren_offset = after_name.len() - after_name.trim_start().len(); + if !after_name[paren_offset..].starts_with('(') { + return Ok(None); + } + let open = name_start + name_len + paren_offset; + let close = matching_paren(input, open).ok_or("Unclosed MongoDB value constructor.")?; + let inner = input[open + 1..close].trim(); + Ok(Some((shell_constructor_to_extended_json(name, inner)?, close + 1))) +} + +fn shell_constructor_to_extended_json(name: &str, inner: &str) -> Result { + let integer = inner.parse::().ok(); + match name { + "ObjectId" if inner.is_empty() => Ok(extended_json("$oid", &ObjectId::new().to_hex())), + "ObjectId" => Ok(extended_json("$oid", &parse_string_arg(inner)?)), + "ISODate" | "Date" if inner.is_empty() => { + Ok(extended_json("$date", &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true))) + } + // `new Date(1735689600000)` takes epoch milliseconds, which extended JSON + // carries as a nested $numberLong rather than a bare number. + "ISODate" | "Date" => match integer { + Some(millis) => Ok(format!(r#"{{"$date":{{"$numberLong":"{millis}"}}}}"#)), + None => Ok(extended_json("$date", &parse_string_arg(inner)?)), + }, + "NumberLong" | "NumberInt" => { + let key = if name == "NumberLong" { "$numberLong" } else { "$numberInt" }; + let value = match integer { + Some(value) => value.to_string(), + None => parse_string_arg(inner)?, + }; + value.parse::().map_err(|_| format!("MongoDB {name}() requires an integer argument."))?; + Ok(extended_json(key, &value)) + } + "NumberDecimal" => { + let value = match parse_json_value(inner) { + Some(Value::Number(number)) => number.to_string(), + _ => parse_string_arg(inner)?, + }; + value.parse::().map_err(|_| "MongoDB NumberDecimal() requires a numeric argument.".to_string())?; + Ok(extended_json("$numberDecimal", &value)) + } + _ => Err(format!("Unsupported MongoDB value constructor {name}().")), + } +} + +fn extended_json(key: &str, value: &str) -> String { + format!("{{\"{key}\":{}}}", serde_json::to_string(value).unwrap_or_else(|_| "null".to_string())) +} + fn parse_json_value(value: &str) -> Option { serde_json::from_str(value).ok() } @@ -1377,6 +1441,53 @@ mod tests { assert_eq!(validate_safety(&command, true, true, false), Ok(())); } + #[test] + fn fills_in_zero_argument_value_constructors() { + let MongoCommand::Update { update, filter, options, .. } = parse( + r#"db.reports.updateOne( + {phone_number: "+84905421172", code: "VN"}, + {$set: {type: ObjectId("5bbc28701f3fd80f00e0211c"), created_at: new Date(), updated_at: ISODate()}}, + {upsert: true} + )"#, + ) + .unwrap() else { + panic!("expected an update command"); + }; + assert_eq!(filter, r#"{"phone_number":"+84905421172","code":"VN"}"#); + assert_eq!(options.as_deref(), Some(r#"{"upsert":true}"#)); + + let set = &parse_json_value(&update).unwrap()["$set"]; + assert_eq!(set["type"]["$oid"], "5bbc28701f3fd80f00e0211c"); + for field in ["created_at", "updated_at"] { + let filled = set[field]["$date"].as_str().expect("filled-in date"); + let filled = chrono::DateTime::parse_from_rfc3339(filled).expect("rfc3339 date"); + assert!( + (Utc::now() - filled.with_timezone(&Utc)).num_seconds().abs() < 60, + "{field} is not the current time" + ); + } + + let MongoCommand::Find { filter, .. } = parse("db.reports.find({_id: ObjectId()})").unwrap() else { + panic!("expected a find command"); + }; + let oid = parse_json_value(&filter).unwrap()["_id"]["$oid"].as_str().unwrap().to_string(); + assert!(oid.len() == 24 && oid.chars().all(|ch| ch.is_ascii_hexdigit()), "{oid}"); + } + + #[test] + fn rewrites_epoch_millisecond_and_numeric_value_constructors() { + let MongoCommand::Find { filter, .. } = + parse(r#"db.orders.find({at: new Date(1735689600000), qty: NumberInt(3), sequence: NumberLong("9223372036854775807"), total: NumberDecimal("12.34")})"#) + .unwrap() + else { + panic!("expected a find command"); + }; + assert_eq!( + filter, + r#"{"at":{"$date":{"$numberLong":"1735689600000"}},"qty":{"$numberInt":"3"},"sequence":{"$numberLong":"9223372036854775807"},"total":{"$numberDecimal":"12.34"}}"# + ); + } + #[test] fn parses_show_databases_aliases_as_read_only() { for source in [ diff --git a/packages/app-tests/mongoShellCommand.test.ts b/packages/app-tests/mongoShellCommand.test.ts index ace6cbb9bf..dc53e7c409 100644 --- a/packages/app-tests/mongoShellCommand.test.ts +++ b/packages/app-tests/mongoShellCommand.test.ts @@ -313,6 +313,50 @@ test("parseMongoFindCommand rewrites NumberLong into extended JSON", () => { assert.deepEqual(JSON.parse(unquoted.filter), { snowflake: { $numberLong: "9007199254740993" } }); }); +test("parseMongoWriteCommand fills in zero-argument new Date, ISODate and ObjectId", () => { + const before = Date.now(); + const command = parseMongoWriteCommand(`db.reports.updateOne( + { phone_number: "+84905421172", code: "VN" }, + { $set: { type: ObjectId("5bbc28701f3fd80f00e0211c"), created_at: new Date(), updated_at: ISODate() } }, + { upsert: true } + )`); + assert.ok(command); + assert.equal(command.kind, "update"); + const update = JSON.parse(command.update) as { $set: Record }; + assert.deepEqual(update.$set.type, { $oid: "5bbc28701f3fd80f00e0211c" }); + for (const field of ["created_at", "updated_at"] as const) { + const filled = Date.parse(update.$set[field]!.$date!); + assert.ok(filled >= before && filled <= Date.now(), `${field} is not the current time`); + } + assert.match(JSON.parse(parseMongoFindCommand("db.reports.find({_id: ObjectId()})")!.filter)._id.$oid, /^[0-9a-f]{24}$/); +}); + +test("parseMongoFindCommand rewrites epoch-millisecond and numeric BSON constructors", () => { + const command = parseMongoFindCommand('db.orders.find({at: new Date(1735689600000), qty: NumberInt(3), total: NumberDecimal("12.34")})'); + assert.ok(command); + assert.deepEqual(JSON.parse(command.filter), { + at: { $date: { $numberLong: "1735689600000" } }, + qty: { $numberInt: "3" }, + total: { $numberDecimal: "12.34" }, + }); +}); + +test("parseMongoFindCommand rejects out-of-range integer constructor arguments", () => { + assert.equal(parseMongoFindCommand("db.orders.find({a: NumberLong(9223372036854775808)})"), null); + assert.equal(parseMongoFindCommand("db.orders.find({a: NumberLong(-9223372036854775809)})"), null); + assert.equal(parseMongoFindCommand("db.orders.find({a: NumberInt(2147483648)})"), null); + assert.equal(parseMongoFindCommand("db.orders.find({a: new Date(99999999999999999999)})"), null); + const atBounds = parseMongoFindCommand("db.orders.find({a: NumberLong(9223372036854775807), b: NumberInt(-2147483648)})"); + assert.ok(atBounds); + assert.deepEqual(JSON.parse(atBounds.filter), { a: { $numberLong: "9223372036854775807" }, b: { $numberInt: "-2147483648" } }); +}); + +test("parseMongoFindCommand does not rewrite constructor text inside strings", () => { + const command = parseMongoFindCommand(`db.orders.find({label: "new Date()", note: 'ObjectId()'})`); + assert.ok(command); + assert.deepEqual(JSON.parse(command.filter), { label: "new Date()", note: "ObjectId()" }); +}); + test("parseMongoFindCommand accepts single-quoted string values and unquoted sort keys", () => { const command = parseMongoFindCommand("db.products.find({category: 'Electronics'}).sort({price: -1}).limit(2)"); assert.ok(command); diff --git a/packages/mongo-shell/src/json.ts b/packages/mongo-shell/src/json.ts index 7914465f25..5f15498301 100644 --- a/packages/mongo-shell/src/json.ts +++ b/packages/mongo-shell/src/json.ts @@ -13,7 +13,8 @@ export function normalizeJsonArgument(value: string): string | null { if (!withoutComments) return "{}"; const withoutEjsonDeserialize = replaceMongoEjsonDeserialize(withoutComments); // Rewrite mongo shell constructors that are not valid JSON into extended JSON - // (mongo_driver::json_value_to_bson): ObjectId / NumberLong / ISODate / new Date. + // (mongo_driver::json_value_to_bson): ObjectId / ISODate / new Date / NumberLong / + // NumberInt / NumberDecimal. const withExtendedJson = replaceMongoShellConstructors(withoutEjsonDeserialize); const preprocessed = quoteUnquotedObjectKeys(convertSingleQuotedStrings(withExtendedJson)); try { @@ -535,8 +536,31 @@ function replaceMongoEjsonDeserialize(source: string): string { return result; } +/** + * Shell value constructors rewritten to extended JSON. + * `Date` is only recognised after `new`, matching the shell where a bare `Date()` + * returns a string rather than a date. + */ +const SHELL_CONSTRUCTORS = new Set(["ObjectId", "ISODate", "Date", "NumberLong", "NumberInt", "NumberDecimal"]); +const CONSTRUCTOR_CALL = /^(new\s+)?([A-Za-z_$][\w$]*)\s*\(/; +const QUOTED_ARGUMENT = /^(["'])([^\\]*)\1$/; +const INTEGER_ARGUMENT = /^-?\d+$/; +const DECIMAL_ARGUMENT = /^-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?$/; + +const INT64_BOUNDS = [-9223372036854775808n, 9223372036854775807n] as const; +const INT32_BOUNDS = [-2147483648n, 2147483647n] as const; + +/** Bounds parity with the Rust parser, which range-checks before emitting $numberLong / $numberInt. */ +function fitsIntegerBounds(value: string, bounds: readonly [bigint, bigint]): boolean { + try { + const parsed = BigInt(value); + return parsed >= bounds[0] && parsed <= bounds[1]; + } catch { + return false; + } +} + function replaceMongoShellConstructors(source: string): string { - const constructor = /^(ObjectId|NumberLong|ISODate)\s*\(\s*["']([^"']+)["']\s*\)|^(ObjectId|NumberLong)\s*\(\s*(-?\d+)\s*\)|^(?:new\s+Date)\s*\(\s*["']([^"']+)["']\s*\)/; let result = ""; let index = 0; while (index < source.length) { @@ -553,23 +577,89 @@ function replaceMongoShellConstructors(source: string): string { result += source.slice(start, index); continue; } - const match = source.slice(index).match(constructor); - if (!match) { + const call = matchShellConstructorCall(source, index); + if (!call) { result += source[index++]!; continue; } - if (match[1]) { - result += match[1] === "ObjectId" ? `{"$oid":"${match[2]}"}` : match[1] === "NumberLong" ? `{"$numberLong":"${match[2]}"}` : `{"$date":"${match[2]}"}`; - } else if (match[3]) { - result += match[3] === "NumberLong" ? `{"$numberLong":"${match[4]}"}` : `{"$oid":"${match[4]}"}`; - } else { - result += `{"$date":"${match[5]}"}`; - } - index += match[0].length; + result += call.json; + index = call.end; } return result; } +/** Rewrite one `Name(...)` / `new Name(...)` call at {@link index}, or null when it is not a known constructor. */ +function matchShellConstructorCall(source: string, index: number): { json: string; end: number } | null { + const match = CONSTRUCTOR_CALL.exec(source.slice(index)); + if (!match) return null; + const name = match[2]!; + if (!SHELL_CONSTRUCTORS.has(name)) return null; + if (name === "Date" && !match[1]) return null; + + const openIndex = index + match[0].length - 1; + const closeIndex = findMatchingParen(source, openIndex); + if (closeIndex < 0) return null; + + const inner = source.slice(openIndex + 1, closeIndex).trim(); + const args = inner ? splitTopLevel(inner) : []; + const json = shellConstructorToExtendedJson(name, args); + return json === null ? null : { json, end: closeIndex + 1 }; +} + +function shellConstructorToExtendedJson(name: string, args: string[]): string | null { + if (args.length > 1) return null; + const arg = args[0]?.trim(); + const literal = arg ? (QUOTED_ARGUMENT.exec(arg)?.[2] ?? null) : null; + + switch (name) { + case "ObjectId": + if (!arg) return wrap("$oid", generateObjectIdHex()); + return literal !== null || INTEGER_ARGUMENT.test(arg) ? wrap("$oid", literal ?? arg) : null; + case "ISODate": + case "Date": + if (!arg) return wrap("$date", new Date().toISOString()); + if (literal !== null) return wrap("$date", literal); + // `new Date(1735689600000)` takes epoch milliseconds, which extended JSON + // carries as a nested $numberLong rather than a bare number. + return INTEGER_ARGUMENT.test(arg) && fitsIntegerBounds(arg, INT64_BOUNDS) ? `{"$date":{"$numberLong":${JSON.stringify(arg)}}}` : null; + case "NumberLong": + case "NumberInt": { + const value = literal ?? arg; + if (value === undefined || !INTEGER_ARGUMENT.test(value)) return null; + if (!fitsIntegerBounds(value, name === "NumberLong" ? INT64_BOUNDS : INT32_BOUNDS)) return null; + return wrap(name === "NumberLong" ? "$numberLong" : "$numberInt", value); + } + case "NumberDecimal": { + const value = literal ?? arg; + if (value === undefined || !DECIMAL_ARGUMENT.test(value)) return null; + return wrap("$numberDecimal", value); + } + default: + return null; + } +} + +function wrap(key: string, value: string): string { + return `{${JSON.stringify(key)}:${JSON.stringify(value)}}`; +} + +const OBJECT_ID_RANDOM = randomHex(5); +let objectIdCounter = Math.floor(Math.random() * 0xffffff); + +/** Client-side ObjectId for a bare `ObjectId()`, mirroring how the shell fills one in. */ +function generateObjectIdHex(): string { + objectIdCounter = (objectIdCounter + 1) % 0x1000000; + const seconds = Math.floor(Date.now() / 1000) % 0x100000000; + return seconds.toString(16).padStart(8, "0") + OBJECT_ID_RANDOM + objectIdCounter.toString(16).padStart(6, "0"); +} + +function randomHex(bytes: number): string { + const values = new Uint8Array(bytes); + if (typeof globalThis.crypto?.getRandomValues === "function") globalThis.crypto.getRandomValues(values); + else for (let i = 0; i < bytes; i += 1) values[i] = Math.floor(Math.random() * 256); + return Array.from(values, (value) => value.toString(16).padStart(2, "0")).join(""); +} + function convertSingleQuotedStrings(source: string): string { let result = ""; let copiedUntil = 0;