SERVER-28436 Implement KeysCollectionManager

5 files changed, 822 insertions, 0 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 1615f604df2..cdf8c5e6464 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -984,6 +984,7 @@ env.Library(
 env.Library(
     target='keys_collection_manager',
     source=[
+        'keys_collection_manager.cpp',
         'keys_collection_cache_reader.cpp',
         'keys_collection_cache_reader_and_updater.cpp',
     ],
@@ -1108,6 +1109,17 @@ env.CppUnitTest(
     ],
 )
 
+env.CppUnitTest(
+    target='keys_collection_manager_test',
+    source=[
+        'keys_collection_manager_test.cpp',
+    ],
+    LIBDEPS=[
+        'keys_collection_manager',
+        '$BUILD_DIR/mongo/s/config_server_test_fixture',
+    ],
+)
+
 env.Library(
     target= 'op_observer_noop',
     source= [
diff --git a/src/mongo/db/keys_collection_cache_reader_and_updater_test.cpp b/src/mongo/db/keys_collection_cache_reader_and_updater_test.cpp
index 187cebd444a..540b4af8898 100644
--- a/src/mongo/db/keys_collection_cache_reader_and_updater_test.cpp
+++ b/src/mongo/db/keys_collection_cache_reader_and_updater_test.cpp
@@ -42,6 +42,7 @@
 #include "mongo/stdx/memory.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/clock_source_mock.h"
+#include "mongo/util/fail_point_service.h"
 
 namespace mongo {
 
@@ -94,6 +95,18 @@ TEST_F(CacheUpdaterTest, ShouldCreate2KeysFromEmpty) {
     ASSERT_NE(key1.getKey(), key2.getKey());
 }
 
+TEST_F(CacheUpdaterTest, ShouldPropagateWriteError) {
+    KeysCollectionCacheReaderAndUpdater updater("dummy", Seconds(5));
+
+    const LogicalTime currentTime(LogicalTime(Timestamp(100, 2)));
+    LogicalClock::get(operationContext())->initClusterTimeFromTrustedSource(currentTime);
+
+    FailPointEnableBlock failWriteBlock("failCollectionInserts");
+
+    auto keyStatus = updater.refresh(operationContext());
+    ASSERT_EQ(ErrorCodes::FailPointEnabled, keyStatus.getStatus());
+}
+
 TEST_F(CacheUpdaterTest, ShouldCreateAnotherKeyIfOnlyOneKeyExists) {
     KeysCollectionCacheReaderAndUpdater updater("dummy", Seconds(5));
 
diff --git a/src/mongo/db/keys_collection_manager.cpp b/src/mongo/db/keys_collection_manager.cpp
new file mode 100644
index 00000000000..f65e1163e8f
--- /dev/null
+++ b/src/mongo/db/keys_collection_manager.cpp
@@ -0,0 +1,324 @@
+/**
+ *    Copyright (C) 2017 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.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/keys_collection_manager.h"
+
+#include "mongo/db/keys_collection_cache_reader.h"
+#include "mongo/db/keys_collection_cache_reader_and_updater.h"
+#include "mongo/db/logical_clock.h"
+#include "mongo/db/logical_time.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+#include "mongo/stdx/memory.h"
+#include "mongo/util/concurrency/idle_thread_block.h"
+#include "mongo/util/mongoutils/str.h"
+#include "mongo/util/time_support.h"
+
+namespace mongo {
+
+namespace {
+
+Milliseconds kDefaultRefreshWaitTime(30 * 1000);
+Milliseconds kRefreshIntervalIfErrored(200);
+
+/**
+ * Returns the amount of time to wait until the monitoring thread should attempt to refresh again.
+ */
+Milliseconds howMuchSleepNeedFor(const LogicalTime& currentTime,
+                                 const LogicalTime& latestExpiredAt,
+                                 const Milliseconds& interval) {
+    auto currentSecs = currentTime.asTimestamp().getSecs();
+    auto expiredSecs = latestExpiredAt.asTimestamp().getSecs();
+
+    if (currentSecs >= expiredSecs) {
+        // This means that the last round didn't generate a usable key for the current time.
+        // However, we don't want to poll too hard as well, so use a low interval.
+        return kRefreshIntervalIfErrored;
+    }
+
+    auto millisBeforeExpire = 1000 * (expiredSecs - currentSecs);
+
+    if (interval.count() <= millisBeforeExpire) {
+        return interval;
+    }
+
+    return kRefreshIntervalIfErrored;
+}
+
+}  // unnamed namespace
+
+KeysCollectionManager::KeysCollectionManager(std::string purpose, Seconds keyValidForInterval)
+    : _purpose(std::move(purpose)),
+      _keyValidForInterval(keyValidForInterval),
+      _keysCache(_purpose) {}
+
+StatusWith<KeysCollectionDocument> KeysCollectionManager::getKeyForValidation(
+    OperationContext* opCtx, long long keyId, const LogicalTime& forThisTime) {
+    auto keyStatus = _getKeyWithKeyIdCheck(keyId, forThisTime);
+
+    if (keyStatus != ErrorCodes::KeyNotFound) {
+        return keyStatus;
+    }
+
+    _refresher.refreshNow(opCtx);
+
+    return _getKeyWithKeyIdCheck(keyId, forThisTime);
+}
+
+StatusWith<KeysCollectionDocument> KeysCollectionManager::getKeyForSigning(
+    OperationContext* opCtx, const LogicalTime& forThisTime) {
+
+    auto keyStatusWith = _getKey(forThisTime);
+    auto keyStatus = keyStatusWith.getStatus();
+
+    if (keyStatus != ErrorCodes::KeyNotFound) {
+        return keyStatusWith;
+    }
+
+    do {
+        _refresher.refreshNow(opCtx);
+
+        keyStatusWith = _getKey(forThisTime);
+        keyStatus = keyStatusWith.getStatus();
+
+        if (keyStatus == ErrorCodes::KeyNotFound) {
+            sleepFor(kRefreshIntervalIfErrored);
+        }
+
+    } while (keyStatus == ErrorCodes::KeyNotFound);
+
+    return keyStatusWith;
+}
+
+StatusWith<KeysCollectionDocument> KeysCollectionManager::_getKeyWithKeyIdCheck(
+    long long keyId, const LogicalTime& forThisTime) {
+    auto keyStatus = _getKey(forThisTime);
+
+    if (!keyStatus.isOK()) {
+        return keyStatus;
+    }
+
+    auto key = keyStatus.getValue();
+
+    if (keyId == key.getKeyId()) {
+        return key;
+    }
+
+    // Key not expired but keyId does not match!
+    return {ErrorCodes::KeyNotFound,
+            str::stream() << "No keys found for " << _purpose << " that is valid for time: "
+                          << forThisTime.toString()
+                          << " with id: "
+                          << keyId};
+}
+
+StatusWith<KeysCollectionDocument> KeysCollectionManager::_getKey(const LogicalTime& forThisTime) {
+    auto keyStatus = _keysCache.getKey(forThisTime);
+
+    if (!keyStatus.isOK()) {
+        return keyStatus;
+    }
+
+    const auto& key = keyStatus.getValue();
+
+    if (key.getExpiresAt() < forThisTime) {
+        return {ErrorCodes::KeyNotFound,
+                str::stream() << "No keys found for " << _purpose << " that is valid for "
+                              << forThisTime.toString()};
+    }
+
+    return key;
+}
+
+void KeysCollectionManager::startMonitoring(ServiceContext* service) {
+    _refresher.setFunc([this](OperationContext* opCtx) { return _keysCache.refresh(opCtx); });
+    _refresher.start(
+        service, str::stream() << "monitoring keys for " << _purpose, _keyValidForInterval);
+}
+
+void KeysCollectionManager::stopMonitoring() {
+    _refresher.stop();
+}
+
+void KeysCollectionManager::enableKeyGenerator(OperationContext* opCtx, bool doEnable) {
+    if (doEnable) {
+        _refresher.switchFunc(opCtx, [this](OperationContext* opCtx) {
+            KeysCollectionCacheReaderAndUpdater keyGenerator(_purpose, _keyValidForInterval);
+            auto keyGenerationStatus = keyGenerator.refresh(opCtx);
+
+            if (ErrorCodes::isShutdownError(keyGenerationStatus.getStatus().code())) {
+                return keyGenerationStatus;
+            }
+
+            // An error encountered by the keyGenerator should not prevent refreshing the cache
+            auto cacheRefreshStatus = _keysCache.refresh(opCtx);
+
+            if (!keyGenerationStatus.isOK()) {
+                return keyGenerationStatus;
+            }
+
+            return cacheRefreshStatus;
+        });
+    } else {
+        _refresher.switchFunc(
+            opCtx, [this](OperationContext* opCtx) { return _keysCache.refresh(opCtx); });
+    }
+}
+
+void KeysCollectionManager::PeriodicRunner::refreshNow(OperationContext* opCtx) {
+    auto refreshRequest = [this]() {
+        stdx::lock_guard<stdx::mutex> lk(_mutex);
+
+        if (_inShutdown) {
+            throw DBException("aborting keys cache refresh because node is shutting down",
+                              ErrorCodes::ShutdownInProgress);
+        }
+
+        if (_refreshRequest) {
+            return _refreshRequest;
+        }
+
+        _refreshNeededCV.notify_all();
+        _refreshRequest = std::make_shared<Notification<void>>();
+        return _refreshRequest;
+    }();
+
+    // note: waitFor waits min(maxTimeMS, kDefaultRefreshWaitTime).
+    // waitFor also throws if timeout, so also throw when notification was not satisfied after
+    // waiting.
+    if (!refreshRequest->waitFor(opCtx, kDefaultRefreshWaitTime)) {
+        throw DBException("timed out waiting for refresh", ErrorCodes::ExceededTimeLimit);
+    }
+}
+
+void KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh(ServiceContext* service,
+                                                               std::string threadName,
+                                                               Milliseconds refreshInterval) {
+    Client::initThreadIfNotAlready(threadName);
+
+    while (true) {
+        auto opCtx = cc().makeOperationContext();
+
+        std::shared_ptr<RefreshFunc> doRefresh;
+        {
+            stdx::lock_guard<stdx::mutex> lock(_mutex);
+
+            if (_inShutdown) {
+                break;
+            }
+
+            invariant(_doRefresh.get() != nullptr);
+            doRefresh = _doRefresh;
+        }
+
+        Milliseconds nextWakeup = kRefreshIntervalIfErrored;
+
+        auto latestKeyStatusWith = (*doRefresh)(opCtx.get());
+        if (latestKeyStatusWith.getStatus().isOK()) {
+            const auto& latestKey = latestKeyStatusWith.getValue();
+            auto currentTime = LogicalClock::get(service)->getClusterTime().getTime();
+
+            nextWakeup =
+                howMuchSleepNeedFor(currentTime, latestKey.getExpiresAt(), refreshInterval);
+        }
+
+        // TODO: Add backoff to nextWakeup if it has a very small value in a row to avoid spinning.
+
+        stdx::unique_lock<stdx::mutex> lock(_mutex);
+
+        if (_refreshRequest) {
+            _refreshRequest->set();
+            _refreshRequest.reset();
+        }
+
+        if (_inShutdown) {
+            break;
+        }
+
+        MONGO_IDLE_THREAD_BLOCK;
+        auto sleepStatus = opCtx->waitForConditionOrInterruptNoAssertUntil(
+            _refreshNeededCV, lock, Date_t::now() + nextWakeup);
+
+        if (ErrorCodes::isShutdownError(sleepStatus.getStatus().code())) {
+            break;
+        }
+    }
+
+    stdx::unique_lock<stdx::mutex> lock(_mutex);
+    if (_refreshRequest) {
+        _refreshRequest->set();
+        _refreshRequest.reset();
+    }
+}
+
+void KeysCollectionManager::PeriodicRunner::setFunc(RefreshFunc newRefreshStrategy) {
+    stdx::lock_guard<stdx::mutex> lock(_mutex);
+    _doRefresh = std::make_shared<RefreshFunc>(std::move(newRefreshStrategy));
+}
+
+void KeysCollectionManager::PeriodicRunner::switchFunc(OperationContext* opCtx,
+                                                       RefreshFunc newRefreshStrategy) {
+    setFunc(newRefreshStrategy);
+
+    // Note: calling refreshNow will ensure that if there is an ongoing method call to the original
+    // refreshStrategy, it will be finished after this.
+    refreshNow(opCtx);
+}
+
+void KeysCollectionManager::PeriodicRunner::start(ServiceContext* service,
+                                                  const std::string& threadName,
+                                                  Milliseconds refreshInterval) {
+    stdx::lock_guard<stdx::mutex> lock(_mutex);
+    invariant(!_backgroundThread.joinable());
+    invariant(!_inShutdown);
+
+    _backgroundThread =
+        stdx::thread(stdx::bind(&KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh,
+                                this,
+                                service,
+                                threadName,
+                                refreshInterval));
+}
+
+void KeysCollectionManager::PeriodicRunner::stop() {
+    {
+        stdx::lock_guard<stdx::mutex> lock(_mutex);
+        if (!_backgroundThread.joinable()) {
+            return;
+        }
+
+        _inShutdown = true;
+        _refreshNeededCV.notify_all();
+    }
+
+    _backgroundThread.join();
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/keys_collection_manager.h b/src/mongo/db/keys_collection_manager.h
new file mode 100644
index 00000000000..57ebed1f1da
--- /dev/null
+++ b/src/mongo/db/keys_collection_manager.h
@@ -0,0 +1,173 @@
+/**
+ *    Copyright (C) 2017 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 <memory>
+
+#include "mongo/base/status_with.h"
+#include "mongo/db/keys_collection_cache_reader.h"
+#include "mongo/db/keys_collection_cache_reader_and_updater.h"
+#include "mongo/db/keys_collection_document.h"
+#include "mongo/stdx/functional.h"
+#include "mongo/stdx/mutex.h"
+#include "mongo/stdx/thread.h"
+#include "mongo/util/concurrency/notification.h"
+#include "mongo/util/duration.h"
+
+namespace mongo {
+
+class OperationContext;
+class LogicalTime;
+class ServiceContext;
+
+/**
+ * This is responsible for providing keys that can be used for HMAC computation. This also supports
+ * automatic key rotation that happens on a configurable interval.
+ */
+class KeysCollectionManager {
+public:
+    KeysCollectionManager(std::string purpose, Seconds keyValidForInterval);
+
+    /**
+     * Return a key that is valid for the given time and also matches the keyId. Note that this call
+     * can block if it will need to do a refresh.
+     *
+     * Throws ErrorCode::ExceededTimeLimit if it times out.
+     */
+    StatusWith<KeysCollectionDocument> getKeyForValidation(OperationContext* opCtx,
+                                                           long long keyId,
+                                                           const LogicalTime& forThisTime);
+
+    /**
+     * Return a key that is valid for the given time. Note that this call can block if it will need
+     * to do a refresh.
+     *
+     * Throws ErrorCode::ExceededTimeLimit if it times out.
+     */
+    StatusWith<KeysCollectionDocument> getKeyForSigning(OperationContext* opCtx,
+                                                        const LogicalTime& forThisTime);
+
+    /**
+     * Starts a background thread that will constantly update the internal cache of keys.
+     *
+     * Cannot call this after stopMonitoring was called at least once.
+     */
+    void startMonitoring(ServiceContext* service);
+
+    /**
+     * Stops the background thread updating the cache.
+     */
+    void stopMonitoring();
+
+    /**
+     * Enable writing new keys to the config server primary. Should only be called if current node
+     * is the config primary.
+     */
+    void enableKeyGenerator(OperationContext* opCtx, bool doEnable);
+
+private:
+    /**
+     * This is responsible for periodically performing refresh in the background.
+     */
+    class PeriodicRunner {
+    public:
+        using RefreshFunc = stdx::function<StatusWith<KeysCollectionDocument>(OperationContext*)>;
+
+        /**
+         * Preemptively inform the monitoring thread it needs to perform a refresh. Returns an
+         * object
+         * that gets notified after the current round of refresh is over. Note that being notified
+         * can
+         * mean either of these things:
+         *
+         * 1. An error occurred and refresh was not performed.
+         * 2. No error occurred but no new key was found.
+         * 3. No error occurred and new keys were found.
+         */
+        void refreshNow(OperationContext* opCtx);
+
+        /**
+         * Sets the refresh function to use.
+         * Should only be used to bootstrap this refresher with initial strategy.
+         */
+        void setFunc(RefreshFunc newRefreshStrategy);
+
+        /**
+         * Switches the current strategy with a new one. This also waits to make sure that the old
+         * strategy is not being used and will no longer be used after this call.
+         */
+        void switchFunc(OperationContext* opCtx, RefreshFunc newRefreshStrategy);
+
+        /**
+         * Starts the refresh thread.
+         */
+        void start(ServiceContext* service,
+                   const std::string& threadName,
+                   Milliseconds refreshInterval);
+
+        /**
+         * Stops the refresh thread.
+         */
+        void stop();
+
+    private:
+        void _doPeriodicRefresh(ServiceContext* service,
+                                std::string threadName,
+                                Milliseconds refreshInterval);
+
+        stdx::mutex _mutex;  // protects all the member variables below.
+        std::shared_ptr<Notification<void>> _refreshRequest;
+        stdx::condition_variable _refreshNeededCV;
+
+        stdx::thread _backgroundThread;
+        std::shared_ptr<RefreshFunc> _doRefresh;
+
+        bool _inShutdown = false;
+    };
+
+    /**
+     * Return a key that is valid for the given time and also matches the keyId.
+     */
+    StatusWith<KeysCollectionDocument> _getKeyWithKeyIdCheck(long long keyId,
+                                                             const LogicalTime& forThisTime);
+
+    /**
+     * Return a key that is valid for the given time.
+     */
+    StatusWith<KeysCollectionDocument> _getKey(const LogicalTime& forThisTime);
+
+    const std::string _purpose;
+    const Seconds _keyValidForInterval;
+
+    // No mutex needed since the members below have their own mutexes.
+    KeysCollectionCacheReader _keysCache;
+    PeriodicRunner _refresher;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/keys_collection_manager_test.cpp b/src/mongo/db/keys_collection_manager_test.cpp
new file mode 100644
index 00000000000..f2c19f6a6c5
--- /dev/null
+++ b/src/mongo/db/keys_collection_manager_test.cpp
@@ -0,0 +1,300 @@
+/**
+ *    Copyright (C) 2017 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.
+ */
+
+#include "mongo/platform/basic.h"
+
+#include <set>
+#include <string>
+
+#include "mongo/db/jsobj.h"
+#include "mongo/db/keys_collection_document.h"
+#include "mongo/db/keys_collection_manager.h"
+#include "mongo/db/logical_clock.h"
+#include "mongo/db/namespace_string.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/s/catalog/dist_lock_manager_mock.h"
+#include "mongo/s/config_server_test_fixture.h"
+#include "mongo/stdx/memory.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/clock_source_mock.h"
+#include "mongo/util/fail_point_service.h"
+
+namespace mongo {
+
+class KeysManagerTest : public ConfigServerTestFixture {
+public:
+    KeysCollectionManager* keyManager() {
+        return _keyManager.get();
+    }
+
+protected:
+    void setUp() override {
+        ConfigServerTestFixture::setUp();
+
+        auto clockSource = stdx::make_unique<ClockSourceMock>();
+        operationContext()->getServiceContext()->setFastClockSource(std::move(clockSource));
+        _keyManager = stdx::make_unique<KeysCollectionManager>("dummy", Seconds(1));
+    }
+
+    void tearDown() override {
+        _keyManager->stopMonitoring();
+
+        ConfigServerTestFixture::tearDown();
+    }
+
+    std::unique_ptr<DistLockManager> makeDistLockManager(
+        std::unique_ptr<DistLockCatalog> distLockCatalog) override {
+        invariant(distLockCatalog);
+        return stdx::make_unique<DistLockManagerMock>(std::move(distLockCatalog));
+    }
+
+private:
+    std::unique_ptr<KeysCollectionManager> _keyManager;
+};
+
+TEST_F(KeysManagerTest, GetKeyForValidationTimesOutIfRefresherIsNotRunning) {
+    operationContext()->setDeadlineAfterNowBy(Microseconds(250 * 1000));
+
+    ASSERT_THROWS(
+        keyManager()->getKeyForValidation(operationContext(), 1, LogicalTime(Timestamp(100, 0))),
+        DBException);
+}
+
+TEST_F(KeysManagerTest, GetKeyForValidationErrorsIfKeyDoesntExist) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    auto keyStatus =
+        keyManager()->getKeyForValidation(operationContext(), 1, LogicalTime(Timestamp(100, 0)));
+    ASSERT_EQ(ErrorCodes::KeyNotFound, keyStatus.getStatus());
+}
+
+TEST_F(KeysManagerTest, GetKeyShouldReturnRightKey) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+
+    auto keyStatus =
+        keyManager()->getKeyForValidation(operationContext(), 1, LogicalTime(Timestamp(100, 0)));
+    ASSERT_OK(keyStatus.getStatus());
+
+    auto key = keyStatus.getValue();
+    ASSERT_EQ(1, key.getKeyId());
+    ASSERT_EQ(origKey1.getKey(), key.getKey());
+    ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
+}
+
+TEST_F(KeysManagerTest, GetKeyShouldErrorIfKeyIdMismatchKey) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+
+    auto keyStatus =
+        keyManager()->getKeyForValidation(operationContext(), 2, LogicalTime(Timestamp(100, 0)));
+    ASSERT_EQ(ErrorCodes::KeyNotFound, keyStatus.getStatus());
+}
+
+TEST_F(KeysManagerTest, GetKeyWithoutRefreshShouldReturnRightKey) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+    KeysCollectionDocument origKey2(
+        2, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(110, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey2.toBSON()));
+
+    {
+        auto keyStatus = keyManager()->getKeyForValidation(
+            operationContext(), 1, LogicalTime(Timestamp(100, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+
+        auto key = keyStatus.getValue();
+        ASSERT_EQ(1, key.getKeyId());
+        ASSERT_EQ(origKey1.getKey(), key.getKey());
+        ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
+    }
+
+    {
+        auto keyStatus = keyManager()->getKeyForValidation(
+            operationContext(), 2, LogicalTime(Timestamp(105, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+
+        auto key = keyStatus.getValue();
+        ASSERT_EQ(2, key.getKeyId());
+        ASSERT_EQ(origKey2.getKey(), key.getKey());
+        ASSERT_EQ(Timestamp(110, 0), key.getExpiresAt().asTimestamp());
+    }
+}
+
+TEST_F(KeysManagerTest, GetKeyForSigningTimesOutIfRefresherIsNotRunning) {
+    operationContext()->setDeadlineAfterNowBy(Microseconds(250 * 1000));
+
+    ASSERT_THROWS(
+        keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 0))),
+        DBException);
+}
+
+TEST_F(KeysManagerTest, GetKeyForSigningTimesOutIfKeyDoesntExist) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    operationContext()->setDeadlineAfterNowBy(Microseconds(250 * 1000));
+
+    ASSERT_THROWS(
+        keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 0))),
+        DBException);
+}
+
+TEST_F(KeysManagerTest, GetKeyForSigningShouldReturnRightKey) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+
+    auto keyStatus =
+        keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 0)));
+    ASSERT_OK(keyStatus.getStatus());
+
+    auto key = keyStatus.getValue();
+    ASSERT_EQ(1, key.getKeyId());
+    ASSERT_EQ(origKey1.getKey(), key.getKey());
+    ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
+}
+
+TEST_F(KeysManagerTest, GetKeyForSigningWithoutRefreshShouldReturnRightKey) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+    KeysCollectionDocument origKey2(
+        2, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(110, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey2.toBSON()));
+
+    {
+        auto keyStatus =
+            keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+
+        auto key = keyStatus.getValue();
+        ASSERT_EQ(1, key.getKeyId());
+        ASSERT_EQ(origKey1.getKey(), key.getKey());
+        ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
+    }
+
+    {
+        auto keyStatus =
+            keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(105, 0)));
+        ASSERT_OK(keyStatus.getStatus());
+
+        auto key = keyStatus.getValue();
+        ASSERT_EQ(2, key.getKeyId());
+        ASSERT_EQ(origKey2.getKey(), key.getKey());
+        ASSERT_EQ(Timestamp(110, 0), key.getExpiresAt().asTimestamp());
+    }
+}
+
+TEST_F(KeysManagerTest, ShouldCreateKeysIfKeyGeneratorEnabled) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    const LogicalTime currentTime(LogicalTime(Timestamp(100, 0)));
+    LogicalClock::get(operationContext())->initClusterTimeFromTrustedSource(currentTime);
+
+    keyManager()->enableKeyGenerator(operationContext(), true);
+
+    auto keyStatus =
+        keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 100)));
+    ASSERT_OK(keyStatus.getStatus());
+
+    auto key = keyStatus.getValue();
+    ASSERT_EQ(Timestamp(101, 0), key.getExpiresAt().asTimestamp());
+}
+
+TEST_F(KeysManagerTest, EnableModeFlipFlopStressTest) {
+    keyManager()->startMonitoring(getServiceContext());
+
+    const LogicalTime currentTime(LogicalTime(Timestamp(100, 0)));
+    LogicalClock::get(operationContext())->initClusterTimeFromTrustedSource(currentTime);
+
+    bool doEnable = true;
+
+    for (int x = 0; x < 10; x++) {
+        keyManager()->enableKeyGenerator(operationContext(), doEnable);
+
+        auto keyStatus =
+            keyManager()->getKeyForSigning(operationContext(), LogicalTime(Timestamp(100, 100)));
+        ASSERT_OK(keyStatus.getStatus());
+
+        auto key = keyStatus.getValue();
+        ASSERT_EQ(Timestamp(101, 0), key.getExpiresAt().asTimestamp());
+
+        doEnable = !doEnable;
+    }
+}
+
+TEST_F(KeysManagerTest, ShouldStillBeAbleToUpdateCacheEvenIfItCantCreateKeys) {
+    KeysCollectionDocument origKey1(
+        1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
+    ASSERT_OK(insertToConfigCollection(
+        operationContext(), NamespaceString(KeysCollectionDocument::ConfigNS), origKey1.toBSON()));
+
+    // Set the time to be very ahead so the updater will be forced to create new keys.
+    const LogicalTime currentTime(LogicalTime(Timestamp(20000, 0)));
+    const SignedLogicalTime fakeTime(currentTime, 2);
+    ASSERT_OK(LogicalClock::get(operationContext())->advanceClusterTimeFromTrustedSource(fakeTime));
+
+    FailPointEnableBlock failWriteBlock("failCollectionInserts");
+
+    {
+        FailPointEnableBlock failQueryBlock("planExecutorAlwaysFails");
+        keyManager()->startMonitoring(getServiceContext());
+        keyManager()->enableKeyGenerator(operationContext(), true);
+    }
+
+    auto keyStatus =
+        keyManager()->getKeyForValidation(operationContext(), 1, LogicalTime(Timestamp(100, 0)));
+    ASSERT_OK(keyStatus.getStatus());
+
+    auto key = keyStatus.getValue();
+    ASSERT_EQ(1, key.getKeyId());
+    ASSERT_EQ(origKey1.getKey(), key.getKey());
+    ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
+}
+
+}  // namespace mongo