SERVER-48709 Fix overflow in key manager wake up calculation

(cherry picked from commit 0cb70e9577c46257798d0385b15ec6bff8dbd28d)

3 files changed, 50 insertions, 11 deletions

diff --git a/src/mongo/db/keys_collection_manager_sharding.cpp b/src/mongo/db/keys_collection_manager_sharding.cpp
index e149e14cd67..ccb9c56f75d 100644
--- a/src/mongo/db/keys_collection_manager_sharding.cpp
+++ b/src/mongo/db/keys_collection_manager_sharding.cpp
@@ -60,14 +60,15 @@ Milliseconds kMaxRefreshWaitTime(10 * 60 * 1000);
 // a successful refresh.
 MONGO_FP_DECLARE(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.
@@ -75,16 +76,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
 
 KeysCollectionManagerSharding::KeysCollectionManagerSharding(
     std::string purpose, std::unique_ptr<KeysCollectionClient> client, Seconds keyValidForInterval)
@@ -250,8 +251,8 @@ void KeysCollectionManagerSharding::PeriodicRunner::_doPeriodicRefresh(
                     _hasSeenKeys = 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_sharding.h b/src/mongo/db/keys_collection_manager_sharding.h
index 63ff3cdbcc0..fab92a83a32 100644
--- a/src/mongo/db/keys_collection_manager_sharding.h
+++ b/src/mongo/db/keys_collection_manager_sharding.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
+
 /**
  * This implementation of the KeysCollectionManager queries the config servers for keys.
  * It maintains in internal background thread that is used to periodically refresh
diff --git a/src/mongo/db/keys_collection_manager_sharding_test.cpp b/src/mongo/db/keys_collection_manager_sharding_test.cpp
index 14b37758089..8bde5344ea6 100644
--- a/src/mongo/db/keys_collection_manager_sharding_test.cpp
+++ b/src/mongo/db/keys_collection_manager_sharding_test.cpp
@@ -395,4 +395,31 @@ TEST_F(KeysManagerShardedTest, ShouldNotReturnKeysInFeatureCompatibilityVersion3
     ASSERT_EQ(ErrorCodes::KeyNotFound, keyStatus.getStatus());
 }
 
+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 mongo