SERVER-39650 Ensure internal options cannot be specified in a raw aggregate command sent to mongos

(cherry picked from commit 3b4c635f2a3a65a8804232c80a48cdefa3c26b65)

5 files changed, 154 insertions, 6 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 d8ecc2e22dc..eb5a187f792 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
@@ -13,6 +13,7 @@ selector:
   - jstests/sharding/database_versioning_upgrade_downgrade.js
   - jstests/sharding/shard_collection_cache_upgrade_downgrade.js
   #### Enable when 4.0 becomes last-stable.
+  - jstests/sharding/aggregation_internal_parameters.js
   - jstests/sharding/change_stream_no_shards.js
   - jstests/sharding/clone_catalog_data.js
   - jstests/sharding/create_database.js
diff --git a/jstests/sharding/aggregation_internal_parameters.js b/jstests/sharding/aggregation_internal_parameters.js
new file mode 100644
index 00000000000..7ac3943c901
--- /dev/null
+++ b/jstests/sharding/aggregation_internal_parameters.js
@@ -0,0 +1,97 @@
+/**
+ * Tests that mongoS rejects 'aggregate' commands which explicitly set any of the
+ * parameters that mongoS uses internally when communicating with the shards.
+ */
+(function() {
+    "use strict";
+
+    const st = new ShardingTest({shards: 2, rs: {nodes: 1, enableMajorityReadConcern: ''}});
+
+    const mongosDB = st.s0.getDB(jsTestName());
+    const mongosColl = mongosDB[jsTestName()];
+
+    assert.commandWorked(mongosDB.dropDatabase());
+
+    // Enable sharding on the test DB and ensure its primary is st.shard0.shardName.
+    assert.commandWorked(mongosDB.adminCommand({enableSharding: mongosDB.getName()}));
+    st.ensurePrimaryShard(mongosDB.getName(), st.rs0.getURL());
+
+    // Test that command succeeds when no internal options have been specified.
+    assert.commandWorked(
+        mongosDB.runCommand({aggregate: mongosColl.getName(), pipeline: [], cursor: {}}));
+
+    // Test that the command fails if we have 'needsMerge: false' without 'fromMongos'.
+    assert.commandFailedWithCode(
+        mongosDB.runCommand(
+            {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, needsMerge: false}),
+        ErrorCodes.FailedToParse);
+
+    // Test that the command fails if we have 'needsMerge: true' without 'fromMongos'.
+    assert.commandFailedWithCode(
+        mongosDB.runCommand(
+            {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, needsMerge: true}),
+        ErrorCodes.FailedToParse);
+
+    // Test that 'fromMongos: true' cannot be specified in a command sent to mongoS.
+    assert.commandFailedWithCode(
+        mongosDB.runCommand(
+            {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, fromMongos: true}),
+        51089);
+
+    // Test that 'fromMongos: false' can be specified in a command sent to mongoS.
+    assert.commandWorked(mongosDB.runCommand(
+        {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, fromMongos: false}));
+
+    // Test that the command fails if we have 'needsMerge: true' with 'fromMongos: false'.
+    assert.commandFailedWithCode(mongosDB.runCommand({
+        aggregate: mongosColl.getName(),
+        pipeline: [],
+        cursor: {},
+        needsMerge: true,
+        fromMongos: false
+    }),
+                                 51089);
+
+    // Test that the command fails if we have 'needsMerge: true' with 'fromMongos: true'.
+    assert.commandFailedWithCode(mongosDB.runCommand({
+        aggregate: mongosColl.getName(),
+        pipeline: [],
+        cursor: {},
+        needsMerge: true,
+        fromMongos: true
+    }),
+                                 51089);
+
+    // Test that 'needsMerge: false' can be specified in a command sent to mongoS along with
+    // 'fromMongos: false'.
+    assert.commandWorked(mongosDB.runCommand({
+        aggregate: mongosColl.getName(),
+        pipeline: [],
+        cursor: {},
+        needsMerge: false,
+        fromMongos: false
+    }));
+
+    // Test that 'mergeByPBRT: true' cannot be specified in a command sent to mongoS.
+    assert.commandFailedWithCode(
+        mongosDB.runCommand(
+            {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, mergeByPBRT: true}),
+        51089);
+
+    // Test that 'mergeByPBRT: false' can be specified in a command sent to mongoS.
+    assert.commandWorked(mongosDB.runCommand(
+        {aggregate: mongosColl.getName(), pipeline: [], cursor: {}, mergeByPBRT: false}));
+
+    // Test that the command fails when all internal parameters have been specified.
+    assert.commandFailedWithCode(mongosDB.runCommand({
+        aggregate: mongosColl.getName(),
+        pipeline: [],
+        cursor: {},
+        needsMerge: true,
+        fromMongos: true,
+        mergeByPBRT: true
+    }),
+                                 51089);
+
+    st.stop();
+})();
diff --git a/src/mongo/db/pipeline/aggregation_request.cpp b/src/mongo/db/pipeline/aggregation_request.cpp
index fdbf900bae1..11b49f23dc0 100644
--- a/src/mongo/db/pipeline/aggregation_request.cpp
+++ b/src/mongo/db/pipeline/aggregation_request.cpp
@@ -327,5 +327,4 @@ Document AggregationRequest::serializeToCommandObj() const {
         {QueryRequest::cmdOptionMaxTimeMS,
          _maxTimeMS == 0 ? Value() : Value(static_cast<int>(_maxTimeMS))}};
 }
-
 }  // namespace mongo
