feat: Implement rate limiting with custom throttling classes for different user types

apps/common/drf/throttling.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+from rest_framework.throttling import SimpleRateThrottle
+
+
+class RateThrottle(SimpleRateThrottle):
+
+    def __init__(self):
+        # Override the usual SimpleRateThrottle, because we can't determine
+        # the rate until called by the view.
+        pass
+
+    def allow_request(self, request, view):
+        if getattr(request, "user", None) and request.user.is_authenticated:
+            if getattr(request.user, "is_service_account", False):
+                self.scope = "service_account"
+            else:
+                self.scope = "user"
+        else:
+            self.scope = "anon"
+
+        self.rate = self.get_rate()
+        self.num_requests, self.duration = self.parse_rate(self.rate)
+        return super().allow_request(request, view)
+
+    def get_cache_key(self, request, view):
+        if request.user and request.user.is_authenticated:
+            ident = request.user.pk
+        else:
+            ident = self.get_ident(request)
+
+        return self.cache_format % {
+            'scope': self.scope,
+            'ident': ident
+        }

apps/jumpserver/settings/libs.py

@@ -38,6 +38,14 @@ REST_FRAMEWORK = {
         "oauth2_provider.contrib.rest_framework.OAuth2Authentication",
         'authentication.backends.drf.SessionAuthentication',
     ),
+    'DEFAULT_THROTTLE_CLASSES': (
+        'common.drf.throttling.RateThrottle',
+    ),
+    'DEFAULT_THROTTLE_RATES': {
+        'anon': '60/min',
+        'user': '180/min',
+        'service_account': '300/min',
+    },
     'DEFAULT_FILTER_BACKENDS': (
         'django_filters.rest_framework.DjangoFilterBackend',
         'common.drf.filters.SearchFilter',