SERVER-47436 Make shards validate shardKey in dataSize command

(cherry picked from commit 59e005fea0e1ca575083ded8c02c518048fb8af0)
(cherry picked from commit 20364840b8f1af16917e4c23c1b5f5efd8b352f8)

4 files changed, 137 insertions, 26 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
index ff964dd230e..4e8b3dbce46 100644
--- a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
@@ -82,6 +82,8 @@ selector:
   - jstests/sharding/clear_jumbo.js
   # Requires count command to be accurate on sharded clusters, introduced in v4.0.
   - jstests/sharding/accurate_count_with_predicate.js
+  # Requires that dataSize command uses shard versioning, introduced in v4.0.
+  - jstests/sharding/mongos_dataSize.js
   # The 3.6 mongos does not support using localOps to find failpoints.
   - jstests/sharding/verify_sessions_expiration_sharded.js
   # Enable when SERVER-33538 is backported.
diff --git a/jstests/sharding/mongos_dataSize.js b/jstests/sharding/mongos_dataSize.js
new file mode 100644
index 00000000000..2a707df707c
--- /dev/null
+++ b/jstests/sharding/mongos_dataSize.js
@@ -0,0 +1,112 @@
+/*
+ * Tests the dataSize command on mongos.
+ */
+(function() {
+    'use strict';
+
+    const kDbName = "foo";
+    const kCollName = "bar";
+    const kNs = kDbName + "." + kCollName;
+    const kNumDocs = 100;
+
+    /*
+     * Returns the global min and max key for the given key pattern.
+     */
+    function getGlobalMinMaxKey(keyPattern) {
+        let globalMin = {};
+        let globalMax = {};
+        for (let field in keyPattern) {
+            globalMin[field] = MinKey;
+            globalMax[field] = MaxKey;
+        }
+        return {globalMin, globalMax};
+    }
+
+    /*
+     * Runs a dataSize command with the given key pattern on the given connection, and
+     * asserts that command works and that the returned numObjects is equal to expectedNumObjects.
+     */
+    function assertDataSizeCmdWorked(conn, keyPattern, expectedNumObjects) {
+        let res = assert.commandWorked(conn.adminCommand({dataSize: kNs, keyPattern: keyPattern}));
+        assert.eq(res.numObjects, expectedNumObjects);
+
+        const {globalMin, globalMax} = getGlobalMinMaxKey(keyPattern);
+        res = assert.commandWorked(conn.adminCommand(
+            {dataSize: kNs, keyPattern: keyPattern, min: globalMin, max: globalMax}));
+        assert.eq(res.numObjects, expectedNumObjects);
+    }
+
+    /*
+     * Runs a dataSize command with the given key pattern on the given connection, and
+     * asserts that command failed with BadValue error.
+     */
+    function assertDataSizeCmdFailedWithBadValue(conn, keyPattern) {
+        assert.commandFailedWithCode(conn.adminCommand({dataSize: kNs, keyPattern: keyPattern}),
+                                     ErrorCodes.BadValue);
+    }
+
+    /*
+     * Runs dataSize commands on the given connection. Asserts the command fails if run with an
+     * invalid
+     * namespace or with the min and max as given by each range in invalidRanges. Asserts that the
+     * command succeeds if run with valid min and max and returns the expected numObjects.
+     */
+    function testDataSizeCmd(conn, keyPattern, invalidRanges, numObjects) {
+        assert.commandFailedWithCode(conn.adminCommand({dataSize: kCollName}),
+                                     ErrorCodes.InvalidNamespace);
+
+        for (let {
+                 min, max, errorCode
+             } of invalidRanges) {
+            const cmdObj = {dataSize: kNs, keyPattern: keyPattern, min: min, max: max};
+            assert.commandFailedWithCode(conn.adminCommand(cmdObj), errorCode);
+        }
+
+        assertDataSizeCmdWorked(conn, {}, numObjects);
+        assertDataSizeCmdWorked(conn, keyPattern, numObjects);
+    }
+
+    const st = new ShardingTest({mongos: 3, shards: 2});
+    assert.commandWorked(st.s.adminCommand({enableSharding: kDbName}));
+    st.ensurePrimaryShard(kDbName, st.shard0.shardName);
+
+    const shardKey1 = {x: 1};
+    jsTest.log(`Sharding the collection with key ${tojson(shardKey1)}`);
+    assert.commandWorked(st.s0.adminCommand({shardCollection: kNs, key: shardKey1}));
+
+    let bulk = st.s0.getCollection(kNs).initializeUnorderedBulkOp();
+    for (let i = 0; i < kNumDocs; ++i) {
+        bulk.insert({_id: i, x: i, y: -i});
+    }
+    assert.commandWorked(bulk.execute());
+
+    jsTest.log("Verify that keyPattern and key range validation works");
+    const invalidRanges1 = [
+        {min: {y: MinKey}, max: {y: MaxKey}, errorCode: ErrorCodes.BadValue},
+        {min: {x: MinKey, y: MinKey}, max: {x: MaxKey, y: MaxKey}, errorCode: ErrorCodes.BadValue},
+        // The command does not throw any particular error when only one of min or max is specified.
+        {min: {}, max: {x: MaxKey}, errorCode: ErrorCodes.UnknownError},
+        {min: {x: MinKey}, max: {}, errorCode: ErrorCodes.UnknownError},
+    ];
+    testDataSizeCmd(st.s0, shardKey1, invalidRanges1, kNumDocs);
+    testDataSizeCmd(st.s1, shardKey1, invalidRanges1, kNumDocs);
+    testDataSizeCmd(st.s2, shardKey1, invalidRanges1, kNumDocs);
+
+    jsTest.log("Dropping the collection");
+    st.s0.getCollection(kNs).drop();
+
+    const shardKey2 = {y: 1};
+    jsTest.log(`Resharding the collection with key ${tojson(shardKey2)}`);
+    assert.commandWorked(st.s0.adminCommand({shardCollection: kNs, key: shardKey2}));
+
+    jsTest.log("Verify that validation occurs on shards not on mongos");
+    // If validation occurs on mongos, the command would fail with BadValue as st.s1 is stale so
+    // it thinks that shardKey1 is the shard key.
+    assertDataSizeCmdWorked(st.s1, shardKey2, 0);
+
+    // If validation occurs on mongos, the command would succeed as st.s2 is stale so it thinks
+    // that shardKey1 is the shard key.
+    assertDataSizeCmdFailedWithBadValue(st.s2, shardKey1);
+
+    st.stop();
+})();
diff --git a/src/mongo/db/commands/dbcommands.cpp b/src/mongo/db/commands/dbcommands.cpp
index 6203a030672..5f9628912c6 100644
--- a/src/mongo/db/commands/dbcommands.cpp
+++ b/src/mongo/db/commands/dbcommands.cpp
@@ -489,9 +489,30 @@ public:
         BSONObj keyPattern = jsobj.getObjectField("keyPattern");
         bool estimate = jsobj["estimate"].trueValue();
 
