SERVER-23570 Make SecureAllocator a real allocator

Also adds SecureT, a wrapper proxy for newing objects on the secure heap.

8 files changed, 305 insertions, 23 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index 2bb03f107d2..5939c9ec499 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -221,6 +221,7 @@ config_header_substs = (
     ('@mongo_config_have_std_is_trivially_copyable@', 'MONGO_CONFIG_HAVE_STD_IS_TRIVIALLY_COPYABLE'),
     ('@mongo_config_have_std_enable_if_t@', 'MONGO_CONFIG_HAVE_STD_ENABLE_IF_T'),
     ('@mongo_config_have_std_make_unique@', 'MONGO_CONFIG_HAVE_STD_MAKE_UNIQUE'),
+    ('@mongo_config_have_std_align@', 'MONGO_CONFIG_HAVE_STD_ALIGN'),
     ('@mongo_config_have_strnlen@', 'MONGO_CONFIG_HAVE_STRNLEN'),
     ('@mongo_config_optimized_build@', 'MONGO_CONFIG_OPTIMIZED_BUILD'),
     ('@mongo_config_ssl@', 'MONGO_CONFIG_SSL'),
diff --git a/src/mongo/base/SConscript b/src/mongo/base/SConscript
index 07d696d48f6..ad7ab191d97 100644
--- a/src/mongo/base/SConscript
+++ b/src/mongo/base/SConscript
@@ -62,6 +62,7 @@ env.Library(
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/util/processinfo',
         '$BUILD_DIR/mongo/util/secure_zero_memory',
     ],
 )
diff --git a/src/mongo/base/secure_allocator.cpp b/src/mongo/base/secure_allocator.cpp
index b6ac5d14eaf..336edbf6364 100644
--- a/src/mongo/base/secure_allocator.cpp
+++ b/src/mongo/base/secure_allocator.cpp
@@ -31,6 +31,9 @@
 
 #include "mongo/base/secure_allocator.h"
 
+#include <memory>
+#include <unordered_map>
+
 #ifdef _WIN32
 #include <windows.h>
 #else
@@ -39,12 +42,19 @@
 #include <sys/types.h>
 #endif
 
