SERVER-41277 Create FLE Data Framing Classes

8 files changed, 377 insertions, 190 deletions

diff --git a/src/mongo/base/data_range.h b/src/mongo/base/data_range.h
index 9cf850838c4..d8b50807fbd 100644
--- a/src/mongo/base/data_range.h
+++ b/src/mongo/base/data_range.h
@@ -117,8 +117,9 @@ public:
     ConstDataRange(const ByteLike (&arr)[N], std::ptrdiff_t debug_offset = 0)
         : ConstDataRange(arr, N, debug_offset) {}
 
-    const byte_type* data() const noexcept {
-        return _begin;
+    template <typename ByteLike = byte_type>
+    const ByteLike* data() const noexcept {
+        return reinterpret_cast<const ByteLike*>(_begin);
     }
 
     size_t length() const noexcept {
diff --git a/src/mongo/crypto/aead_encryption.cpp b/src/mongo/crypto/aead_encryption.cpp
index b5e0ae4ce1c..27d3c0e5d77 100644
--- a/src/mongo/crypto/aead_encryption.cpp
+++ b/src/mongo/crypto/aead_encryption.cpp
@@ -41,7 +41,6 @@ namespace mongo {
 namespace crypto {
 
 namespace {
-constexpr size_t kHmacOutSize = 32;
 constexpr size_t kIVSize = 16;
 
 // AssociatedData can be 2^24 bytes but since there needs to be room for the ciphertext in the
@@ -183,23 +182,41 @@ size_t aeadCipherOutputLength(size_t plainTextLen) {
     return aesOutLen + kHmacOutSize;
 }
 
-StatusWith<size_t> aeadGetMaximumPlainTextLength(size_t cipherTextLen) {
-    if (cipherTextLen > (aesCBCIVSize + kHmacOutSize)) {
-        return cipherTextLen - aesCBCIVSize - kHmacOutSize;
+Status aeadEncryptLocalKMS(const SymmetricKey& key,
+                           const ConstDataRange in,
+                           uint8_t* out,
+                           size_t outLen) {
+    if (key.getKeySize() != kFieldLevelEncryptionKeySize) {
+        return Status(ErrorCodes::BadValue,
+                      "AEAD encryption key is the incorrect length. "
+                      "Must be 96 bytes.");
     }
 
-    return Status(ErrorCodes::BadValue, "Invalid cipher text length");
-}
+    // According to the rfc on AES encryption, the associatedDataLength is defined as the
+    // number of bits in associatedData in BigEndian format. This is what the code segment
+    // below describes.
+    // RFC: (https://tools.ietf.org/html/draft-mcgrew-aead-aes-cbc-hmac-sha2-01#section-2.1)
+    std::array<uint8_t, sizeof(uint64_t)> dataLenBitsEncodedStorage;
+    DataRange dataLenBitsEncoded(dataLenBitsEncodedStorage);
+    dataLenBitsEncoded.write<BigEndian<uint64_t>>(static_cast<uint64_t>(0));
 
-Status aeadEncrypt(const SymmetricKey& key,
-                   const uint8_t* in,
-                   const size_t inLen,
-                   const uint8_t* associatedData,
-                   const uint64_t associatedDataLen,
-                   uint8_t* out,
-                   size_t outLen) {
+    ConstDataRange keyCDR(key.getKey(), kAeadAesHmacKeySize);
 
-    if (associatedDataLen >= kMaxAssociatedDataLength) {
+    return aeadEncryptWithIV(keyCDR,
+                             in.data<uint8_t>(),
+                             in.length(),
+                             nullptr,
+                             0,
+                             nullptr,
+                             0,
+                             dataLenBitsEncoded,
+                             out,
+                             outLen);
+}
+
+Status aeadEncryptDataFrame(FLEEncryptionFrame& dataframe) {
+    auto associatedData = dataframe.getAssociatedData();
+    if (associatedData.length() >= kMaxAssociatedDataLength) {
         return Status(ErrorCodes::BadValue,
                       str::stream()
                           << "AssociatedData for encryption is too large. Cannot be larger than "
@@ -212,49 +229,51 @@ Status aeadEncrypt(const SymmetricKey& key,
     // RFC: (https://tools.ietf.org/html/draft-mcgrew-aead-aes-cbc-hmac-sha2-01#section-2.1)
     std::array<uint8_t, sizeof(uint64_t)> dataLenBitsEncodedStorage;
     DataRange dataLenBitsEncoded(dataLenBitsEncodedStorage);
-    dataLenBitsEncoded.write<BigEndian<uint64_t>>(associatedDataLen * 8);
+    dataLenBitsEncoded.write<BigEndian<uint64_t>>(static_cast<uint64_t>(associatedData.length()) *
+                                                  8);
 
-    ConstDataRange aeadKey(key.getKey(), kAeadAesHmacKeySize);
 
-    if (key.getKeySize() != kFieldLevelEncryptionKeySize) {
+    auto key = dataframe.getKey();
+
+    auto plaintext = dataframe.getPlaintext();
+
+    if (key->getKeySize() != kFieldLevelEncryptionKeySize) {
         return Status(ErrorCodes::BadValue, "Invalid key size.");
     }
 
-    if (in == nullptr || !in) {
+    if (plaintext.data() == nullptr) {
         return Status(ErrorCodes::BadValue, "Invalid AEAD plaintext input.");
     }
 
-    if (key.getAlgorithm() != aesAlgorithm) {
+    if (key->getAlgorithm() != aesAlgorithm) {
         return Status(ErrorCodes::BadValue, "Invalid algorithm for key.");
     }
 
-    ConstDataRange hmacCDR(nullptr, 0);
+    ConstDataRange iv(nullptr, 0);
     SHA512Block hmacOutput;
-    if (associatedData != nullptr &&
-        static_cast<int>(associatedData[0]) ==
-            FleAlgorithmInt_serializer(FleAlgorithmInt::kDeterministic)) {
-        const uint8_t* ivKey = key.getKey() + kAeadAesHmacKeySize;
-        hmacOutput = SHA512Block::computeHmac(ivKey,
-                                              sym256KeySize,
-                                              {ConstDataRange(associatedData, associatedDataLen),
-                                               dataLenBitsEncoded,
-                                               ConstDataRange(in, inLen)});
+
+    if (dataframe.getFLEAlgorithmType() == FleAlgorithmInt::kDeterministic) {
+        const uint8_t* ivKey = key->getKey() + kAeadAesHmacKeySize;
+        hmacOutput = SHA512Block::computeHmac(
+            ivKey, sym256KeySize, {associatedData, dataLenBitsEncoded, plaintext});
 
         static_assert(SHA512Block::kHashLength >= kIVSize,
                       "Invalid AEAD parameters. Generated IV too short.");
 
-        hmacCDR = ConstDataRange(hmacOutput.data(), kIVSize);
+        iv = ConstDataRange(hmacOutput.data(), kIVSize);
     }
+
+    ConstDataRange aeadKey(key->getKey(), kAeadAesHmacKeySize);
     return aeadEncryptWithIV(aeadKey,
-                             in,
-                             inLen,
-                             reinterpret_cast<const uint8_t*>(hmacCDR.data()),
-                             hmacCDR.length(),
-                             associatedData,
-                             associatedDataLen,
+                             plaintext.data<uint8_t>(),
+                             plaintext.length(),
+                             iv.data<uint8_t>(),
+                             iv.length(),
+                             associatedData.data<uint8_t>(),
+                             associatedData.length(),
                              dataLenBitsEncoded,
-                             out,
-                             outLen);
+                             dataframe.getCiphertextMutable(),
+                             dataframe.getDataLength());
 }
 
 Status aeadEncryptWithIV(ConstDataRange key,
@@ -320,8 +339,7 @@ Status aeadEncryptWithIV(ConstDataRange key,
 }
 
 Status aeadDecrypt(const SymmetricKey& key,
-                   const uint8_t* cipherText,
-                   const size_t cipherLen,
+                   ConstDataRange ciphertext,
                    const uint8_t* associatedData,
                    const uint64_t associatedDataLen,
                    uint8_t* out,
@@ -330,15 +348,16 @@ Status aeadDecrypt(const SymmetricKey& key,
         return Status(ErrorCodes::BadValue, "Invalid key size.");
     }
 
-    if (!(cipherText && out)) {
+    if (!out) {
         return Status(ErrorCodes::BadValue, "Invalid AEAD parameters.");
     }
 
-    if (cipherLen < kHmacOutSize) {
+    if (ciphertext.length() < kHmacOutSize) {
         return Status(ErrorCodes::BadValue, "Ciphertext is not long enough.");
     }
 
-    size_t expectedMaximumPlainTextSize = uassertStatusOK(aeadGetMaximumPlainTextLength(cipherLen));
+    size_t expectedMaximumPlainTextSize =
+        uassertStatusOK(aeadGetMaximumPlainTextLength(ciphertext.length()));
     if ((*outLen) != expectedMaximumPlainTextSize) {
         return Status(ErrorCodes::BadValue, "Output buffer must be as long as the cipherText.");
     }
@@ -353,7 +372,7 @@ Status aeadDecrypt(const SymmetricKey& key,
     const uint8_t* macKey = key.getKey();
     const uint8_t* encKey = key.getKey() + sym256KeySize;
 
-    size_t aesLen = cipherLen - kHmacOutSize;
+    size_t aesLen = ciphertext.length() - kHmacOutSize;
 
     // According to the rfc on AES encryption, the associatedDataLength is defined as the
     // number of bits in associatedData in BigEndian format. This is what the code segment
@@ -366,11 +385,11 @@ Status aeadDecrypt(const SymmetricKey& key,
         SHA512Block::computeHmac(macKey,
                                  sym256KeySize,
                                  {ConstDataRange(associatedData, associatedDataLen),
-                                  ConstDataRange(cipherText, aesLen),
+                                  ConstDataRange(ciphertext.data(), aesLen),
                                   dataLenBitsEncoded});
 
     if (consttimeMemEqual(reinterpret_cast<const unsigned char*>(hmacOutput.data()),
-                          reinterpret_cast<const unsigned char*>(cipherText + aesLen),
+                          ciphertext.data<const unsigned char>() + aesLen,
                           kHmacOutSize) == false) {
         return Status(ErrorCodes::BadValue, "HMAC data authentication failed.");
     }
@@ -378,7 +397,7 @@ Status aeadDecrypt(const SymmetricKey& key,
     SymmetricKey symEncKey(encKey, sym256KeySize, aesAlgorithm, key.getKeyId(), 1);
 
     auto sDecrypt =
-        _aesDecrypt(symEncKey, ConstDataRange(cipherText, aesLen), out, *outLen, outLen);
+        _aesDecrypt(symEncKey, ConstDataRange(ciphertext.data(), aesLen), out, aesLen, outLen);
     if (!sDecrypt.isOK()) {
         return sDecrypt;
     }
@@ -386,5 +405,27 @@ Status aeadDecrypt(const SymmetricKey& key,
     return Status::OK();
 }
 
+Status aeadDecryptDataFrame(FLEDecryptionFrame& dataframe) {
+    auto ciphertext = dataframe.getCiphertext();
+    auto associatedData = dataframe.getAssociatedData();
+    auto& plaintext = dataframe.getPlaintextMutable();
+    size_t outLen = plaintext.size();
+    uassertStatusOK(aeadDecrypt(*dataframe.getKey(),
+                                ciphertext,
+                                associatedData.data<uint8_t>(),
+                                associatedData.length(),
+                                plaintext.data(),
+                                &outLen));
+    plaintext.resize(outLen);
+    return Status::OK();
+}
+
+Status aeadDecryptLocalKMS(const SymmetricKey& key,
+                           const ConstDataRange cipher,
+                           uint8_t* out,
+                           size_t* outLen) {
+    return aeadDecrypt(key, cipher, nullptr, 0, out, outLen);
+}
+
 }  // namespace crypto
 }  // namespace mongo
diff --git a/src/mongo/crypto/aead_encryption.h b/src/mongo/crypto/aead_encryption.h
index 134e83bff08..0b20291dd88 100644
--- a/src/mongo/crypto/aead_encryption.h
+++ b/src/mongo/crypto/aead_encryption.h
@@ -32,6 +32,9 @@
 #include <cstddef>
 #include <cstdint>
 
+#include "mongo/crypto/fle_data_frames.h"
+#include "mongo/shell/kms_gen.h"
+
 #include "mongo/base/data_view.h"
 #include "mongo/base/status.h"
 #include "mongo/crypto/symmetric_key.h"
@@ -51,26 +54,27 @@ constexpr size_t kAeadAesHmacKeySize = 64;
  */
 size_t aeadCipherOutputLength(size_t plainTextLen);
 
-
 /**
- * Returns the length of the plaintext output given the ciphertext length. Only for AEAD.
+ * Encrypts a dataframe object following the AEAD_AES_256_CBC_HMAC_SHA_512 encryption
+ * algorithm. Used for field level encryption.
  */
-StatusWith<size_t> aeadGetMaximumPlainTextLength(size_t cipherTextLen);
-
+Status aeadEncryptDataFrame(FLEEncryptionFrame& dataframe);
 
 /**
- * Encrypts the plaintext using following the AEAD_AES_256_CBC_HMAC_SHA_512 encryption
- * algorithm. Writes output to out.
+ * Decrypts a dataframe object following the AEAD_AES_256_CBC_HMAC_SHA_512 decryption
+ * algorithm. Used for field level encryption.
  */
-Status aeadEncrypt(const SymmetricKey& key,
-                   const uint8_t* in,
-                   const size_t inLen,
-                   const uint8_t* associatedData,
-                   const uint64_t associatedDataLen,
-                   uint8_t* out,
-                   size_t outLen);
+Status aeadDecryptDataFrame(FLEDecryptionFrame& dataframe);
 
 /**
+ * Uses AEAD_AES_256_CBC_HMAC_SHA_512 encryption to encrypt a local datakey.
+ * Writes output to out.
+ */
+Status aeadEncryptLocalKMS(const SymmetricKey& key,
+                           const ConstDataRange in,
+                           uint8_t* out,
+                           size_t outLen);
+/**
  * Internal calls for the aeadEncryption algorithm. Only used for testing.
  */
 Status aeadEncryptWithIV(ConstDataRange key,
@@ -85,16 +89,23 @@ Status aeadEncryptWithIV(ConstDataRange key,
                          size_t outLen);
 
 /**
- * Decrypts the cipherText using AEAD_AES_256_CBC_HMAC_SHA_512 decryption. Writes output
- * to out.
+ * Internal call for the aeadDecryption algorithm. Only used for testing.
  */
 Status aeadDecrypt(const SymmetricKey& key,
-                   const uint8_t* cipherText,
-                   const size_t cipherLen,
+                   ConstDataRange ciphertext,
                    const uint8_t* associatedData,
                    const uint64_t associatedDataLen,
                    uint8_t* out,
                    size_t* outLen);
 
+/**
+ * Decrypts the cipherText using AEAD_AES_256_CBC_HMAC_SHA_512 decryption. Writes output
+ * to out.
+ */
+Status aeadDecryptLocalKMS(const SymmetricKey& key,
+                           const ConstDataRange cipher,
+                           uint8_t* out,
+                           size_t* outLen);
+
 }  // namespace crypto
 }  // namespace mongo
diff --git a/src/mongo/crypto/aead_encryption_test.cpp b/src/mongo/crypto/aead_encryption_test.cpp
index f0417cec23e..fc7b4e6de9e 100644
--- a/src/mongo/crypto/aead_encryption_test.cpp
+++ b/src/mongo/crypto/aead_encryption_test.cpp
@@ -29,6 +29,8 @@
 
 #include <algorithm>
 
+#include "mongo/base/data_range.h"
+
 #include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 
@@ -139,8 +141,7 @@ TEST(AEAD, EncryptAndDecrypt) {
     std::array<uint8_t, 144> plainText = {};
     size_t plainTextDecryptLen = 144;
     ASSERT_OK(crypto::aeadDecrypt(key,
-                                  cryptoBuffer.data(),
-                                  cryptoBuffer.size(),
+                                  ConstDataRange(cryptoBuffer),
                                   associatedData.data(),
                                   dataLen,
                                   plainText.data(),
@@ -152,8 +153,7 @@ TEST(AEAD, EncryptAndDecrypt) {
     (*aesVector)[0] ^= 1;
     key = SymmetricKey(aesVector, aesAlgorithm, "aeadEncryptDecryptTest");
     ASSERT_NOT_OK(crypto::aeadDecrypt(key,
-                                      cryptoBuffer.data(),
-                                      cryptoBuffer.size(),
+                                      ConstDataRange(cryptoBuffer),
                                       associatedData.data(),
                                       dataLen,
                                       plainText.data(),
diff --git a/src/mongo/crypto/fle_data_frames.h b/src/mongo/crypto/fle_data_frames.h
new file mode 100644
index 00000000000..a72b31dceb1
--- /dev/null
+++ b/src/mongo/crypto/fle_data_frames.h
@@ -0,0 +1,191 @@
+/**
+ *    Copyright (C) 2019-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/base/data_range.h"
+#include "mongo/crypto/symmetric_crypto.h"
+#include "mongo/db/matcher/schema/encrypt_schema_gen.h"
+#include "mongo/shell/kms_gen.h"
+#include "mongo/util/uuid.h"
+
+namespace mongo {
+constexpr int kAssociatedDataLength = 18;
+constexpr size_t kHmacOutSize = 32;
+
+/**
+ * Returns the length of the plaintext output given the ciphertext length. Only for AEAD.
+ */
+inline StatusWith<size_t> aeadGetMaximumPlainTextLength(size_t cipherTextLen) {
+    if (cipherTextLen > (crypto::aesCBCIVSize + kHmacOutSize)) {
+        return cipherTextLen - crypto::aesCBCIVSize - kHmacOutSize;
+    }
+
+    return Status(ErrorCodes::BadValue, "Invalid cipher text length");
+}
+
+/**
+ * This class is a helper for encryption. It holds a ConstDataRange over the
+ * plaintext to be encrypted and owns a buffer where the BinData subtype 6 is
+ * written out to. The encrypt function can only be called after the constructor
+ * to initialize the plaintext, the associated data, and the key has been called.
+ */
+class FLEEncryptionFrame {
+
+public:
+    FLEEncryptionFrame(std::shared_ptr<SymmetricKey> key,
+                       FleAlgorithmInt algorithm,
+                       UUID uuid,
+                       BSONType type,
+                       ConstDataRange plaintext,
+                       size_t cipherLength)
+        : _key(std::move(key)), _plaintext(plaintext) {
+        // As per the description of the encryption algorithm for FLE, the
+        // associated data is constructed of the following -
+        // associatedData[0] = the FleAlgorithmEnum
+        //      - either a 1 or a 2 depending on whether the iv is provided.
+        // associatedData[1-16] = the uuid in bytes
+        // associatedData[17] = the bson type
+        _data.resize(kAssociatedDataLength + cipherLength);
+        _data[0] = FleAlgorithmInt_serializer(algorithm);
+        auto uuidCDR = uuid.toCDR();
+        invariant(uuidCDR.length() == 16);
+        std::copy(
+            uuidCDR.data<uint8_t>(), uuidCDR.data<uint8_t>() + uuidCDR.length(), _data.begin() + 1);
+        _data[17] = static_cast<uint8_t>(type);
+    }
+
+    FLEEncryptionFrame() : _plaintext(ConstDataRange(nullptr, 0)){};
+
+    const ConstDataRange get() const& {
+        return ConstDataRange(_data);
+    }
+
+    std::shared_ptr<SymmetricKey> getKey() const {
+        return _key;
+    }
+
+    const ConstDataRange getPlaintext() const {
+        return _plaintext;
+    }
+
+    const ConstDataRange getAssociatedData() const {
+        return ConstDataRange(_data.data(), kAssociatedDataLength);
+    }
+
+    uint8_t* getCiphertextMutable() & {
+        invariant(_data.size() > kAssociatedDataLength);
+        return _data.data() + kAssociatedDataLength;
+    }
+
+    uint8_t* getCiphertextMutable() && = delete;
+
+    FleAlgorithmInt getFLEAlgorithmType() {
+        return FleAlgorithmInt_parse(IDLParserErrorContext("root"), _data[0]);
+    }
+
+    size_t getDataLength() const {
+        return _data.size() - kAssociatedDataLength;
+    }
+
+private:
+    std::shared_ptr<SymmetricKey> _key;
+    std::vector<uint8_t> _data;
+    ConstDataRange _plaintext;
+};
+
+/**
+ * This is a helper class for decryption. It has a ConstDataRange over a
+ * vector owned by the instantiator of this class and allows
+ * containing the plaintext object after it has been decrypted.
+ */
+class FLEDecryptionFrame {
+public:
+    FLEDecryptionFrame(ConstDataRange data) : _data(data) {
+        uassert(ErrorCodes::BadValue,
+                "Ciphertext blob too small",
+                _data.length() > kAssociatedDataLength);
+        uassert(ErrorCodes::BadValue,
+                "Ciphertext blob algorithm unknown",
+                (getFLEAlgorithmType() == FleAlgorithmInt::kDeterministic ||
+                 getFLEAlgorithmType() == FleAlgorithmInt::kRandom));
+
+        _plaintext.resize(
+            uassertStatusOK(aeadGetMaximumPlainTextLength(data.length() - kAssociatedDataLength)));
+    }
+
+    void setKey(std::shared_ptr<SymmetricKey> key) {
+        _key = key;
+    }
+
+    std::shared_ptr<SymmetricKey> getKey() {
+        return _key;
+    }
+
+    const ConstDataRange getAssociatedData() const {
+        return ConstDataRange(_data.data(), kAssociatedDataLength);
+    }
+
+    UUID getUUID() const {
+        auto uuid = UUID::fromCDR(ConstDataRange(_data.data() + 1, UUID::kNumBytes));
+        return uuid;
+    }
+
+    ConstDataRange getCiphertext() const {
+        return ConstDataRange(_data.data() + kAssociatedDataLength, getDataLength());
+    }
+
+    ConstDataRange getPlaintext() const {
+        return ConstDataRange(_plaintext);
+    }
+
+    std::vector<uint8_t>& getPlaintextMutable() {
+        return _plaintext;
+    }
+
+    uint8_t getBSONType() {
+        return *(_data.data<uint8_t>() + 17);
+    }
+
+private:
+    FleAlgorithmInt getFLEAlgorithmType() const {
+        return FleAlgorithmInt_parse(IDLParserErrorContext("root"), *_data.data<uint8_t>());
+    }
+
+    size_t getDataLength() const {
+        return _data.length() - kAssociatedDataLength;
+    }
+
+
+    std::shared_ptr<SymmetricKey> _key;
+    ConstDataRange _data;
+    std::vector<uint8_t> _plaintext;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/shell/encrypted_dbclient_base.cpp b/src/mongo/shell/encrypted_dbclient_base.cpp
index be82f6b97bb..d8530c49ec5 100644
--- a/src/mongo/shell/encrypted_dbclient_base.cpp
+++ b/src/mongo/shell/encrypted_dbclient_base.cpp
@@ -36,6 +36,7 @@
 #include "mongo/bson/bson_depth.h"
 #include "mongo/client/dbclient_base.h"
 #include "mongo/crypto/aead_encryption.h"
+#include "mongo/crypto/fle_data_frames.h"
 #include "mongo/crypto/symmetric_crypto.h"
 #include "mongo/db/client.h"
 #include "mongo/db/commands.h"
@@ -183,27 +184,15 @@ void EncryptedDBClientBase::encryptMarking(const BSONObj& elem,
 void EncryptedDBClientBase::decryptPayload(ConstDataRange data,
                                            BSONObjBuilder* builder,
                                            StringData elemName) {
+    invariant(builder);
     uassert(ErrorCodes::BadValue, "Invalid decryption blob", data.length() > kAssociatedDataLength);
-    ConstDataRange uuidCdr = ConstDataRange(data.data() + 1, 16);
-    UUID uuid = UUID::fromCDR(uuidCdr);
-
-    auto key = getDataKey(uuid);
-    std::vector<uint8_t> out(uassertStatusOK(
-        crypto::aeadGetMaximumPlainTextLength(data.length() - kAssociatedDataLength)));
-    size_t outLen = out.size();
-
-    uassertStatusOK(
-        crypto::aeadDecrypt(*key,
-                            reinterpret_cast<const uint8_t*>(data.data() + kAssociatedDataLength),
-                            data.length() - kAssociatedDataLength,
-                            reinterpret_cast<const uint8_t*>(data.data()),
-                            kAssociatedDataLength,
-                            out.data(),
-                            &outLen));
+
+    FLEDecryptionFrame dataFrame = createDecryptionFrame(data);
+    auto plaintext = dataFrame.getPlaintext();
 
     // extract type byte
-    const uint8_t bsonType = static_cast<const uint8_t>(*(data.data() + 17));
-    BSONObj decryptedObj = validateBSONElement(ConstDataRange(out.data(), outLen), bsonType);
+    const uint8_t bsonType = dataFrame.getBSONType();
+    BSONObj decryptedObj = validateBSONElement(plaintext, bsonType);
     if (bsonType == BSONType::Object) {
         builder->append(elemName, decryptedObj);
     } else {
@@ -361,7 +350,7 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
     UUID uuid = UUID::fromCDR(ConstDataRange(binData.data(), binData.size()));
     BSONType bsonType = BSONType::EOO;
 
-    BufBuilder plaintext;
+    BufBuilder plaintextBuilder;
     if (args.get(1).isObject()) {
         JS::RootedObject rootedObj(cx, &args.get(1).toObject());
         auto jsclass = JS_GetClass(rootedObj);
@@ -374,7 +363,7 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
             // If it is a JS Object, then we can extract all the information by simply calling
             // ValueWriter.toBSON and setting the type bit, which is what is happening below.
             BSONObj valueObj = mozjs::ValueWriter(cx, args.get(1)).toBSON();
-            plaintext.appendBuf(valueObj.objdata(), valueObj.objsize());
+            plaintextBuilder.appendBuf(valueObj.objdata(), valueObj.objsize());
             if (strcmp(jsclass->name, "Array") == 0) {
                 bsonType = BSONType::Array;
             } else {
@@ -408,11 +397,9 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
             // If it is one of our Mongo defined types, then we have to use the ValueWriter
             // writeThis function, which takes in a set of WriteFieldRecursionFrames (setting
             // a limit on how many times we can recursively dig into an object's nested
-            // structure)
-            // and writes the value out to a BSONObjBuilder. We can then extract that
-            // information
-            // from the object by building it and pulling out the first element, which is the
-            // object we are trying to get.
+            // structure) and writes the value out to a BSONObjBuilder. We can then extract
+            // that information from the object by building it and pulling out the first
+            // element, which is the object we are trying to get.
             mozjs::ObjectWrapper::WriteFieldRecursionFrames frames;
             frames.emplace(cx, rootedObj.get(), nullptr, StringData{});
             BSONObjBuilder builder;
@@ -421,7 +408,7 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
             BSONObj object = builder.obj();
             auto elem = object.getField("value"_sd);
 
-            plaintext.appendBuf(elem.value(), elem.valuesize());
+            plaintextBuilder.appendBuf(elem.value(), elem.valuesize());
             bsonType = elem.type();
         }
 
@@ -431,8 +418,8 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
             uasserted(ErrorCodes::BadValue, "Plaintext string to encrypt too long.");
         }
 
-        plaintext.appendNum(static_cast<uint32_t>(valueStr.size() + 1));
-        plaintext.appendStr(valueStr, true);
+        plaintextBuilder.appendNum(static_cast<uint32_t>(valueStr.size() + 1));
+        plaintextBuilder.appendStr(valueStr, true);
         bsonType = BSONType::String;
 
     } else if (args.get(1).isNumber()) {
@@ -441,7 +428,7 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
                 algorithm != FleAlgorithmInt::kDeterministic);
 
         double valueNum = mozjs::ValueWriter(cx, args.get(1)).toNumber();
-        plaintext.appendNum(valueNum);
+        plaintextBuilder.appendNum(valueNum);
         bsonType = BSONType::NumberDouble;
     } else if (args.get(1).isBoolean()) {
         uassert(ErrorCodes::BadValue,
@@ -450,22 +437,23 @@ void EncryptedDBClientBase::encrypt(mozjs::MozJSImplScope* scope,
 
         bool boolean = mozjs::ValueWriter(cx, args.get(1)).toBoolean();
         if (boolean) {
-            plaintext.appendChar(0x01);
+            plaintextBuilder.appendChar(0x01);
         } else {
-            plaintext.appendChar(0x00);
+            plaintextBuilder.appendChar(0x00);
         }
         bsonType = BSONType::Bool;
     } else {
         uasserted(ErrorCodes::BadValue, "Cannot encrypt valuetype provided.");
     }
-    ConstDataRange plaintextRange(plaintext.buf(), plaintext.len());
 
-    auto key = getDataKey(uuid);
-    std::vector<uint8_t> fleBlob =
-        encryptWithKey(uuid, key, plaintextRange, bsonType, FleAlgorithmInt_serializer(algorithm));
+    ConstDataRange plaintext(plaintextBuilder.buf(), plaintextBuilder.len());
+
+    FLEEncryptionFrame encryptionFrame =
+        createEncryptionFrame(getDataKey(uuid), algorithm, uuid, bsonType, plaintext);
 
     // Prepare the return value
-    std::string blobStr = base64::encode(reinterpret_cast<char*>(fleBlob.data()), fleBlob.size());
+    ConstDataRange ciphertextBlob(encryptionFrame.get());
+    std::string blobStr = base64::encode(ciphertextBlob.data(), ciphertextBlob.length());
     JS::AutoValueArray<2> arr(cx);
 
     arr[0].setInt32(BinDataType::Encrypt);
@@ -481,43 +469,20 @@ void EncryptedDBClientBase::decrypt(mozjs::MozJSImplScope* scope,
             "decrypt argument must be a BinData subtype Encrypt object",
             args.get(0).isObject());
 
-    if (!scope->getProto<mozjs::BinDataInfo>().instanceOf(args.get(0))) {
-        uasserted(ErrorCodes::BadValue,
-                  "decrypt argument must be a BinData subtype Encrypt object");
-    }
+    uassert(ErrorCodes::BadValue,
+            "decrypt argument must be a BinData subtype Encrypt object",
+            scope->getProto<mozjs::BinDataInfo>().instanceOf(args.get(0)));
 
     JS::RootedObject obj(cx, &args.get(0).get().toObject());
-    std::vector<uint8_t> binData = getBinDataArg(scope, cx, args, 0, BinDataType::Encrypt);
+    std::vector<uint8_t> data = getBinDataArg(scope, cx, args, 0, BinDataType::Encrypt);
 
-    uassert(
-        ErrorCodes::BadValue, "Ciphertext blob too small", binData.size() > kAssociatedDataLength);
-    uassert(ErrorCodes::BadValue,
-            "Ciphertext blob algorithm unknown",
-            (FleAlgorithmInt(binData[0]) == FleAlgorithmInt::kDeterministic ||
-             FleAlgorithmInt(binData[0]) == FleAlgorithmInt::kRandom));
-
-    ConstDataRange uuidCdr = ConstDataRange(&binData[1], UUID::kNumBytes);
-    UUID uuid = UUID::fromCDR(uuidCdr);
-
-    auto key = getDataKey(uuid);
-    std::vector<uint8_t> out(uassertStatusOK(
-        crypto::aeadGetMaximumPlainTextLength(binData.size() - kAssociatedDataLength)));
-    size_t outLen = out.size();
-
-    auto decryptStatus = crypto::aeadDecrypt(*key,
-                                             &binData[kAssociatedDataLength],
-                                             binData.size() - kAssociatedDataLength,
-                                             &binData[0],
-                                             kAssociatedDataLength,
-                                             out.data(),
-                                             &outLen);
-    if (!decryptStatus.isOK()) {
-        uasserted(decryptStatus.code(), decryptStatus.reason());
-    }
+    ConstDataRange ciphertextBlob(data);
 
-    uint8_t bsonType = binData[17];
+    FLEDecryptionFrame dataFrame = createDecryptionFrame(ciphertextBlob);
+
+    const uint8_t bsonType = dataFrame.getBSONType();
     BSONObj parent;
-    BSONObj decryptedObj = validateBSONElement(ConstDataRange(out.data(), outLen), bsonType);
+    BSONObj decryptedObj = validateBSONElement(dataFrame.getPlaintext(), bsonType);
     if (bsonType == BSONType::Object) {
         mozjs::ValueReader(cx, args.rval()).fromBSON(decryptedObj, &parent, true);
     } else {
@@ -570,6 +535,26 @@ bool EncryptedDBClientBase::isMongos() const {
     return _conn->isMongos();
 }
 
+FLEEncryptionFrame EncryptedDBClientBase::createEncryptionFrame(std::shared_ptr<SymmetricKey> key,
+                                                                FleAlgorithmInt algorithm,
+                                                                UUID uuid,
+                                                                BSONType type,
+                                                                ConstDataRange plaintext) {
+
+    auto cipherLength = crypto::aeadCipherOutputLength(plaintext.length());
+    FLEEncryptionFrame dataframe(key, algorithm, uuid, type, plaintext, cipherLength);
+    uassertStatusOK(crypto::aeadEncryptDataFrame(dataframe));
+    return dataframe;
+}
+
+FLEDecryptionFrame EncryptedDBClientBase::createDecryptionFrame(ConstDataRange data) {
+    auto frame = FLEDecryptionFrame(data);
+    auto key = getDataKey(frame.getUUID());
+    frame.setKey(key);
+    uassertStatusOK(crypto::aeadDecryptDataFrame(frame));
+    return frame;
+}
+
 NamespaceString EncryptedDBClientBase::getCollectionNS() {
     JS::RootedValue fullNameRooted(_cx);
     JS::RootedObject collectionRooted(_cx, &_collection.get().toObject());
@@ -652,36 +637,6 @@ std::shared_ptr<SymmetricKey> EncryptedDBClientBase::getDataKeyFromDisk(const UU
         std::move(decryptedKey), crypto::aesAlgorithm, "kms_encryption");
 }
 
-std::vector<uint8_t> EncryptedDBClientBase::encryptWithKey(UUID uuid,
-                                                           const std::shared_ptr<SymmetricKey>& key,
-                                                           ConstDataRange plaintext,
-                                                           BSONType bsonType,
-                                                           int32_t algorithm) {
-    // As per the description of the encryption algorithm for FLE, the
-    // associated data is constructed of the following -
-    // associatedData[0] = the FleAlgorithmEnum
-    //      - either a 1 or a 2 depending on whether the iv is provided.
-    // associatedData[1-16] = the uuid in bytes
-    // associatedData[17] = the bson type
-
-    ConstDataRange uuidCdr = uuid.toCDR();
-    uint64_t outputLength = crypto::aeadCipherOutputLength(plaintext.length());
-    std::vector<uint8_t> outputBuffer(kAssociatedDataLength + outputLength);
-    outputBuffer[0] = static_cast<uint8_t>(algorithm);
-    std::memcpy(&outputBuffer[1], uuidCdr.data(), uuidCdr.length());
-    outputBuffer[17] = static_cast<uint8_t>(bsonType);
-    uassertStatusOK(crypto::aeadEncrypt(*key,
-                                        reinterpret_cast<const uint8_t*>(plaintext.data()),
-                                        plaintext.length(),
-                                        outputBuffer.data(),
-                                        18,
-                                        // The ciphertext starts 18 bytes into the output
-                                        // buffer, as described above.
-                                        outputBuffer.data() + 18,
-                                        outputLength));
-    return outputBuffer;
-}
-
 namespace {
 
 /**
diff --git a/src/mongo/shell/encrypted_dbclient_base.h b/src/mongo/shell/encrypted_dbclient_base.h
index 1cd05dca2cf..518c3631a5f 100644
--- a/src/mongo/shell/encrypted_dbclient_base.h
+++ b/src/mongo/shell/encrypted_dbclient_base.h
@@ -59,7 +59,6 @@ namespace mongo {
 
 constexpr std::size_t kEncryptedDBCacheSize = 50;
 
-constexpr int kAssociatedDataLength = 18;
 constexpr uint8_t kIntentToEncryptBit = 0x00;
 constexpr uint8_t kDeterministicEncryptionBit = 0x01;
 constexpr uint8_t kRandomEncryptionBit = 0x02;
@@ -155,11 +154,13 @@ protected:
 
     std::shared_ptr<SymmetricKey> getDataKey(const UUID& uuid);
 
-    std::vector<uint8_t> encryptWithKey(UUID uuid,
-                                        const std::shared_ptr<SymmetricKey>& key,
-                                        ConstDataRange plaintext,
-                                        BSONType bsonType,
-                                        int32_t algorithm);
+    FLEEncryptionFrame createEncryptionFrame(std::shared_ptr<SymmetricKey> key,
+                                             FleAlgorithmInt algorithm,
+                                             UUID uuid,
+                                             BSONType type,
+                                             ConstDataRange plaintext);
+
+    FLEDecryptionFrame createDecryptionFrame(ConstDataRange data);
 
 private:
     virtual void encryptMarking(const BSONObj& elem, BSONObjBuilder* builder, StringData elemName);
diff --git a/src/mongo/shell/kms_local.cpp b/src/mongo/shell/kms_local.cpp
index 32d5f760383..8e91326b706 100644
--- a/src/mongo/shell/kms_local.cpp
+++ b/src/mongo/shell/kms_local.cpp
@@ -69,13 +69,7 @@ private:
 std::vector<uint8_t> LocalKMSService::encrypt(ConstDataRange cdr, StringData kmsKeyId) {
     std::vector<std::uint8_t> ciphertext(crypto::aeadCipherOutputLength(cdr.length()));
 
-    uassertStatusOK(crypto::aeadEncrypt(_key,
-                                        reinterpret_cast<const uint8_t*>(cdr.data()),
-                                        cdr.length(),
-                                        nullptr,
-                                        0,
-                                        ciphertext.data(),
-                                        ciphertext.size()));
+    uassertStatusOK(crypto::aeadEncryptLocalKMS(_key, cdr, ciphertext.data(), ciphertext.size()));
 
     return ciphertext;
 }
@@ -93,17 +87,10 @@ BSONObj LocalKMSService::encryptDataKey(ConstDataRange cdr, StringData keyId) {
 }
 
 SecureVector<uint8_t> LocalKMSService::decrypt(ConstDataRange cdr, BSONObj masterKey) {
-    SecureVector<uint8_t> plaintext(
-        uassertStatusOK(crypto::aeadGetMaximumPlainTextLength(cdr.length())));
+    SecureVector<uint8_t> plaintext(uassertStatusOK(aeadGetMaximumPlainTextLength(cdr.length())));
 
     size_t outLen = plaintext->size();
-    uassertStatusOK(crypto::aeadDecrypt(_key,
-                                        reinterpret_cast<const uint8_t*>(cdr.data()),
-                                        cdr.length(),
-                                        nullptr,
-                                        0,
-                                        plaintext->data(),
-                                        &outLen));
+    uassertStatusOK(crypto::aeadDecryptLocalKMS(_key, cdr, plaintext->data(), &outLen));
     plaintext->resize(outLen);
 
     return plaintext;