Enable throttling

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/tests/test_fix_api.py

@@ -259,12 +259,13 @@ def test_api_with_single_vulnerability_and_vulnerable_package(self):
         }
 
     def test_api_with_all_vulnerable_packages(self):
-        with self.assertNumQueries(4):
+        with self.assertNumQueries(5):
             # There are 4 queries:
             # 1. SAVEPOINT
             # 2. Authenticating user
-            # 3. Get all vulnerable packages
-            # 4. RELEASE SAVEPOINT
+            # 3. Checking if user is staff user for throttling purposes
+            # 4. Get all vulnerable packages
+            # 5. RELEASE SAVEPOINT
             response = self.csrf_client.get(f"/api/packages/all", format="json").data
             assert len(response) == 11
             assert response == [

vulnerabilities/throttling.py

@@ -0,0 +1,52 @@
+#
+# 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 ExceptionalUserRateThrottle(UserRateThrottle):
+    def allow_request(self, request, view):
+        """
+        Give special access to a few special accounts.
+
+        Mirrors code in super class with minor tweaks.
+        """
+        if self.rate is None:
+            return True
+
+        self.key = self.get_cache_key(request, view)
+        if self.key is None:
+            return True
+
+        self.history = self.cache.get(self.key, [])
+        self.now = self.timer()
+
+        # Adjust if user has special privileges.
+
+        user = User.objects.get(username=request.user.username)
+
+        if user:
+            if user.is_superuser or user.is_staff:
+                # No throttling for superusers or staff.
+                return True
+
+            else:
+                self.num_requests = self.num_requests
+                self.duration = self.duration
+
+        # Drop any requests from the history which have now passed the
+        # throttle duration
+        while self.history and self.history[-1] <= self.now - self.duration:
+            self.history.pop()
+        if len(self.history) >= self.num_requests:
+            return self.throttle_failure()
+        return self.throttle_success()

vulnerablecode/settings.py

@@ -184,6 +184,10 @@
         "django_filters.rest_framework.DjangoFilterBackend",
         "rest_framework.filters.SearchFilter",
     ),
+    "DEFAULT_THROTTLE_CLASSES": [
+        "vulnerabilities.throttling.ExceptionalUserRateThrottle",
+    ],
+    "DEFAULT_THROTTLE_RATES": {"user": "1000/hour"},
     "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,