SERVER-67406 Ensure we don't remove user $fields when removing metadata

(cherry picked from commit e4ffe4b63da0ef16e129ac693dda79b64047a469)

3 files changed, 107 insertions, 2 deletions

diff --git a/jstests/sharding/query/metadata_removal.js b/jstests/sharding/query/metadata_removal.js
new file mode 100644
index 00000000000..7caba6433bc
--- /dev/null
+++ b/jstests/sharding/query/metadata_removal.js
@@ -0,0 +1,46 @@
+// Ensure that metadata fields (and only metadata) are removed before documents are returned to
+// customers. As we can't insert documents with metadata, this test focuses on queries that might
+// inject metadata when communication results to mongos for merging. In addition, we want to ensure
+// that other fields beginning with a '$' are not modified before returning them to customers.
+// @tags: [
+//   requires_fcv_60
+// ]
+
+(function() {
+"use strict";
+
+function runTest(coll, keyword) {
+    const document = ({_id: 5, x: 1, $set: {$inc: {x: 5}}});
+    assert.commandWorked(coll.insert(document));
+
+    // Leaves $set intact without shard merging.
+    assert.eq([document], coll.find().hint({$natural: 1}).toArray());
+    assert.eq([document], coll.find().hint({_id: 1}).toArray());
+
+    // Leaves $set intact when shard merging on agg query.
+    const docWithY = ({_id: 5, x: 1, $set: {$inc: {x: 5}}, y: 1});
+    assert.eq([docWithY], coll.aggregate([{$addFields: {y: 1}}]).toArray());
+
+    // Leaves $set intact when shard merging on sort query.
+    assert.eq([document], coll.find().sort({_id: 1}).toArray());
+
+    // Leaves added fields intact when user adds $ prefixed field in an agg pipeline.
+    const docWithDollarY = ({_id: 5, x: 1, $set: {$inc: {x: 5}}, $y: 1});
+    assert.eq(
+        [docWithDollarY],
+        coll.aggregate(
+                [{$replaceWith: {$setField: {field: {$literal: "$y"}, input: "$$ROOT", value: 1}}}])
+            .toArray());
+}
+
+const st = new ShardingTest({shards: 2});
+try {
+    assert.commandWorked(st.s0.adminCommand({enableSharding: 'test'}));
+    st.ensurePrimaryShard('test', st.shard1.shardName);
+    assert.commandWorked(st.s0.adminCommand({shardCollection: 'test.coll', key: {x: 'hashed'}}));
+
+    runTest(st.getDB("test").coll);
+} finally {
+    st.stop();
+}
+})();
diff --git a/src/mongo/s/query/router_stage_remove_metadata_fields.cpp b/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
index 320c441ef46..7f3d89beef1 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
@@ -53,9 +53,14 @@ StatusWith<ClusterQueryResult> RouterStageRemoveMetadataFields::next() {
     }
 
     BSONObjIterator iterator(*childResult.getValue().getResult());
+
     // Find the first field that we need to remove.
-    while (iterator.more() && (*iterator).fieldName()[0] != '$') {
-        ++iterator;
+    for (; iterator.more(); ++iterator) {
+        // To save some time, we ensure that the current field name starts with a $
+        // before checking if it's actually a metadata field in the map.
+        if ((*iterator).fieldName()[0] == '$' && _metaFields.contains((*iterator).fieldName())) {
+            break;
+        }
     }
 
     if (!iterator.more()) {
diff --git a/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp b/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
index a18aa0cbb31..9fa199a2ac9 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields_test.cpp
@@ -194,6 +194,60 @@ TEST(RouterStageRemoveMetadataFieldsTest, ForwardsAwaitDataTimeout) {
     ASSERT_EQ(789, durationCount<Milliseconds>(awaitDataTimeout.getValue()));
 }
 
+// Grabs the next document from a stage and ensure it matches expectedDoc.
+// Assumes that remotes are exhausted after one document.
+void verifyNextDocument(RouterExecStage* stage, const BSONObj& expectedDoc) {
+    auto result = stage->next();
+    ASSERT_OK(result.getStatus());
+    ASSERT(result.getValue().getResult());
+    ASSERT_BSONOBJ_EQ(*result.getValue().getResult(), expectedDoc);
+    ASSERT_TRUE(stage->remotesExhausted());
+}
+
+TEST(RouterStageRemoveMetadataFieldsTest, AllowsNonMetaDataDollars) {
+    auto mockStage = std::make_unique<RouterStageMock>(opCtx);
+    mockStage->queueResult(BSON("$a" << 1 << "$sortKey" << 1 << "b" << 1));
+    mockStage->queueResult(BSON("a" << 2 << "$sortKey" << 1 << "$b" << 2));
+    mockStage->markRemotesExhausted();
+
+    auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
+        opCtx, std::move(mockStage), StringDataSet{"$sortKey"_sd});
+    ASSERT_TRUE(sortKeyStage->remotesExhausted());
+
+    verifyNextDocument(sortKeyStage.get(), BSON("$a" << 1 << "b" << 1));
+    verifyNextDocument(sortKeyStage.get(), BSON("a" << 2 << "$b" << 2));
+
+    auto endResult = sortKeyStage->next();
+    ASSERT_OK(endResult.getStatus());
+    ASSERT(endResult.getValue().isEOF());
+    ASSERT_TRUE(sortKeyStage->remotesExhausted());
+}
+
+// For every keyword, ensure it's removed if it's in the first, middle, or
+// last position, and that the remainder of the document is undisturbed.
+TEST(RouterStageRemoveMetadataFieldsTest, RemovesAllMetaDataDollars) {
+    for (auto& keyword : Document::allMetadataFieldNames) {
+        auto mockStage = std::make_unique<RouterStageMock>(opCtx);
+        mockStage->queueResult(BSON(keyword << 1 << "$a" << 1));
+        mockStage->queueResult(BSON("$a" << 1 << keyword << 1));
+        mockStage->queueResult(BSON("$a" << 1 << keyword << 1 << "$b" << 1));
+        mockStage->markRemotesExhausted();
+
+        auto sortKeyStage = std::make_unique<RouterStageRemoveMetadataFields>(
+            opCtx, std::move(mockStage), Document::allMetadataFieldNames);
+        ASSERT_TRUE(sortKeyStage->remotesExhausted());
+
+        verifyNextDocument(sortKeyStage.get(), BSON("$a" << 1));
+        verifyNextDocument(sortKeyStage.get(), BSON("$a" << 1));
+        verifyNextDocument(sortKeyStage.get(), BSON("$a" << 1 << "$b" << 1));
+
+        auto endResult = sortKeyStage->next();
+        ASSERT_OK(endResult.getStatus());
+        ASSERT(endResult.getValue().isEOF());
+        ASSERT_TRUE(sortKeyStage->remotesExhausted());
+    }
+}
+
 }  // namespace
 
 }  // namespace mongo