SERVER-43840 fix invariant failure in exclusion projection parser

3 files changed, 27 insertions, 12 deletions

diff --git a/jstests/core/dbref3.js b/jstests/core/dbref3.js
index 5bf0470442d..6a91a56da2a 100644
--- a/jstests/core/dbref3.js
+++ b/jstests/core/dbref3.js
@@ -42,4 +42,18 @@ assert.eq(1, distinctDBs.length);
 t.insert({sub: {$ref: "foo", $id: [{x: 1, y: 1}, {x: 2, y: 2}, {x: 3, y: 3}]}});
 var k = t.findOne({'sub.$id': {$elemMatch: {x: 2}}}, {_id: 0, 'sub.$id.$': 1});
 print('k = ' + tojson(k));
-assert.eq({sub: {$id: [{x: 2, y: 2}]}}, k);
\ No newline at end of file
+assert.eq({sub: {$id: [{x: 2, y: 2}]}}, k);
+
+// Check that DBRef fields can be excluded
+assert.commandWorked(
+    t.insert({_id: 0, shouldExclude: 1, sub: {$ref: "foo", $id: 10, $db: "someDb"}}));
+assert.eq(t.find({shouldExclude: 1}, {"sub.$ref": 0}).toArray(),
+          [{_id: 0, shouldExclude: 1, sub: {$id: 10, $db: "someDb"}}]);
+assert.eq(t.find({shouldExclude: 1}, {"sub.$id": 0}).toArray(),
+          [{_id: 0, shouldExclude: 1, sub: {$ref: "foo", $db: "someDb"}}]);
+assert.eq(t.find({shouldExclude: 1}, {"sub.$id": 0, "sub.$ref": 0, "sub.$db": 0}).toArray(),
+          [{_id: 0, shouldExclude: 1, sub: {}}]);
+
+// It should be legal to exclude a DBRef field anywhere, even at the top layer.
+assert.eq(t.aggregate([{$match: {shouldExclude: 1}}, {$project: {"$id": 0, sub: 0}}]).toArray(),
+          [{_id: 0, shouldExclude: 1}]);
diff --git a/src/mongo/db/pipeline/parsed_exclusion_projection.cpp b/src/mongo/db/pipeline/parsed_exclusion_projection.cpp
index e91bff4884c..de31a6cb7e4 100644
--- a/src/mongo/db/pipeline/parsed_exclusion_projection.cpp
+++ b/src/mongo/db/pipeline/parsed_exclusion_projection.cpp
@@ -56,9 +56,6 @@ void ParsedExclusionProjection::parse(const BSONObj& spec, ExclusionNode* node,
     for (auto elem : spec) {
         const auto fieldName = elem.fieldNameStringData();
 
-        // A $ should have been detected by ParsedAggregationProjection before we get here.
-        invariant(fieldName[0] != '$');
-
         // Track whether the projection spec specifies a desired behavior for the _id field.
         idSpecified = idSpecified || fieldName == "_id"_sd || fieldName.startsWith("_id."_sd);
 
diff --git a/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp b/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp
index cc7271be016..485691472a5 100644
--- a/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp
+++ b/src/mongo/db/pipeline/parsed_exclusion_projection_test.cpp
@@ -84,14 +84,6 @@ ParsedExclusionProjection makeExclusionProjectionWithNoArrayRecursion() {
 //
 
 DEATH_TEST(ExclusionProjectionExecutionTest,
-           ShouldRejectComputedField,
-           "Invariant failure fieldName[0] != '$'") {
-    // Top-level expression.
-    auto exclusion = makeExclusionProjectionWithDefaultPolicies();
-    exclusion.parse(BSON("a" << false << "b" << BSON("$literal" << 1)));
-}
-
-DEATH_TEST(ExclusionProjectionExecutionTest,
            ShouldFailWhenGivenIncludedNonIdField,
            "Invariant failure !elem.trueValue() || elem.fieldNameStringData() == \"_id\"_sd") {
     auto exclusion = makeExclusionProjectionWithDefaultPolicies();
@@ -417,6 +409,18 @@ TEST(ExclusionProjectionExecutionTest, ShouldAllowExclusionOfIdSubfieldWithDefau
     ASSERT_DOCUMENT_EQ(result, expectedResult);
 }
 
+TEST(ExclusionProjectionExecutionTest, ShouldAllowLimitedDollarPrefixedFields) {
+    auto exclusion = makeExclusionProjectionWithDefaultIdExclusion();
+    exclusion.parse(
+        BSON("$id" << false << "$db" << false << "$ref" << false << "$sortKey" << false));
+
+    auto result = exclusion.applyProjection(
+        Document{{"$id", 5}, {"$db", 3}, {"$ref", 4}, {"$sortKey", 5}, {"someField", 6}});
+    auto expectedResult = Document{{"someField", 6}};
+
+    ASSERT_DOCUMENT_EQ(result, expectedResult);
+}
+
 //
 // Nested array recursion.
 //