From 1ab5aa42a735a2e4b1926303a37f531af7db6a48 Mon Sep 17 00:00:00 2001
From: wangruidong <940853815@qq.com>
Date: Wed, 24 Dec 2025 15:37:57 +0800
Subject: [PATCH] feat: Implement rate limiting with custom throttling classes
 for different user types

---
 apps/common/drf/throttling.py    | 34 ++++++++++++++++++++++++++++++++
 apps/jumpserver/settings/libs.py |  8 ++++++++
 2 files changed, 42 insertions(+)
 create mode 100644 apps/common/drf/throttling.py

diff --git a/apps/common/drf/throttling.py b/apps/common/drf/throttling.py
new file mode 100644
index 000000000..3483c6de4
--- /dev/null
+++ b/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
+        }
diff --git a/apps/jumpserver/settings/libs.py b/apps/jumpserver/settings/libs.py
index a5ec30c9a..21b59a93d 100644
--- a/apps/jumpserver/settings/libs.py
+++ b/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',