From 0a1a88755daa504c67e13dd49667a06c3c606ee8 Mon Sep 17 00:00:00 2001
From: Davis Haupt <davis.haupt@mongodb.com>
Date: Thu, 6 Jan 2022 22:17:46 +0000
Subject: SERVER-58358 validate arguments passed to dataSize command

---
 jstests/core/datasize_validation.js  | 53 ++++++++++++++++++++++++++++++++++++
 src/mongo/db/commands/dbcommands.cpp | 18 +++++++++++-
 2 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 jstests/core/datasize_validation.js

diff --git a/jstests/core/datasize_validation.js b/jstests/core/datasize_validation.js
new file mode 100644
index 00000000000..c0f9c9d67b2
--- /dev/null
+++ b/jstests/core/datasize_validation.js
@@ -0,0 +1,53 @@
+// Cannot implicitly shard accessed collections because the "dataSize" command returns an
+// "keyPattern must equal shard key" error response.
+// @tags: [assumes_unsharded_collection, requires_fcv_53]
+
+//
+// Test argument validation for dataSize command
+//
+
+(function() {
+let coll = db[jsTestName()];
+coll.drop();
+coll.insertOne({_id: 1});
+coll.insertOne({_id: 2});
+coll.insertOne({_id: 3});
+coll.insertOne({_id: 4});
+
+assert.commandFailed(db.runCommand({
+    'dataSize': coll.getFullName(),
+    min: NumberLong("1"),
+    max: NumberLong("2"),
+    estimate: false
+}),
+                     "min and max should be objects");
+
+assert.commandFailed(
+    db.runCommand({'dataSize': coll.getFullName(), min: {_id: NumberLong("1")}, estimate: false}),
+    "min and max should both be present");
+assert.commandFailed(
+    db.runCommand({'dataSize': coll.getFullName(), max: {_id: NumberLong("2")}, estimate: false}),
+    "min and max should both be present");
+
+let resultWithKey = assert.commandWorked(db.runCommand({
+    'dataSize': coll.getFullName(),
+    keyPattern: {_id: 1},
+    min: {_id: NumberLong("1")},
+    max: {_id: NumberLong("2")},
+    estimate: false
+}));
+assert.eq(1, resultWithKey.numObjects, "only 1 object should be inspected between min/max bounds.");
+
+let result = assert.commandWorked(db.runCommand({
+    'dataSize': coll.getFullName(),
+    min: {_id: NumberLong("1")},
+    max: {_id: NumberLong("2")},
+    estimate: false
+}));
+assert.eq(result.size,
+          resultWithKey.size,
+          "measured size should be equal after keyPattern properly inferred from min/max bounds.");
+assert.eq(result.numObjects,
+          resultWithKey.numObjects,
+          "numObjects should be equal after keyPattern properly inferred from min/max bounds.");
+})();
diff --git a/src/mongo/db/commands/dbcommands.cpp b/src/mongo/db/commands/dbcommands.cpp
index 9167dc0d18e..1addfe77936 100644
--- a/src/mongo/db/commands/dbcommands.cpp
+++ b/src/mongo/db/commands/dbcommands.cpp
@@ -304,8 +304,24 @@ public:
                    const BSONObj& jsobj,
                    std::string& errmsg,
                    BSONObjBuilder& result) override {
-        Timer timer;
+        auto hasMin = jsobj.hasField("min");
+        auto hasMax = jsobj.hasField("max");
+
+        uassert(ErrorCodes::BadValue,
+                hasMin ? "max must be set if min is set" : "min must be set if max is set",
+                hasMin == hasMax);
+
+        if (hasMin) {
+            uassert(ErrorCodes::BadValue,
+                    "min key must be an object",
+                    jsobj["min"].type() == BSONType::Object);
 
+            uassert(ErrorCodes::BadValue,
+                    "max key must be an object",
+                    jsobj["max"].type() == BSONType::Object);
+        }
+
+        Timer timer;
         std::string ns = jsobj.firstElement().String();
         BSONObj min = jsobj.getObjectField("min");
         BSONObj max = jsobj.getObjectField("max");
-- 
cgit v1.2.1