Skip to content

Commit 7cad40f

Browse files
TG1999TG1999
committed
Remove endpoint based throttling
Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>
1 parent 5f15a26 commit 7cad40f

5 files changed

Lines changed: 72 additions & 182 deletions

File tree

‎vulnerabilities/api.py‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616
from rest_framework import viewsets
1717
from rest_framework.decorators import action
1818
from rest_framework.response import Response
19+
from rest_framework.throttling import AnonRateThrottle
20+
from rest_framework.throttling import UserRateThrottle
1921

2022
from vulnerabilities.models import Alias
2123
from vulnerabilities.models import Package
@@ -231,11 +233,10 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
231233
serializer_class = PackageSerializer
232234
filter_backends = (filters.DjangoFilterBackend,)
233235
filterset_class = PackageFilterSet
234-
throttle_classes = [StaffUserRateThrottle]
235-
throttle_scope = "packages"
236+
throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
236237

237238
# TODO: Fix the swagger documentation for this endpoint
238-
@action(detail=False, methods=["post"], throttle_scope="bulk_search_packages")
239+
@action(detail=False, methods=["post"])
239240
def bulk_search(self, request):
240241
"""
241242
Lookup for vulnerable packages using many Package URLs at once.
@@ -289,7 +290,7 @@ def bulk_search(self, request):
289290
vulnerable_purls = [str(package.package_url) for package in vulnerable_purls]
290291
return Response(data=vulnerable_purls)
291292

292-
@action(detail=False, methods=["get"], throttle_scope="vulnerable_packages")
293+
@action(detail=False, methods=["get"])
293294
def all(self, request):
294295
"""
295296
Return the Package URLs of all packages known to be vulnerable.
@@ -341,8 +342,7 @@ def get_queryset(self):
341342
serializer_class = VulnerabilitySerializer
342343
filter_backends = (filters.DjangoFilterBackend,)
343344
filterset_class = VulnerabilityFilterSet
344-
throttle_classes = [StaffUserRateThrottle]
345-
throttle_scope = "vulnerabilities"
345+
throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
346346

347347

348348
class CPEFilterSet(filters.FilterSet):
@@ -363,11 +363,10 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
363363
).distinct()
364364
serializer_class = VulnerabilitySerializer
365365
filter_backends = (filters.DjangoFilterBackend,)
366-
throttle_classes = [StaffUserRateThrottle]
366+
throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]
367367
filterset_class = CPEFilterSet
368-
throttle_scope = "cpes"
369368

