SERVER-67206 Generate tags from a FLE2 Range Payload

4 files changed, 39 insertions, 34 deletions

diff --git a/src/mongo/db/query/fle/encrypted_predicate.cpp b/src/mongo/db/query/fle/encrypted_predicate.cpp
index c23e1a2a5ae..aeed43cc884 100644
--- a/src/mongo/db/query/fle/encrypted_predicate.cpp
+++ b/src/mongo/db/query/fle/encrypted_predicate.cpp
@@ -46,6 +46,17 @@ void logTagsExceeded(const ExceptionFor<ErrorCodes::FLEMaxTagLimitExceeded>& ex)
         6672410, 2, "FLE Max tag limit hit during query rewrite", "__error__"_attr = ex.what());
 }
 
+std::unique_ptr<MatchExpression> makeTagDisjunction(BSONArray&& tagArray) {
+    auto tagElems = std::vector<BSONElement>();
+    tagArray.elems(tagElems);
+
+    auto newExpr = std::make_unique<InMatchExpression>(kSafeContent);
+    newExpr->setBackingBSON(std::move(tagArray));
+    uassertStatusOK(newExpr->setEqualities(std::move(tagElems)));
+
+    return newExpr;
+}
+
 BSONArray toBSONArray(std::vector<PrfBlock>&& vec) {
     auto bab = BSONArrayBuilder();
     for (auto& elt : vec) {
diff --git a/src/mongo/db/query/fle/encrypted_predicate.h b/src/mongo/db/query/fle/encrypted_predicate.h
index b2797fd1eb9..ac4d55b811a 100644
--- a/src/mongo/db/query/fle/encrypted_predicate.h
+++ b/src/mongo/db/query/fle/encrypted_predicate.h
@@ -82,6 +82,8 @@ BSONArray toBSONArray(std::vector<PrfBlock>&& vec);
  */
 std::vector<Value> toValues(std::vector<PrfBlock>&& vec);
 
+std::unique_ptr<MatchExpression> makeTagDisjunction(BSONArray&& tagArray);
+
 void logTagsExceeded(const ExceptionFor<ErrorCodes::FLEMaxTagLimitExceeded>& ex);
 /**
  * Interface for implementing a server rewrite for an encrypted index. Each type of predicate
diff --git a/src/mongo/db/query/fle/equality_predicate.cpp b/src/mongo/db/query/fle/equality_predicate.cpp
index fc46bcbfe3a..a349962145b 100644
--- a/src/mongo/db/query/fle/equality_predicate.cpp
+++ b/src/mongo/db/query/fle/equality_predicate.cpp
@@ -29,6 +29,7 @@
 
 #include "equality_predicate.h"
 
+#include "mongo/bson/bsontypes.h"
 #include "mongo/crypto/fle_crypto.h"
 #include "mongo/crypto/fle_tags.h"
 #include "mongo/db/matcher/expression_expr.h"
@@ -63,16 +64,7 @@ std::unique_ptr<MatchExpression> EqualityPredicate::rewriteToTagDisjunction(
             if (!isPayload(payload)) {
                 return nullptr;
             }
-            auto obj = toBSONArray(generateTags(payload));
-
-            auto tags = std::vector<BSONElement>();
-            obj.elems(tags);
-
-            auto inExpr = std::make_unique<InMatchExpression>(kSafeContent);
-            inExpr->setBackingBSON(std::move(obj));
-            auto status = inExpr->setEqualities(std::move(tags));
-            uassertStatusOK(status);
-            return inExpr;
+            return makeTagDisjunction(toBSONArray(generateTags(payload)));
         }
         case MatchExpression::MATCH_IN: {
             auto inExpr = static_cast<InMatchExpression*>(expr);
@@ -93,22 +85,14 @@ std::unique_ptr<MatchExpression> EqualityPredicate::rewriteToTagDisjunction(
             auto backingBSONBuilder = BSONArrayBuilder();
 
             for (auto& eq : inExpr->getEqualities()) {
-                auto obj = toBSONArray(generateTags(eq));
+                auto obj = generateTags(eq);
                 for (auto&& elt : obj) {
-                    backingBSONBuilder.append(elt);
+                    backingBSONBuilder.appendBinData(elt.size(), BinDataGeneral, elt.data());
                 }
             }
 
             auto backingBSON = backingBSONBuilder.arr();
-            auto allTags = std::vector<BSONElement>();
-            backingBSON.elems(allTags);
-
-            auto newExpr = std::make_unique<InMatchExpression>(kSafeContent);
-            newExpr->setBackingBSON(std::move(backingBSON));
-            auto status = newExpr->setEqualities(std::move(allTags));
-            uassertStatusOK(status);
-
-            return newExpr;
+            return makeTagDisjunction(std::move(backingBSON));
         }
         default:
             MONGO_UNREACHABLE_TASSERT(6911300);
diff --git a/src/mongo/db/query/fle/range_predicate.cpp b/src/mongo/db/query/fle/range_predicate.cpp
index 7c9be5c90fa..0ebf33790fc 100644
--- a/src/mongo/db/query/fle/range_predicate.cpp
+++ b/src/mongo/db/query/fle/range_predicate.cpp
@@ -29,7 +29,11 @@
 
 #include "range_predicate.h"
 
+#include <iterator>
+
 #include "mongo/crypto/encryption_fields_gen.h"
+#include "mongo/crypto/fle_crypto.h"
+#include "mongo/crypto/fle_tags.h"
 #include "mongo/db/matcher/expression_leaf.h"
 #include "mongo/db/query/fle/encrypted_predicate.h"
 
@@ -39,29 +43,33 @@ REGISTER_ENCRYPTED_MATCH_PREDICATE_REWRITE_WITH_FLAG(BETWEEN,
                                                      RangePredicate,
                                                      gFeatureFlagFLE2Range);
 
-// TODO: SERVER-67206 Generate tags for range payload.
 std::vector<PrfBlock> RangePredicate::generateTags(BSONValue payload) const {
-    return {};
+    auto parsedPayload = parseFindPayload<ParsedFindRangePayload>(payload);
+    std::vector<PrfBlock> tags;
+    for (auto& edge : parsedPayload.edges) {
+        auto tagsForEdge = readTags(*_rewriter->getEscReader(),
+                                    *_rewriter->getEccReader(),
+                                    edge.esc,
+                                    edge.ecc,
+                                    edge.edc,
+                                    parsedPayload.maxCounter);
+        tags.insert(tags.end(),
+                    std::make_move_iterator(tagsForEdge.begin()),
+                    std::make_move_iterator(tagsForEdge.end()));
+    }
+    return tags;
 }
 
 std::unique_ptr<MatchExpression> RangePredicate::rewriteToTagDisjunction(
     MatchExpression* expr) const {
     invariant(expr->matchType() == MatchExpression::BETWEEN);
     auto betExpr = static_cast<BetweenMatchExpression*>(expr);
-    auto ffp = betExpr->rhs();
+    auto payload = betExpr->rhs();
 
-    if (!isPayload(ffp)) {
+    if (!isPayload(payload)) {
         return nullptr;
     }
-
-    auto obj = toBSONArray(generateTags(ffp));
-    auto tags = std::vector<BSONElement>();
-    obj.elems(tags);
-    auto inExpr = std::make_unique<InMatchExpression>(kSafeContent);
-    inExpr->setBackingBSON(std::move(obj));
-    auto status = inExpr->setEqualities(std::move(tags));
-    uassertStatusOK(status);
-    return inExpr;
+    return makeTagDisjunction(toBSONArray(generateTags(payload)));
 }
 
 // TODO: SERVER-67209 Server-side rewrite for agg expressions with $between.