feat: Add JWT key-set size custom attribute to aid in key rotations

Part of edx/edx-arch-experiments#261

CHANGELOG.rst

@@ -12,6 +12,14 @@ Change Log
 Unreleased
 ----------
 
+[8.5.0] - 2023-04-05
+--------------------
+
+Added
+~~~~~
+
+* Added ``jwt_auth_verify_keys_count`` custom attribute to aid in key rotations
+
 [8.4.1] - 2022-12-18
 --------------------
 

edx_rest_framework_extensions/__init__.py

@@ -1,3 +1,3 @@
 """ edx Django REST Framework extensions. """
 
-__version__ = '8.4.1'  # pragma: no cover
+__version__ = '8.5.0'  # pragma: no cover

edx_rest_framework_extensions/auth/jwt/decoder.py

@@ -10,6 +10,7 @@
 
 import jwt
 from django.conf import settings
+from edx_django_utils.monitoring import set_custom_attribute
 from jwkest.jwk import KEYS
 from jwkest.jws import JWS
 from rest_framework_jwt.settings import api_settings
@@ -173,6 +174,9 @@ def _set_filters(token):
 
 def _verify_jwt_signature(token, jwt_issuer, decode_symmetric_token):
     key_set = _get_signing_jwk_key_set(jwt_issuer, add_symmetric_keys=decode_symmetric_token)
+    # Reporting the size of the public key list will aid in key rotations: Once all
+    # servers are reporting a larger list, the private key can be changed.
+    set_custom_attribute('jwt_auth_verify_keys_count', len(key_set))
 
     try:
         _ = JWS().verify_compact(token, key_set)

edx_rest_framework_extensions/auth/jwt/tests/test_decoder.py

@@ -199,6 +199,23 @@ def test_success_asymmetric_jwt_decode(self):
         token = generate_asymmetric_jwt_token(self.payload)
         self.assertEqual(get_asymmetric_only_jwt_decode_handler(token), self.payload)
 
+    @mock.patch('edx_rest_framework_extensions.auth.jwt.decoder.set_custom_attribute')
+    def test_keyset_size_monitoring(self, mock_set_custom_attribute):
+        """
+        Validates that a custom attribute is recorded for the keyset size.
+        """
+        token = generate_asymmetric_jwt_token(self.payload)
+
+        # The secret key is included by default making a list of length 2, but for
+        # asymmetric-only there is only 1 key in the keyset.
+        self.assertEqual(jwt_decode_handler(token), self.payload)
+        self.assertEqual(get_asymmetric_only_jwt_decode_handler(token), self.payload)
+
+        assert mock_set_custom_attribute.call_args_list == [
+            mock.call('jwt_auth_verify_keys_count', 2),
+            mock.call('jwt_auth_verify_keys_count', 1),
+        ]
+
 
 def _jwt_decode_handler_with_defaults(token):  # pylint: disable=unused-argument
     """