From eb2b72cf9c0269f086223d499ac9be8a270d268c Mon Sep 17 00:00:00 2001
From: Benety Goh <benety@mongodb.com>
Date: Tue, 29 Sep 2020 18:29:57 -0400
Subject: SERVER-50131 make EncryptionHooks check in Sorter work with tests
 without service context

---
 src/mongo/db/sorter/SConscript      |  1 -
 src/mongo/db/sorter/sorter.cpp      | 24 ++++++++++++++++++++----
 src/mongo/db/sorter/sorter_test.cpp |  5 ++---
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/src/mongo/db/sorter/SConscript b/src/mongo/db/sorter/SConscript
index 1f8a51a7a7a..b3d6a25db0f 100644
--- a/src/mongo/db/sorter/SConscript
+++ b/src/mongo/db/sorter/SConscript
@@ -12,7 +12,6 @@ sorterEnv.CppUnitTest(
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/db/service_context',
-        '$BUILD_DIR/mongo/db/service_context_test_fixture',
         '$BUILD_DIR/mongo/db/storage/encryption_hooks',
         '$BUILD_DIR/mongo/db/storage/storage_options',
         '$BUILD_DIR/mongo/s/is_mongos',
diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index d41738d2778..ef4fe7773b7 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -87,6 +87,24 @@ void checkNoExternalSortOnMongos(const SortOptions& opts) {
             !(isMongos() && opts.extSortAllowed));
 }
 
+/**
+ * Returns the current EncryptionHooks registered with the global service context.
+ * Returns nullptr if the service context is not available; or if the EncyptionHooks
+ * registered is not enabled.
+ */
+EncryptionHooks* getEncryptionHooksIfEnabled() {
+    // Some tests may not run with a global service context.
+    if (!hasGlobalServiceContext()) {
+        return nullptr;
+    }
+    auto service = getGlobalServiceContext();
+    auto encryptionHooks = EncryptionHooks::get(service);
+    if (!encryptionHooks->enabled()) {
+        return nullptr;
+    }
+    return encryptionHooks;
+}
+
 }  // namespace
 
 namespace sorter {
@@ -287,8 +305,7 @@ private:
         read(_buffer.get(), blockSize);
         uassert(16816, "file too short?", !_done);
 
-        auto encryptionHooks = EncryptionHooks::get(getGlobalServiceContext());
-        if (encryptionHooks->enabled()) {
+        if (auto encryptionHooks = getEncryptionHooksIfEnabled()) {
             std::unique_ptr<char[]> out(new char[blockSize]);
             size_t outLen;
             Status status =
@@ -1078,8 +1095,7 @@ void SortedFileWriter<Key, Value>::spill() {
     }
 
     std::unique_ptr<char[]> out;
-    auto encryptionHooks = EncryptionHooks::get(getGlobalServiceContext());
-    if (encryptionHooks->enabled()) {
+    if (auto encryptionHooks = getEncryptionHooksIfEnabled()) {
         size_t protectedSizeMax = size + encryptionHooks->additionalBytesForProtectedBuffer();
         out.reset(new char[protectedSizeMax]);
         size_t resultLen;
diff --git a/src/mongo/db/sorter/sorter_test.cpp b/src/mongo/db/sorter/sorter_test.cpp
index 781475c522a..1fe6cec7563 100644
--- a/src/mongo/db/sorter/sorter_test.cpp
+++ b/src/mongo/db/sorter/sorter_test.cpp
@@ -37,7 +37,6 @@
 #include "mongo/base/data_type_endian.h"
 #include "mongo/base/static_assert.h"
 #include "mongo/config.h"
-#include "mongo/db/service_context_test_fixture.h"
 #include "mongo/db/sorter/sorter.h"
 #include "mongo/logv2/log.h"
 #include "mongo/platform/random.h"
@@ -275,7 +274,7 @@ public:
     }
 };
 
-class SortedFileWriterAndFileIteratorTests : public ScopedGlobalServiceContextForTest {
+class SortedFileWriterAndFileIteratorTests {
 public:
     void run() {
         unittest::TempDir tempDir("sortedFileWriterTests");
@@ -362,7 +361,7 @@ public:
 };
 
 namespace SorterTests {
-class Basic : public ScopedGlobalServiceContextForTest {
+class Basic {
 public:
     virtual ~Basic() {}
 
-- 
cgit v1.2.1