feat: add JWT roles based throttling mechanism (#4506)

---------
Co-authored-by: Adam Stankiewicz <[email protected]>

course_discovery/apps/api/tests/jwt_utils.py

@@ -5,18 +5,21 @@
 from django.conf import settings
 
 
-def generate_jwt_payload(user):
+def generate_jwt_payload(user, payload=None):
     """Generate a valid JWT payload given a user."""
     now = int(time())
     ttl = 5
-    return {
+    jwt_payload = {
         'iss': settings.JWT_AUTH['JWT_ISSUER'],
         'aud': settings.JWT_AUTH['JWT_AUDIENCE'],
         'username': user.username,
         'email': user.email,
         'iat': now,
         'exp': now + ttl
     }
+    if payload:
+        jwt_payload.update(payload)
+    return jwt_payload
 
 
 def generate_jwt_token(payload):
@@ -29,8 +32,8 @@ def generate_jwt_header(token):
     return f'JWT {token}'
 
 
-def generate_jwt_header_for_user(user):
-    payload = generate_jwt_payload(user)
+def generate_jwt_header_for_user(user, payload=None):
+    payload = generate_jwt_payload(user, payload)
     token = generate_jwt_token(payload)
 
     return generate_jwt_header(token)

course_discovery/apps/core/tests/test_throttles.py

@@ -1,15 +1,19 @@
 from unittest.mock import patch
 
+import ddt
+from django.test.utils import override_settings
 from django.urls import reverse
 from rest_framework.test import APITestCase
 
+from course_discovery.apps.api.tests.jwt_utils import generate_jwt_header_for_user
 from course_discovery.apps.api.tests.mixins import SiteMixin
 from course_discovery.apps.core.models import UserThrottleRate
 from course_discovery.apps.core.tests.factories import USER_PASSWORD, UserFactory
 from course_discovery.apps.core.throttles import OverridableUserRateThrottle, throttling_cache
 from course_discovery.apps.publisher.tests.factories import GroupFactory
 
 
+@ddt.ddt
 class RateLimitingExceededTest(SiteMixin, APITestCase):
     """
     Testing rate limiting of API calls.
@@ -26,7 +30,15 @@ def tearDown(self):
         super().tearDown()
         throttling_cache().clear()
 
-    def _make_requests(self, count=None):
+    def _build_jwt_headers(self, user, payload=None):
+        """
+        Helper function for creating headers for the JWT authentication.
+        """
+        token = generate_jwt_header_for_user(user, payload)
+        headers = {'HTTP_AUTHORIZATION': token}
+        return headers
+
+    def _make_requests(self, count=None, **headers):
         """ Make multiple requests until the throttle's limit is exceeded.
 
         Returns
@@ -37,19 +49,19 @@ def _make_requests(self, count=None):
         with patch('rest_framework.views.APIView.throttle_classes', (OverridableUserRateThrottle,)):
             with patch.object(OverridableUserRateThrottle, 'THROTTLE_RATES', user_rates):
                 for __ in range(count - 1):
-                    response = self.client.get(self.url)
+                    response = self.client.get(self.url, **headers)
                     assert response.status_code == 200
-                response = self.client.get(self.url)
+                response = self.client.get(self.url, **headers)
         return response
 
-    def assert_rate_limit_successfully_exceeded(self):
+    def assert_rate_limit_successfully_exceeded(self, **headers):
         """ Asserts that the throttle's rate limit can be exceeded without encountering an error. """
-        response = self._make_requests()
+        response = self._make_requests(**headers)
         assert response.status_code == 200
 
-    def assert_rate_limited(self, count=None):
+    def assert_rate_limited(self, count=None, **headers):
         """ Asserts that the throttle's rate limit was exceeded and we were denied. """
-        response = self._make_requests(count)
+        response = self._make_requests(count, **headers)
         assert response.status_code == 429
 
     def test_rate_limiting(self):
@@ -84,3 +96,23 @@ def test_staff_with_user_throttle_rate(self):
         self.user.save()
         UserThrottleRate.objects.create(user=self.user, rate='10/hour')
         self.assert_rate_limited(11)
+
+    @ddt.data(
+        ([], True),
+        (['enterprise_learner:*'], False),
+        (['enterprise_admin:*'], False),
+        (['enterprise_openedx_operator:*'], False),
+    )
+    @ddt.unpack
+    @override_settings(ENHANCED_THROTTLE_LIMIT='10/hour')
+    def test_enterprise_user_throttling_with_jwt_authentication(self, jwt_roles, is_rate_limited):
+        """ Verify enterprise users are throttled at a higher rate. """
+        payload = {
+            'roles': jwt_roles,
+        }
+        headers = self._build_jwt_headers(self.user, payload)
+        if is_rate_limited:
+            self.assert_rate_limited(**headers)
+        else:
+            self.assert_rate_limit_successfully_exceeded(count=5, **headers)
+            self.assert_rate_limited(**headers)

course_discovery/apps/core/throttles.py

@@ -1,5 +1,7 @@
 """Custom API throttles."""
+from django.conf import settings
 from django.core.cache import InvalidCacheBackendError, caches
+from edx_rest_framework_extensions.auth.jwt.decoder import configured_jwt_decode_handler
 from rest_framework.throttling import UserRateThrottle
 
 from course_discovery.apps.core.models import UserThrottleRate
@@ -16,6 +18,23 @@ def throttling_cache():
         return caches['default']
 
 
+def is_priviledged_user_with_enhanced_throttle(request):
+    """
+    Determine whether a JWT-authenticated user has increased throttling rates
+    based on the `roles` in the decoded JWT token associated with the request.
+    """
+    jwt_token = request.auth
+    if not jwt_token:
+        return False
+    decoded_jwt = configured_jwt_decode_handler(jwt_token)
+    roles = decoded_jwt.get('roles', [])
+    return any(
+        privileged_role_keyword in role
+        for privileged_role_keyword in settings.ENHANCED_THROTTLE_JWT_ROLE_KEYWORDS
+        for role in roles
+    )
+
+
 class OverridableUserRateThrottle(UserRateThrottle):
     """Rate throttling of requests, overridable on a per-user basis."""
     cache = throttling_cache()
@@ -28,10 +47,15 @@ def allow_request(self, request, view):
                 # Override this throttle's rate if applicable
                 user_throttle = UserThrottleRate.objects.get(user=user)
                 self.rate = user_throttle.rate
-                self.num_requests, self.duration = self.parse_rate(self.rate)
             except UserThrottleRate.DoesNotExist:
                 # If we don't have a custom user override, skip throttling if they are a privileged user
                 if user.is_superuser or user.is_staff or is_publisher_user(user):
                     return True
 
+                # If the user has privileged throttling limits, increase the rate
+                if is_priviledged_user_with_enhanced_throttle(request):
+                    self.rate = settings.ENHANCED_THROTTLE_LIMIT
+
+        self.num_requests, self.duration = self.parse_rate(self.rate)
+
         return super().allow_request(request, view)

course_discovery/settings/base.py

@@ -790,3 +790,8 @@
 )
 
 ENABLE_COURSE_REVIEWS_ADMIN = False
+
+# Keywords that will be searched for inside the `roles` key of the JWT in case a user uses JWT authentication.
+# If the keyword is found, the user has more lenient throttling limits.
+ENHANCED_THROTTLE_JWT_ROLE_KEYWORDS = []
+ENHANCED_THROTTLE_LIMIT = '400/hour'

course_discovery/settings/test.py

@@ -161,3 +161,5 @@
 ELASTICSEARCH_DSL_QUERYSET_PAGINATION = 10000
 
 ELASTICSEARCH_DSL_LOAD_PER_QUERY = 10000
+
+ENHANCED_THROTTLE_JWT_ROLE_KEYWORDS = ['enterprise']

docs/decisions/0029-jwt-based-throttling.rst

@@ -0,0 +1,41 @@
+29. JWT roles based Throttling
+==============================
+
+Status
+--------
+Accepted (Dec 2024)
+
+Context
+---------
+Course Discovery APIs are used by a number of consumers like external partners and marketers, the LMS, Publisher, and the
+enterprise-catalog service. One of these consumers that stands out is the Enterprise Learner Portal. Just like Publisher, this
+portal is an MFE that queries discovery at runtime. However, unlike Publisher, where our users are mostly staff and partners, the
+learner portal is intended to be used by regular enterprise learners and admins i.e those without any staff or privileged access.
+
+As of late, our regular throttling limits have proven to be a little too aggressive for some of these learners. As this is clearly 
+non-malicious traffic, we would like to find a way to set more lenient throttling limits for these learners.
+
+Decision
+----------
+Since these learners authenticate with Discovery by using JWT tokens issued by the LMS, we have decided to use the `roles` key
+in the JWT to identify them. Enterprise customers are guaranteed to have one of a small number of fixed roles assigned to them. 
+Once we identify that an incoming request's user has one of these roles in their JWT, we will enable higher rate limits for them.
+
+Two new settings, `ENHANCED_THROTTLE_JWT_ROLE_KEYWORDS` AND `ENHANCED_THROTTLE_LIMIT` will be added. The `ENHANCED_THROTTLE_JWT_ROLE_KEYWORDS`
+setting accepts a list of strings representing keywords that identify privileged JWT roles. Each keyword is matched against the
+roles in a JWT token; if a keyword is found in any role, the user's throttle limits are set to the value of `ENHANCED_THROTTLE_LIMIT`.  
+
+Consequences
+--------------
+- Some regular users (i.e without staff or any django group privileges) will have privileged rate limits
+- We will need to be vigilant to ensure that the list of roles we consider privileged is accurate i.e its entries should only
+  be associated to enterprise learners (i.e the users we need higher limits for). Furthermore, it should be impossible for
+  "other" users to attain those roles.
+
+Alternatives Considered
+-------------------------
+- Raise the throttle limits for everyone. However, that increases the possibility of scraping and service degradation attacks
+  by malicious actors
+- Extend the user model in discovery to store identification information for users with enhanced throttle limits. However, that
+  conflicts with the responsibility of the LMS as the sole authentication provider. Furthermore, adding this information in 
+  discovery for each enterprise user would needlessly complicate the enterprise onboarding and offboarding process.