diff --git a/src/mongo/s/commands/cluster_aggregate.cpp b/src/mongo/s/commands/cluster_aggregate.cpp
index 38e01a6ab73..d58ce76220a 100644
--- a/src/mongo/s/commands/cluster_aggregate.cpp
+++ b/src/mongo/s/commands/cluster_aggregate.cpp
@@ -930,6 +930,14 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
                                       const AggregationRequest& request,
                                       BSONObj cmdObj,
                                       BSONObjBuilder* result) {
+    uassert(51089,
+            str::stream() << "Internal parameter(s) [" << AggregationRequest::kNeedsMergeName
+                          << ", "
+                          << AggregationRequest::kFromMongosName
+                          << ", "
+                          << AggregationRequest::kMergeByPBRTName
+                          << "] cannot be set to 'true' when sent to mongos",
+            !request.needsMerge() && !request.isFromMongos() && !request.mergeByPBRT());
     auto executionNsRoutingInfoStatus = getExecutionNsRoutingInfo(opCtx, namespaces.executionNss);
     boost::optional<CachedCollectionRoutingInfo> routingInfo;
     LiteParsedPipeline litePipe(request);
diff --git a/src/mongo/s/commands/cluster_aggregate_test.cpp b/src/mongo/s/commands/cluster_aggregate_test.cpp
index db4d2981090..59ffa830242 100644
--- a/src/mongo/s/commands/cluster_aggregate_test.cpp
+++ b/src/mongo/s/commands/cluster_aggregate_test.cpp
@@ -59,12 +59,12 @@ class ClusterAggregateTest : public CatalogCacheTestFixture {
 protected:
     const BSONObj kAggregateCmdTargeted{
         fromjson("{aggregate: 'coll', pipeline: [{$match: {_id: 0}}], explain: false, "
-                 "allowDiskUse: false, fromMongos: true, "
-                 "cursor: {batchSize: 10}, maxTimeMS: 100, readConcern: {level: 'snapshot'}}")};
+                 "allowDiskUse: false, cursor: {batchSize: 10}, maxTimeMS: 100, readConcern: "
+                 "{level: 'snapshot'}}")};
 
-    const BSONObj kAggregateCmdScatterGather{fromjson(
-        "{aggregate: 'coll', pipeline: [], explain: false, allowDiskUse: false, fromMongos: true, "
-        "cursor: {batchSize: 10}, readConcern: {level: 'snapshot'}}")};
+    const BSONObj kAggregateCmdScatterGather{
+        fromjson("{aggregate: 'coll', pipeline: [], explain: false, allowDiskUse: false, cursor: "
+                 "{batchSize: 10}, readConcern: {level: 'snapshot'}}")};
 
     void setUp() {
         CatalogCacheTestFixture::setUp();
@@ -188,6 +188,27 @@ protected:
 
         future.timed_get(kFutureTimeout);
     }
+
+    /**
+     * This method should only be used to test early exits from Cluster::runAggregate, before
+     * a request is sent to the shards. Otherwise the call would get blocked as no expect* hooks
+     * are provided in this method.
+     */
+    Status testRunAggregateEarlyExit(const BSONObj& inputBson) {
+        BSONObjBuilder result;
+        NamespaceString nss{"a.collection"};
+        auto client = getServiceContext()->makeClient("ClusterCmdClient");
+        auto opCtx = client->makeOperationContext();
+        auto request = AggregationRequest::parseFromBSON(nss, inputBson);
+        if (request.getStatus() != Status::OK()) {
+            return request.getStatus();
+        }
+        return ClusterAggregate::runAggregate(opCtx.get(),
+                                              ClusterAggregate::Namespaces{nss, nss},
+                                              request.getValue(),
+                                              inputBson,
+                                              &result);
+    }
 };
 
 TEST_F(ClusterAggregateTest, NoErrors) {
@@ -242,5 +263,27 @@ TEST_F(ClusterAggregateTest, MaxRetriesSnapshotErrors) {
     runAggCommandMaxErrors(kAggregateCmdScatterGather, ErrorCodes::SnapshotTooOld, false);
 }
 
+TEST_F(ClusterAggregateTest, ShouldFailWhenFromMongosIsTrue) {
+    const BSONObj inputBson = fromjson("{pipeline: [], cursor: {}, fromMongos: true}");
+    ASSERT_THROWS_CODE(testRunAggregateEarlyExit(inputBson).ignore(), AssertionException, 51089);
+}
+
+TEST_F(ClusterAggregateTest, ShouldFailWhenNeedsMergeIstrueAndFromMongosIsFalse) {
+    const BSONObj inputBson =
+        fromjson("{pipeline: [], cursor: {}, needsMerge: true, fromMongos: false}");
+    ASSERT_THROWS_CODE(testRunAggregateEarlyExit(inputBson).ignore(), AssertionException, 51089);
+}
+
+TEST_F(ClusterAggregateTest, ShouldFailWhenNeedsMergeIstrueAndFromMongosIsTrue) {
+    const BSONObj inputBson =
+        fromjson("{pipeline: [], cursor: {}, needsMerge: true, fromMongos: true}");
+    ASSERT_THROWS_CODE(testRunAggregateEarlyExit(inputBson).ignore(), AssertionException, 51089);
+}
+
+TEST_F(ClusterAggregateTest, ShouldFailWhenMergeByPBRTIsTrue) {
+    const BSONObj inputBson = fromjson("{pipeline: [], cursor: {}, mergeByPBRT: true}");
+    ASSERT_THROWS_CODE(testRunAggregateEarlyExit(inputBson).ignore(), AssertionException, 51089);
+}
+
 }  // namespace
 }  // namespace mongo