From 604906790b567563bc311feeaa9e667cf182e533 Mon Sep 17 00:00:00 2001
From: Tim McCormack <tmccormack@edx.org>
Date: Wed, 5 Apr 2023 14:04:12 +0000
Subject: [PATCH] feat: Add JWT key-set size custom attribute to aid in key
 rotations (#317)

Part of https://github.com/edx/edx-arch-experiments/issues/261
---
 CHANGELOG.rst                                   |  8 ++++++++
 edx_rest_framework_extensions/__init__.py       |  2 +-
 .../auth/jwt/decoder.py                         |  8 ++++++++
 .../auth/jwt/tests/test_decoder.py              | 17 +++++++++++++++++
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 33b2287f..f40a2c73 100644
--- a/CHANGELOG.rst
+++ b/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
 --------------------
 
diff --git a/edx_rest_framework_extensions/__init__.py b/edx_rest_framework_extensions/__init__.py
index 736ee77e..caa68a28 100644
--- a/edx_rest_framework_extensions/__init__.py
+++ b/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
diff --git a/edx_rest_framework_extensions/auth/jwt/decoder.py b/edx_rest_framework_extensions/auth/jwt/decoder.py
index 18929c35..77f4cdfa 100644
--- a/edx_rest_framework_extensions/auth/jwt/decoder.py
+++ b/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,13 @@ 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)
+    # .. custom_attribute_name: jwt_auth_verify_keys_count
+    # .. custom_attribute_description: Number of JWT verification keys in use for this
+    #   verification. Should be same as number of asymmetric public keys, plus one if
+    #   a symmetric key secret is set. This is intended to aid in key rotations; once
+    #   the average count stabilizes at a higher number after adding a public key, it
+    #   should be safe to change the secret key.
+    set_custom_attribute('jwt_auth_verify_keys_count', len(key_set))
 
     try:
         _ = JWS().verify_compact(token, key_set)
diff --git a/edx_rest_framework_extensions/auth/jwt/tests/test_decoder.py b/edx_rest_framework_extensions/auth/jwt/tests/test_decoder.py
index 73dd030f..b16301e3 100644
--- a/edx_rest_framework_extensions/auth/jwt/tests/test_decoder.py
+++ b/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
     """