+#include "mongo/base/disallow_copying.h"
+#include "mongo/base/init.h"
+#include "mongo/stdx/memory.h"
+#include "mongo/stdx/mutex.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/log.h"
+#include "mongo/util/processinfo.h"
 #include "mongo/util/secure_zero_memory.h"
 
 namespace mongo {
 
+namespace {
+
 /**
  * NOTE(jcarey): Why not new/delete?
  *
@@ -53,11 +63,9 @@ namespace mongo {
  * platforms do offer those semantics, they're not available globally, so we
  * have to flow all allocations through page based allocations.
  */
-namespace secure_allocator_details {
-
 #ifdef _WIN32
 
-void* allocate(std::size_t bytes) {
+void* systemAllocate(std::size_t bytes) {
     // Flags:
     //
     // MEM_COMMIT - allocates the memory charges and zeros the underlying
@@ -85,9 +93,7 @@ void* allocate(std::size_t bytes) {
     return ptr;
 }
 
-void deallocate(void* ptr, std::size_t bytes) {
-    secureZeroMemory(ptr, bytes);
-
+void systemDeallocate(void* ptr, std::size_t bytes) {
     if (VirtualUnlock(ptr, bytes) == 0) {
         auto str = errnoWithPrefix("Failed to VirtualUnlock");
         severe() << str;
@@ -121,7 +127,7 @@ void deallocate(void* ptr, std::size_t bytes) {
 #error "Could not determine a way to map anonymous memory, required for secure allocation"
 #endif
 
-void* allocate(std::size_t bytes) {
+void* systemAllocate(std::size_t bytes) {
     // Flags:
     //
     // PROT_READ | PROT_WRITE - allows read write access to the page
@@ -151,9 +157,7 @@ void* allocate(std::size_t bytes) {
     return ptr;
 }
 
-void deallocate(void* ptr, std::size_t bytes) {
-    secureZeroMemory(ptr, bytes);
-
+void systemDeallocate(void* ptr, std::size_t bytes) {
     if (munlock(ptr, bytes) != 0) {
         severe() << errnoWithPrefix("Failed to munlock");
         fassertFailed(28833);
@@ -167,5 +171,105 @@ void deallocate(void* ptr, std::size_t bytes) {
 
 #endif
 
+/**
+ * Each page represent one call to mmap+mlock / VirtualAlloc+VirtualLock on construction and will
+ * match with an equivalent unlock and unmap on destruction.  Pages are rounded up to the nearest
+ * page size and allocate returns aligned pointers.
+ */
+class Allocation {
+    MONGO_DISALLOW_COPYING(Allocation);
+
+public:
+    explicit Allocation(std::size_t initialAllocation) {
+        auto pageSize = ProcessInfo::getPageSize();
+        std::size_t remainder = initialAllocation % pageSize;
+
+        _size = _remaining =
+            remainder ? initialAllocation + pageSize - remainder : initialAllocation;
+        _start = _ptr = systemAllocate(_size);
+    }
+
+    ~Allocation() {
+        systemDeallocate(_start, _size);
+    }
+
+    /**
+     * Allocates an aligned pointer of given size from the locked page we have a pointer to.
+     *
+     * Returns null if the request can't be satisifed.
+     */
+    void* allocate(std::size_t size, std::size_t alignOf) {
+        if (stdx::align(alignOf, size, _ptr, _remaining)) {
+            auto result = _ptr;
+            _ptr = static_cast<char*>(_ptr) + size;
+            _remaining -= size;
+
+            return result;
+        }
+
+        // We'll make a new allocation if this one doesn't have enough room
+        return nullptr;
+    }
+
+private:
+    void* _start;            // Start of the allocation
+    void* _ptr;              // Start of unallocated bytes
+    std::size_t _size;       // Total size
+    std::size_t _remaining;  // Remaining bytes
+};
+
+// See secure_allocator_details::allocate for a more detailed comment on what these are used for
+stdx::mutex allocatorMutex;  // Protects the values below
+std::unordered_map<void*, std::shared_ptr<Allocation>> secureTable;
+std::shared_ptr<Allocation> lastAllocation = nullptr;
+
+}  // namespace
+
+MONGO_INITIALIZER_GENERAL(SecureAllocator,
+                          ("SystemInfo"),
+                          MONGO_NO_DEPENDENTS)(InitializerContext* context) {
+    return Status::OK();
+}
+
+namespace secure_allocator_details {
+
+/**
+ * To save on allocations, we try to serve multiple requests out of the same mlocked page where
+ * possible.  We do this by invoking the system allocator in multiples of a full page, and keep the
+ * last page around, giving out pointers from that page if its possible to do so.  We also keep an
+ * unordered_map of all the allocations we've handed out, which hold shared_ptrs that get rid of
+ * pages when we're not using them anymore.
+ */
+void* allocate(std::size_t bytes, std::size_t alignOf) {
+    stdx::lock_guard<stdx::mutex> lk(allocatorMutex);
+
+    if (lastAllocation) {
+        auto out = lastAllocation->allocate(bytes, alignOf);
+
+        if (out) {
+            secureTable[out] = lastAllocation;
+            return out;
+        }
+    }
+
+    lastAllocation = std::make_shared<Allocation>(bytes);
+    auto out = lastAllocation->allocate(bytes, alignOf);
+    secureTable[out] = lastAllocation;
+    return out;
+}
+
+/**
+ * Deallocates a secure allocation.
+ *
+ * We zero memory before derefing the associated allocation.
+ */
+void deallocate(void* ptr, std::size_t bytes) {
+    secureZeroMemory(ptr, bytes);
+
+    stdx::lock_guard<stdx::mutex> lk(allocatorMutex);
+
+    secureTable.erase(ptr);
+}
+
 }  // namespace secure_allocator_details
 }  // namespace mongo
diff --git a/src/mongo/base/secure_allocator.h b/src/mongo/base/secure_allocator.h
index ed10d075968..416ff12ebfa 100644
--- a/src/mongo/base/secure_allocator.h
+++ b/src/mongo/base/secure_allocator.h
@@ -31,15 +31,19 @@
 
 #include <cstddef>
 #include <limits>
+#include <memory>
 #include <string>
 #include <type_traits>
 #include <vector>
 
+#include "mongo/stdx/type_traits.h"
+#include "mongo/util/assert_util.h"
+
 namespace mongo {
 
 namespace secure_allocator_details {
 
-void* allocate(std::size_t bytes);
+void* allocate(std::size_t bytes, std::size_t alignOf);
 void deallocate(void* ptr, std::size_t bytes);
 
 }  // namespace secure_allocator_details
@@ -107,7 +111,8 @@ struct SecureAllocator {
     };
 
     pointer allocate(size_type n) {
-        return static_cast<pointer>(secure_allocator_details::allocate(sizeof(value_type) * n));
+        return static_cast<pointer>(secure_allocator_details::allocate(
+            sizeof(value_type) * n, std::alignment_of<T>::value));
     }
 
     pointer allocate(size_type n, const_void_pointer) {
@@ -173,9 +178,114 @@ bool operator!=(const SecureAllocator<T>& lhs, const SecureAllocator<U>& rhs) {
     return !(lhs == rhs);
 }
 
+/**
+ * SecureHandle offers a smart pointer-ish interface to a type.
+ *
+ * It attempts to solve the problem of using container types with small object optimizations that
+ * might accidentally leave important data on the stack if they're too small to spill to the heap.
+ * It uses the secure allocator for allocations and deallocations.
+ *
+ * This type is meant to offer more value like semantics than a unique_ptr:
+ *   - SecureHandle's are only default constructible if their T is
+ *   - You can construct a SecureHandle like the underlying value:
+ *       SecureHandle<Mytype> type(foo, bar baz);
+ *   - SecureHandle's are copyable if their T's are
+ *   - only moving can produce a non-usable T
+ *
+ * While still being cheap and convenient like a unique_ptr:
+ *   - SecureHandle's move by pointer to their secure storage
+ *   - T& operator*()
+ *   - T* operator->()
+ *
+ * In a moved from state, SecureHandle's may be copy or move assigned and the destructor may run,
+ * but all other operations will invariant.
+ */
+template <typename T>
+class SecureHandle {
+public:
+    // NOTE: (jcarey)
+    //
+    // We have the default ctor and the perfect forwarding ctor because msvc 2013 ice's on some
+    // default constructed types without it (sfinae was falling over for some reason).
+    //
+    // For non-default constructible t's, we'll fail to substitute the forwarded call to new
+    SecureHandle() : _t(_new()) {}
+
+    // Generic constructor that forwards to the underlying T if the first arg isn't a SecureHandle
+    template <
+        typename Arg,
+        typename... Args,
+        stdx::enable_if_t<!std::is_same<SecureHandle<T>, typename std::decay<Arg>::type>::value,
+                          int> = 0>
+    SecureHandle(Arg&& arg, Args&&... args)
+        : _t(_new(std::forward<Arg>(arg), std::forward<Args>(args)...)) {}
+
+    SecureHandle(const SecureHandle& other) : _t(_new(*other)) {}
+
+    SecureHandle& operator=(const SecureHandle& other) {
+        if (_t) {
+            *_t = *other;
+        } else {
+            _t = _new(*other);
+        }
+
+        return *this;
+    }
+
+    SecureHandle(SecureHandle&& other) : _t(other._t) {
+        other._t = nullptr;
+    }
+
+    SecureHandle& operator=(SecureHandle&& other) {
+        if (&other == this) {
+            return *this;
+        }
+
+        _delete();
+
+        _t = other._t;
+        other._t = nullptr;
+
+        return *this;
+    }
+
+    ~SecureHandle() {
+        _delete();
+    }
+
+    T& operator*() const {
+        invariant(_t);
+
+        return *_t;
+    }
+
+    T* operator->() const {
+        invariant(_t);
+
+        return _t;
+    }
+
+private:
+    template <typename... Args>
+    static T* _new(Args&&... args) {
+        return ::new (secure_allocator_details::allocate(sizeof(T), std::alignment_of<T>::value))
+            T(std::forward<Args>(args)...);
+    }
+
+    void _delete() {
+        if (_t) {
+            _t->~T();
+            secure_allocator_details::deallocate(_t, sizeof(T));
+        }
+    }
+
+    T* _t;
+};
+
 template <typename T>
-using SecureVector = std::vector<T, SecureAllocator<T>>;
+using SecureVector = SecureHandle<std::vector<T, SecureAllocator<T>>>;
 
-using SecureString = std::basic_string<char, std::char_traits<char>, SecureAllocator<char>>;
+using SecureString =
+    SecureHandle<std::basic_string<char, std::char_traits<char>, SecureAllocator<char>>>;
 
 }  // namespace mongo
diff --git a/src/mongo/base/secure_allocator_test.cpp b/src/mongo/base/secure_allocator_test.cpp
index f04b286ef2b..750dc8c5c78 100644
--- a/src/mongo/base/secure_allocator_test.cpp
+++ b/src/mongo/base/secure_allocator_test.cpp
@@ -30,6 +30,8 @@
 
 #include "mongo/base/secure_allocator.h"
 
+#include <array>
+
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
@@ -37,21 +39,50 @@ namespace mongo {
 TEST(SecureAllocator, SecureVector) {
     SecureVector<int> vec;
 
-    vec.push_back(1);
-    vec.push_back(2);
+    vec->push_back(1);
+    vec->push_back(2);
 
-    ASSERT_EQUALS(1, vec[0]);
-    ASSERT_EQUALS(2, vec[1]);
+    ASSERT_EQUALS(1, (*vec)[0]);
+    ASSERT_EQUALS(2, (*vec)[1]);
 
-    vec.resize(2000, 3);
-    ASSERT_EQUALS(3, vec[2]);
+    vec->resize(2000, 3);
+    ASSERT_EQUALS(3, (*vec)[2]);
 }
 
 TEST(SecureAllocator, SecureString) {
     SecureString str;
 
-    str.resize(2000, 'x');
-    ASSERT_EQUALS(0, str.compare(SecureString(2000, 'x')));
+    str->resize(2000, 'x');
+    ASSERT_EQUALS(0, str->compare(*SecureString(2000, 'x')));
+
+    SecureString str2(str);
+    ASSERT_NOT_EQUALS(&*str, &*str2);
+    str2 = str;
+    ASSERT_NOT_EQUALS(&*str, &*str2);
+
+    auto strPtr = &*str;
+    auto str2Ptr = &*str2;
+    SecureString str3(std::move(str));
+    ASSERT_EQUALS(strPtr, &*str3);
+    str3 = std::move(str2);
+    ASSERT_EQUALS(str2Ptr, &*str3);
+}
+
+// Verify that we can make a good number of secure objects.  Under the initial secure allocator
+// design (page per object), you couldn't make more than 8-50 objects before running out of lockable
+// pages.
+TEST(SecureAllocator, ManySecureBytes) {
+    std::array<SecureHandle<char>, 4096> chars;
+    std::vector<SecureHandle<char>> e_chars(4096, 'e');
+}
+
+TEST(SecureAllocator, NonDefaultConstructibleWorks) {
+    struct Foo {
+        Foo(int) {}
+        Foo() = delete;
+    };
+
+    SecureHandle<Foo> foo(10);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/config.h.in b/src/mongo/config.h.in
index c94fa484177..e3274e09615 100644
--- a/src/mongo/config.h.in
+++ b/src/mongo/config.h.in
@@ -67,6 +67,9 @@
 // Defined if std::make_unique is available
 @mongo_config_have_std_make_unique@
 
+// Defined if std::align is available
+@mongo_config_have_std_align@
+
 // Defined if strnlen is available
 @mongo_config_have_strnlen@
 
diff --git a/src/mongo/stdx/memory.h b/src/mongo/stdx/memory.h
index 83793f86c5a..8d04d728263 100644
--- a/src/mongo/stdx/memory.h
+++ b/src/mongo/stdx/memory.h
@@ -55,3 +55,29 @@ using boost::make_unique;  // NOLINT
 }  // namespace mongo
 
 #endif
+
+#if defined(MONGO_CONFIG_HAVE_STD_ALIGN)
+
+#include <memory>
+
+namespace mongo {
+namespace stdx {
+
+using ::std::align;  // NOLINT
+
+}  // namespace stdx
+}  // namespace mongo
+
+#else
+
+#include <boost/align/align.hpp>
+
+namespace mongo {
+namespace stdx {
+
+using boost::alignment::align;  // NOLINT
+
+}  // namespace stdx
+}  // namespace mongo
+
+#endif
diff --git a/src/mongo/util/processinfo.cpp b/src/mongo/util/processinfo.cpp
index 8afcb2f0eeb..baaa874eaff 100644
--- a/src/mongo/util/processinfo.cpp
+++ b/src/mongo/util/processinfo.cpp
@@ -85,7 +85,13 @@ void ProcessInfo::initializeSystemInfo() {
     }
 }
 
-MONGO_INITIALIZER(SystemInfo)(InitializerContext* context) {
+/**
+ * We need this get the system page size for the secure allocator, which the enterprise modules need
+ * for storage for command line parameters.
+ */
+MONGO_INITIALIZER_GENERAL(SystemInfo,
+                          MONGO_NO_PREREQUISITES,
+                          MONGO_NO_DEPENDENTS)(InitializerContext* context) {
     ProcessInfo::initializeSystemInfo();
     return Status::OK();
 }