SERVER-63385 Add support for ESC documents

1 files changed, 435 insertions, 0 deletions

diff --git a/src/mongo/crypto/fle_crypto.cpp b/src/mongo/crypto/fle_crypto.cpp
index b2c894bb4fa..adaaf994221 100644
--- a/src/mongo/crypto/fle_crypto.cpp
+++ b/src/mongo/crypto/fle_crypto.cpp
@@ -133,6 +133,316 @@ PrfBlock prf(ConstDataRange key, uint64_t value) {
     return prf(key, bufValue);
 }
 
+PrfBlock prf(ConstDataRange key, uint64_t value, int64_t value2) {
+    SHA256Block block;
+
+    std::array<char, sizeof(uint64_t)> bufValue;
+    DataView(bufValue.data()).write<LittleEndian<uint64_t>>(value);
+
+
+    std::array<char, sizeof(uint64_t)> bufValue2;
+    DataView(bufValue2.data()).write<LittleEndian<uint64_t>>(value2);
+
+    SHA256Block::computeHmac(key.data<uint8_t>(),
+                             key.length(),
+                             {
+                                 ConstDataRange{bufValue},
+                                 ConstDataRange{bufValue2},
+                             },
+                             &block);
+    return blockToArray(block);
+}
+
+ConstDataRange binDataToCDR(const BSONElement element) {
+    uassert(6338501, "Expected binData BSON element", element.type() == BinData);
+
+    int len;
+    const char* data = element.binData(len);
+    return ConstDataRange(data, data + len);
+}
+
+template <typename T>
+void toBinData(StringData field, T t, BSONObjBuilder* builder) {
+    BSONObj obj = t.toBSON();
+
+    builder->appendBinData(field, obj.objsize(), BinDataType::BinDataGeneral, obj.objdata());
+}
+
+void toBinData(StringData field, PrfBlock block, BSONObjBuilder* builder) {
+    builder->appendBinData(field, block.size(), BinDataType::BinDataGeneral, block.data());
+}
+
+void toBinData(StringData field, ConstDataRange block, BSONObjBuilder* builder) {
+    builder->appendBinData(field, block.length(), BinDataType::BinDataGeneral, block.data());
+}
+
+void toBinData(StringData field, std::vector<uint8_t>& block, BSONObjBuilder* builder) {
+    builder->appendBinData(field, block.size(), BinDataType::BinDataGeneral, block.data());
+}
+
+/**
+ * AEAD AES + SHA256
+ * Block size = 16 bytes
+ * SHA-256 - block size = 256 bits = 32 bytes
+ */
+// TODO (SERVER-63382) - replace with call to CTR AEAD algorithm
+StatusWith<std::vector<uint8_t>> encryptDataWithAssociatedData(ConstDataRange key,
+                                                               ConstDataRange associatedData,
+                                                               ConstDataRange plainText) {
+
+    std::array<uint8_t, sizeof(uint64_t)> dataLenBitsEncodedStorage;
+    DataRange dataLenBitsEncoded(dataLenBitsEncodedStorage);
+    dataLenBitsEncoded.write<BigEndian<uint64_t>>(
+        static_cast<uint64_t>(associatedData.length() * 8));
+
+    std::vector<uint8_t> out;
+    out.resize(crypto::aeadCipherOutputLength(plainText.length()));
+
+    // TODO - key is too short, we have 32, need 64. The new API should only 32 bytes and this can
+    // be removed
+    std::array<uint8_t, 64> bigToken;
+    std::copy(key.data(), key.data() + key.length(), bigToken.data());
+    std::copy(key.data(), key.data() + key.length(), bigToken.data() + key.length());
+
+    auto s = crypto::aeadEncryptWithIV(
+        bigToken, plainText, ConstDataRange(0, 0), associatedData, dataLenBitsEncoded, out);
+    if (!s.isOK()) {
+        return s;
+    }
+
+    return {out};
+}
+
+StatusWith<std::vector<uint8_t>> encryptData(ConstDataRange key, ConstDataRange plainText) {
+
+    return encryptDataWithAssociatedData(key, ConstDataRange(0, 0), plainText);
+}
+
+StatusWith<std::vector<uint8_t>> encryptData(ConstDataRange key, uint64_t value) {
+
+    std::array<char, sizeof(uint64_t)> bufValue;
+    DataView(bufValue.data()).write<LittleEndian<uint64_t>>(value);
+
+    return encryptData(key, bufValue);
+}
+
+// TODO (SERVER-63382) - replace with call to CTR AEAD algorithm
+StatusWith<std::vector<uint8_t>> decryptDataWithAssociatedData(ConstDataRange key,
+                                                               ConstDataRange associatedData,
+                                                               ConstDataRange cipherText) {
+    // TODO - key is too short, we have 32, need 64. The new API should only 32 bytes and this can
+    // be removed
+    std::array<uint8_t, 64> bigToken;
+    std::copy(key.data(), key.data() + key.length(), bigToken.data());
+    std::copy(key.data(), key.data() + key.length(), bigToken.data() + key.length());
+
+    SymmetricKey sk(reinterpret_cast<const uint8_t*>(bigToken.data()),
+                    bigToken.size(),
+                    0,
+                    SymmetricKeyId("ignore"),
+                    0);
+
+    auto swLen = aeadGetMaximumPlainTextLength(cipherText.length());
+    if (!swLen.isOK()) {
+        return swLen.getStatus();
+    }
+    std::vector<uint8_t> out;
+    out.resize(swLen.getValue());
+
+    auto swOutLen = crypto::aeadDecrypt(sk, cipherText, associatedData, out);
+    if (!swOutLen.isOK()) {
+        return swOutLen.getStatus();
+    }
+    out.resize(swOutLen.getValue());
+    return out;
+}
+
+
+StatusWith<std::vector<uint8_t>> decryptData(ConstDataRange key, ConstDataRange cipherText) {
+    return decryptDataWithAssociatedData(key, ConstDataRange(0, 0), cipherText);
+}
+
+
+template <typename T>
+struct FLEStoragePackTypeHelper;
+
+template <>
+struct FLEStoragePackTypeHelper<uint64_t> {
+    using Type = LittleEndian<uint64_t>;
+};
+
+template <>
+struct FLEStoragePackTypeHelper<PrfBlock> {
+    using Type = PrfBlock;
+};
+
+template <typename T>
+struct FLEStoragePackType {
+    // Note: the reference must be removed before the const
+    using Type =
+        typename FLEStoragePackTypeHelper<std::remove_const_t<std::remove_reference_t<T>>>::Type;
+};
+
+template <typename T1, typename T2, FLETokenType TokenT>
+StatusWith<std::vector<uint8_t>> packAndEncrypt(std::tuple<T1, T2> tuple, FLEToken<TokenT> token) {
+    DataBuilder builder(sizeof(T1) + sizeof(T2));
+    Status s = builder.writeAndAdvance<typename FLEStoragePackType<T1>::Type>(std::get<0>(tuple));
+    if (!s.isOK()) {
+        return s;
+    }
+
+    s = builder.writeAndAdvance<typename FLEStoragePackType<T2>::Type>(std::get<1>(tuple));
+    if (!s.isOK()) {
+        return s;
+    }
+
+    dassert(builder.getCursor().length() == (sizeof(T1) + sizeof(T2)));
+    return encryptData(token.toCDR(), builder.getCursor());
+}
+
+
+template <typename T1, typename T2, FLETokenType TokenT>
+StatusWith<std::tuple<T1, T2>> decryptAndUnpack(ConstDataRange cdr, FLEToken<TokenT> token) {
+    auto swVec = decryptData(token.toCDR(), cdr);
+    if (!swVec.isOK()) {
+        return swVec.getStatus();
+    }
+
+    auto& data = swVec.getValue();
+
+    ConstDataRangeCursor cdrc(data);
+
+    auto swt1 = cdrc.readAndAdvanceNoThrow<typename FLEStoragePackType<T1>::Type>();
+    if (!swt1.isOK()) {
+        return swt1.getStatus();
+    }
+
+    auto swt2 = cdrc.readAndAdvanceNoThrow<typename FLEStoragePackType<T2>::Type>();
+    if (!swt2.isOK()) {
+        return swt2.getStatus();
+    }
+
+    return std::tie(swt1.getValue(), swt2.getValue());
+}
+
+
+//#define DEBUG_ENUM_BINARY 1
+
+template <typename collectionT, typename tagTokenT, typename valueTokenT>
+uint64_t emuBinaryCommon(FLEStateCollectionReader* reader,
+                         tagTokenT tagToken,
+                         valueTokenT valueToken) {
+
+    // Default search parameters
+    uint64_t lambda = 0;
+    uint64_t i = 0;
+
+    // Step 2:
+    // Search for null record
+    PrfBlock nullRecordId = collectionT::generateId(tagToken, boost::none);
+
+    BSONObj nullDoc = reader->getById(nullRecordId);
+
+    if (!nullDoc.isEmpty()) {
+        auto swNullEscDoc = collectionT::decryptNullDocument(valueToken, nullDoc);
+        uassertStatusOK(swNullEscDoc.getStatus());
+        lambda = swNullEscDoc.getValue().pos + 1;
+#ifdef DEBUG_ENUM_BINARY
+        std::cout << fmt::format("start: null_document: lambda {}, i: {}", lambda, i) << std::endl;
+#endif
+    }
+
+    // step 4, 5: get document count
+    uint64_t rho = reader->getDocumentCount();
+
+#ifdef DEBUG_ENUM_BINARY
+    std::cout << fmt::format("start: lambda: {}, i: {}, rho: {}", lambda, i, rho) << std::endl;
+#endif
+
+    // step 6
+    bool flag = true;
+
+    // step 7
+    // TODO - this loop never terminates unless it finds a document, need to add a terminating
+    // condition
+    while (flag) {
+        // 7 a
+        BSONObj doc = reader->getById(collectionT::generateId(tagToken, rho + lambda));
+
+#ifdef DEBUG_ENUM_BINARY
+        std::cout << fmt::format("search1: rho: {},  doc: {}", rho, doc.toString()) << std::endl;
+#endif
+
+        // 7 b
+        if (!doc.isEmpty()) {
+            rho = 2 * rho;
+        } else {
+            flag = false;
+        }
+    }
+
+    // Step 8:
+    uint64_t median = 0, min = 1, max = rho;
+
+    // Step 9
+    uint64_t maxIterations = ceil(log2(rho));
+
+#ifdef DEBUG_ENUM_BINARY
+    std::cout << fmt::format("start2: maxIterations {}", maxIterations) << std::endl;
+#endif
+
+    for (uint64_t j = 1; j <= maxIterations; j++) {
+        // 9a
+        median = ceil(static_cast<double>(max - min) / 2) + min;
+
+
+        // 9b
+        BSONObj doc = reader->getById(collectionT::generateId(tagToken, median + lambda));
+
+#ifdef DEBUG_ENUM_BINARY
+        std::cout << fmt::format("search_stat: min: {}, median: {}, max: {}, i: {}, doc: {}",
+                                 min,
+                                 median,
+                                 max,
+                                 i,
+                                 doc.toString())
+                  << std::endl;
+#endif
+
+        // 9c
+        if (!doc.isEmpty()) {
+            // 9 c i
+            min = median;
+
+            // 9 c ii
+            if (j == maxIterations) {
+                i = min + lambda;
+            }
+            // 9d
+        } else {
+            // 9 d i
+            max = median;
+
+            // 9 d ii
+            // Binary search has ended without finding a document, check for the first document
+            // explicitly
+            if (j == maxIterations && min == 1) {
+                // 9 d ii A
+                BSONObj doc = reader->getById(collectionT::generateId(tagToken, 1 + lambda));
+                // 9 d ii B
+                if (!doc.isEmpty()) {
+                    i = 1 + lambda;
+                }
+            } else if (j == maxIterations && min != 1) {
+                i = min + lambda;
+            }
+        }
+    }
+
+    return i;
+}
+
+
 }  // namespace
 
 
