dev: rest framework throttling (makeplane#4534)

wangxukun · May 21, 2024 · 410f04c · 410f04c
1 parent 4feec35
commit 410f04c
Show file tree

Hide file tree

Showing 7 changed files with 67 additions and 3 deletions.
diff --git a/apiserver/plane/authentication/adapter/error.py b/apiserver/plane/authentication/adapter/error.py
@@ -55,6 +55,8 @@
     "ADMIN_USER_ALREADY_EXIST": 5180,
     "ADMIN_USER_DOES_NOT_EXIST": 5185,
     "ADMIN_USER_DEACTIVATED": 5190,
+    # Rate limit
+    "RATE_LIMIT_EXCEEDED": 5900,
 }
 
 

diff --git a/apiserver/plane/authentication/adapter/exception.py b/apiserver/plane/authentication/adapter/exception.py
@@ -1,5 +1,10 @@
+# Third party imports
 from rest_framework.views import exception_handler
 from rest_framework.exceptions import NotAuthenticated
+from rest_framework.exceptions import Throttled
+
+# Module imports
+from plane.authentication.adapter.error import AuthenticationException, AUTHENTICATION_ERROR_CODES
 
 
 def auth_exception_handler(exc, context):
@@ -9,4 +14,14 @@ def auth_exception_handler(exc, context):
     if isinstance(exc, NotAuthenticated):
         response.status_code = 401
 
+    # Check if an Throttled exception is raised.
+    if isinstance(exc, Throttled):
+        exc = AuthenticationException(
+            error_code=AUTHENTICATION_ERROR_CODES["RATE_LIMIT_EXCEEDED"],
+            error_message="RATE_LIMIT_EXCEEDED",
+        )
+        response.data = exc.get_error_dict()
+        response.status_code = 429
+
+    # Return the response that is generated by the default exception handler.
     return response
diff --git a/apiserver/plane/authentication/rate_limit.py b/apiserver/plane/authentication/rate_limit.py
@@ -0,0 +1,26 @@
+# Third party imports
+from rest_framework.throttling import AnonRateThrottle
+from rest_framework import status
+from rest_framework.response import Response
+
+# Module imports
+from plane.authentication.adapter.error import (
+    AuthenticationException,
+    AUTHENTICATION_ERROR_CODES,
+)
+
+
+class AuthenticationThrottle(AnonRateThrottle):
+    rate = "30/minute"
+    scope = "authentication"
+
+    def throttle_failure_view(self, request, *args, **kwargs):
+        try:
+            raise AuthenticationException(
+                error_code=AUTHENTICATION_ERROR_CODES["RATE_LIMIT_EXCEEDED"],
+                error_message="RATE_LIMIT_EXCEEDED",
+            )
+        except AuthenticationException as e:
+            return Response(
+                e.get_error_dict(), status=status.HTTP_429_TOO_MANY_REQUESTS
+            )
diff --git a/apiserver/plane/authentication/views/app/check.py b/apiserver/plane/authentication/views/app/check.py
@@ -15,14 +15,18 @@
     AuthenticationException,
     AUTHENTICATION_ERROR_CODES,
 )
-
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 class EmailCheckSignUpEndpoint(APIView):
 
     permission_classes = [
         AllowAny,
     ]
 
+    throttle_classes = [
+        AuthenticationThrottle,
+    ]
+
     def post(self, request):
         try:
             # Check instance configuration
@@ -86,6 +90,10 @@ class EmailCheckSignInEndpoint(APIView):
         AllowAny,
     ]
 
+    throttle_classes = [
+        AuthenticationThrottle,
+    ]
+
     def post(self, request):
         try:
             # Check instance configuration

diff --git a/apiserver/plane/authentication/views/app/password_management.py b/apiserver/plane/authentication/views/app/password_management.py
@@ -32,7 +32,7 @@
     AuthenticationException,
     AUTHENTICATION_ERROR_CODES,
 )
-
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 def generate_password_token(user):
     uidb64 = urlsafe_base64_encode(smart_bytes(user.id))
@@ -46,6 +46,10 @@ class ForgotPasswordEndpoint(APIView):
         AllowAny,
     ]
 
+    throttle_classes = [
+        AuthenticationThrottle,
+    ]
+
     def post(self, request):
         email = request.data.get("email")
 

diff --git a/apiserver/plane/authentication/views/space/check.py b/apiserver/plane/authentication/views/space/check.py
@@ -15,14 +15,18 @@
     AUTHENTICATION_ERROR_CODES,
     AuthenticationException,
 )
-
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 class EmailCheckEndpoint(APIView):
 
     permission_classes = [
         AllowAny,
     ]
 
+    throttle_classes = [
+        AuthenticationThrottle,
+    ]
+
     def post(self, request):
         # Check instance configuration
         instance = Instance.objects.first()

diff --git a/apiserver/plane/authentication/views/space/password_management.py b/apiserver/plane/authentication/views/space/password_management.py
@@ -32,6 +32,7 @@
     AuthenticationException,
     AUTHENTICATION_ERROR_CODES,
 )
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 
 def generate_password_token(user):
@@ -46,6 +47,10 @@ class ForgotPasswordSpaceEndpoint(APIView):
         AllowAny,
     ]
 
+    throttle_classes = [
+        AuthenticationThrottle,
+    ]
+
     def post(self, request):
         email = request.data.get("email")