SERVER-47733 SymmetricEncryptorWindows shouldn't pad when update is called

(cherry picked from commit 2f6e5d0f94c06fde943ed6a25a9b7ecf6f774ce5)

3 files changed, 129 insertions, 15 deletions

diff --git a/src/mongo/crypto/symmetric_crypto_windows.cpp b/src/mongo/crypto/symmetric_crypto_windows.cpp
index 25dd5f304b8..0c90c3b4a1b 100644
--- a/src/mongo/crypto/symmetric_crypto_windows.cpp
+++ b/src/mongo/crypto/symmetric_crypto_windows.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/platform/basic.h"
 
+#include <bcrypt.h>
 #include <memory>
 #include <vector>
 
@@ -206,16 +207,99 @@ protected:
     std::vector<unsigned char> _iv;
 };
 
+/**
+ * Like other symmetric encryptors, this class encrypts block-by-block with update and then only
+ * pads once finalize is called. However, the Windows's BCrypt implementation does not natively
+ * implement this functionality (see SERVER-47733), and will either require block aligned inputs or
+ * will attempt to pad every input. This class bulks together inputs in a local buffer which is
+ * flushed to BCrypt whenever a full block is accumulated via update invocations. Data provided to
+ * update may be encrypted immediately, on a subsequent call to update, or on the call to finalize.
+ */
 class SymmetricEncryptorWindows : public SymmetricImplWindows<SymmetricEncryptor> {
 public:
     using SymmetricImplWindows::SymmetricImplWindows;
 
+    SymmetricEncryptorWindows(const SymmetricKey& key,
+                              aesMode mode,
+                              const uint8_t* iv,
+                              size_t ivLen)
+        : _blockData(_blockBuffer->data(), _blockBuffer->size()),
+          _blockCursor(_blockData),
+          SymmetricImplWindows<SymmetricEncryptor>(key, mode, iv, ivLen) {}
+
     StatusWith<size_t> update(const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen) final {
-        ULONG len = 0;
+        ULONG blockBufferEncryptLen = 0;
+        ULONG inputEncryptLen = 0;
+        ConstDataRange inData(in, inLen);
+        ConstDataRangeCursor inCursor(inData);
+
+        // If we have an incomplete block, we need to fill it before encrypting.
+        // If the total amount of input bytes will not fill the blockBuffer, just add it all to
+        // the buffer.
+        if (inLen < _blockCursor.length()) {
+            _blockCursor.writeAndAdvance(inCursor);
+            return 0;
+        } else if (_blockCursor.length() < _blockData.length() && _blockCursor.length() > 0) {
+            // Entering this code path means that we had data left over from the last time update
+            // was called. What we do below is fill the buffer with new input data until it is full.
+            // We then encrypt that buffer. We skip this step when the buffer is empty.
+            uint8_t bytesToFill = _blockCursor.length();
+            ConstDataRange bytesToFillRange(inCursor.data(), bytesToFill);
+            _blockCursor.writeAndAdvance(bytesToFillRange);
+            inCursor.advance(bytesToFill);
+            // We now encrypt the full buffer.
+            NTSTATUS status = BCryptEncrypt(_keyHandle,
+                                            const_cast<PUCHAR>(_blockBuffer->data()),
+                                            _blockBuffer->size(),
+                                            NULL,
+                                            _iv.data(),
+                                            _iv.size(),
+                                            out,
+                                            outLen,
+                                            &blockBufferEncryptLen,
+                                            0);
+            if (status != STATUS_SUCCESS) {
+                return Status{ErrorCodes::OperationFailed,
+                              str::stream() << "Encrypt failed: " << statusWithDescription(status)};
+            }
+            _blockCursor = DataRangeCursor(_blockData);
+        }
+
+        // we will attempt to encrypt as much of the remaining data as we can (i.e. the largest
+        // available size that is a multiple of the block length)
+        size_t remainingBytes = inCursor.length();
+        ULONG bytesToEncrypt = remainingBytes - (remainingBytes % aesBlockSize);
 
         NTSTATUS status = BCryptEncrypt(_keyHandle,
-                                        const_cast<PUCHAR>(in),
-                                        inLen,
+                                        const_cast<PUCHAR>(inCursor.data<UCHAR>()),
+                                        bytesToEncrypt,
+                                        NULL,
+                                        _iv.data(),
+                                        _iv.size(),
+                                        out + blockBufferEncryptLen,
+                                        outLen - blockBufferEncryptLen,
+                                        &inputEncryptLen,
+                                        0);
+
+        if (status != STATUS_SUCCESS) {
+            return Status{ErrorCodes::OperationFailed,
+                          str::stream() << "Encrypt failed: " << statusWithDescription(status)};
+        }
+
+        inCursor.advance(bytesToEncrypt);
+
+        // we now have to store what is left of the input in the block buffer
+        _blockCursor.writeAndAdvance(inCursor);
+
+        return static_cast<size_t>(blockBufferEncryptLen + inputEncryptLen);
+    }
+
+    StatusWith<size_t> finalize(uint8_t* out, size_t outLen) final {
+        // if there is any data left over in the block buffer, we will encrypt it with padding
+        ULONG len = 0;
+        NTSTATUS status = BCryptEncrypt(_keyHandle,
+                                        const_cast<PUCHAR>(_blockBuffer->data()),
+                                        _blockBuffer->size() - _blockCursor.length(),
                                         NULL,
                                         _iv.data(),
                                         _iv.size(),
@@ -229,18 +313,23 @@ public:
                           str::stream() << "Encrypt failed: " << statusWithDescription(status)};
         }
 
-        return static_cast<size_t>(len);
-    }
+        // we will now start a new block
+        _blockCursor = DataRangeCursor(_blockData);
 
