SERVER-64101shreyaskalyan/SERVER-64101

3 files changed, 96 insertions, 68 deletions

diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index 7c7e15dc073..7dedeeefb28 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -632,32 +632,26 @@ StatusWith<std::vector<uint8_t>> KeyIdAndValue::decrypt(FLEUserKey userKey,
  */
 class EDCClientPayload {
 public:
-    static FLE2InsertUpdatePayload parseIUPayload(ConstDataRange cdr);
-    static FLE2FindEqualityPayload parseFindPayload(ConstDataRange cdr);
+    static FLE2InsertUpdatePayload parseInsertUpdatePayload(ConstDataRange cdr);
 
-    static FLE2InsertUpdatePayload serializeIUPayload(FLEIndexKeyAndId indexKey,
+    static FLE2InsertUpdatePayload serializeInsertUpdatePayload(FLEIndexKeyAndId indexKey,
                                                       FLEUserKeyAndId userKey,
                                                       BSONElement element,
                                                       uint64_t maxContentionFactor);
-
-    static FLE2FindEqualityPayload serializeFindPayload(FLEIndexKeyAndId indexKey,
-                                                        FLEUserKeyAndId userKey,
-                                                        BSONElement element,
-                                                        uint64_t maxContentionFactor);
 };
 
 
-FLE2InsertUpdatePayload EDCClientPayload::parseIUPayload(ConstDataRange cdr) {
+FLE2InsertUpdatePayload EDCClientPayload::parseInsertUpdatePayload(ConstDataRange cdr) {
     return parseFromCDR<FLE2InsertUpdatePayload>(cdr);
 }
 
 
-FLE2FindEqualityPayload EDCClientPayload::parseFindPayload(ConstDataRange cdr) {
+FLE2FindEqualityPayload FLEClientCrypto::parseFindPayload(ConstDataRange cdr) {
     return parseFromCDR<FLE2FindEqualityPayload>(cdr);
 }
 
 
-FLE2InsertUpdatePayload EDCClientPayload::serializeIUPayload(FLEIndexKeyAndId indexKey,
+FLE2InsertUpdatePayload EDCClientPayload::serializeInsertUpdatePayload(FLEIndexKeyAndId indexKey,
                                                              FLEUserKeyAndId userKey,
                                                              BSONElement element,
                                                              uint64_t maxContentionFactor) {
@@ -719,20 +713,17 @@ FLE2InsertUpdatePayload EDCClientPayload::serializeIUPayload(FLEIndexKeyAndId in
 }
 
 
-FLE2FindEqualityPayload EDCClientPayload::serializeFindPayload(FLEIndexKeyAndId indexKey,
+FLE2FindEqualityPayload FLEClientCrypto::serializeFindPayload(FLEIndexKeyAndId indexKey,
                                                                FLEUserKeyAndId userKey,
                                                                BSONElement element,
                                                                uint64_t maxContentionFactor) {
     auto value = ConstDataRange(element.value(), element.value() + element.valuesize());
 
     auto collectionToken = FLELevel1TokenGenerator::generateCollectionsLevel1Token(indexKey.key);
-    auto serverEncryptToken =
-        FLELevel1TokenGenerator::generateServerDataEncryptionLevel1Token(indexKey.key);
 
     auto edcToken = FLECollectionTokenGenerator::generateEDCToken(collectionToken);
     auto escToken = FLECollectionTokenGenerator::generateESCToken(collectionToken);
     auto eccToken = FLECollectionTokenGenerator::generateECCToken(collectionToken);
-    auto ecocToken = FLECollectionTokenGenerator::generateECOCToken(collectionToken);
 
     EDCDerivedFromDataToken edcDatakey =
         FLEDerivedFromDataTokenGenerator::generateEDCDerivedFromDataToken(edcToken, value);
@@ -953,7 +944,7 @@ void convertToFLE2Payload(FLEKeyVault* keyVault,
 
             if (ep.getType() == Fle2PlaceholderType::kInsert) {
                 auto iupayload =
-                    EDCClientPayload::serializeIUPayload(indexKey, userKey, el, ep.getMaxContentionCounter());
+                    EDCClientPayload::serializeInsertUpdatePayload(indexKey, userKey, el, ep.getMaxContentionCounter());
 
                 toEncryptedBinData(fieldNameToSerialize,
                                    EncryptedBinDataType::kFLE2InsertUpdatePayload,
@@ -961,12 +952,14 @@ void convertToFLE2Payload(FLEKeyVault* keyVault,
                                    builder);
             } else if (ep.getType() == Fle2PlaceholderType::kFind) {
                 auto findpayload =
-                    EDCClientPayload::serializeFindPayload(indexKey, userKey, el, ep.getMaxContentionCounter());
+                    FLEClientCrypto::serializeFindPayload(indexKey, userKey, el, ep.getMaxContentionCounter());
 
                 toEncryptedBinData(fieldNameToSerialize,
                                    EncryptedBinDataType::kFLE2FindEqualityPayload,
                                    findpayload,
                                    builder);
+            } else {
+                uasserted(6410100, "No other FLE2 placeholders supported at this time.");
             }
         } else {
             uasserted(6338603, "Only FLE 2 style encryption placeholders are supported");
@@ -980,8 +973,8 @@ void convertToFLE2Payload(FLEKeyVault* keyVault,
     }
 }
 
-void parseAndVerifyIUPayload(std::vector<EDCServerPayloadInfo>* pFields, StringData fieldPath, EncryptedBinDataType type, ConstDataRange subCdr) {
-    auto iupayload = EDCClientPayload::parseIUPayload(subCdr);
+void parseAndVerifyInsertUpdatePayload(std::vector<EDCServerPayloadInfo>* pFields, StringData fieldPath, EncryptedBinDataType type, ConstDataRange subCdr) {
+    auto iupayload = EDCClientPayload::parseInsertUpdatePayload(subCdr);
 
     uassert(6373504,
             str::stream() << "Type '" << typeName(static_cast<BSONType>(iupayload.getType()))
@@ -992,11 +985,6 @@ void parseAndVerifyIUPayload(std::vector<EDCServerPayloadInfo>* pFields, StringD
     pFields->push_back({std::move(iupayload), fieldPath.toString(), 0});
 }
 
-void parseAndVerifyFindPayload(std::vector<EDCServerPayloadInfo>* pFields, StringData fieldPath, EncryptedBinDataType type, ConstDataRange subCdr) {
-    // No-op;
-    return;
-}
-
 void collectEDCServerInfo(std::vector<EDCServerPayloadInfo>* pFields,
                           ConstDataRange cdr,
                           StringData fieldPath) {
@@ -1010,7 +998,7 @@ void collectEDCServerInfo(std::vector<EDCServerPayloadInfo>* pFields,
     if (encryptedTypeBinding == EncryptedBinDataType::kFLE2InsertUpdatePayload) {
         parseAndVerifyIUPayload(pFields, fieldPath, encryptedTypeBinding, subCdr);
     } else if (encryptedTypeBinding == EncryptedBinDataType::kFLE2FindEqualityPayload) {
-        parseAndVerifyFindPayload(pFields, fieldPath, encryptedTypeBinding, subCdr);
+        // No-op
     } else {
         uasserted(6373503,
                   str::stream() << "Unexpected encrypted payload type: "
@@ -1288,7 +1276,7 @@ std::vector<uint8_t> FLEClientCrypto::encrypt(BSONElement element,
                                               FLEUserKeyAndId userKey,
                                               FLECounter counter) {
 
-    auto iupayload = EDCClientPayload::serializeIUPayload(indexKey, userKey, element, counter);
+    auto iupayload = EDCClientPayload::serializeInsertUpdatePayload(indexKey, userKey, element, counter);
 
     return toEncryptedVector(EncryptedBinDataType::kFLE2InsertUpdatePayload, iupayload);
 }
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index a64f567219b..e259a7848f0 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -714,6 +714,14 @@ public:
     static std::pair<BSONType, std::vector<uint8_t>> decrypt(BSONElement element,
                                                              FLEKeyVault* keyVault);
 
+    static FLE2FindEqualityPayload parseFindPayload(ConstDataRange cdr);
+
+    static FLE2FindEqualityPayload serializeFindPayload(FLEIndexKeyAndId indexKey,
+                                                        FLEUserKeyAndId userKey,
+                                                        BSONElement element,
+                                                        uint64_t maxContentionFactor);
+
+
     /**
      * Generates a client-side payload that is sent to the server.
      *
diff --git a/src/mongo/crypto/fle_crypto_test.cpp b/src/mongo/crypto/fle_crypto_test.cpp
index 23de9e0ff5a..9235430e453 100644
--- a/src/mongo/crypto/fle_crypto_test.cpp
+++ b/src/mongo/crypto/fle_crypto_test.cpp
@@ -600,42 +600,33 @@ BSONObj encryptDocumentForInsert(BSONObj obj, FLEKeyVault* keyVault) {
     return finalDoc;
 }
 
-
-BSONObj generateTokensForFind(BSONObj obj, FLEKeyVault* keyVault) {
-    auto result = FLEClientCrypto::transformPlaceholders(obj, keyVault);
-
-    // TODO: once query work for find is done, we can return finalDoc. Until then
-    // we should just return result to see the object with placeholders.
-
-    // // Start Server Side
-    // auto serverPayload = EDCServerCollection::getEncryptedFieldInfo(result);
-
-    // for (auto& payload : serverPayload) {
-    //     payload.count = 1;
-    // }
-
-    // // Finalize document for find
-    // auto finalDoc = EDCServerCollection::finalizeForUpdate(result, serverPayload);
-
-    return result;
-}
-
 BSONObj encryptDocument(BSONObj obj, FLEKeyVault* keyVault, Operation operation) {
     if (operation == Operation::kInsert) {
         return encryptDocumentForInsert(obj, keyVault);
     } else if (operation == Operation::kFind) {
-        return generateTokensForFind(obj, keyVault);
+        return FLEClientCrypto::transformPlaceholders(obj, keyVault);
     } else {
         return BSONObj();
     }
 }
 
-void roundTripTest(BSONObj doc, BSONType type, Operation operation) {
+void assertPayload(BSONElement elem, Operation operation) {
+    int len;
+    const char* data(elem.binData(len));
+    ConstDataRange cdr(data, len);
 
+    const auto& [encryptedType, subCdr] = fromEncryptedConstDataRange(cdr);
+    if (operation == Operation::kFind) {
+        ASSERT_TRUE(encryptedType == EncryptedBinDataType::kFLE2FindEqualityPayload);
+    } else {
+        ASSERT_TRUE(encryptedType == EncryptedBinDataType::kFLE2InsertUpdatePayload);
+    }
+}
+
+void roundTripTest(BSONObj doc, BSONType type, Operation operation) {
     auto element = doc.firstElement();
     ASSERT_EQ(element.type(), type);
 
-
     TestKeyVault keyVault;
 
     auto inputDoc = BSON("plainText"
@@ -656,6 +647,7 @@ void roundTripTest(BSONObj doc, BSONType type, Operation operation) {
     // TODO : when query enables server side work for Find, remove this
     // if statement.
     if (operation == Operation::kFind) {
+        assertPayload(finalDoc["encrypted"], operation);
         return;
     }
 
@@ -668,28 +660,68 @@ void roundTripTest(BSONObj doc, BSONType type, Operation operation) {
     ASSERT_BSONOBJ_EQ(inputDoc, decryptedDoc);
 }
 
+void roundTripMultiencrypted(BSONObj doc1, BSONObj doc2, Operation operation1, Operation operation2) {
+    auto element1 = doc1.firstElement();
+    auto element2 = doc2.firstElement();
+
+    TestKeyVault keyVault;
+
+    auto inputDoc = BSON("plainText"
+                         << "sample"
+                         << "encrypted1" << element1
+                         << "encrypted2" << element2);
+
+    auto buf1 = generatePlaceholder(element1, operation1);
+    auto buf2 = generatePlaceholder(element2, operation2);
+
+    BSONObjBuilder builder;
+    builder.append("plaintext", "sample");
+    builder.appendBinData("encrypted1", buf1.size(), BinDataType::Encrypt, buf1.data());
+    builder.appendBinData("encrypted2", buf2.size(), BinDataType::Encrypt, buf2.data());
+
+    auto finalDoc = FLEClientCrypto::transformPlaceholders(builder.obj(), &keyVault);
+
+    ASSERT_EQ(finalDoc["encrypted1"].type(), BinData);
+    ASSERT_TRUE(finalDoc["encrypted1"].isBinData(BinDataType::Encrypt));
+
+    ASSERT_EQ(finalDoc["encrypted2"].type(), BinData);
+    ASSERT_TRUE(finalDoc["encrypted2"].isBinData(BinDataType::Encrypt));
+
+    assertPayload(finalDoc["encrypted1"], operation1);
+    assertPayload(finalDoc["encrypted2"], operation2);
+}
+
+const std::vector<std::pair<BSONObj, BSONType>> objects{
+    { BSON("sample" << "value123"), String },
+    { BSON("sample" << BSONBinData(testValue.data(), testValue.size(), BinDataType::BinDataGeneral)), BinData },
+    { BSON("sample" << OID()), jstOID },
+    { BSON("sample" << false), Bool },
+    { BSON("sample" << true), Bool },
+    { BSON("sample" << Date_t()), Date },
+    { BSON("sample" << BSONRegEx("value1", "value2")), RegEx },
+    { BSON("sample" << 123456), NumberInt },
+    { BSON("sample" << Timestamp()), bsonTimestamp },
+    { BSON("sample" << 12345678901234567LL), NumberLong },
+    { BSON("sample" << BSONCode("value")), Code }
+};
+
 TEST(FLE_EDC, Allowed_Types) {
     std::vector<Operation> opTypes{Operation::kInsert, Operation::kFind};
-    
-    for (const auto& type : opTypes) {
-        roundTripTest(BSON("sample"
-                        << "value123"),
-                    String, type);
-        roundTripTest(BSON("sample" << BSONBinData(
-                            testValue.data(), testValue.size(), BinDataType::BinDataGeneral)),
-                    BinData, type);
-        roundTripTest(BSON("sample" << OID()), jstOID, type);
-
-
-        roundTripTest(BSON("sample" << false), Bool, type);
-        roundTripTest(BSON("sample" << true), Bool, type);
-        roundTripTest(BSON("sample" << Date_t()), Date, type);
-        roundTripTest(BSON("sample" << BSONRegEx("value1", "value2")), RegEx, type);
-        roundTripTest(BSON("sample" << 123456), NumberInt, type);
-        roundTripTest(BSON("sample" << Timestamp()), bsonTimestamp, type);
-        roundTripTest(BSON("sample" << 12345678901234567LL), NumberLong, type);
-        roundTripTest(BSON("sample" << BSONCode("value")), Code, type);
+
+    for (const auto& opType : opTypes) {
+        for (const auto& [obj, objType] : objects) {
+            roundTripTest(obj, objType, opType);
+        }
     };
+
+    for (const auto& [obj1, _] : objects) {
+        for (const auto& [obj2, _] : objects) {
+            roundTripMultiencrypted(obj1, obj2, Operation::kInsert, Operation::kInsert);
+            roundTripMultiencrypted(obj1, obj2, Operation::kInsert, Operation::kFind);
+            roundTripMultiencrypted(obj1, obj2, Operation::kFind, Operation::kInsert);
+            roundTripMultiencrypted(obj1, obj2, Operation::kFind, Operation::kFind);
+        }
+    }
 }
 
 void illegalBSONType(BSONObj doc, BSONType type) {