SERVER-34644 Only explicitly do $-prefix check when writing sharded metadata

1 files changed, 22 insertions, 8 deletions

diff --git a/src/mongo/s/shard_key_pattern.cpp b/src/mongo/s/shard_key_pattern.cpp
index b2ce514472e..d37df0cceb2 100644
--- a/src/mongo/s/shard_key_pattern.cpp
+++ b/src/mongo/s/shard_key_pattern.cpp
@@ -102,12 +102,15 @@ std::vector<std::unique_ptr<FieldRef>> parseShardKeyPattern(const BSONObj& keyPa
     return parsedPaths;
 }
 
-bool isShardKeyElement(const BSONElement& element, bool allowRegex) {
-    if (element.eoo() || element.type() == Array)
+bool isValidShardKeyElement(const BSONElement& element) {
+    return !element.eoo() && element.type() != Array;
+}
+
+bool isValidShardKeyElementForStorage(const BSONElement& element) {
+    if (!isValidShardKeyElement(element))
         return false;
 
-    // TODO: Disallow regex all the time
-    if (!allowRegex && element.type() == RegEx)
+    if (element.type() == RegEx)
         return false;
 
     if (element.type() == Object && !element.embeddedObject().storageValidEmbedded().isOK())
@@ -166,6 +169,17 @@ Status ShardKeyPattern::checkShardKeySize(const BSONObj& shardKey) {
                           << " bytes"};
 }
 
+Status ShardKeyPattern::checkShardKeyIsValidForMetadataStorage(const BSONObj& shardKey) {
+    for (const auto& elem : shardKey) {
+        if (!isValidShardKeyElementForStorage(elem)) {
+            return {ErrorCodes::BadValue,
+                    str::stream() << "Shard key element " << elem << " is not valid for storage"};
+        }
+    }
+
+    return Status::OK();
+}
+
 ShardKeyPattern::ShardKeyPattern(const BSONObj& keyPattern)
     : _keyPattern(keyPattern),
       _keyPatternPaths(parseShardKeyPattern(keyPattern)),
@@ -200,7 +214,7 @@ bool ShardKeyPattern::isShardKey(const BSONObj& shardKey) const {
     for (const auto& patternEl : keyPatternBSON) {
         BSONElement keyEl = shardKey[patternEl.fieldNameStringData()];
 
-        if (!isShardKeyElement(keyEl, true))
+        if (!isValidShardKeyElement(keyEl))
             return false;
     }
 
@@ -219,7 +233,7 @@ BSONObj ShardKeyPattern::normalizeShardKey(const BSONObj& shardKey) const {
 
         BSONElement keyEl = shardKey[patternEl.fieldNameStringData()];
 
-        if (!isShardKeyElement(keyEl, true))
+        if (!isValidShardKeyElement(keyEl))
             return BSONObj();
 
         keyBuilder.appendAs(keyEl, patternEl.fieldName());
@@ -238,7 +252,7 @@ BSONObj ShardKeyPattern::extractShardKeyFromMatchable(const MatchableDocument& m
         BSONElement matchEl =
             extractKeyElementFromMatchable(matchable, patternEl.fieldNameStringData());
 
-        if (!isShardKeyElement(matchEl, true))
+        if (!isValidShardKeyElement(matchEl))
             return BSONObj();
 
         if (isHashedPatternEl(patternEl)) {
@@ -304,7 +318,7 @@ BSONObj ShardKeyPattern::extractShardKeyFromQuery(const CanonicalQuery& query) c
         const FieldRef& patternPath = **it;
         BSONElement equalEl = findEqualityElement(equalities, patternPath);
 
-        if (!isShardKeyElement(equalEl, false))
+        if (!isValidShardKeyElementForStorage(equalEl))
             return BSONObj();
 
         if (isHashedPattern()) {