SERVER-28435 Implement getNewKey for catalog client

11 files changed, 183 insertions, 1 deletions

diff --git a/src/mongo/crypto/sha1_block.cpp b/src/mongo/crypto/sha1_block.cpp
index 9debee81fcc..3f36ef6611f 100644
--- a/src/mongo/crypto/sha1_block.cpp
+++ b/src/mongo/crypto/sha1_block.cpp
@@ -90,4 +90,8 @@ bool SHA1Block::operator!=(const SHA1Block& other) const {
     return !(*this == other);
 }
 
+std::ostream& operator<<(std::ostream& os, const SHA1Block& sha1) {
+    return os << sha1.toString();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/crypto/sha1_block.h b/src/mongo/crypto/sha1_block.h
index a228540bcb0..3cf92c9e187 100644
--- a/src/mongo/crypto/sha1_block.h
+++ b/src/mongo/crypto/sha1_block.h
@@ -97,4 +97,6 @@ private:
     HashType _hash;
 };
 
+std::ostream& operator<<(std::ostream& os, const SHA1Block& sha1);
+
 }  // namespace mongo
diff --git a/src/mongo/db/keys_collection_document.cpp b/src/mongo/db/keys_collection_document.cpp
index 24392fdd9d5..155eb730a31 100644
--- a/src/mongo/db/keys_collection_document.cpp
+++ b/src/mongo/db/keys_collection_document.cpp
@@ -47,6 +47,8 @@ const char kExpiresAtFieldName[] = "expiresAt";
 
 }  // namespace
 
