SERVER-48709 Fix overflow in key manager wake up calculation

3 files changed, 50 insertions, 11 deletions

diff --git a/src/mongo/db/keys_collection_manager.cpp b/src/mongo/db/keys_collection_manager.cpp
index dd3277bdd1b..ff11c1779c3 100644
--- a/src/mongo/db/keys_collection_manager.cpp
+++ b/src/mongo/db/keys_collection_manager.cpp
@@ -61,14 +61,15 @@ Milliseconds kMaxRefreshWaitTime(10 * 60 * 1000);
 // a successful refresh.
 MONGO_FAIL_POINT_DEFINE(maxKeyRefreshWaitTimeOverrideMS);
 
-/**
- * Returns the amount of time to wait until the monitoring thread should attempt to refresh again.
- */
+}  // unnamed namespace
+
+namespace keys_collection_manager_util {
+
 Milliseconds howMuchSleepNeedFor(const LogicalTime& currentTime,
                                  const LogicalTime& latestExpiredAt,
                                  const Milliseconds& interval) {
-    auto currentSecs = currentTime.asTimestamp().getSecs();
-    auto expiredSecs = latestExpiredAt.asTimestamp().getSecs();
+    auto currentSecs = Seconds(currentTime.asTimestamp().getSecs());
+    auto expiredSecs = Seconds(latestExpiredAt.asTimestamp().getSecs());
 
     if (currentSecs >= expiredSecs) {
         // This means that the last round didn't generate a usable key for the current time.
@@ -76,16 +77,16 @@ Milliseconds howMuchSleepNeedFor(const LogicalTime& currentTime,
         return kRefreshIntervalIfErrored;
     }
 
-    auto millisBeforeExpire = 1000 * (expiredSecs - currentSecs);
+    Milliseconds millisBeforeExpire = Milliseconds(expiredSecs) - Milliseconds(currentSecs);
 
-    if (interval.count() <= millisBeforeExpire) {
+    if (interval <= millisBeforeExpire) {
         return interval;
     }
 
-    return Milliseconds(millisBeforeExpire);
+    return millisBeforeExpire;
 }
 
-}  // unnamed namespace
+}  // namespace keys_collection_manager_util
 
 KeysCollectionManager::KeysCollectionManager(std::string purpose,
                                              std::unique_ptr<KeysCollectionClient> client,
@@ -255,8 +256,8 @@ void KeysCollectionManager::PeriodicRunner::_doPeriodicRefresh(ServiceContext* s
 
                 _hasSeenKeys.store(true);
 
-                nextWakeup =
-                    howMuchSleepNeedFor(currentTime, latestKey.getExpiresAt(), refreshInterval);
+                nextWakeup = keys_collection_manager_util::howMuchSleepNeedFor(
+                    currentTime, latestKey.getExpiresAt(), refreshInterval);
             } else {
                 errorCount += 1;
                 nextWakeup = Milliseconds(kRefreshIntervalIfErrored.count() * errorCount);
diff --git a/src/mongo/db/keys_collection_manager.h b/src/mongo/db/keys_collection_manager.h
index 751de346ef0..0fe5ee6f732 100644
--- a/src/mongo/db/keys_collection_manager.h
+++ b/src/mongo/db/keys_collection_manager.h
@@ -50,6 +50,17 @@ class LogicalTime;
 class ServiceContext;
 class KeysCollectionClient;
 
+namespace keys_collection_manager_util {
+
+/**
+ * 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);
+
+}  // namespace keys_collection_manager_util
+
 /**
  * The KeysCollectionManager queries the config servers for keys that can be used for
  * HMAC computation. It maintains an internal background thread that is used to periodically
diff --git a/src/mongo/db/keys_collection_manager_sharding_test.cpp b/src/mongo/db/keys_collection_manager_sharding_test.cpp
index 97fd233b29e..d1190bb5efb 100644
--- a/src/mongo/db/keys_collection_manager_sharding_test.cpp
+++ b/src/mongo/db/keys_collection_manager_sharding_test.cpp
@@ -376,5 +376,32 @@ TEST_F(KeysManagerShardedTest, HasSeenKeysIsFalseUntilKeysAreFound) {
     ASSERT_EQ(true, keyManager()->hasSeenKeys());
 }
 
+LogicalTime addSeconds(const LogicalTime& logicalTime, const Seconds& seconds) {
+    auto asTimestamp = logicalTime.asTimestamp();
+    return LogicalTime(Timestamp(asTimestamp.getSecs() + seconds.count(), asTimestamp.getInc()));
+}
+
+TEST(KeysCollectionManagerUtilTest, HowMuchSleepNeededForCalculationDoesNotOverflow) {
+    auto secondsSinceEpoch = durationCount<Seconds>(Date_t::now().toDurationSinceEpoch());
+    auto defaultKeysIntervalSeconds = Seconds(KeysRotationIntervalSec);
+
+    // Mock inputs that would have caused an overflow without the changes from SERVER-48709.
+    // "currentTime" is the current logical time in the LogicalClock, which will typically be close
+    // to a timestamp constructed from the number of seconds since the unix epoch. "latestExpiredAt"
+    // is the highest expiration logical time of any key, which will at most be currentTime +
+    // (default key rotation interval * 2) because two valid keys are kept at a time. "interval" is
+    // the duration a key is valid for, which defaults to 90 days = 7,776,000 seconds.
+    auto currentTime = LogicalTime(Timestamp(secondsSinceEpoch, 0));
+    auto latestExpiredAt = addSeconds(currentTime, defaultKeysIntervalSeconds * 2);
+    auto interval = Milliseconds(defaultKeysIntervalSeconds);
+
+    // Despite the default rotation interval seconds * 1000 not fitting in a 32 bit unsigned
+    // integer (7,776,000,000 vs. 4,294,967,295), the calculation should not overflow, and the next
+    // wakeup should correctly be the default interval.
+    auto nextWakeupMillis =
+        keys_collection_manager_util::howMuchSleepNeedFor(currentTime, latestExpiredAt, interval);
+    ASSERT_EQ(nextWakeupMillis, interval);
+}
+
 }  // namespace
 }  // namespace mongo