@@ -228,4 +538,129 @@ ECCTwiceDerivedValueToken FLETwiceDerivedTokenGenerator::generateECCTwiceDerived
 }
 
 
+PrfBlock ESCCollection::generateId(ESCTwiceDerivedTagToken tagToken,
+                                   boost::optional<uint64_t> index) {
+    if (index.has_value()) {
+        return prf(tagToken.data, kESCNonNullId, index.value());
+    } else {
+        return prf(tagToken.data, kESCNullId, 0);
+    }
+}
+
+BSONObj ESCCollection::generateNullDocument(ESCTwiceDerivedTagToken tagToken,
+                                            ESCTwiceDerivedValueToken valueToken,
+                                            uint64_t pos,
+                                            uint64_t count) {
+    auto block = ESCCollection::generateId(tagToken, boost::none);
+
+    auto swCipherText = packAndEncrypt(std::tie(pos, count), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+
+BSONObj ESCCollection::generateInsertDocument(ESCTwiceDerivedTagToken tagToken,
+                                              ESCTwiceDerivedValueToken valueToken,
+                                              uint64_t index,
+                                              uint64_t count) {
+    auto block = ESCCollection::generateId(tagToken, index);
+
+    auto swCipherText = packAndEncrypt(std::tie(KESCInsertRecordValue, count), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+
+BSONObj ESCCollection::generatePositionalDocument(ESCTwiceDerivedTagToken tagToken,
+                                                  ESCTwiceDerivedValueToken valueToken,
+                                                  uint64_t index,
+                                                  uint64_t pos,
+                                                  uint64_t count) {
+    auto block = ESCCollection::generateId(tagToken, index);
+
+    auto swCipherText = packAndEncrypt(std::tie(pos, count), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+
+BSONObj ESCCollection::generateCompactionPlaceholderDocument(ESCTwiceDerivedTagToken tagToken,
+                                                             ESCTwiceDerivedValueToken valueToken,
+                                                             uint64_t index) {
+    auto block = ESCCollection::generateId(tagToken, index);
+
+    auto swCipherText = packAndEncrypt(
+        std::tie(kESCompactionRecordValue, kESCompactionRecordCountPlaceholder), valueToken);
+    uassertStatusOK(swCipherText);
+
+    BSONObjBuilder builder;
+    toBinData(kId, block, &builder);
+    toBinData(kValue, swCipherText.getValue(), &builder);
+
+    return builder.obj();
+}
+
+StatusWith<ESCNullDocument> ESCCollection::decryptNullDocument(ESCTwiceDerivedValueToken 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 ESCNullDocument{std::get<0>(value), std::get<1>(value)};
+}
+
+
+StatusWith<ESCDocument> ESCCollection::decryptDocument(ESCTwiceDerivedValueToken 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 ESCDocument{
+        std::get<0>(value) == kESCompactionRecordValue, std::get<0>(value), std::get<1>(value)};
+}
+
+
+uint64_t ESCCollection::emuBinary(FLEStateCollectionReader* reader,
+                                  ESCTwiceDerivedTagToken tagToken,
+                                  ESCTwiceDerivedValueToken valueToken) {
+    return emuBinaryCommon<ESCCollection, ESCTwiceDerivedTagToken, ESCTwiceDerivedValueToken>(
+        reader, tagToken, valueToken);
+}
+
 }  // namespace mongo