SERVER-73165 Improve JWKManager's just in-time refresh

3 files changed, 41 insertions, 21 deletions

diff --git a/src/mongo/crypto/jwk_manager.cpp b/src/mongo/crypto/jwk_manager.cpp
index 2465cf63f2d..32234719608 100644
--- a/src/mongo/crypto/jwk_manager.cpp
+++ b/src/mongo/crypto/jwk_manager.cpp
@@ -56,11 +56,11 @@ StringData reduceInt(StringData value) {
 
 }  // namespace
 
-JWKManager::JWKManager(StringData source) : _keyURI(source) {
+JWKManager::JWKManager(StringData source) : _keyURI(source), _isKeyModified(false) {
     _loadKeysFromUri(true /* isInitialLoad */);
 }
 
-JWKManager::JWKManager(BSONObj keys) {
+JWKManager::JWKManager(BSONObj keys) : _isKeyModified(false) {
     _setAndValidateKeys(keys, true /* isInitialLoad */);
 }
 
@@ -117,12 +117,6 @@ void JWKManager::_setAndValidateKeys(const BSONObj& keys, bool isInitialLoad) {
                 str::stream() << "Key IDs must be unique, duplicate '" << keyId << "'",
                 newKeyMaterial->find(keyId) == newKeyMaterial->end());
 
-        // Does not need to be loaded atomically because isInitialLoad is only true when invoked
-        // from the constructor, so there will not be any concurrent reads.
-        if (isInitialLoad) {
-            _initialKeyMaterial.insert({keyId, key.copy()});
-        }
-
         newKeyMaterial->insert({keyId, key.copy()});
 
         auto swValidator = JWSValidator::create(JWK.getType(), key);
@@ -132,6 +126,9 @@ void JWKManager::_setAndValidateKeys(const BSONObj& keys, bool isInitialLoad) {
         newValidators->insert({keyId, shValidator});
         LOGV2_DEBUG(7070202, 3, "Loaded JWK key", "kid"_attr = keyId, "typ"_attr = JWK.getType());
     }
+    // We compare the old keys from the new ones and return if any old key is not present in the new
+    // set of keys.
+    _isKeyModified |= _haveKeysBeenModified(*newKeyMaterial);
 
     // A mutex is not used here because no single thread consumes both _validators and _keyMaterial.
     // _keyMaterial is used purely for reflection of current key state while _validators is used
@@ -177,4 +174,23 @@ void JWKManager::serialize(BSONObjBuilder* bob) const {
     jwks.serialize(bob);
 }
 
+bool JWKManager::_haveKeysBeenModified(const KeyMap& newKeyMaterial) const {
+    if (!_keyMaterial) {
+        return false;
+    }
+    const bool hasBeenModified =
+        std::any_of((*_keyMaterial).cbegin(), (*_keyMaterial).cend(), [&](const auto& entry) {
+            auto newKey = newKeyMaterial.find(entry.first);
+            if (newKey == newKeyMaterial.end()) {
+                // Key no longer exists in this JWKS.
+                return true;
+            }
+
+            // Key still exists.
+            return entry.second.woCompare(newKey->second) != 0;
+        });
+
+    return hasBeenModified;
+}
+
 }  // namespace mongo::crypto
diff --git a/src/mongo/crypto/jwk_manager.h b/src/mongo/crypto/jwk_manager.h
index 03a05cb1ee8..2752f99235f 100644
--- a/src/mongo/crypto/jwk_manager.h
+++ b/src/mongo/crypto/jwk_manager.h
@@ -68,13 +68,14 @@ public:
     }
 
     /**
-     * Get a snapshot of the keys the key manager was initialized with. It is not
-     * modified after initialization and can be used to determine whether
-     * invalidation is necessary.
-     * TODO SERVER-73165 Generalize this to better handle invalidation.
+     * Get current keys.
      */
-    const KeyMap& getInitialKeys() const {
-        return _initialKeyMaterial;
+    const KeyMap& getKeys() const {
+        return *_keyMaterial;
+    }
+
+    bool getIsKeyModified() const {
+        return _isKeyModified;
     }
 
     /**
@@ -87,15 +88,18 @@ private:
 
     void _loadKeysFromUri(bool isInitialLoad);
 
-private:
-    // Stores the key material that the manager was initialized with.
-    KeyMap _initialKeyMaterial;
+    bool _haveKeysBeenModified(const KeyMap& newKeyMaterial) const;
 
+private:
     // Stores the current key material of the manager, which may have been refreshed.
     std::shared_ptr<KeyMap> _keyMaterial;
 
     boost::optional<std::string> _keyURI;
     std::shared_ptr<std::map<std::string, SharedValidator>> _validators;
+
+    // If an existing key got deleted or modified while doing a just in time refresh, we activate
+    // this flag to indicate that a refresh occurred during this JWKManager's lifetime.
+    bool _isKeyModified;
 };
 
 }  // namespace mongo::crypto
diff --git a/src/mongo/crypto/jwt_test.cpp b/src/mongo/crypto/jwt_test.cpp
index ae414ae17f0..6acccaeae44 100644
--- a/src/mongo/crypto/jwt_test.cpp
+++ b/src/mongo/crypto/jwt_test.cpp
@@ -57,11 +57,11 @@ TEST(JWKManager, parseJWKSetBasicFromSource) {
     manager.serialize(&bob);
     ASSERT_BSONOBJ_EQ(bob.obj(), data);
 
-    const auto& initialKeys = manager.getInitialKeys();
+    const auto& currentKeys = manager.getKeys();
     for (const auto& key : data["keys"_sd].Obj()) {
-        auto initialKey = initialKeys.find(key["kid"_sd].str());
-        ASSERT(initialKey != initialKeys.end());
-        ASSERT_BSONOBJ_EQ(key.Obj(), initialKey->second);
+        auto currentKey = currentKeys.find(key["kid"_sd].str());
+        ASSERT(currentKey != currentKeys.end());
+        ASSERT_BSONOBJ_EQ(key.Obj(), currentKey->second);
     }
 
     for (const auto& key : data["keys"_sd].Obj()) {