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

4 files changed, 131 insertions, 92 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,