SERVER-41277 Create FLE Data Framing Classes

4 files changed, 311 insertions, 68 deletions

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