SERVER-59869 Robustify sharded distinct command response handling

2 files changed, 32 insertions, 31 deletions

diff --git a/src/mongo/s/cluster_commands_helpers.cpp b/src/mongo/s/cluster_commands_helpers.cpp
index 8dfb3f20b4c..577e3ab2a3e 100644
--- a/src/mongo/s/cluster_commands_helpers.cpp
+++ b/src/mongo/s/cluster_commands_helpers.cpp
@@ -657,31 +657,6 @@ RawResponsesResult appendRawResponses(
     return {false, shardsWithSuccessResponses, successARSResponses, firstStaleConfigErrorReceived};
 }
 
-BSONObj extractQuery(const BSONObj& cmdObj) {
-    auto queryElem = cmdObj["query"];
-    if (!queryElem)
-        queryElem = cmdObj["q"];
-
-    if (!queryElem || queryElem.isNull())
-        return BSONObj();
-
-    uassert(ErrorCodes::TypeMismatch,
-            str::stream() << "Illegal type specified for query " << queryElem,
-            queryElem.type() == Object);
-    return queryElem.embeddedObject();
-}
-
-BSONObj extractCollation(const BSONObj& cmdObj) {
-    const auto collationElem = cmdObj["collation"];
-    if (!collationElem)
-        return BSONObj();
-
-    uassert(ErrorCodes::TypeMismatch,
-            str::stream() << "Illegal type specified for collation " << collationElem,
-            collationElem.type() == Object);
-    return collationElem.embeddedObject();
-}
-
 bool appendEmptyResultSet(OperationContext* opCtx,
                           BSONObjBuilder& result,
                           Status status,
diff --git a/src/mongo/s/commands/cluster_distinct_cmd.cpp b/src/mongo/s/commands/cluster_distinct_cmd.cpp
index 31741e57bba..68a4d889144 100644
--- a/src/mongo/s/commands/cluster_distinct_cmd.cpp
+++ b/src/mongo/s/commands/cluster_distinct_cmd.cpp
@@ -99,8 +99,20 @@ public:
         const BSONObj& cmdObj = opMsgRequest.body;
         const NamespaceString nss(parseNs(dbname, cmdObj));
 
-        auto targetingQuery = extractQuery(cmdObj);
-        auto targetingCollation = extractCollation(cmdObj);
+        auto parsedDistinctCmd =
+            ParsedDistinct::parse(opCtx, nss, cmdObj, ExtensionsCallbackNoop(), true);
+        uassertStatusOK(parsedDistinctCmd.getStatus());
+
+        auto distinctCanonicalQuery = parsedDistinctCmd.getValue().releaseQuery();
+        auto targetingQuery = distinctCanonicalQuery->getQueryObj();
+        auto targetingCollation = distinctCanonicalQuery->getFindCommandRequest().getCollation();
+
+        // Construct collator for deduping.
+        std::unique_ptr<CollatorInterface> collator;
+        if (!targetingCollation.isEmpty()) {
+            collator = uassertStatusOK(CollatorFactoryInterface::get(opCtx->getServiceContext())
+                                           ->makeFromBSON(targetingCollation));
+        }
 
         const auto explainCmd = ClusterExplain::wrapAsExplain(cmdObj, verbosity);
 
@@ -169,8 +181,13 @@ public:
         CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
         const NamespaceString nss(parseNs(dbName, cmdObj));
 
-        auto query = extractQuery(cmdObj);
-        auto collation = extractCollation(cmdObj);
+        auto parsedDistinctCmd =
+            ParsedDistinct::parse(opCtx, nss, cmdObj, ExtensionsCallbackNoop(), false);
+        uassertStatusOK(parsedDistinctCmd.getStatus());
+
+        auto distinctCanonicalQuery = parsedDistinctCmd.getValue().releaseQuery();
+        auto query = distinctCanonicalQuery->getQueryObj();
+        auto collation = distinctCanonicalQuery->getFindCommandRequest().getCollation();
 
         // Construct collator for deduping.
         std::unique_ptr<CollatorInterface> collator;
@@ -247,8 +264,17 @@ public:
                 : response.swResponse.getStatus();
             uassertStatusOK(status);
 
-            BSONObj res = std::move(response.swResponse.getValue().data);
-            BSONObjIterator it(res["values"].embeddedObject());
+            BSONObj res = response.swResponse.getValue().data;
+            auto values = res["values"];
+            uassert(5986900,
+                    str::stream() << "No 'values' field in distinct command response: "
+                                  << res.toString() << ". Original command: " << cmdObj.toString(),
+                    !values.eoo());
+            uassert(5986901,
+                    str::stream() << "Expected 'values' field to be of type Array, but found "
+                                  << typeName(values.type()),
+                    values.type() == BSONType::Array);
+            BSONObjIterator it(values.embeddedObject());
             while (it.more()) {
                 BSONElement nxt = it.next();
                 BSONObjBuilder temp(32);