370-
@action(detail=False, methods=["post"], throttle_scope="bulk_search_cpes")
369+
@action(detail=False, methods=["post"])
371370
def bulk_search(self, request):
372371
"""
373372
Lookup for vulnerabilities using many CPEs at once.
@@ -409,5 +408,4 @@ class AliasViewSet(viewsets.ReadOnlyModelViewSet):
409408
serializer_class = VulnerabilitySerializer
410409
filter_backends = (filters.DjangoFilterBackend,)
411410
filterset_class = AliasFilterSet
412-
throttle_classes = [StaffUserRateThrottle]
413-
throttle_scope = "aliases"
411+
throttle_classes = [StaffUserRateThrottle, AnonRateThrottle]

‎vulnerabilities/tests/test_auth.py‎

Lines changed: 0 additions & 33 deletions
This file was deleted.

‎vulnerabilities/tests/test_throttling.py‎

Lines changed: 39 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,10 @@ def setUp(self):
3030
self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
3131

3232
self.csrf_client_anon = APIClient(enforce_csrf_checks=True)
33+
self.csrf_client_anon_1 = APIClient(enforce_csrf_checks=True)
3334

34-
def test_packages_endpoint_throttling(self):
35-
36-
# A basic user can only access /packages endpoint 10 times a day
37-
for i in range(0, 10):
35+
def test_package_endpoint_throttling(self):
36+
for i in range(0, 20):
3837
response = self.csrf_client.get("/api/packages")
3938
self.assertEqual(response.status_code, 200)
4039
response = self.staff_csrf_client.get("/api/packages")
@@ -50,86 +49,48 @@ def test_packages_endpoint_throttling(self):
5049

5150
# A anonymous user can only access /packages endpoint 10 times a day
5251
for i in range(0, 10):
52+
print(i)
5353
response = self.csrf_client_anon.get("/api/packages")
5454
self.assertEqual(response.status_code, 200)
5555

5656
response = self.csrf_client_anon.get("/api/packages")
5757
# 429 - too many requests for anon user
5858
self.assertEqual(response.status_code, 429)
59+
self.assertEqual(
60+
response.data.get("message"),
61+
"Your request has been throttled. Please contact support@nexb.com",
62+
)
5963

60-
def test_cpes_endpoint_throttling(self):
61-
62-
# A basic user can only access /cpes endpoint 4 times a day
63-
for i in range(0, 4):
64-
response = self.csrf_client.get("/api/cpes")
65-
self.assertEqual(response.status_code, 200)
66-
response = self.staff_csrf_client.get("/api/cpes")
67-
self.assertEqual(response.status_code, 200)
68-
69-
response = self.csrf_client.get("/api/cpes")
70-
# 429 - too many requests for basic user
71-
self.assertEqual(response.status_code, 429)
72-
73-
response = self.staff_csrf_client.get("/api/cpes", format="json")
74-
# 200 - staff user can access API unlimited times
75-
self.assertEqual(response.status_code, 200)
76-
77-
def test_all_vulnerable_packages_endpoint_throttling(self):
78-
79-
# A basic user can only access /packages/all 1 time a day
80-
for i in range(0, 1):
81-
response = self.csrf_client.get("/api/packages/all")
82-
self.assertEqual(response.status_code, 200)
83-
response = self.staff_csrf_client.get("/api/packages/all")
84-
self.assertEqual(response.status_code, 200)
85-
86-
response = self.csrf_client.get("/api/packages/all")
87-
# 429 - too many requests for basic user
88-
self.assertEqual(response.status_code, 429)
89-
90-
response = self.staff_csrf_client.get("/api/packages/all", format="json")
91-
# 200 - staff user can access API unlimited times
92-
self.assertEqual(response.status_code, 200)
93-
94-
def test_vulnerabilities_endpoint_throttling(self):
95-
96-
# A basic user can only access /vulnerabilities 8 times a day
97-
for i in range(0, 8):
98-
response = self.csrf_client.get("/api/vulnerabilities")
99-
self.assertEqual(response.status_code, 200)
100-
response = self.staff_csrf_client.get("/api/vulnerabilities")
101-
self.assertEqual(response.status_code, 200)
102-
103-
response = self.csrf_client.get("/api/vulnerabilities")
104-
# 429 - too many requests for basic user
64+
response = self.csrf_client_anon.get("/api/vulnerabilties")
65+
# 429 - too many requests for anon user
10566
self.assertEqual(response.status_code, 429)
67+
self.assertEqual(
68+
response.data.get("message"),
69+
"Your request has been throttled. Please contact support@nexb.com",
70+
)
10671

107-
response = self.staff_csrf_client.get("/api/vulnerabilities", format="json")
108-
# 200 - staff user can access API unlimited times
109-
self.assertEqual(response.status_code, 200)
110-
111-
def test_aliases_endpoint_throttling(self):
11272

113-
# A basic user can only access /alias 2 times a day
114-
for i in range(0, 2):
115-
response = self.csrf_client.get("/api/aliases")
116-
self.assertEqual(response.status_code, 200)
117-
response = self.staff_csrf_client.get("/api/aliases")
118-
self.assertEqual(response.status_code, 200)
73+
class ThrottleApiTestsForPostRequest(APITestCase):
74+
def setUp(self):
75+
# create a basic user
76+
self.user = ApiUser.objects.create_api_user(username="e@mail.com")
77+
self.auth = f"Token {self.user.auth_token.key}"
78+
self.csrf_client = APIClient(enforce_csrf_checks=True)
79+
self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
11980

120-
response = self.csrf_client.get("/api/aliases")
121-
# 429 - too many requests for basic user
122-
self.assertEqual(response.status_code, 429)
81+
# create a staff user
82+
self.staff_user = ApiUser.objects.create_api_user(username="staff@mail.com", is_staff=True)
83+
self.staff_auth = f"Token {self.staff_user.auth_token.key}"
84+
self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
85+
self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
12386

124-
response = self.staff_csrf_client.get("/api/aliases", format="json")
125-
# 200 - staff user can access API unlimited times
126-
self.assertEqual(response.status_code, 200)
87+
self.csrf_client_anon = APIClient(enforce_csrf_checks=True)
12788

12889
def test_bulk_search_packages_endpoint_throttling(self):
12990
data = json.dumps({"purls": ["pkg:foo/bar"]})
13091

131-
# A basic user can only access /packages/bulk_search 6 times a day
132-
for i in range(0, 6):
92+
# A basic user can only access /packages/bulk_search 20 times a day
93+
for i in range(0, 20):
13394
response = self.csrf_client.post(
13495
"/api/packages/bulk_search", data=data, content_type="application/json"
13596
)
@@ -151,28 +112,19 @@ def test_bulk_search_packages_endpoint_throttling(self):
151112
# 200 - staff user can access API unlimited times
152113
self.assertEqual(response.status_code, 200)
153114

154-
def test_bulk_search_cpes_endpoint_throttling(self):
155-
data = json.dumps({"cpes": ["cpe:foo/bar"]})
156-
157-
# A basic user can only access /cpes/bulk_search 5 times a day
158-
for i in range(0, 5):
159-
response = self.csrf_client.post(
160-
"/api/cpes/bulk_search", data=data, content_type="application/json"
161-
)
162-
self.assertEqual(response.status_code, 200)
163-
response = self.staff_csrf_client.post(
164-
"/api/cpes/bulk_search", data=data, content_type="application/json"
115+
# A anonymous user can only access /packages endpoint 10 times a day
116+
for i in range(0, 10):
117+
response = self.csrf_client_anon.post(
118+
"/api/packages/bulk_search", data=data, content_type="application/json"
165119
)
166120
self.assertEqual(response.status_code, 200)
167121

168-
response = self.csrf_client.post(
169-
"/api/cpes/bulk_search", data=data, content_type="application/json"
122+
response = self.csrf_client_anon.post(
123+
"/api/packages/bulk_search", data=data, content_type="application/json"
170124
)
171-
# 429 - too many requests for basic user
125+
# 429 - too many requests for anon user
172126
self.assertEqual(response.status_code, 429)
173-
174-
response = self.staff_csrf_client.post(
175-
"/api/cpes/bulk_search", data=data, content_type="application/json"
127+
self.assertEqual(
128+
response.data.get("message"),
129+
"Your request has been throttled. Please contact support@nexb.com",
176130
)
177-
# 200 - staff user can access API unlimited times
178-
self.assertEqual(response.status_code, 200)

‎vulnerabilities/throttling.py‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,23 +6,33 @@
66
# See https://github.com/nexB/vulnerablecode for support or download.
77
# See https://aboutcode.org for more information about nexB OSS projects.
88
#
9-
from rest_framework.throttling import ScopedRateThrottle
9+
from rest_framework.exceptions import Throttled
10+
from rest_framework.throttling import UserRateThrottle
11+
from rest_framework.views import exception_handler
1012

1113

12-
class StaffUserRateThrottle(ScopedRateThrottle):
14+
class StaffUserRateThrottle(UserRateThrottle):
1315
def allow_request(self, request, view):
1416
"""
1517
Do not apply throttling for superusers and admins.
1618
"""
1719
if request.user.is_superuser or request.user.is_staff:
1820
return True
1921

20-
scope = view.throttle_scope
21-
if request.user.is_authenticated:
22-
scope = f"auth_{scope}"
23-
else:
24-
scope = f"anon_{scope}"
22+
return super().allow_request(request, view)
2523

26-
view.throttle_scope = scope
2724

28-
return super().allow_request(request, view)
25+
def throttled_exception_handler(exception, context):
26+
"""
27+
Return this response whenever a request has been throttled
28+
"""
29+
30+
response = exception_handler(exception, context)
31+
32+
if isinstance(exception, Throttled):
33+
response_data = {
34+
"message": "Your request has been throttled. Please contact support@nexb.com"
35+
}
36+
response.data = response_data
37+
38+
return response

‎vulnerablecode/settings.py‎

Lines changed: 5 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -171,51 +171,11 @@
171171
LOGIN_REDIRECT_URL = "/"
172172
LOGOUT_REDIRECT_URL = "/"
173173

174-
REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
175-
"anon_vulnerable_packages": env.str(
176-
"VULNERABLECODE_ALL_VULNERABLE_PACKAGES_THROTTLING_RATE", default="1/hour"
177-
),
178-
"anon_bulk_search_packages": env.str(
179-
"VULNERABLECODE_BULK_SEARCH_PACKAGE_THROTTLING_RATE", default="5/hour"
180-
),
181-
"anon_packages": env.str("VULNERABLECODE_PACKAGES_SEARCH_THROTTLING_RATE", default="10/minute"),
182-
"anon_vulnerabilities": env.str(
183-
"VULNERABLECODE_VULNERABILITIES_SEARCH_THROTTLING_RATE", default="10/minute"
184-
),
185-
"anon_aliases": env.str("VULNERABLECODE_ALIASES_SEARCH_THROTTLING_RATE", default="5/minute"),
186-
"anon_cpes": env.str("VULNERABLECODE_CPE_SEARCH_THROTTLING_RATE", default="5/minute"),
187-
"anon_bulk_search_cpes": env.str(
188-
"VULNERABLECODE_BULK_SEARCH_CPE_THROTTLING_RATE", default="5/minute"
189-
),
190-
"auth_vulnerable_packages": env.str(
191-
"VULNERABLECODE_ALL_VULNERABLE_PACKAGES_THROTTLING_RATE", default="1/hour"
192-
),
193-
"auth_bulk_search_packages": env.str(
194-
"VULNERABLECODE_BULK_SEARCH_PACKAGE_THROTTLING_RATE", default="5/hour"
195-
),
196-
"auth_packages": env.str("VULNERABLECODE_PACKAGES_SEARCH_THROTTLING_RATE", default="10/minute"),
197-
"auth_vulnerabilities": env.str(
198-
"VULNERABLECODE_VULNERABILITIES_SEARCH_THROTTLING_RATE", default="10/minute"
199-
),
200-
"auth_aliases": env.str("VULNERABLECODE_ALIASES_SEARCH_THROTTLING_RATE", default="5/minute"),
201-
"auth_cpes": env.str("VULNERABLECODE_CPE_SEARCH_THROTTLING_RATE", default="5/minute"),
202-
"auth_bulk_search_cpes": env.str(
203-
"VULNERABLECODE_BULK_SEARCH_CPE_THROTTLING_RATE", default="5/minute"
204-
),
205-
}
174+
REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {"anon": "3600/hour", "user": "10800/hour"}
206175

207176
if IS_TESTS:
208177
VULNERABLECODEIO_REQUIRE_AUTHENTICATION = False
209-
REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {
210-
"auth_vulnerable_packages": "1/day",
211-
"auth_bulk_search_packages": "6/day",
212-
"auth_packages": "10/day",
213-
"anon_packages": "10/day",
214-
"auth_vulnerabilities": "8/day",
215-
"auth_aliases": "2/day",
216-
"auth_cpes": "4/day",
217-
"auth_bulk_search_cpes": "5/day",
218-
}
178+
REST_FRAMEWORK_DEFAULT_THROTTLE_RATES = {"anon": "10/day", "user": "20/day"}
219179

220180

221181
USE_L10N = True
@@ -253,8 +213,11 @@
253213
),
254214
"DEFAULT_THROTTLE_CLASSES": [
255215
"vulnerabilities.throttling.StaffUserRateThrottle",
216+
"rest_framework.throttling.AnonRateThrottle",
217+
"rest_framework.throttling.UserRateThrottle",
256218
],
257219
"DEFAULT_THROTTLE_RATES": REST_FRAMEWORK_DEFAULT_THROTTLE_RATES,
220+
"EXCEPTION_HANDLER": "vulnerabilities.throttling.throttled_exception_handler",
258221
"DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",
259222
# Limit the load on the Database returning a small number of records by default. https://github.com/nexB/vulnerablecode/issues/819
260223
"PAGE_SIZE": 10,

0 commit comments

Comments
 (0)