SERVER-72926 Implement v2 changes to Queryable Encryption finds

1 files changed, 165 insertions, 7 deletions

diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index 03b8f260354..64d3cede3ee 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -745,12 +745,18 @@ protected:
         FleCrudTest::tearDown();
     }
     std::vector<PrfBlock> readTagsWithContention(BSONObj obj, uint64_t contention = 0) {
+        if (!gFeatureFlagFLE2ProtocolVersion2.isEnabledAndIgnoreFCV()) {
+            auto s = getTestESCDataToken(obj);
+            auto c = getTestECCDataToken(obj);
+            auto d = getTestEDCDataToken(obj);
+            auto esc = CollectionReader("test.esc", *_queryImpl);
+            auto ecc = CollectionReader("test.ecc", *_queryImpl);
+            return mongo::fle::readTagsWithContention(esc, ecc, s, c, d, contention, 100, {});
+        }
         auto s = getTestESCDataToken(obj);
-        auto c = getTestECCDataToken(obj);
         auto d = getTestEDCDataToken(obj);
         auto esc = CollectionReader("test.esc", *_queryImpl);
-        auto ecc = CollectionReader("test.ecc", *_queryImpl);
-        return mongo::fle::readTagsWithContention(esc, ecc, s, c, d, contention, 100, {});
+        return mongo::fle::readTagsWithContentionV2(esc, s, d, contention, 100, {});
     }
     std::vector<PrfBlock> readTags(BSONObj obj, uint64_t cm = 0) {
         auto s = getTestESCDataToken(obj);
@@ -777,6 +783,7 @@ TEST_F(FleCrudTest, InsertOne) {
                      .isEmpty());
 }
 
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, InsertOneV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
     auto doc = BSON("encrypted"
@@ -793,6 +800,7 @@ TEST_F(FleCrudTest, InsertOneV2) {
             .isEmpty());
 }
 
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, InsertOneRangeV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
     auto doc = BSON("encrypted" << 5);
