SERVER-64749 Test tags reader with non-trivial contention factors

5 files changed, 213 insertions, 96 deletions

diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 9303bcbb61f..de40b6d736f 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -636,7 +636,7 @@ public:
     static FLE2InsertUpdatePayload serializeInsertUpdatePayload(FLEIndexKeyAndId indexKey,
                                                                 FLEUserKeyAndId userKey,
                                                                 BSONElement element,
-                                                                uint64_t maxContentionFactor);
+                                                                uint64_t contentionFactor);
 };
 
 
@@ -644,12 +644,10 @@ FLE2InsertUpdatePayload EDCClientPayload::parseInsertUpdatePayload(ConstDataRang
     return parseFromCDR<FLE2InsertUpdatePayload>(cdr);
 }
 
-
-FLE2InsertUpdatePayload EDCClientPayload::serializeInsertUpdatePayload(
-    FLEIndexKeyAndId indexKey,
-    FLEUserKeyAndId userKey,
-    BSONElement element,
-    uint64_t maxContentionFactor) {
+FLE2InsertUpdatePayload EDCClientPayload::serializeInsertUpdatePayload(FLEIndexKeyAndId indexKey,
+                                                                       FLEUserKeyAndId userKey,
+                                                                       BSONElement element,
+                                                                       uint64_t contentionFactor) {
     auto value = ConstDataRange(element.value(), element.value() + element.valuesize());
 
     auto collectionToken = FLELevel1TokenGenerator::generateCollectionsLevel1Token(indexKey.key);
@@ -667,12 +665,6 @@ FLE2InsertUpdatePayload EDCClientPayload::serializeInsertUpdatePayload(
     ECCDerivedFromDataToken eccDatakey =
         FLEDerivedFromDataTokenGenerator::generateECCDerivedFromDataToken(eccToken, value);
 
-    uint64_t contentionFactor = 0;
-    if (maxContentionFactor > 0) {
-        // Generate a number between [1,maxContentionFactor]
-        contentionFactor = SecureRandom().nextInt64(maxContentionFactor) + 1;
-    }
-
     EDCDerivedFromDataTokenAndContentionFactorToken edcDataCounterkey =
         FLEDerivedFromDataTokenAndContentionFactorTokenGenerator::
             generateEDCDerivedFromDataTokenAndContentionFactorToken(edcDatakey, contentionFactor);
@@ -888,7 +880,8 @@ void visitEncryptedBSON(const BSONObj& object,
 void convertToFLE2Payload(FLEKeyVault* keyVault,
                           ConstDataRange cdr,
                           BSONObjBuilder* builder,
-                          StringData fieldNameToSerialize) {
+                          StringData fieldNameToSerialize,
+                          const ContentionFactorFn& contentionFactor) {
     auto [encryptedType, subCdr] = fromEncryptedConstDataRange(cdr);
 
     if (encryptedType == EncryptedBinDataType::kFLE2Placeholder) {
@@ -908,7 +901,7 @@ void convertToFLE2Payload(FLEKeyVault* keyVault,
 
             if (ep.getType() == Fle2PlaceholderType::kInsert) {
                 auto iupayload = EDCClientPayload::serializeInsertUpdatePayload(
-                    indexKey, userKey, el, ep.getMaxContentionCounter());
+                    indexKey, userKey, el, contentionFactor(ep));
 
                 toEncryptedBinData(fieldNameToSerialize,
                                    EncryptedBinDataType::kFLE2InsertUpdatePayload,
@@ -1146,6 +1139,10 @@ stdx::unordered_map<std::string, EncryptedField> toFieldMap(const EncryptedField
     return fields;
 }
 
+uint64_t generateRandomContention(uint64_t cm) {
+    return cm > 0 ? SecureRandom().nextInt64(cm) + 1 : 0;
+}
+
 }  // namespace
 
 
@@ -1265,18 +1262,25 @@ std::vector<uint8_t> FLEClientCrypto::encrypt(BSONElement element,
                                               FLEIndexKeyAndId indexKey,
                                               FLEUserKeyAndId userKey,
                                               FLECounter counter) {
-
-    auto iupayload =
-        EDCClientPayload::serializeInsertUpdatePayload(indexKey, userKey, element, counter);
-
-    return toEncryptedVector(EncryptedBinDataType::kFLE2InsertUpdatePayload, iupayload);
+    return toEncryptedVector(EncryptedBinDataType::kFLE2InsertUpdatePayload,
+                             EDCClientPayload::serializeInsertUpdatePayload(
+                                 indexKey, userKey, element, generateRandomContention(counter)));
 }
 
 
 BSONObj FLEClientCrypto::transformPlaceholders(const BSONObj& obj, FLEKeyVault* keyVault) {
+    return transformPlaceholders(obj, keyVault, [](const FLE2EncryptionPlaceholder& ep) {
+        // Generate a number between [1,maxContentionFactor]
+        return generateRandomContention(ep.getMaxContentionCounter());
+    });
+}
+
+BSONObj FLEClientCrypto::transformPlaceholders(const BSONObj& obj,
+                                               FLEKeyVault* keyVault,
+                                               const ContentionFactorFn& cf) {
     auto ret = transformBSON(
-        obj, [keyVault](ConstDataRange cdr, BSONObjBuilder* builder, StringData field) {
-            convertToFLE2Payload(keyVault, cdr, builder, field);
+        obj, [keyVault, cf](ConstDataRange cdr, BSONObjBuilder* builder, StringData field) {
+            convertToFLE2Payload(keyVault, cdr, builder, field, cf);
         });
 
     return ret;
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index 1092a817d82..f72204f941e 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -721,6 +721,7 @@ private:
     }
 };
 
+using ContentionFactorFn = std::function<uint64_t(const FLE2EncryptionPlaceholder&)>;
 
 class FLEClientCrypto {
 public:
@@ -772,6 +773,14 @@ public:
      */
     static BSONObj transformPlaceholders(const BSONObj& obj, FLEKeyVault* keyVault);
 
+    /**
+     * Generates a client-side payload that is sent to the server. Contention factor is given
+     * explicitly as a lambda expression.
+     */
+    static BSONObj transformPlaceholders(const BSONObj& obj,
+                                         FLEKeyVault* keyVault,
+                                         const ContentionFactorFn& contentionFactor);
+
 
     /**
      * For every encrypted field path in the EncryptedFieldConfig, this generates
diff --git a/src/mongo/crypto/fle_tags.cpp b/src/mongo/crypto/fle_tags.cpp
index 4422182d5dc..29dd3f86f8c 100644
--- a/src/mongo/crypto/fle_tags.cpp
+++ b/src/mongo/crypto/fle_tags.cpp
@@ -70,82 +70,93 @@ using TwiceDerived = FLETwiceDerivedTokenGenerator;
 //          }
 //          pos++
 //      return [EDC::encrypt(i) | i in tags]
-//
-// Note, a positive contention factor (cm) means we must repeat the above process (cm) times.
+void readTagsWithContention(const FLEStateCollectionReader& esc,
+                            const FLEStateCollectionReader& ecc,
+                            ESCDerivedFromDataToken s,
+                            ECCDerivedFromDataToken c,
+                            EDCDerivedFromDataToken d,
+                            uint64_t cf,
+                            std::vector<PrfBlock>& binaryTags) {
+
+    auto escTok = DerivedToken::generateESCDerivedFromDataTokenAndContentionFactorToken(s, cf);
+    auto escTag = TwiceDerived::generateESCTwiceDerivedTagToken(escTok);
+    auto escVal = TwiceDerived::generateESCTwiceDerivedValueToken(escTok);
+
+    auto eccTok = DerivedToken::generateECCDerivedFromDataTokenAndContentionFactorToken(c, cf);
+    auto eccTag = TwiceDerived::generateECCTwiceDerivedTagToken(eccTok);
+    auto eccVal = TwiceDerived::generateECCTwiceDerivedValueToken(eccTok);
+
+    auto edcTok = DerivedToken::generateEDCDerivedFromDataTokenAndContentionFactorToken(d, cf);
+    auto edcTag = TwiceDerived::generateEDCTwiceDerivedToken(edcTok);
+
+    // (1) Query ESC for the counter value after the most recent insert.
+    //     0 => 0 inserts for this field value pair.
+    //     n => n inserts for this field value pair.
+    //     none => compaction => query ESC for null document to find # of inserts.
+    auto insertCounter = ESCCollection::emuBinary(esc, escTag, escVal);
+    if (insertCounter && insertCounter.get() == 0) {
+        return;
+    }
+
+    stdx::unordered_set<int64_t> tags;
+
+    auto numInserts = insertCounter
+        ? uassertStatusOK(
+              ESCCollection::decryptDocument(
+                  escVal, esc.getById(ESCCollection::generateId(escTag, insertCounter))))
+              .count
+        : uassertStatusOK(ESCCollection::decryptNullDocument(
+                              escVal, esc.getById(ESCCollection::generateId(escTag, boost::none))))
+              .count;
+
+    // (2) Set the initial set of tags to 1..n inclusive - a superset of the true tag set.
+    for (uint64_t i = 1; i <= numInserts; i++) {
+        tags.insert(i);
+    }
+
+    // (3) Query ECC for a null document.
+    auto eccNullDoc = ecc.getById(ECCCollection::generateId(eccTag, boost::none));
+    auto pos = eccNullDoc.isEmpty()
+        ? 1
+        : uassertStatusOK(ECCCollection::decryptNullDocument(eccVal, eccNullDoc)).position + 2;
+
+    // (3) Search ECC for deleted tag(counter) values.
+    while (true) {
+        auto eccObj = ecc.getById(ECCCollection::generateId(eccTag, pos));
+        if (eccObj.isEmpty()) {
+            break;
+        }
+        auto eccDoc = uassertStatusOK(ECCCollection::decryptDocument(eccVal, eccObj));
+        // Compaction placeholders only present for positive contention factors (cm).
+        if (eccDoc.valueType == ECCValueType::kCompactionPlaceholder) {
+            break;
+        }
+        for (uint64_t i = eccDoc.start; i <= eccDoc.end; i++) {
+            tags.erase(i);
+        }
+        pos++;
+    }
+
+    for (auto counter : tags) {
+        // (4) Derive binary tag values (encrypted) from the set of counter values (tags).
+        binaryTags.emplace_back(EDCServerCollection::generateTag(edcTag, counter));
+    }
+}
+
+// A positive contention factor (cm) means we must run the above algorithm (cm) times.
 std::vector<PrfBlock> readTags(const FLEStateCollectionReader& esc,
                                const FLEStateCollectionReader& ecc,
                                ESCDerivedFromDataToken s,
                                ECCDerivedFromDataToken c,
                                EDCDerivedFromDataToken d,
                                boost::optional<int64_t> cm) {
-
     std::vector<PrfBlock> binaryTags;
-
-    for (auto i = 0; cm && i <= cm.get(); i++) {
-        auto escTok = DerivedToken::generateESCDerivedFromDataTokenAndContentionFactorToken(s, i);
-        auto escTag = TwiceDerived::generateESCTwiceDerivedTagToken(escTok);
-        auto escVal = TwiceDerived::generateESCTwiceDerivedValueToken(escTok);
-
-        auto eccTok = DerivedToken::generateECCDerivedFromDataTokenAndContentionFactorToken(c, i);
-        auto eccTag = TwiceDerived::generateECCTwiceDerivedTagToken(eccTok);
-        auto eccVal = TwiceDerived::generateECCTwiceDerivedValueToken(eccTok);
-
-        auto edcTok = DerivedToken::generateEDCDerivedFromDataTokenAndContentionFactorToken(d, i);
-        auto edcTag = TwiceDerived::generateEDCTwiceDerivedToken(edcTok);
-
-        // (1) Query ESC for the counter value after the most recent insert.
-        //     0 => 0 inserts for this field value pair.
-        //     n => n inserts for this field value pair.
-        //     none => compaction => query ESC for null document to find # of inserts.
-        auto insertCounter = ESCCollection::emuBinary(esc, escTag, escVal);
-        if (insertCounter && insertCounter.get() == 0) {
-            continue;
-        }
-
-        stdx::unordered_set<int64_t> tags;
-
-        auto numInserts = insertCounter
-            ? uassertStatusOK(
-                  ESCCollection::decryptDocument(
-                      escVal, esc.getById(ESCCollection::generateId(escTag, insertCounter))))
-                  .count
-            : uassertStatusOK(
-                  ESCCollection::decryptNullDocument(
-                      escVal, esc.getById(ESCCollection::generateId(escTag, boost::none))))
-                  .count;
-
-        // (2) Set the initial set of tags to 1..n inclusive - a superset of the true tag set.
-        for (uint64_t i = 1; i <= numInserts; i++) {
-            tags.insert(i);
-        }
-
-        // (3) Query ECC for a null document.
-        auto eccNullDoc = ecc.getById(ECCCollection::generateId(eccTag, boost::none));
-        auto pos = eccNullDoc.isEmpty()
-            ? 1
-            : uassertStatusOK(ECCCollection::decryptNullDocument(eccVal, eccNullDoc)).position + 2;
-
-        // (3) Search ECC for deleted tag(counter) values.
-        while (true) {
-            auto eccObj = ecc.getById(ECCCollection::generateId(eccTag, pos));
-            if (eccObj.isEmpty()) {
-                break;
-            }
-            auto eccDoc = uassertStatusOK(ECCCollection::decryptDocument(eccVal, eccObj));
-            // Compaction placeholders only present for positive contention factors (cm).
-            if (eccDoc.valueType == ECCValueType::kCompactionPlaceholder) {
-                break;
-            }
-            for (uint64_t i = eccDoc.start; i <= eccDoc.end; i++) {
-                tags.erase(i);
-            }
-            pos++;
-        }
-
-        for (auto counter : tags) {
-            // (4) Derive binary tag values (encrypted) from the set of counter values (tags).
-            binaryTags.emplace_back(EDCServerCollection::generateTag(edcTag, counter));
-        }
+    if (!cm || cm.get() == 0) {
+        readTagsWithContention(esc, ecc, s, c, d, 0, binaryTags);
+        return binaryTags;
+    }
+    for (auto i = 1; i <= cm.get(); i++) {
+        readTagsWithContention(esc, ecc, s, c, d, i, binaryTags);
     }
     return binaryTags;
 }
diff --git a/src/mongo/crypto/fle_tags.h b/src/mongo/crypto/fle_tags.h
index 08cd7b1e392..83144341366 100644
--- a/src/mongo/crypto/fle_tags.h
+++ b/src/mongo/crypto/fle_tags.h
@@ -35,7 +35,22 @@
 
 namespace mongo::fle {
 
-// Read a list of binary tags given ESC, ECC, and EDC derived tokens.
+/**
+ * Read a list of binary tags given ESC, ECC, and EDC derived tokens and a specific contention
+ * factor.
+ */
+void readTagsWithContention(const FLEStateCollectionReader& esc,
+                            const FLEStateCollectionReader& ecc,
+                            ESCDerivedFromDataToken s,
+                            ECCDerivedFromDataToken c,
+                            EDCDerivedFromDataToken d,
+                            uint64_t contentionFactor,
+                            std::vector<PrfBlock>& binaryTags);
+
+/**
+ * Read a list of binary tags given ESC, ECC, and EDC derived tokens and a maximum contention
+ * factor.
+ */
 std::vector<PrfBlock> readTags(const FLEStateCollectionReader& esc,
                                const FLEStateCollectionReader& ecc,
                                ESCDerivedFromDataToken s,
diff --git a/src/mongo/db/fle_crud_test.cpp b/src/mongo/db/fle_crud_test.cpp
index 8b89b500c3b..1567ec1b84b 100644
--- a/src/mongo/db/fle_crud_test.cpp
+++ b/src/mongo/db/fle_crud_test.cpp
@@ -151,6 +151,11 @@ protected:
     void doSingleInsert(int id, BSONElement element);
     void doSingleInsert(int id, BSONObj obj);
 
+    void doSingleInsertWithContention(
+        int id, BSONElement element, int64_t cm, uint64_t cf, EncryptedFieldConfig efc);
+    void doSingleInsertWithContention(
+        int id, BSONObj obj, int64_t cm, uint64_t cf, EncryptedFieldConfig efc);
+
     void doSingleDelete(int id);
 
     void doSingleUpdate(int id, BSONElement element);
@@ -422,7 +427,7 @@ void FleCrudTest::validateDocument(int id, boost::optional<BSONObj> doc) {
 }
 
 // Use different keys for index and user
-std::vector<char> generateSinglePlaceholder(BSONElement value) {
+std::vector<char> generateSinglePlaceholder(BSONElement value, int64_t cm = 0) {
     FLE2EncryptionPlaceholder ep;
 
     ep.setAlgorithm(mongo::Fle2AlgorithmInt::kEquality);
@@ -430,7 +435,7 @@ std::vector<char> generateSinglePlaceholder(BSONElement value) {
     ep.setIndexKeyId(indexKeyId);
     ep.setValue(value);
     ep.setType(mongo::Fle2PlaceholderType::kInsert);
-    ep.setMaxContentionCounter(0);
+    ep.setMaxContentionCounter(cm);
 
     BSONObj obj = ep.toBSON();
 
@@ -464,6 +469,30 @@ void FleCrudTest::doSingleInsert(int id, BSONObj obj) {
     doSingleInsert(id, obj.firstElement());
 }
 
+void FleCrudTest::doSingleInsertWithContention(
+    int id, BSONElement element, int64_t cm, uint64_t cf, EncryptedFieldConfig efc) {
+    auto buf = generateSinglePlaceholder(element, cm);
+    BSONObjBuilder builder;
+    builder.append("_id", id);
+    builder.append("counter", 1);
+    builder.append("plainText", "sample");
+    builder.appendBinData("encrypted", buf.size(), BinDataType::Encrypt, buf.data());
+
+    auto clientDoc = builder.obj();
+
+    auto result = FLEClientCrypto::transformPlaceholders(
+        clientDoc, &_keyVault, [cf](const FLE2EncryptionPlaceholder&) { return cf; });
+
+    auto serverPayload = EDCServerCollection::getEncryptedFieldInfo(result);
+
+    uassertStatusOK(processInsert(_queryImpl.get(), _edcNs, serverPayload, efc, 0, result));
+}
+
+void FleCrudTest::doSingleInsertWithContention(
+    int id, BSONObj obj, int64_t cm, uint64_t cf, EncryptedFieldConfig efc) {
+    doSingleInsertWithContention(id, obj.firstElement(), cm, cf, efc);
+}
+
 void FleCrudTest::doSingleUpdate(int id, BSONObj obj) {
     doSingleUpdate(id, obj.firstElement());
 }
@@ -583,13 +612,23 @@ protected:
     void tearDown() {
         FleCrudTest::tearDown();
     }
-    std::vector<PrfBlock> readTags(BSONObj obj) {
+    std::vector<PrfBlock> readTagsWithContention(BSONObj obj, uint64_t contention = 0) {
         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::readTags(esc, ecc, s, c, d, 0);
+        std::vector<PrfBlock> binaryTags;
+        mongo::fle::readTagsWithContention(esc, ecc, s, c, d, contention, binaryTags);
+        return binaryTags;
+    }
+    std::vector<PrfBlock> readTags(BSONObj obj, uint64_t cm = 0) {
+        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::readTags(esc, ecc, s, c, d, cm);
     }
 };
 
@@ -1097,5 +1136,44 @@ TEST_F(FleTagsTest, InsertAndUpdate) {
     ASSERT_EQ(0, readTags(doc1).size());
     ASSERT_EQ(1, readTags(doc2).size());
 }
+
+TEST_F(FleTagsTest, ContentionFactor) {
+    auto efc = EncryptedFieldConfig::parse(IDLParserErrorContext("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 1 and once with a contention factor or 4.
+    doSingleInsertWithContention(1, doc1, 4, 1, efc);
+    doSingleInsertWithContention(4, doc1, 4, 4, efc);
+    doSingleInsertWithContention(5, doc1, 4, 1, efc);
+
+    // Insert doc2 once with a contention factor of 3 and once with a contention factor of 4.
+    doSingleInsertWithContention(7, doc2, 4, 3, efc);
+    doSingleInsertWithContention(8, doc2, 4, 4, efc);
+
+    ASSERT_EQ(2, readTagsWithContention(doc1, 1).size());
+    ASSERT_EQ(0, readTagsWithContention(doc2, 1).size());
+    ASSERT_EQ(0, readTagsWithContention(doc1, 2).size());
+    ASSERT_EQ(0, readTagsWithContention(doc2, 2).size());
+    ASSERT_EQ(0, readTagsWithContention(doc1, 3).size());
+    ASSERT_EQ(1, readTagsWithContention(doc2, 3).size());
+    ASSERT_EQ(1, readTagsWithContention(doc1, 4).size());
+    ASSERT_EQ(1, readTagsWithContention(doc2, 4).size());
+    ASSERT_EQ(3, readTags(doc1, 4).size());
+    ASSERT_EQ(2, readTags(doc2, 4).size());
+}
 }  // namespace
 }  // namespace mongo