Merge pull request #988 from TG1999/throttling

Enable throttling

Author: TG1999

CHANGELOG.rst

@@ -3,6 +3,13 @@ Release notes
 
 
 
+Version v30.2.2
+----------------
+
+- We enabled API throttling for a basic user and for a staff user 
+  they can have unlimited access on API.
+
+
 Version v30.2.1
 ----------------
 

vulnerabilities/tests/test_throttling.py

@@ -0,0 +1,48 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/nexB/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+from django.contrib.auth import get_user_model
+from rest_framework.test import APIClient
+from rest_framework.test import APITestCase
+
+User = get_user_model()
+
+
+class ThrottleApiTests(APITestCase):
+    def setUp(self):
+        # create a basic user
+        self.user = User.objects.create_user("username", "[email protected]", "secret")
+        self.auth = f"Token {self.user.auth_token.key}"
+        self.csrf_client = APIClient(enforce_csrf_checks=True)
+        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
+
+        # create a staff user
+        self.staff_user = User.objects.create_user(
+            "staff", "[email protected]", "secret", is_staff=True
+        )
+        self.staff_auth = f"Token {self.staff_user.auth_token.key}"
+        self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
+        self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
+
+    def test_api_throttling(self):
+
+        # A basic user can only access API 5 times a day
+        for i in range(0, 5):
+            response = self.csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+            response = self.staff_csrf_client.get("/api/packages")
+            self.assertEqual(response.status_code, 200)
+
+        response = self.csrf_client.get("/api/packages")
+        # 429 - too many requests for basic user
+        self.assertEqual(response.status_code, 429)
+
+        response = self.staff_csrf_client.get("/api/packages", format="json")
+        # 200 - staff user can access API unlimited times
+        self.assertEqual(response.status_code, 200)

vulnerabilities/throttling.py

@@ -0,0 +1,24 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/nexB/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+from django.contrib.auth import get_user_model
+from rest_framework.throttling import UserRateThrottle
+
+User = get_user_model()
+
+
+class StaffUserRateThrottle(UserRateThrottle):
+    def allow_request(self, request, view):
+        """
+        Do not apply throttling for superusers and admins.
+        """
+        if request.user.is_superuser or request.user.is_staff:
+            return True
+
+        return super().allow_request(request, view)

vulnerablecode/settings.py

@@ -150,9 +150,12 @@
 
 LOGIN_REDIRECT_URL = "/"
 LOGOUT_REDIRECT_URL = "/"
+THROTTLING_RATE = env.str("THROTTLING_RATE", default="1000/day")
 
 if IS_TESTS:
     VULNERABLECODEIO_REQUIRE_AUTHENTICATION = True
+    THROTTLING_RATE = "5/day"
+
 
 USE_L10N = True
 
@@ -184,6 +187,10 @@
         "django_filters.rest_framework.DjangoFilterBackend",
         "rest_framework.filters.SearchFilter",
     ),
+    "DEFAULT_THROTTLE_CLASSES": [
+        "vulnerabilities.throttling.StaffUserRateThrottle",
+    ],
+    "DEFAULT_THROTTLE_RATES": {"user": THROTTLING_RATE},
     "DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",
     # Limit the load on the Database returning a small number of records by default. https://github.com/nexB/vulnerablecode/issues/819
     "PAGE_SIZE": 10,