SERVER-21404: Add rollback file coverage to ESE

5 files changed, 182 insertions, 1 deletions

diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp
index f0cc920452f..9e3951c2045 100644
--- a/src/mongo/db/dbhelpers.cpp
+++ b/src/mongo/db/dbhelpers.cpp
@@ -57,7 +57,9 @@
 #include "mongo/db/range_arithmetic.h"
 #include "mongo/db/repl/repl_client_info.h"
 #include "mongo/db/repl/replication_coordinator_global.h"
+#include "mongo/db/storage/data_protector.h"
 #include "mongo/db/storage/storage_options.h"
+#include "mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h"
 #include "mongo/db/write_concern.h"
 #include "mongo/db/write_concern_options.h"
 #include "mongo/db/s/collection_metadata.h"
@@ -582,6 +584,57 @@ Helpers::RemoveSaver::RemoveSaver(const string& a, const string& b, const string
     stringstream ss;
     ss << why << "." << terseCurrentTime(false) << "." << NUM++ << ".bson";
     _file /= ss.str();
+
+    auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
+    if (hooks->enabled()) {
+        _protector = hooks->getDataProtector();
+    }
+}
+
+Helpers::RemoveSaver::~RemoveSaver() {
+    if (_protector && _out) {
+        auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
+        invariant(hooks->enabled());
+
+        size_t protectedSizeMax = hooks->additionalBytesForProtectedBuffer();
+        std::unique_ptr<uint8_t[]> protectedBuffer(new uint8_t[protectedSizeMax]);
+
+        size_t resultLen;
+        Status status = _protector->finalize(protectedBuffer.get(), protectedSizeMax, &resultLen);
+        if (!status.isOK()) {
+            severe() << "Unable to finalize DataProtector while closing RemoveSaver: "
+                     << status.reason();
+            fassertFailed(34350);
+        }
+
+        _out->write(reinterpret_cast<const char*>(protectedBuffer.get()), resultLen);
+        if (_out->fail()) {
+            severe() << "Couldn't write finalized DataProtector data to: " << _file.string()
+                     << " for remove saving: " << errnoWithDescription();
+            fassertFailed(34351);
+        }
+
+        protectedBuffer.reset(new uint8_t[protectedSizeMax]);
+        status = _protector->finalizeTag(protectedBuffer.get(), protectedSizeMax, &resultLen);
+        if (!status.isOK()) {
+            severe() << "Unable to get finalizeTag from DataProtector while closing RemoveSaver: "
+                     << status.reason();
+            fassertFailed(34352);
+        }
+        if (resultLen != _protector->getNumberOfBytesReservedForTag()) {
+            severe() << "Attempted to write tag of size " << resultLen
+                     << " when DataProtector only reserved "
+                     << _protector->getNumberOfBytesReservedForTag() << " bytes";
+            fassertFailed(34353);
+        }
+        _out->seekp(0);
+        _out->write(reinterpret_cast<const char*>(protectedBuffer.get()), resultLen);
+        if (_out->fail()) {
+            severe() << "Couldn't write finalizeTag from DataProtector to: " << _file.string()
+                     << " for remove saving: " << errnoWithDescription();
+            fassertFailed(34354);
+        }
+    }
 }
 
 Status Helpers::RemoveSaver::goingToDelete(const BSONObj& o) {
@@ -597,7 +650,30 @@ Status Helpers::RemoveSaver::goingToDelete(const BSONObj& o) {
             return Status(ErrorCodes::FileNotOpen, msg);
         }
     }
