TEST_F(KeyGeneratorUpdateTest, ShouldCreateAnotherKeyIfOnlyOneKeyExists) {
KeyGenerator generator("dummy", catalogClient(), Seconds(5));
LogicalClock::get(operationContext())
->setClusterTimeFromTrustedSource(LogicalTime(Timestamp(100, 2)));
KeysCollectionDocument origKey1(
1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
ASSERT_OK(insertToConfigCollection(
operationContext(), KeysCollectionDocument::ConfigNS, origKey1.toBSON()));
{
auto allKeys = getKeys(operationContext());
ASSERT_EQ(1u, allKeys.size());
const auto& key1 = allKeys.front();
ASSERT_EQ(1, key1.getKeyId());
ASSERT_EQ("dummy", key1.getPurpose());
ASSERT_EQ(Timestamp(105, 0), key1.getExpiresAt().asTimestamp());
}
auto currentTime = LogicalClock::get(operationContext())->getClusterTime();
auto generateStatus = generator.generateNewKeysIfNeeded(operationContext());
ASSERT_OK(generateStatus);
{
auto allKeys = getKeys(operationContext());
ASSERT_EQ(2u, allKeys.size());
const auto& key1 = allKeys.front();
ASSERT_EQ(1, key1.getKeyId());
ASSERT_EQ("dummy", key1.getPurpose());
ASSERT_EQ(origKey1.getKey(), key1.getKey());
ASSERT_EQ(Timestamp(105, 0), key1.getExpiresAt().asTimestamp());
const auto& key2 = allKeys.back();
ASSERT_EQ(currentTime.asTimestamp().asLL(), key2.getKeyId());
ASSERT_EQ("dummy", key2.getPurpose());
ASSERT_EQ(Timestamp(110, 0), key2.getExpiresAt().asTimestamp());
ASSERT_NE(key1.getKey(), key2.getKey());
}
}
TEST_F(KeyGeneratorUpdateTest, ShouldCreate2KeysIfAllKeysAreExpired) {
KeyGenerator generator("dummy", catalogClient(), Seconds(5));
LogicalClock::get(operationContext())
->setClusterTimeFromTrustedSource(LogicalTime(Timestamp(120, 2)));
KeysCollectionDocument origKey1(
1, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(105, 0)));
ASSERT_OK(insertToConfigCollection(
operationContext(), KeysCollectionDocument::ConfigNS, origKey1.toBSON()));
KeysCollectionDocument origKey2(
2, "dummy", TimeProofService::generateRandomKey(), LogicalTime(Timestamp(110, 0)));
ASSERT_OK(insertToConfigCollection(
operationContext(), KeysCollectionDocument::ConfigNS, origKey2.toBSON()));
{
auto allKeys = getKeys(operationContext());
ASSERT_EQ(2u, allKeys.size());
const auto& key1 = allKeys.front();
ASSERT_EQ(1, key1.getKeyId());
ASSERT_EQ("dummy", key1.getPurpose());
ASSERT_EQ(Timestamp(105, 0), key1.getExpiresAt().asTimestamp());
const auto& key2 = allKeys.back();
ASSERT_EQ(2, key2.getKeyId());
ASSERT_EQ("dummy", key2.getPurpose());
ASSERT_EQ(Timestamp(110, 0), key2.getExpiresAt().asTimestamp());
}
auto currentTime = LogicalClock::get(operationContext())->getClusterTime();
auto generateStatus = generator.generateNewKeysIfNeeded(operationContext());
ASSERT_OK(generateStatus);
auto allKeys = getKeys(operationContext());
ASSERT_EQ(4u, allKeys.size());
auto citer = allKeys.cbegin();
std::set<std::string> seenKeys;
{
const auto& key = *citer;
ASSERT_EQ(1, key.getKeyId());
ASSERT_EQ("dummy", key.getPurpose());
ASSERT_EQ(origKey1.getKey(), key.getKey());
ASSERT_EQ(Timestamp(105, 0), key.getExpiresAt().asTimestamp());
bool inserted = false;
std::tie(std::ignore, inserted) = seenKeys.insert(key.getKey().toString());
ASSERT_TRUE(inserted);
}
{
++citer;
const auto& key = *citer;
ASSERT_EQ(2, key.getKeyId());
ASSERT_EQ("dummy", key.getPurpose());
ASSERT_EQ(origKey2.getKey(), key.getKey());
ASSERT_EQ(Timestamp(110, 0), key.getExpiresAt().asTimestamp());
bool inserted = false;
std::tie(std::ignore, inserted) = seenKeys.insert(key.getKey().toString());
ASSERT_TRUE(inserted);
}
{
++citer;
const auto& key = *citer;
ASSERT_EQ(currentTime.asTimestamp().asLL(), key.getKeyId());
ASSERT_EQ("dummy", key.getPurpose());
ASSERT_EQ(Timestamp(125, 0), key.getExpiresAt().asTimestamp());
bool inserted = false;
std::tie(std::ignore, inserted) = seenKeys.insert(key.getKey().toString());
ASSERT_TRUE(inserted);
}
{
++citer;
const auto& key = *citer;
ASSERT_EQ(currentTime.asTimestamp().asLL() + 1, key.getKeyId());
ASSERT_EQ("dummy", key.getPurpose());
ASSERT_NE(origKey1.getKey(), key.getKey());
ASSERT_NE(origKey2.getKey(), key.getKey());
ASSERT_EQ(Timestamp(130, 0), key.getExpiresAt().asTimestamp());
bool inserted = false;
std::tie(std::ignore, inserted) = seenKeys.insert(key.getKey().toString());
ASSERT_TRUE(inserted);
}
}
Status KeyGenerator::generateNewKeysIfNeeded(OperationContext* opCtx) {
if (MONGO_FAIL_POINT(disableKeyGeneration)) {
return {ErrorCodes::FailPointEnabled, "key generation disabled"};
}
auto currentTime = LogicalClock::get(opCtx)->getClusterTime();
auto keyStatus = _client->getNewKeys(opCtx, _purpose, currentTime);
if (!keyStatus.isOK()) {
return keyStatus.getStatus();
}
const auto& newKeys = keyStatus.getValue();
auto keyIter = newKeys.cbegin();
LogicalTime currentKeyExpiresAt;
long long keyId = currentTime.asTimestamp().asLL();
if (keyIter == newKeys.cend()) {
currentKeyExpiresAt = addSeconds(currentTime, _keyValidForInterval);
auto status = insertNewKey(opCtx, _client, keyId, _purpose, currentKeyExpiresAt);
if (!status.isOK()) {
return status;
}
keyId++;
} else if (keyIter->getExpiresAt() < currentTime) {
currentKeyExpiresAt = addSeconds(currentTime, _keyValidForInterval);
auto status = insertNewKey(opCtx, _client, keyId, _purpose, currentKeyExpiresAt);
if (!status.isOK()) {
return status;
}
keyId++;
++keyIter;
} else {
currentKeyExpiresAt = keyIter->getExpiresAt();
++keyIter;
}
// Create a new key in advance if we don't have a key on standby after the current one
// expires.
// Note: Convert this block into a loop if more reserved keys are desired.
if (keyIter == newKeys.cend()) {
auto reserveKeyExpiresAt = addSeconds(currentKeyExpiresAt, _keyValidForInterval);
auto status = insertNewKey(opCtx, _client, keyId, _purpose, reserveKeyExpiresAt);
if (!status.isOK()) {
return status;
}
} else if (keyIter->getExpiresAt() < currentTime) {
currentKeyExpiresAt = addSeconds(currentKeyExpiresAt, _keyValidForInterval);
auto status = insertNewKey(opCtx, _client, keyId, _purpose, currentKeyExpiresAt);
if (!status.isOK()) {
return status;
}
}
return Status::OK();
}