+const std::string KeysCollectionDocument::ConfigNS = "admin.system.keys";
+
 StatusWith<KeysCollectionDocument> KeysCollectionDocument::fromBSON(const BSONObj& source) {
     long long keyId;
     Status status = bsonExtractIntegerField(source, kKeyIdFieldName, &keyId);
diff --git a/src/mongo/db/keys_collection_document.h b/src/mongo/db/keys_collection_document.h
index b1eef99343b..5571982fdac 100644
--- a/src/mongo/db/keys_collection_document.h
+++ b/src/mongo/db/keys_collection_document.h
@@ -47,6 +47,8 @@ namespace mongo {
  */
 class KeysCollectionDocument {
 public:
+    static const std::string ConfigNS;
+
     KeysCollectionDocument(long long keyId,
                            std::string purpose,
                            TimeProofService::Key key,
@@ -75,7 +77,7 @@ public:
     const LogicalTime& getExpiresAt() const;
 
 private:
-    long long _keyId;
+    long long _keyId = 0;
     std::string _purpose;
     TimeProofService::Key _key;
     LogicalTime _expiresAt;
diff --git a/src/mongo/s/catalog/SConscript b/src/mongo/s/catalog/SConscript
index 879d3080806..f32a443534a 100644
--- a/src/mongo/s/catalog/SConscript
+++ b/src/mongo/s/catalog/SConscript
@@ -11,6 +11,7 @@ env.Library(
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/db/write_concern_options',
+        '$BUILD_DIR/mongo/db/keys_collection_document',
     ],
 )
 
diff --git a/src/mongo/s/catalog/sharding_catalog_client.h b/src/mongo/s/catalog/sharding_catalog_client.h
index 9a37abad3f3..f0e5a7090ce 100644
--- a/src/mongo/s/catalog/sharding_catalog_client.h
+++ b/src/mongo/s/catalog/sharding_catalog_client.h
@@ -34,6 +34,7 @@
 #include <vector>
 
 #include "mongo/base/disallow_copying.h"
+#include "mongo/db/keys_collection_document.h"
 #include "mongo/db/repl/optime_with.h"
 #include "mongo/db/write_concern_options.h"
 #include "mongo/s/catalog/dist_lock_manager.h"
@@ -52,6 +53,7 @@ struct ChunkVersion;
 class CollectionType;
 class ConnectionString;
 class DatabaseType;
+class LogicalTime;
 class NamespaceString;
 class OperationContext;
 class ShardKeyPattern;
@@ -357,6 +359,15 @@ public:
                                                      repl::ReadConcernLevel readConcern) = 0;
 
     /**
+     * Returns keys for the given purpose and with an expiresAt value greater than newerThanThis.
+     */
+    virtual StatusWith<std::vector<KeysCollectionDocument>> getNewKeys(
+        OperationContext* opCtx,
+        StringData purpose,
+        const LogicalTime& newerThanThis,
+        repl::ReadConcernLevel readConcernLevel) = 0;
+
+    /**
      * Directly sends the specified command to the config server and returns the response.
      *
      * NOTE: Usage of this function is disallowed in new code, which should instead go through
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
index 8c36fccae77..8495c5f6b07 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.cpp
@@ -1838,4 +1838,42 @@ Status ShardingCatalogClientImpl::appendInfoForConfigServerDatabases(
     return Status::OK();
 }
 
+StatusWith<std::vector<KeysCollectionDocument>> ShardingCatalogClientImpl::getNewKeys(
+    OperationContext* opCtx,
+    StringData purpose,
+    const LogicalTime& newerThanThis,
+    repl::ReadConcernLevel readConcernLevel) {
+    auto config = Grid::get(opCtx)->shardRegistry()->getConfigShard();
+
+    BSONObjBuilder queryBuilder;
+    queryBuilder.append("purpose", purpose);
+    queryBuilder.append("expiresAt", BSON("$gt" << newerThanThis.asTimestamp()));
+
+    auto findStatus =
+        config->exhaustiveFindOnConfig(opCtx,
+                                       kConfigReadSelector,
+                                       readConcernLevel,
+                                       NamespaceString(KeysCollectionDocument::ConfigNS),
+                                       queryBuilder.obj(),
+                                       BSON("expiresAt" << 1),
+                                       boost::none);
+
+    if (!findStatus.isOK()) {
+        return findStatus.getStatus();
+    }
+
+    const auto& keyDocs = findStatus.getValue().docs;
+    std::vector<KeysCollectionDocument> keys;
+    for (auto&& keyDoc : keyDocs) {
+        auto parseStatus = KeysCollectionDocument::fromBSON(keyDoc);
+        if (!parseStatus.isOK()) {
+            return parseStatus.getStatus();
+        }
+
+        keys.push_back(std::move(parseStatus.getValue()));
+    }
+
+    return keys;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/catalog/sharding_catalog_client_impl.h b/src/mongo/s/catalog/sharding_catalog_client_impl.h
index 0a94a3a18eb..38f5fe3c77c 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_impl.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_impl.h
@@ -185,6 +185,12 @@ public:
                                const BSONObj& cmdObj,
                                BSONObjBuilder* result);
 
+    StatusWith<std::vector<KeysCollectionDocument>> getNewKeys(
+        OperationContext* opCtx,
+        StringData purpose,
+        const LogicalTime& newerThanThis,
+        repl::ReadConcernLevel readConcernLevel) override;
+
 private:
     /**
      * Selects an optimal shard on which to place a newly created database from the set of
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
index 730a411af29..7261d6d7093 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.cpp
@@ -235,4 +235,12 @@ Status ShardingCatalogClientMock::appendInfoForConfigServerDatabases(
     return Status::OK();
 }
 
+StatusWith<std::vector<KeysCollectionDocument>> ShardingCatalogClientMock::getNewKeys(
+    OperationContext* opCtx,
+    StringData purpose,
+    const LogicalTime& newerThanThis,
+    repl::ReadConcernLevel readConcernLevel) {
+    return {ErrorCodes::InternalError, "Method not implemented"};
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/catalog/sharding_catalog_client_mock.h b/src/mongo/s/catalog/sharding_catalog_client_mock.h
index a2d223f2384..1be61153b10 100644
--- a/src/mongo/s/catalog/sharding_catalog_client_mock.h
+++ b/src/mongo/s/catalog/sharding_catalog_client_mock.h
@@ -162,6 +162,12 @@ public:
                                               const BSONObj& listDatabasesCmd,
                                               BSONArrayBuilder* builder) override;
 
+    StatusWith<std::vector<KeysCollectionDocument>> getNewKeys(
+        OperationContext* opCtx,
+        StringData purpose,
+        const LogicalTime& newerThanThis,
+        repl::ReadConcernLevel readConcernLevel) override;
+
 private:
     std::unique_ptr<DistLockManager> _distLockManager;
 };
diff --git a/src/mongo/s/catalog/sharding_catalog_test.cpp b/src/mongo/s/catalog/sharding_catalog_test.cpp
index d062d59023e..7b21c2429f4 100644
--- a/src/mongo/s/catalog/sharding_catalog_test.cpp
+++ b/src/mongo/s/catalog/sharding_catalog_test.cpp
@@ -2435,5 +2435,107 @@ TEST_F(ShardingCatalogClientTest, RetryOnFindCommandNetworkErrorSucceedsAtMaxRet
     future.timed_get(kFutureTimeout);
 }
 
+TEST_F(ShardingCatalogClientTest, GetNewKeys) {
+    configTargeter()->setFindHostReturnValue(HostAndPort("config:123"));
+
+    std::string purpose("none");
+    LogicalTime currentTime(Timestamp(1234, 5678));
+    repl::ReadConcernLevel readConcernLevel(repl::ReadConcernLevel::kMajorityReadConcern);
+
+    auto future = launchAsync([this, purpose, currentTime, readConcernLevel] {
+        auto status =
+            catalogClient()->getNewKeys(operationContext(), purpose, currentTime, readConcernLevel);
+        ASSERT_OK(status.getStatus());
+        return status.getValue();
+    });
+
+    LogicalTime dummyTime(Timestamp(9876, 5432));
+    auto randomKey1 = TimeProofService::generateRandomKey();
+    KeysCollectionDocument key1(1, "none", randomKey1, dummyTime);
+
+    LogicalTime dummyTime2(Timestamp(123456, 789));
+    auto randomKey2 = TimeProofService::generateRandomKey();
+    KeysCollectionDocument key2(2, "none", randomKey2, dummyTime2);
+
+    onFindCommand([this, key1, key2](const RemoteCommandRequest& request) {
+        ASSERT_EQ("config:123", request.target.toString());
+        ASSERT_EQ("admin", request.dbname);
+
+        const NamespaceString nss(request.dbname, request.cmdObj.firstElement().String());
+        ASSERT_EQ(KeysCollectionDocument::ConfigNS, nss.ns());
+
+        auto query = assertGet(QueryRequest::makeFromFindCommand(nss, request.cmdObj, false));
+
+        BSONObj expectedQuery(
+            fromjson("{purpose: 'none',"
+                     "expiresAt: {$gt: {$timestamp: {t: 1234, i: 5678}}}}"));
+
+        ASSERT_EQ(KeysCollectionDocument::ConfigNS, query->ns());
+        ASSERT_BSONOBJ_EQ(expectedQuery, query->getFilter());
+        ASSERT_BSONOBJ_EQ(BSON("expiresAt" << 1), query->getSort());
+        ASSERT_FALSE(query->getLimit().is_initialized());
+
+        checkReadConcern(request.cmdObj, Timestamp(0, 0), repl::OpTime::kUninitializedTerm);
+
+        return vector<BSONObj>{key1.toBSON(), key2.toBSON()};
+    });
+
+    const auto keyDocs = future.timed_get(kFutureTimeout);
+    ASSERT_EQ(2u, keyDocs.size());
+
+    const auto& key1Result = keyDocs.front();
+    ASSERT_EQ(1, key1Result.getKeyId());
+    ASSERT_EQ("none", key1Result.getPurpose());
+    ASSERT_EQ(randomKey1, key1Result.getKey());
+    ASSERT_EQ(Timestamp(9876, 5432), key1Result.getExpiresAt().asTimestamp());
+
+    const auto& key2Result = keyDocs.back();
+    ASSERT_EQ(2, key2Result.getKeyId());
+    ASSERT_EQ("none", key2Result.getPurpose());
+    ASSERT_EQ(randomKey2, key2Result.getKey());
+    ASSERT_EQ(Timestamp(123456, 789), key2Result.getExpiresAt().asTimestamp());
+}
+
+TEST_F(ShardingCatalogClientTest, GetNewKeysWithEmptyCollection) {
+    configTargeter()->setFindHostReturnValue(HostAndPort("config:123"));
+
+    std::string purpose("none");
+    LogicalTime currentTime(Timestamp(1234, 5678));
+    repl::ReadConcernLevel readConcernLevel(repl::ReadConcernLevel::kMajorityReadConcern);
+
+    auto future = launchAsync([this, purpose, currentTime, readConcernLevel] {
+        auto status =
+            catalogClient()->getNewKeys(operationContext(), purpose, currentTime, readConcernLevel);
+        ASSERT_OK(status.getStatus());
+        return status.getValue();
+    });
+
+    onFindCommand([this](const RemoteCommandRequest& request) {
+        ASSERT_EQ("config:123", request.target.toString());
+        ASSERT_EQ("admin", request.dbname);
+
+        const NamespaceString nss(request.dbname, request.cmdObj.firstElement().String());
+        ASSERT_EQ(KeysCollectionDocument::ConfigNS, nss.ns());
+
+        auto query = assertGet(QueryRequest::makeFromFindCommand(nss, request.cmdObj, false));
+
+        BSONObj expectedQuery(
+            fromjson("{purpose: 'none',"
+                     "expiresAt: {$gt: {$timestamp: {t: 1234, i: 5678}}}}"));
+
+        ASSERT_EQ(KeysCollectionDocument::ConfigNS, query->ns());
+        ASSERT_BSONOBJ_EQ(expectedQuery, query->getFilter());
+        ASSERT_BSONOBJ_EQ(BSON("expiresAt" << 1), query->getSort());
+        ASSERT_FALSE(query->getLimit().is_initialized());
+
+        checkReadConcern(request.cmdObj, Timestamp(0, 0), repl::OpTime::kUninitializedTerm);
+
+        return vector<BSONObj>{};
+    });
+
+    const auto keyDocs = future.timed_get(kFutureTimeout);
+    ASSERT_EQ(0u, keyDocs.size());
+}
+
 }  // namespace
 }  // namespace mongo