-    _out->write(o.objdata(), o.objsize());
+
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(o.objdata());
+    size_t dataSize = o.objsize();
+
+    std::unique_ptr<uint8_t[]> protectedBuffer;
+    if (_protector) {
+        auto hooks = WiredTigerCustomizationHooks::get(getGlobalServiceContext());
+        invariant(hooks->enabled());
+
+        size_t protectedSizeMax = dataSize + hooks->additionalBytesForProtectedBuffer();
+        protectedBuffer.reset(new uint8_t[protectedSizeMax]);
+
+        size_t resultLen;
+        Status status = _protector->protect(
+            data, dataSize, protectedBuffer.get(), protectedSizeMax, &resultLen);
+        if (!status.isOK()) {
+            return status;
+        }
+
+        data = protectedBuffer.get();
+        dataSize = resultLen;
+    }
+
+    _out->write(reinterpret_cast<const char*>(data), dataSize);
     if (_out->fail()) {
         string msg = str::stream() << "couldn't write document to file: " << _file.string()
                                    << " for remove saving: " << errnoWithDescription();
diff --git a/src/mongo/db/dbhelpers.h b/src/mongo/db/dbhelpers.h
index 2c77bf8ca2d..5be0ba37185 100644
--- a/src/mongo/db/dbhelpers.h
+++ b/src/mongo/db/dbhelpers.h
@@ -33,11 +33,13 @@
 
 #include "mongo/db/db.h"
 #include "mongo/db/record_id.h"
+#include "mongo/db/storage/data_protector.h"
 
 namespace mongo {
 
 class Collection;
 class Cursor;
+class DataProtector;
 class OperationContext;
 struct KeyRange;
 struct WriteConcernOptions;
@@ -217,6 +219,7 @@ struct Helpers {
 
     public:
         RemoveSaver(const std::string& type, const std::string& ns, const std::string& why);
+        ~RemoveSaver();
 
         /**
          * Writes document to file. File is created lazily before writing the first document.
@@ -228,6 +231,7 @@ struct Helpers {
     private:
         boost::filesystem::path _root;
         boost::filesystem::path _file;
+        std::unique_ptr<DataProtector> _protector;
         std::unique_ptr<std::ostream> _out;
     };
 };
diff --git a/src/mongo/db/storage/data_protector.h b/src/mongo/db/storage/data_protector.h
new file mode 100644
index 00000000000..7eac67cc508
--- /dev/null
+++ b/src/mongo/db/storage/data_protector.h
@@ -0,0 +1,87 @@
+/**
+ *    Copyright (C) 2015 MongoDB Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    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
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ *    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 GNU Affero General 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 <cstddef>
+#include <cstdint>
+
+namespace mongo {
+
+class Status;
+
+/**
+ * Performs an implementation specific transformation on a series of input buffers to
+ * produce a protected form of their concatenated contents.
+ *
+ * protect() must be called on each input buffer, to produce a series of output buffers.
+ * The caller must ensure that the output buffers are large enough to contain the protected
+ * data resulting from the call. The caller must concatenate the output buffers together, in order.
+ * Once all input buffers have been protect()ed, finalize() must be called, and its output
+ * appended to the end of the protected data.
+ * The caller may then call finalizeTag() to get implementation defined metadata.
+ */
+class DataProtector {
+public:
+    virtual ~DataProtector() = default;
+
+    /**
+     * Copy `inLen` bytes from `in`, process them, and write the processed bytes into `out`.
+     * As processing may produce more or fewer bytes than were provided, the actual number
+     * of bytes written will placed in `bytesWritten`.
+     */
+    virtual Status protect(const std::uint8_t* in,
+                           std::size_t inLen,
+                           std::uint8_t* out,
+                           std::size_t outLen,
+                           std::size_t* bytesWritten) = 0;
+
+    /**
+     * Declares that this DataProtector will be provided no more data to protect.
+     * Fills `out` with any leftover state that needs serialization.
+     */
+    virtual Status finalize(std::uint8_t* out, std::size_t outLen, std::size_t* bytesWritten) = 0;
+
+    /**
+     * Returns the number of bytes reserved for metadata at the beginning of the first output
+     * buffer.
+     * Not all implementations will choose to reserve this space. They will return 0.
+     */
+    virtual std::size_t getNumberOfBytesReservedForTag() const = 0;
+
+    /**
+     * Fills buffer `out` of size `outLen`, with implementation defined metadata that had to be
+     * calculated after finalization.
+     * `bytesWritten` is filled with the number of bytes written into `out`.
+     */
+    virtual Status finalizeTag(std::uint8_t* out,
+                               std::size_t outLen,
+                               std::size_t* bytesWritten) = 0;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
index 48f135e1bae..42c37314ea3 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.cpp
@@ -34,6 +34,7 @@
 #include "mongo/base/init.h"
 #include "mongo/base/string_data.h"
 #include "mongo/db/service_context.h"
+#include "mongo/db/storage/data_protector.h"
 #include "mongo/stdx/memory.h"
 
 namespace mongo {
@@ -77,6 +78,11 @@ std::string EmptyWiredTigerCustomizationHooks::getOpenConfig(StringData tableNam
     return "";
 }
 
+
+std::unique_ptr<DataProtector> EmptyWiredTigerCustomizationHooks::getDataProtector() {
+    return std::unique_ptr<DataProtector>();
+}
+
 Status EmptyWiredTigerCustomizationHooks::protectTmpData(
     const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) {
     return Status(ErrorCodes::InternalError,
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
index c3bc0e1d08b..16bf22f2dbe 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_customization_hooks.h
@@ -36,6 +36,7 @@
 #include "mongo/db/jsobj.h"
 
 namespace mongo {
+class DataProtector;
 class StringData;
 class ServiceContext;
 
@@ -75,6 +76,11 @@ public:
     }
 
     /**
+     * Get the data protector object
+     */
+    virtual std::unique_ptr<DataProtector> getDataProtector() = 0;
+
+    /**
      * Transform temp data to non-readable form before writing it to disk.
      */
     virtual Status protectTmpData(
@@ -98,6 +104,8 @@ public:
 
     std::string getOpenConfig(StringData tableName) override;
 
+    std::unique_ptr<DataProtector> getDataProtector() override;
+
     Status protectTmpData(
         const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) override;