SERVER-63511 Add support for ECC documents

3 files changed, 329 insertions, 0 deletions

diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index adaaf994221..f819fde1db8 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -263,6 +263,18 @@ StatusWith<std::vector<uint8_t>> decryptData(ConstDataRange key, ConstDataRange
 }
 
 
+StatusWith<uint64_t> decryptUInt64(ConstDataRange key, ConstDataRange cipherText) {
+    auto swPlainText = decryptData(key, cipherText);
+    if (!swPlainText.isOK()) {
+        return swPlainText.getStatus();
+    }
+
+    ConstDataRange cdr(swPlainText.getValue());
+
+    return cdr.readNoThrow<LittleEndian<uint64_t>>();
+}
+
+
 template <typename T>
 struct FLEStoragePackTypeHelper;
 
@@ -663,4 +675,120 @@ uint64_t ESCCollection::emuBinary(FLEStateCollectionReader* reader,
         reader, tagToken, valueToken);
 }
 
+
+PrfBlock ECCCollection::generateId(ECCTwiceDerivedTagToken tagToken,
+                                   boost::optional<uint64_t> index) {
+    if (index.has_value()) {
+        return prf(tagToken.data, kECCNonNullId, index.value());
+    } else {
+        return prf(tagToken.data, kECCNullId, 0);
+    }
+}
+
+BSONObj ECCCollection::generateNullDocument(ECCTwiceDerivedTagToken tagToken,
+                                            ECCTwiceDerivedValueToken valueToken,
+                                            uint64_t count) {
+    auto block = ECCCollection::generateId(tagToken, boost::none);
+
+    auto swCipherText = encryptData(valueToken.data, count);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+BSONObj ECCCollection::generateDocument(ECCTwiceDerivedTagToken tagToken,
+                                        ECCTwiceDerivedValueToken valueToken,
+                                        uint64_t index,
+                                        uint64_t start,
+                                        uint64_t end) {
+    auto block = ECCCollection::generateId(tagToken, index);
+
+    auto swCipherText = packAndEncrypt(std::tie(start, end), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+BSONObj ECCCollection::generateDocument(ECCTwiceDerivedTagToken tagToken,
+                                        ECCTwiceDerivedValueToken valueToken,
+                                        uint64_t index,
+                                        uint64_t count) {
+    return generateDocument(tagToken, valueToken, index, count, count);
+}
+
+BSONObj ECCCollection::generateCompactionDocument(ECCTwiceDerivedTagToken tagToken,
+                                                  ECCTwiceDerivedValueToken valueToken,
+                                                  uint64_t index) {
+    auto block = ECCCollection::generateId(tagToken, index);
+
+    auto swCipherText =
+        packAndEncrypt(std::tie(kECCompactionRecordValue, kECCompactionRecordValue), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+
+StatusWith<ECCNullDocument> ECCCollection::decryptNullDocument(ECCTwiceDerivedValueToken valueToken,
+                                                               BSONObj& doc) {
+    BSONElement encryptedValue;
+    auto status = bsonExtractTypedField(doc, kValue, BinData, &encryptedValue);
+    if (!status.isOK()) {
+        return status;
+    }
+
+    auto swUnpack = decryptUInt64(valueToken.data, binDataToCDR(encryptedValue));
+
+    if (!swUnpack.isOK()) {
+        return swUnpack.getStatus();
+    }
+
+    auto& value = swUnpack.getValue();
+
+    return ECCNullDocument{value};
+}
+
+
+StatusWith<ECCDocument> ECCCollection::decryptDocument(ECCTwiceDerivedValueToken valueToken,
+                                                       BSONObj& doc) {
+    BSONElement encryptedValue;
+    auto status = bsonExtractTypedField(doc, kValue, BinData, &encryptedValue);
+    if (!status.isOK()) {
+        return status;
+    }
+
+    auto swUnpack = decryptAndUnpack<uint64_t, uint64_t>(binDataToCDR(encryptedValue), valueToken);
+
+    if (!swUnpack.isOK()) {
+        return swUnpack.getStatus();
+    }
+
+    auto& value = swUnpack.getValue();
+
+    return ECCDocument{std::get<0>(value) != kECCompactionRecordValue
+                           ? ECCValueType::kNormal
+                           : ECCValueType::kCompactionPlaceholder,
+                       std::get<0>(value),
+                       std::get<1>(value)};
+}
+
+uint64_t ECCCollection::emuBinary(FLEStateCollectionReader* reader,
+                                  ECCTwiceDerivedTagToken tagToken,
+                                  ECCTwiceDerivedValueToken valueToken) {
+    return emuBinaryCommon<ECCCollection, ECCTwiceDerivedTagToken, ECCTwiceDerivedValueToken>(
+        reader, tagToken, valueToken);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/crypto/fle_crypto.h b/src/mongo/crypto/fle_crypto.h
index 5a6c0d6ac8a..02e3946c824 100644
--- a/src/mongo/crypto/fle_crypto.h
+++ b/src/mongo/crypto/fle_crypto.h
@@ -498,4 +498,151 @@ public:
                               ESCTwiceDerivedValueToken valueToken);
 };
 
+
+/**
+ * ECC Collection
+ * - a record of deleted documents
+ *
+ * {
+ *    _id : HMAC(ECCTwiceDerivedTagToken, type || pos )
+ *    value : Encrypt(ECCTwiceDerivedValueToken,  count OR start || end)
+ * }
+ *
+ * where
+ *  type = uint64_t
+ *  pos = uint64_t
+ *  value is either:
+ *       count = uint64_t  // Null records
+ *    OR
+ *       start = uint64_t  // Other records
+ *       end = uint64_t
+ *
+ * where type:
+ *   0 - null record
+ *   1 - regular record or compaction record
+ *
+ * where start and end:
+ *   [0..UINT_64_MAX) - regular start and end
+ *   UINT64_MAX - compaction placeholder
+ *
+ * Record types:
+ *
+ * Document Counts
+ * Null: 0 or 1
+ * Regular: 0 or more
+ * Compaction: 0 or 1
+ *
+ * Null record:
+ * {
+ *    _id : HMAC(ECCTwiceDerivedTagToken, null )
+ *    value : Encrypt(ECCTwiceDerivedValueToken,  count)
+ * }
+ *
+ * Regular record:
+ * {
+ *    _id : HMAC(ECCTwiceDerivedTagToken, pos )
+ *    value : Encrypt(ECCTwiceDerivedValueToken,  start || end)
+ * }
+ *
+ * Compaction placeholder record:
+ * {
+ *    _id : HMAC(ECCTwiceDerivedTagToken, pos )
+ *    value : Encrypt(ECCTwiceDerivedValueToken,  UINT64_MAX || UINT64_MAX)
+ * }
+ *
+ * PlainText of tag
+ * struct {
+ *    uint64_t type;
+ *    uint64_t pos;
+ * }
+ *
+ * PlainText of value for null records
+ * struct {
+ *    uint64_t count;
+ * }
+ *
+ * PlainText of value for non-null records
+ * struct {
+ *    uint64_t start;
+ *    uint64_t end;
+ * }
+ */
+enum class ECCValueType : uint64_t {
+    kNormal = 0,
+    kCompactionPlaceholder = 1,
+};
+
+
+struct ECCNullDocument {
+    // Id is not included as it HMAC generated and cannot be reversed
+    uint64_t pos;
+};
+
+
+struct ECCDocument {
+    // Id is not included as it HMAC generated and cannot be reversed
+    ECCValueType valueType;
+    uint64_t start;
+    uint64_t end;
+};
+
+class ECCCollection {
+public:
+    /**
+     * Generate the _id value
+     */
+    static PrfBlock generateId(ECCTwiceDerivedTagToken tagToken, boost::optional<uint64_t> index);
+
+    /**
+     * Generate a null document which will be the "first" document for a given field.
+     */
+    static BSONObj generateNullDocument(ECCTwiceDerivedTagToken tagToken,
+                                        ECCTwiceDerivedValueToken valueToken,
+                                        uint64_t count);
+
+    /**
+     * Generate a regular ECC document for (count).
+     *
+     * Note: it is stored as (count, count)
+     */
+    static BSONObj generateDocument(ECCTwiceDerivedTagToken tagToken,
+                                    ECCTwiceDerivedValueToken valueToken,
+                                    uint64_t index,
+                                    uint64_t count);
+
+    /**
+     * Generate a regular ECC document for (start, end)
+     */
+    static BSONObj generateDocument(ECCTwiceDerivedTagToken tagToken,
+                                    ECCTwiceDerivedValueToken valueToken,
+                                    uint64_t index,
+                                    uint64_t start,
+                                    uint64_t end);
+
+    /**
+     * Generate a compaction ECC document.
+     */
+    static BSONObj generateCompactionDocument(ECCTwiceDerivedTagToken tagToken,
+                                              ECCTwiceDerivedValueToken valueToken,
+                                              uint64_t index);
+
+    /**
+     * Decrypt the null document.
+     */
+    static StatusWith<ECCNullDocument> decryptNullDocument(ECCTwiceDerivedValueToken valueToken,
+                                                           BSONObj& doc);
+
+    /**
+     * Decrypt a regular document.
+     */
+    static StatusWith<ECCDocument> decryptDocument(ECCTwiceDerivedValueToken valueToken,
+                                                   BSONObj& doc);
+
+    /**
+     * Search for the highest document id for a given field/value pair based on the token.
+     */
+    static uint64_t emuBinary(FLEStateCollectionReader* reader,
+                              ECCTwiceDerivedTagToken tagToken,
+                              ECCTwiceDerivedValueToken valueToken);
+};
 }  // namespace mongo
diff --git a/src/mongo/crypto/fle_crypto_test.cpp b/src/mongo/crypto/fle_crypto_test.cpp
index 1425702be45..7c35914fc48 100644
--- a/src/mongo/crypto/fle_crypto_test.cpp
+++ b/src/mongo/crypto/fle_crypto_test.cpp
@@ -296,6 +296,60 @@ TEST(FLE_ESC, RoundTrip) {
 }
 
 
+TEST(FLE_ECC, RoundTrip) {
+
+    ConstDataRange value(testValue);
+
+    auto c1 = FLELevel1TokenGenerator::generateCollectionsLevel1Token(indexKey);
+    auto token = FLECollectionTokenGenerator::generateECCToken(c1);
+
+    ECCDerivedFromDataToken datakey =
+        FLEDerivedFromDataTokenGenerator::generateECCDerivedFromDataToken(token, value);
+
+    ECCDerivedFromDataTokenAndContentionFactorToken dataCounterkey =
+        FLEDerivedFromDataTokenAndContentionFactorTokenGenerator::
+            generateECCDerivedFromDataTokenAndContentionFactorToken(datakey, 0);
+
+    auto twiceTag = FLETwiceDerivedTokenGenerator::generateECCTwiceDerivedTagToken(dataCounterkey);
+    auto twiceValue =
+        FLETwiceDerivedTokenGenerator::generateECCTwiceDerivedValueToken(dataCounterkey);
+
+
+    {
+        BSONObj doc = ECCCollection::generateNullDocument(twiceTag, twiceValue, 123456789);
+        auto swDoc = ECCCollection::decryptNullDocument(twiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT_EQ(swDoc.getValue().pos, 123456789);
+    }
+
+
+    {
+        BSONObj doc = ECCCollection::generateDocument(twiceTag, twiceValue, 123, 123456789);
+        auto swDoc = ECCCollection::decryptDocument(twiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT(swDoc.getValue().valueType == ECCValueType::kNormal);
+        ASSERT_EQ(swDoc.getValue().start, 123456789);
+        ASSERT_EQ(swDoc.getValue().end, 123456789);
+    }
+
+    {
+        BSONObj doc =
+            ECCCollection::generateDocument(twiceTag, twiceValue, 123, 123456789, 983456789);
+        auto swDoc = ECCCollection::decryptDocument(twiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT(swDoc.getValue().valueType == ECCValueType::kNormal);
+        ASSERT_EQ(swDoc.getValue().start, 123456789);
+        ASSERT_EQ(swDoc.getValue().end, 983456789);
+    }
+
+    {
+        BSONObj doc = ECCCollection::generateCompactionDocument(twiceTag, twiceValue, 123456789);
+        auto swDoc = ECCCollection::decryptDocument(twiceValue, doc);
+        ASSERT_OK(swDoc.getStatus());
+        ASSERT(swDoc.getValue().valueType == ECCValueType::kCompactionPlaceholder);
+    }
+}
+
 class TestDocumentCollection : public FLEStateCollectionReader {
 public:
     void insert(BSONObj& obj) {