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

(cherry picked from commit e4ffe4b63da0ef16e129ac693dda79b64047a469)
(cherry picked from commit a7ac8403a5ea9a84b8ba218d7c63396963256c12)

2 files changed, 61 insertions, 2 deletions

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 5ce2b8e677e..4b171b93711 100644
--- a/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
+++ b/src/mongo/s/query/router_stage_remove_metadata_fields.cpp
@@ -54,9 +54,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 1558438600d..c349c9a3c28 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(RouterExecStage::ExecContext::kInitialFind);
+    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(RouterExecStage::ExecContext::kInitialFind);
+    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(RouterExecStage::ExecContext::kInitialFind);
+        ASSERT_OK(endResult.getStatus());
+        ASSERT(endResult.getValue().isEOF());
+        ASSERT_TRUE(sortKeyStage->remotesExhausted());
+    }
+}
+
 }  // namespace
 
 }  // namespace mongo