From 410f04c327b933306dec2810790611583046ff36 Mon Sep 17 00:00:00 2001
From: Nikhil <118773738+pablohashescobar@users.noreply.github.com>
Date: Tue, 21 May 2024 15:03:01 +0530
Subject: [PATCH] dev: rest framework throttling (#4534)

---
 .../plane/authentication/adapter/error.py     |  2 ++
 .../plane/authentication/adapter/exception.py | 15 +++++++++++
 apiserver/plane/authentication/rate_limit.py  | 26 +++++++++++++++++++
 .../plane/authentication/views/app/check.py   | 10 ++++++-
 .../views/app/password_management.py          |  6 ++++-
 .../plane/authentication/views/space/check.py |  6 ++++-
 .../views/space/password_management.py        |  5 ++++
 7 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 apiserver/plane/authentication/rate_limit.py

diff --git a/apiserver/plane/authentication/adapter/error.py b/apiserver/plane/authentication/adapter/error.py
index 47dbc8e8a..457a67f4f 100644
--- a/apiserver/plane/authentication/adapter/error.py
+++ b/apiserver/plane/authentication/adapter/error.py
@@ -55,6 +55,8 @@ AUTHENTICATION_ERROR_CODES = {
     "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
index 12845ea02..a6f7637a9 100644
--- 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
new file mode 100644
index 000000000..744bd38fe
--- /dev/null
+++ 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
index 4f164e308..2448aee55 100644
--- a/apiserver/plane/authentication/views/app/check.py
+++ b/apiserver/plane/authentication/views/app/check.py
@@ -15,7 +15,7 @@ from plane.authentication.adapter.error import (
     AuthenticationException,
     AUTHENTICATION_ERROR_CODES,
 )
-
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 class EmailCheckSignUpEndpoint(APIView):
 
@@ -23,6 +23,10 @@ class EmailCheckSignUpEndpoint(APIView):
         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
index dd14ceb91..43054867e 100644
--- a/apiserver/plane/authentication/views/app/password_management.py
+++ b/apiserver/plane/authentication/views/app/password_management.py
@@ -32,7 +32,7 @@ from plane.authentication.adapter.error import (
     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
index 83f52e28f..1b20d19a2 100644
--- a/apiserver/plane/authentication/views/space/check.py
+++ b/apiserver/plane/authentication/views/space/check.py
@@ -15,7 +15,7 @@ from plane.authentication.adapter.error import (
     AUTHENTICATION_ERROR_CODES,
     AuthenticationException,
 )
-
+from plane.authentication.rate_limit import AuthenticationThrottle
 
 class EmailCheckEndpoint(APIView):
 
@@ -23,6 +23,10 @@ class EmailCheckEndpoint(APIView):
         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
index fa20fa618..3e0379b96 100644
--- a/apiserver/plane/authentication/views/space/password_management.py
+++ b/apiserver/plane/authentication/views/space/password_management.py
@@ -32,6 +32,7 @@ from plane.authentication.adapter.error import (
     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")