-        AutoGetCollectionForReadCommand ctx(opCtx, NamespaceString(ns));
-
+        const NamespaceString nss(ns);
+        AutoGetCollectionForReadCommand ctx(opCtx, nss);
         Collection* collection = ctx.getCollection();
+
+        const auto metadata = CollectionShardingState::get(opCtx, nss)->getCurrentMetadata();
+
+        if (metadata->isSharded()) {
+            const ShardKeyPattern shardKeyPattern(metadata->getKeyPattern());
+            uassert(ErrorCodes::BadValue,
+                    "keyPattern must be empty or must be an object that equals the shard key",
+                    keyPattern.isEmpty() || (SimpleBSONObjComparator::kInstance.evaluate(
+                                                shardKeyPattern.toBSON() == keyPattern)));
+
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << "min value " << min << " does not have shard key",
+                    min.isEmpty() || shardKeyPattern.isShardKey(min));
+            min = shardKeyPattern.normalizeShardKey(min);
+
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << "max value " << max << " does not have shard key",
+                    max.isEmpty() || shardKeyPattern.isShardKey(max));
+            max = shardKeyPattern.normalizeShardKey(max);
+        }
+
         long long numRecords = 0;
         if (collection) {
             numRecords = collection->numRecords(opCtx);
diff --git a/src/mongo/s/commands/cluster_data_size_cmd.cpp b/src/mongo/s/commands/cluster_data_size_cmd.cpp
index 57d6659f26a..a6a18c4386f 100644
--- a/src/mongo/s/commands/cluster_data_size_cmd.cpp
+++ b/src/mongo/s/commands/cluster_data_size_cmd.cpp
@@ -77,30 +77,6 @@ public:
 
         auto routingInfo =
             uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss));
-        const auto cm = routingInfo.cm();
-
-        BSONObj min = cmdObj.getObjectField("min");
-        BSONObj max = cmdObj.getObjectField("max");
-
-        if (cm) {
-            auto keyPattern = cmdObj["keyPattern"];
-
-            uassert(ErrorCodes::BadValue,
-                    "keyPattern must be empty or must be an object that equals the shard key",
-                    !keyPattern || (keyPattern.type() == Object &&
-                                    SimpleBSONObjComparator::kInstance.evaluate(
-                                        cm->getShardKeyPattern().toBSON() == keyPattern.Obj())));
-
-            uassert(ErrorCodes::BadValue,
-                    str::stream() << "min value " << min << " does not have shard key",
-                    cm->getShardKeyPattern().isShardKey(min));
-            min = cm->getShardKeyPattern().normalizeShardKey(min);
-
-            uassert(ErrorCodes::BadValue,
-                    str::stream() << "max value " << max << " does not have shard key",
-                    cm->getShardKeyPattern().isShardKey(max));
-            max = cm->getShardKeyPattern().normalizeShardKey(max);
-        }
 
         auto shardResults = scatterGatherVersionedTargetByRoutingTable(
             opCtx,