@@ -822,6 +830,7 @@ TEST_F(FleCrudTest, InsertTwoSame) {
 }
 
 // Insert two documents with same values
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, InsertTwoSameV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
     auto doc = BSON("encrypted"
@@ -867,6 +876,7 @@ TEST_F(FleCrudTest, InsertTwoDifferent) {
                      .isEmpty());
 }
 
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, InsertTwoDifferentV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
     doSingleInsert(1,
@@ -922,6 +932,7 @@ TEST_F(FleCrudTest, Insert100Fields) {
 }
 
 // Insert 1 document with 100 fields
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, Insert100FieldsV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
 
@@ -985,6 +996,7 @@ TEST_F(FleCrudTest, Insert20Fields50Rows) {
 }
 
 // Insert 100 documents each with 20 fields with 7 distinct values per field
+// TODO: SERVER-73303 delete when v2 is enabled by default
 TEST_F(FleCrudTest, Insert20Fields50RowsV2) {
     RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
     uint64_t fieldCount = 20;
@@ -1556,6 +1568,15 @@ TEST_F(FleTagsTest, InsertOne) {
     ASSERT_EQ(1, readTags(doc).size());
 }
 
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertOneV2) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc = BSON("encrypted"
+                    << "a");
+    doSingleInsert(1, doc);
+    ASSERT_EQ(1, readTags(doc).size());
+}
+
 TEST_F(FleTagsTest, InsertTwoSame) {
     auto doc = BSON("encrypted"
                     << "a");
@@ -1566,6 +1587,16 @@ TEST_F(FleTagsTest, InsertTwoSame) {
     ASSERT_EQ(2, readTags(doc).size());
 }
 
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertTwoSameV2) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc = BSON("encrypted"
+                    << "a");
+    doSingleInsert(1, doc);
+    doSingleInsert(2, doc);
+    ASSERT_EQ(2, readTags(doc).size());
+}
+
 TEST_F(FleTagsTest, InsertTwoDifferent) {
     auto doc1 = BSON("encrypted"
                      << "a");
@@ -1579,7 +1610,23 @@ TEST_F(FleTagsTest, InsertTwoDifferent) {
     ASSERT_EQ(1, readTags(doc2).size());
 }
 
-TEST_F(FleTagsTest, InsertAndDeleteOne) {
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertTwoDifferentV2) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    doSingleInsert(1, doc1);
+    doSingleInsert(2, doc2);
+
+    ASSERT_EQ(1, readTags(doc1).size());
+    ASSERT_EQ(1, readTags(doc2).size());
+}
+
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertAndDeleteOneV1) {
     auto doc = BSON("encrypted"
                     << "a");
 
@@ -1589,7 +1636,19 @@ TEST_F(FleTagsTest, InsertAndDeleteOne) {
     ASSERT_EQ(0, readTags(doc).size());
 }
 
-TEST_F(FleTagsTest, InsertTwoSameAndDeleteOne) {
+TEST_F(FleTagsTest, InsertAndDeleteOne) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+    doSingleDelete(1, Fle2AlgorithmInt::kEquality);
+
+    ASSERT_EQ(1, readTags(doc).size());
+}
+
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertTwoSameAndDeleteOneV1) {
     auto doc = BSON("encrypted"
                     << "a");
 
@@ -1600,7 +1659,20 @@ TEST_F(FleTagsTest, InsertTwoSameAndDeleteOne) {
     ASSERT_EQ(1, readTags(doc).size());
 }
 
-TEST_F(FleTagsTest, InsertTwoDifferentAndDeleteOne) {
+TEST_F(FleTagsTest, InsertTwoSameAndDeleteOne) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    doSingleInsert(1, doc);
+    doSingleInsert(2, doc);
+    doSingleDelete(2, Fle2AlgorithmInt::kEquality);
+
+    ASSERT_EQ(2, readTags(doc).size());
+}
+
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, InsertTwoDifferentAndDeleteOneV1) {
     auto doc1 = BSON("encrypted"
                      << "a");
     auto doc2 = BSON("encrypted"
@@ -1614,6 +1686,21 @@ TEST_F(FleTagsTest, InsertTwoDifferentAndDeleteOne) {
     ASSERT_EQ(1, readTags(doc2).size());
 }
 
+TEST_F(FleTagsTest, InsertTwoDifferentAndDeleteOne) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    doSingleInsert(1, doc1);
+    doSingleInsert(2, doc2);
+    doSingleDelete(1, Fle2AlgorithmInt::kEquality);
+
+    ASSERT_EQ(1, readTags(doc1).size());
+    ASSERT_EQ(1, readTags(doc2).size());
+}
+
 TEST_F(FleTagsTest, InsertAndUpdate) {
     auto doc1 = BSON("encrypted"
                      << "a");
@@ -1666,7 +1753,49 @@ TEST_F(FleTagsTest, ContentionFactor) {
     ASSERT_EQ(2, readTags(doc2, 4).size());
 }
 
-TEST_F(FleTagsTest, MemoryLimit) {
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, ContentionFactorV2) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto efc = EncryptedFieldConfig::parse(IDLParserContext("root"), fromjson(R"({
+        "escCollection": "esc",
+        "eccCollection": "ecc",
+        "ecocCollection": "ecoc",
+        "fields": [{
+            "keyId": { "$uuid": "12345678-1234-9876-1234-123456789012"},
+            "path": "encrypted",
+            "bsonType": "string",
+            "queries": {"queryType": "equality", "contention": NumberLong(4)}
+        }]
+    })"));
+
+    auto doc1 = BSON("encrypted"
+                     << "a");
+    auto doc2 = BSON("encrypted"
+                     << "b");
+
+    // Insert doc1 twice with a contention factor of 0 and once with a contention factor or 3.
+    doSingleInsertWithContention(1, doc1, 4, 0, efc);
+    doSingleInsertWithContention(4, doc1, 4, 3, efc);
+    doSingleInsertWithContention(5, doc1, 4, 0, efc);
+
+    // Insert doc2 once with a contention factor of 2 and once with a contention factor of 3.
+    doSingleInsertWithContention(7, doc2, 4, 2, efc);
+    doSingleInsertWithContention(8, doc2, 4, 3, efc);
+
+    ASSERT_EQ(2, readTagsWithContention(doc1, 0).size());
+    ASSERT_EQ(0, readTagsWithContention(doc2, 0).size());
+    ASSERT_EQ(0, readTagsWithContention(doc1, 1).size());
+    ASSERT_EQ(0, readTagsWithContention(doc2, 1).size());
+    ASSERT_EQ(0, readTagsWithContention(doc1, 2).size());
+    ASSERT_EQ(1, readTagsWithContention(doc2, 2).size());
+    ASSERT_EQ(1, readTagsWithContention(doc1, 3).size());
+    ASSERT_EQ(1, readTagsWithContention(doc2, 3).size());
+    ASSERT_EQ(3, readTags(doc1, 4).size());
+    ASSERT_EQ(2, readTags(doc2, 4).size());
+}
+
+// TODO: SERVER-73303 delete when v2 is enabled by default
+TEST_F(FleTagsTest, MemoryLimitV1) {
     auto doc = BSON("encrypted"
                     << "a");
 
@@ -1694,6 +1823,35 @@ TEST_F(FleTagsTest, MemoryLimit) {
     ASSERT_EQ(tagLimit, readTags(doc).size());
 }
 
+TEST_F(FleTagsTest, MemoryLimit) {
+    RAIIServerParameterControllerForTest controller("featureFlagFLE2ProtocolVersion2", true);
+    auto doc = BSON("encrypted"
+                    << "a");
+
+    const auto tagLimit = 10;
+
+    // Set memory limit to 10 tags * 40 bytes per tag
+    internalQueryFLERewriteMemoryLimit.store(tagLimit * 40);
+
+    // Do 10 inserts
+    for (auto i = 0; i < tagLimit; i++) {
+        doSingleInsert(i, doc);
+    }
+
+    // readTags returns 10 tags which does not exceed memory limit.
+    ASSERT_EQ(tagLimit, readTags(doc).size());
+
+    doSingleInsert(10, doc);
+
+    // readTags returns 11 tags which does exceed memory limit.
+    ASSERT_THROWS_CODE(readTags(doc), DBException, ErrorCodes::FLEMaxTagLimitExceeded);
+
+    doSingleDelete(5, Fle2AlgorithmInt::kEquality);
+
+    // readTags returns 11 tags which does exceed memory limit.
+    ASSERT_THROWS_CODE(readTags(doc), DBException, ErrorCodes::FLEMaxTagLimitExceeded);
+}
+
 TEST_F(FleTagsTest, SampleMemoryLimit) {
 
     struct S {