-    StatusWith<size_t> finalize(uint8_t* out, size_t outLen) final {
-        // No finalize needed
-        return 0;
+        return static_cast<size_t>(len);
     }
 
     StatusWith<size_t> finalizeTag(uint8_t* out, size_t outLen) final {
         // Not a tagged cipher mode, write nothing.
         return 0;
     }
+
+private:
+    // buffer to store a single block of data, to be encrypted by update when filled, or by finalize
+    // with padding. 16 is the block length for AES.
+    SecureAllocatorDefaultDomain::SecureHandle<std::array<uint8_t, aesBlockSize>> _blockBuffer;
+    DataRange _blockData;
+    DataRangeCursor _blockCursor;
 };
 
 class SymmetricDecryptorWindows : public SymmetricImplWindows<SymmetricDecryptor> {
diff --git a/src/mongo/db/storage/remove_saver.cpp b/src/mongo/db/storage/remove_saver.cpp
index e24f33ecadd..440d6139b31 100644
--- a/src/mongo/db/storage/remove_saver.cpp
+++ b/src/mongo/db/storage/remove_saver.cpp
@@ -35,6 +35,7 @@
 
 #include <boost/filesystem/operations.hpp>
 #include <fstream>
+#include <ios>
 
 #include "mongo/db/service_context.h"
 #include "mongo/db/storage/encryption_hooks.h"
@@ -50,7 +51,11 @@ using std::stringstream;
 
 namespace mongo {
 
-RemoveSaver::RemoveSaver(const string& a, const string& b, const string& why) {
+RemoveSaver::RemoveSaver(const string& a,
+                         const string& b,
+                         const string& why,
+                         std::unique_ptr<Storage> storage)
+    : _storage(std::move(storage)) {
     static int NUM = 0;
 
     _root = storageGlobalParams.dbpath;
@@ -119,16 +124,14 @@ RemoveSaver::~RemoveSaver() {
                      << " for remove saving: " << redact(errnoWithDescription());
             fassertFailed(34354);
         }
+
+        _storage->dumpBuffer();
     }
 }
 
 Status RemoveSaver::goingToDelete(const BSONObj& o) {
     if (!_out) {
-        // We don't expect to ever pass "" to create_directories below, but catch
-        // this anyway as per SERVER-26412.
-        invariant(!_root.empty());
-        boost::filesystem::create_directories(_root);
-        _out.reset(new ofstream(_file.string().c_str(), ios_base::out | ios_base::binary));
+        _out = _storage->makeOstream(_file, _root);
 
         if (_out->fail()) {
             string msg = str::stream() << "couldn't create file: " << _file.string()
@@ -175,4 +178,13 @@ Status RemoveSaver::goingToDelete(const BSONObj& o) {
     return Status::OK();
 }
 
+std::unique_ptr<std::ostream> RemoveSaver::Storage::makeOstream(
+    const boost::filesystem::path& file, const boost::filesystem::path& root) {
+    // We don't expect to ever pass "" to create_directories below, but catch
+    // this anyway as per SERVER-26412.
+    invariant(!root.empty());
+    boost::filesystem::create_directories(root);
+    return std::make_unique<std::ofstream>(file.string().c_str(),
+                                           std::ios_base::out | std::ios_base::binary);
+}
 }  // namespace mongo
diff --git a/src/mongo/db/storage/remove_saver.h b/src/mongo/db/storage/remove_saver.h
index 0704ddb4d49..463fe7f8410 100644
--- a/src/mongo/db/storage/remove_saver.h
+++ b/src/mongo/db/storage/remove_saver.h
@@ -51,7 +51,11 @@ class RemoveSaver {
     RemoveSaver& operator=(const RemoveSaver&) = delete;
 
 public:
-    RemoveSaver(const std::string& type, const std::string& ns, const std::string& why);
+    class Storage;
+    RemoveSaver(const std::string& type,
+                const std::string& ns,
+                const std::string& why,
+                std::unique_ptr<Storage> storage = std::make_unique<Storage>());
     ~RemoveSaver();
 
     /**
@@ -77,11 +81,20 @@ public:
     }
     void file() && = delete;
 
+    class Storage {
+    public:
+        virtual ~Storage() = default;
+        virtual std::unique_ptr<std::ostream> makeOstream(const boost::filesystem::path& file,
+                                                          const boost::filesystem::path& root);
+        virtual void dumpBuffer() {}
+    };
+
 private:
     boost::filesystem::path _root;
     boost::filesystem::path _file;
     std::unique_ptr<DataProtector> _protector;
     std::unique_ptr<std::ostream> _out;
+    std::unique_ptr<Storage> _storage;
 };
 
 }  // namespace mongo