SERVER-67629 encryptedBetween agg operator

6 files changed, 254 insertions, 49 deletions

diff --git a/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp b/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
index 0dcf20fc16a..32665a44f78 100644
--- a/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
+++ b/src/mongo/db/pipeline/abt/agg_expression_visitor.cpp
@@ -331,6 +331,10 @@ public:
         unsupportedExpression(expr->getOpName());
     }
 
+    void visit(const ExpressionInternalFLEBetween* expr) override final {
+        unsupportedExpression(expr->getOpName());
+    }
+
     void visit(const ExpressionMap* expr) override final {
         unsupportedExpression("$map");
     }
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index 3325897f1aa..f9697e682dc 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -3894,17 +3894,8 @@ ExpressionInternalFLEEqual::ExpressionInternalFLEEqual(ExpressionContext* const
                                                        int64_t contentionFactor,
                                                        ConstDataRange edcToken)
     : Expression(expCtx, {std::move(field)}),
-      _serverToken(PrfBlockfromCDR(serverToken)),
-      _edcToken(PrfBlockfromCDR(edcToken)),
-      _contentionFactor(contentionFactor) {
+      _evaluator(serverToken, contentionFactor, {edcToken}) {
     expCtx->sbeCompatible = false;
-
-    auto tokens =
-        EDCServerCollection::generateEDCTokens(ConstDataRange(_edcToken), _contentionFactor);
-
-    for (auto& token : tokens) {
-        _cachedEDCTokens.insert(std::move(token.data));
-    }
 }
 
 void ExpressionInternalFLEEqual::_doAddDependencies(DepsTracker* deps) const {
@@ -3926,27 +3917,24 @@ intrusive_ptr<Expression> ExpressionInternalFLEEqual::parse(ExpressionContext* c
 
     auto serverTokenPair = fromEncryptedConstDataRange(fleEq.getServerEncryptionToken());
 
-    uassert(6672405,
+    uassert(6762901,
             "Invalid server token",
             serverTokenPair.first == EncryptedBinDataType::kFLE2TransientRaw &&
                 serverTokenPair.second.length() == sizeof(PrfBlock));
 
     auto edcTokenPair = fromEncryptedConstDataRange(fleEq.getEdcDerivedToken());
 
-    uassert(6672406,
+    uassert(6762902,
             "Invalid edc token",
             edcTokenPair.first == EncryptedBinDataType::kFLE2TransientRaw &&
                 edcTokenPair.second.length() == sizeof(PrfBlock));
 
 
     auto cf = fleEq.getMaxCounter();
-    uassert(6672408, "Contention factor must be between 0 and 10000", cf >= 0 && cf < 10000);
+    uassert(6762903, "Contention factor must be between 0 and 10000", cf >= 0 && cf < 10000);
 
-    return new ExpressionInternalFLEEqual(expCtx,
-                                          std::move(fieldExpr),
-                                          serverTokenPair.second,
-                                          fleEq.getMaxCounter(),
-                                          edcTokenPair.second);
+    return new ExpressionInternalFLEEqual(
+        expCtx, std::move(fieldExpr), serverTokenPair.second, cf, edcTokenPair.second);
 }
 
 Value toValue(const std::array<std::uint8_t, 32>& buf) {
@@ -3955,47 +3943,116 @@ Value toValue(const std::array<std::uint8_t, 32>& buf) {
 }
 
 Value ExpressionInternalFLEEqual::serialize(bool explain) const {
-    return Value(Document{{kInternalFleEq,
-                           Document{{"field", _children[0]->serialize(explain)},
-                                    {"edc", toValue(_edcToken)},
-                                    {"counter", Value(static_cast<long long>(_contentionFactor))},
-                                    {"server", toValue(_serverToken)}}}});
+    return Value(Document{
+        {kInternalFleEq,
+         Document{{"field", _children[0]->serialize(explain)},
+                  {"edc", toValue(_evaluator.edcTokens()[0])},
+                  {"counter", Value(static_cast<long long>(_evaluator.contentionFactor()))},
+                  {"server", toValue(_evaluator.serverToken())}}}});
 }
 
 Value ExpressionInternalFLEEqual::evaluate(const Document& root, Variables* variables) const {
-    // Inputs
-    // 1. Value for FLE2IndexedEqualityEncryptedValue field
-
-    Value fieldValue = _children[0]->evaluate(root, variables);
-
+    auto fieldValue = _children[0]->evaluate(root, variables);
     if (fieldValue.nullish()) {
         return Value(BSONNULL);
     }
+    return Value(_evaluator.evaluate<FLE2IndexedEqualityEncryptedValue>(
+        fieldValue,
+        EncryptedBinDataType::kFLE2EqualityIndexedValue,
+        [](auto token, auto serverValue) {
+            return EDCServerCollection::decryptAndParse(token, serverValue);
+        }));
+}
 
-    if (fieldValue.getType() != BinData) {
-        return Value(false);
+const char* ExpressionInternalFLEEqual::getOpName() const {
+    return kInternalFleEq.rawData();
+}
+
+/* ----------------------- ExpressionInternalFLEBetween ---------------------------- */
+
+constexpr auto kInternalFleBetween = "$_internalFleBetween"_sd;
+
+ExpressionInternalFLEBetween::ExpressionInternalFLEBetween(ExpressionContext* const expCtx,
+                                                           boost::intrusive_ptr<Expression> field,
+                                                           ConstDataRange serverToken,
+                                                           int64_t contentionFactor,
+                                                           std::vector<ConstDataRange> edcTokens)
+    : Expression(expCtx, {std::move(field)}), _evaluator(serverToken, contentionFactor, edcTokens) {
+    expCtx->sbeCompatible = false;
+}
+
+void ExpressionInternalFLEBetween::_doAddDependencies(DepsTracker* deps) const {
+    for (auto&& operand : _children) {
+        operand->addDependencies(deps);
     }
+}
 
-    auto fieldValuePair = fromEncryptedBinData(fieldValue);
+REGISTER_STABLE_EXPRESSION(_internalFleBetween, ExpressionInternalFLEBetween::parse);
 
-    uassert(6672407,
-            "Invalid encrypted indexed field",
-            fieldValuePair.first == EncryptedBinDataType::kFLE2EqualityIndexedValue);
+intrusive_ptr<Expression> ExpressionInternalFLEBetween::parse(ExpressionContext* const expCtx,
+                                                              BSONElement expr,
+                                                              const VariablesParseState& vps) {
+    IDLParserContext ctx(kInternalFleBetween);
+    auto fleBetween = InternalFleBetweenStruct::parse(ctx, expr.Obj());
 
-    // Value matches if
-    // 1. Decrypt field is successful
-    // 2. EDC_u Token is in GenTokens(EDC Token, ContentionFactor)
-    //
-    auto swIndexed =
-        EDCServerCollection::decryptAndParse(ConstDataRange(_serverToken), fieldValuePair.second);
-    uassertStatusOK(swIndexed);
-    auto indexed = swIndexed.getValue();
+    auto fieldExpr = Expression::parseOperand(expCtx, fleBetween.getField().getElement(), vps);
 
-    return Value(_cachedEDCTokens.count(indexed.edc.data) == 1);
+    auto serverTokenPair = fromEncryptedConstDataRange(fleBetween.getServerEncryptionToken());
+
+    uassert(6762904,
+            "Invalid server token",
+            serverTokenPair.first == EncryptedBinDataType::kFLE2TransientRaw &&
+                serverTokenPair.second.length() == sizeof(PrfBlock));
+
+    std::vector<ConstDataRange> edcTokens;
+    for (auto& elem : fleBetween.getEdcDerivedTokens()) {
+        auto [first, second] = fromEncryptedConstDataRange(elem);
+        uassert(6762905,
+                "Invalid edc token",
+                first == EncryptedBinDataType::kFLE2TransientRaw &&
+                    second.length() == sizeof(PrfBlock));
+        edcTokens.push_back(second);
+    }
+
+    auto cf = fleBetween.getMaxCounter();
+    uassert(6762906, "Contention factor must be between 0 and 10000", cf >= 0 && cf < 10000);
+
+    return new ExpressionInternalFLEBetween(
+        expCtx, std::move(fieldExpr), serverTokenPair.second, cf, edcTokens);
 }
 
-const char* ExpressionInternalFLEEqual::getOpName() const {
-    return kInternalFleEq.rawData();
+Value ExpressionInternalFLEBetween::serialize(bool explain) const {
+    std::vector<Value> edcValues;
+    edcValues.reserve(_evaluator.edcTokens().size());
+    for (auto& token : _evaluator.edcTokens()) {
+        edcValues.push_back(toValue(PrfBlockfromCDR(token)));
+    }
+    return Value(Document{
+        {kInternalFleBetween,
+         Document{{"field", _children[0]->serialize(explain)},
+                  {"edc", Value(edcValues)},
+                  {"counter", Value(static_cast<long long>(_evaluator.contentionFactor()))},
+                  {"server", toValue(_evaluator.serverToken())}}}});
+}
+
+Value ExpressionInternalFLEBetween::evaluate(const Document& root, Variables* variables) const {
+    // TODO(SERVER-67627): Uncomment for runtime tag matching.
+    // auto fieldValue = _children[0]->evaluate(root, variables);
+    // if (fieldValue.nullish()) {
+    //     return Value(BSONNULL);
+    // }
+    // return Value(_evaluator.evaluate<FLE2IndexedRangeEncryptedValue>(
+    //     fieldValue,
+    //     EncryptedBinDataType::kFLE2RangeIndexedValue,
+    //     [](auto token, auto serverValue) {
+    //         return EDCServerCollection::decryptAndParseRange(token, serverValue);
+    //     }));
+    uasserted(ErrorCodes::InternalErrorNotSupported,
+              str::stream() << "$_internalFleBetween not supported.");
+}
+
+const char* ExpressionInternalFLEBetween::getOpName() const {
+    return kInternalFleBetween.rawData();
 }
 
 /* ------------------------ ExpressionNary ----------------------------- */
diff --git a/src/mongo/db/pipeline/expression.h b/src/mongo/db/pipeline/expression.h
index fb2c91648e2..8746ca20c6e 100644
--- a/src/mongo/db/pipeline/expression.h
+++ b/src/mongo/db/pipeline/expression.h
@@ -41,6 +41,7 @@
 #include <vector>
 
 #include "mongo/base/init.h"
+#include "mongo/crypto/fle_crypto.h"
 #include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/exec/document_value/document.h"
 #include "mongo/db/exec/document_value/value.h"
@@ -2253,6 +2254,35 @@ public:
     }
 };
 
+class ExpressionInternalFLEBetween final : public Expression {
+public:
+    ExpressionInternalFLEBetween(ExpressionContext* expCtx,
+                                 boost::intrusive_ptr<Expression> field,
+                                 ConstDataRange serverToken,
+                                 int64_t contentionFactor,
+                                 std::vector<ConstDataRange> edcTokens);
+    Value serialize(bool explain) const final;
+
+    Value evaluate(const Document& root, Variables* variables) const final;
+    const char* getOpName() const;
+
+    static boost::intrusive_ptr<Expression> parse(ExpressionContext* expCtx,
+                                                  BSONElement expr,
+                                                  const VariablesParseState& vps);
+    void _doAddDependencies(DepsTracker* deps) const final;
+
+    void acceptVisitor(ExpressionMutableVisitor* visitor) final {
+        return visitor->visit(this);
+    }
+
+    void acceptVisitor(ExpressionConstVisitor* visitor) const final {
+        return visitor->visit(this);
+    }
+
+private:
+    EncryptedPredicateEvaluator _evaluator;
+};
+
 class ExpressionInternalFLEEqual final : public Expression {
 public:
     ExpressionInternalFLEEqual(ExpressionContext* expCtx,
@@ -2279,10 +2309,7 @@ public:
     }
 
 private:
-    std::array<std::uint8_t, 32> _serverToken;
-    std::array<std::uint8_t, 32> _edcToken;
-    int64_t _contentionFactor;
-    stdx::unordered_set<std::array<std::uint8_t, 32>> _cachedEDCTokens;
+    EncryptedPredicateEvaluator _evaluator;
 };
 
 class ExpressionMap final : public Expression {
diff --git a/src/mongo/db/pipeline/expression_parser.idl b/src/mongo/db/pipeline/expression_parser.idl
index 9f1cde70856..fc400bf04a3 100644
--- a/src/mongo/db/pipeline/expression_parser.idl
+++ b/src/mongo/db/pipeline/expression_parser.idl
@@ -54,4 +54,25 @@ structs:
             type: long
             cpp_name: maxCounter
 
+    InternalFleBetweenStruct:
+        description: "Struct for $_internalFleBetween"
+        strict: true
+        fields:
+          field:
+            description: "Expression"
+            type: IDLAnyType
+            cpp_name: field
+          edc:
+            description: "EDCDerivedFromDataTokens"
+            type: array<bindata_encrypt>
+            cpp_name: edcDerivedTokens
+          server:
+            description: "ServerDataEncryptionLevel1Token"
+            type: bindata_encrypt
+            cpp_name: serverEncryptionToken
+          counter:
+            description: "Queryable Encryption max counter"
+            type: long
+            cpp_name: maxCounter
+
 
diff --git a/src/mongo/db/pipeline/expression_test.cpp b/src/mongo/db/pipeline/expression_test.cpp
index 184a5e04bd3..36793467420 100644
--- a/src/mongo/db/pipeline/expression_test.cpp
+++ b/src/mongo/db/pipeline/expression_test.cpp
@@ -4162,5 +4162,98 @@ TEST(ExpressionFLETest, TestBinData_RoundTrip) {
     ASSERT_BSONOBJ_EQ(value.getDocument().toBson(), roundTripExpr);
 }
 
+TEST(ExpressionFLETest, ParseAndSerializeBetween) {
+    auto expCtx = ExpressionContextForTest();
+    auto vps = expCtx.variablesParseState;
+
+    auto expr = fromjson(R"({$_internalFleBetween: {
+    field: {
+        "$binary": {
+            "base64":
+            "BxI0VngSNJh2EjQSNFZ4kBIQ0JE8aMUFkPk5sSTVqfdNNfjqUfQQ1Uoj0BBcthrWoe9wyU3cN6zmWaQBPJ97t0ZPbecnMsU736yXre6cBO4Zdt/wThtY+v5+7vFgNnWpgRP0e+vam6QPmLvbBrO0LdsvAPTGW4yqwnzCIXCoEg7QPGfbfAXKPDTNenBfRlawiblmTOhO/6ljKotWsMp22q/rpHrn9IEIeJmecwuuPIJ7EA+XYQ3hOKVccYf2ogoK73+8xD/Vul83Qvr84Q8afc4QUMVs8A==",
+                "subType": "6"
+        }
+    },
+    server: {
+        "$binary": {
+            "base64": "COuac/eRLYakKX6B0vZ1r3QodOQFfjqJD+xlGiPu4/Ps",
+            "subType": "6"
+        }
+    },
+    counter: {
+        "$numberLong": "3"
+    },
+    edc: [{
+        "$binary": {
+            "base64": "CEWSmQID7SfwyAUI3ZkSFkATKryDQfnxXEOGad5d4Rsg",
+            "subType": "6"
+        }
+    }]   } })");
+
+    auto exprFle = ExpressionInternalFLEBetween::parse(&expCtx, expr.firstElement(), vps);
+    auto value = exprFle->serialize(false);
+
+    auto roundTripExpr = fromjson(R"({$_internalFleBetween: {
+    field: {
+        "$const" : { "$binary": {
+            "base64":
+            "BxI0VngSNJh2EjQSNFZ4kBIQ0JE8aMUFkPk5sSTVqfdNNfjqUfQQ1Uoj0BBcthrWoe9wyU3cN6zmWaQBPJ97t0ZPbecnMsU736yXre6cBO4Zdt/wThtY+v5+7vFgNnWpgRP0e+vam6QPmLvbBrO0LdsvAPTGW4yqwnzCIXCoEg7QPGfbfAXKPDTNenBfRlawiblmTOhO/6ljKotWsMp22q/rpHrn9IEIeJmecwuuPIJ7EA+XYQ3hOKVccYf2ogoK73+8xD/Vul83Qvr84Q8afc4QUMVs8A==",
+                "subType": "6"
+        }}
+    },
+    edc: [{
+        "$binary": {
+            "base64": "CEWSmQID7SfwyAUI3ZkSFkATKryDQfnxXEOGad5d4Rsg",
+            "subType": "6"
+        }
+    }],
+    counter: {
+        "$numberLong": "3"
+    },
+    server: {
+        "$binary": {
+            "base64": "COuac/eRLYakKX6B0vZ1r3QodOQFfjqJD+xlGiPu4/Ps",
+            "subType": "6"
+        }
+    }
+        } })");
+    ASSERT_BSONOBJ_EQ(value.getDocument().toBson(), roundTripExpr);
+}
+
+TEST(ExpressionFLETest, EvalBetween) {
+    auto expCtx = ExpressionContextForTest();
+    auto vps = expCtx.variablesParseState;
+
+    auto expr = fromjson(R"({$_internalFleBetween: {
+    field: {
+        "$binary": {
+            "base64":
+            "BxI0VngSNJh2EjQSNFZ4kBIQ0JE8aMUFkPk5sSTVqfdNNfjqUfQQ1Uoj0BBcthrWoe9wyU3cN6zmWaQBPJ97t0ZPbecnMsU736yXre6cBO4Zdt/wThtY+v5+7vFgNnWpgRP0e+vam6QPmLvbBrO0LdsvAPTGW4yqwnzCIXCoEg7QPGfbfAXKPDTNenBfRlawiblmTOhO/6ljKotWsMp22q/rpHrn9IEIeJmecwuuPIJ7EA+XYQ3hOKVccYf2ogoK73+8xD/Vul83Qvr84Q8afc4QUMVs8A==",
+                "subType": "6"
+        }
+    },
+    server: {
+        "$binary": {
+            "base64": "COuac/eRLYakKX6B0vZ1r3QodOQFfjqJD+xlGiPu4/Ps",
+            "subType": "6"
+        }
+    },
+    counter: {
+        "$numberLong": "3"
+    },
+    edc: [{
+        "$binary": {
+            "base64": "CEWSmQID7SfwyAUI3ZkSFkATKryDQfnxXEOGad5d4Rsg",
+            "subType": "6"
+        }
+    }]   } })");
+
+    auto exprFle = ExpressionInternalFLEBetween::parse(&expCtx, expr.firstElement(), vps);
+    // TODO(SERVER-67627): Remove assertion when runtime tag matching is implemented.
+    ASSERT_THROWS_CODE(exprFle->evaluate({}, &expCtx.variables),
+                       AssertionException,
+                       ErrorCodes::InternalErrorNotSupported);
+}
+
 }  // namespace ExpressionTests
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/expression_visitor.h b/src/mongo/db/pipeline/expression_visitor.h
index 6b7c4fc4cdd..47d0eb9a18e 100644
--- a/src/mongo/db/pipeline/expression_visitor.h
+++ b/src/mongo/db/pipeline/expression_visitor.h
@@ -153,6 +153,7 @@ class ExpressionHyperbolicSine;
 class ExpressionInternalFindSlice;
 class ExpressionInternalFindPositional;
 class ExpressionInternalFindElemMatch;
+class ExpressionInternalFLEBetween;
 class ExpressionInternalFLEEqual;
 class ExpressionInternalJsEmit;
 class ExpressionFunction;
@@ -246,6 +247,7 @@ public:
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionLn>) = 0;
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionLog>) = 0;
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionLog10>) = 0;
+    virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionInternalFLEBetween>) = 0;
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionInternalFLEEqual>) = 0;
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionMap>) = 0;
     virtual void visit(expression_walker::MaybeConstPtr<IsConst, ExpressionMeta>) = 0;
@@ -426,6 +428,7 @@ struct SelectiveConstExpressionVisitorBase : public ExpressionConstVisitor {
     void visit(const ExpressionLn*) override {}
     void visit(const ExpressionLog*) override {}
     void visit(const ExpressionLog10*) override {}
+    void visit(const ExpressionInternalFLEBetween*) override {}
     void visit(const ExpressionInternalFLEEqual*) override {}
     void visit(const ExpressionMap*) override {}
     void visit(const ExpressionMeta*) override {}