Skip to content

Commit f1b53d3

Browse files
TG1999TG1999
committed
Override throttle rate for each endpoint
Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>
1 parent 6f72ecf commit f1b53d3

4 files changed

Lines changed: 215 additions & 9 deletions

File tree

‎vulnerabilities/api.py‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,13 @@
2323
from vulnerabilities.models import VulnerabilityReference
2424
from vulnerabilities.models import VulnerabilitySeverity
2525
from vulnerabilities.models import get_purl_query_lookups
26+
from vulnerabilities.throttling import AliasesAPIThrottle
27+
from vulnerabilities.throttling import BulkSearchCPEAPIThrottle
28+
from vulnerabilities.throttling import BulkSearchPackagesAPIThrottle
29+
from vulnerabilities.throttling import CPEAPIThrottle
30+
from vulnerabilities.throttling import PackagesAPIThrottle
31+
from vulnerabilities.throttling import VulnerabilitiesAPIThrottle
32+
from vulnerabilities.throttling import VulnerablePackagesAPIThrottle
2633

2734

2835
class VulnerabilitySeveritySerializer(serializers.ModelSerializer):
@@ -221,6 +228,15 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
221228
filter_backends = (filters.DjangoFilterBackend,)
222229
filterset_class = PackageFilterSet
223230

231+
def get_throttles(self):
232+
if self.action == "bulk_search":
233+
throttle_classes = [BulkSearchPackagesAPIThrottle]
234+
elif self.action == "all":
235+
throttle_classes = [VulnerablePackagesAPIThrottle]
236+
else:
237+
throttle_classes = [PackagesAPIThrottle]
238+
return [throttle() for throttle in throttle_classes]
239+
224240
# TODO: Fix the swagger documentation for this endpoint
225241
@action(detail=False, methods=["post"])
226242
def bulk_search(self, request):
@@ -246,7 +262,7 @@ def bulk_search(self, request):
246262
if purl_data:
247263
purl_response = PackageSerializer(purl_data[0], context={"request": request}).data
248264
else:
249-
purl_response = purl
265+
purl_response = purl.to_dict()
250266
purl_response["unresolved_vulnerabilities"] = []
251267
purl_response["resolved_vulnerabilities"] = []
252268
purl_response["purl"] = purl_string
@@ -302,6 +318,7 @@ def get_queryset(self):
302318
serializer_class = VulnerabilitySerializer
303319
filter_backends = (filters.DjangoFilterBackend,)
304320
filterset_class = VulnerabilityFilterSet
321+
throttle_classes = [VulnerabilitiesAPIThrottle]
305322

306323

307324
class CPEFilterSet(filters.FilterSet):
@@ -320,6 +337,13 @@ class CPEViewSet(viewsets.ReadOnlyModelViewSet):
320337
filter_backends = (filters.DjangoFilterBackend,)
321338
filterset_class = CPEFilterSet
322339

340+
def get_throttles(self):
341+
if self.action == "bulk_search":
342+
throttle_classes = [BulkSearchCPEAPIThrottle]
343+
else:
344+
throttle_classes = [CPEAPIThrottle]
345+
return [throttle() for throttle in throttle_classes]
346+
323347
@action(detail=False, methods=["post"])
324348
def bulk_search(self, request):
325349
"""
@@ -357,3 +381,4 @@ class AliasViewSet(viewsets.ReadOnlyModelViewSet):
357381
serializer_class = VulnerabilitySerializer
358382
filter_backends = (filters.DjangoFilterBackend,)
359383
filterset_class = AliasFilterSet
384+
throttle_classes = [AliasesAPIThrottle]

‎vulnerabilities/tests/test_throttling.py‎

Lines changed: 125 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
# See https://aboutcode.org for more information about nexB OSS projects.
88
#
99

10+
import json
11+
1012
from django.contrib.auth import get_user_model
1113
from rest_framework.test import APIClient
1214
from rest_framework.test import APITestCase
@@ -30,10 +32,10 @@ def setUp(self):
3032
self.staff_csrf_client = APIClient(enforce_csrf_checks=True)
3133
self.staff_csrf_client.credentials(HTTP_AUTHORIZATION=self.staff_auth)
3234

33-
def test_api_throttling(self):
35+
def test_packages_endpoint_throttling(self):
3436

35-
# A basic user can only access API 5 times a day
36-
for i in range(0, 5):
37+
# A basic user can only access /packages endpoint 10 times a day
38+
for i in range(0, 10):
3739
response = self.csrf_client.get("/api/packages")
3840
self.assertEqual(response.status_code, 200)
3941
response = self.staff_csrf_client.get("/api/packages")
@@ -46,3 +48,123 @@ def test_api_throttling(self):
4648
response = self.staff_csrf_client.get("/api/packages", format="json")
4749
# 200 - staff user can access API unlimited times
4850
self.assertEqual(response.status_code, 200)
51+
52+
def test_cpes_endpoint_throttling(self):
53+
54+
# A basic user can only access /cpes endpoint 5 times a day
55+
for i in range(0, 4):
56+
response = self.csrf_client.get("/api/cpes")
57+
self.assertEqual(response.status_code, 200)
58+
response = self.staff_csrf_client.get("/api/cpes")
59+
self.assertEqual(response.status_code, 200)
60+
61+
response = self.csrf_client.get("/api/cpes")
62+
# 429 - too many requests for basic user
63+
self.assertEqual(response.status_code, 429)
64+
65+
response = self.staff_csrf_client.get("/api/cpes", format="json")
66+
# 200 - staff user can access API unlimited times
67+
self.assertEqual(response.status_code, 200)
68+
69+
def test_all_vulnerable_packages_endpoint_throttling(self):
70+
71+
# A basic user can only access /packages/all 1 time a day
72+
for i in range(0, 1):
73+
response = self.csrf_client.get("/api/packages/all")
74+
self.assertEqual(response.status_code, 200)
75+
response = self.staff_csrf_client.get("/api/packages/all")
76+
self.assertEqual(response.status_code, 200)
77+
78+
response = self.csrf_client.get("/api/packages/all")
79+
# 429 - too many requests for basic user
80+
self.assertEqual(response.status_code, 429)
81+
82+
response = self.staff_csrf_client.get("/api/packages/all", format="json")
83+
# 200 - staff user can access API unlimited times
84+
self.assertEqual(response.status_code, 200)
85+
86+
def test_vulnerabilities_endpoint_throttling(self):
87+
88+
# A basic user can only access /vulnerabilities 10 times a day
89+
for i in range(0, 8):
90+
response = self.csrf_client.get("/api/vulnerabilities")
91+
self.assertEqual(response.status_code, 200)
92+
response = self.staff_csrf_client.get("/api/vulnerabilities")
93+
self.assertEqual(response.status_code, 200)
94+
95+
response = self.csrf_client.get("/api/vulnerabilities")
96+
# 429 - too many requests for basic user
97+
self.assertEqual(response.status_code, 429)
98+
99+
response = self.staff_csrf_client.get("/api/vulnerabilities", format="json")
100+
# 200 - staff user can access API unlimited times
101+
self.assertEqual(response.status_code, 200)
102+
103+
def test_aliases_endpoint_throttling(self):
104+
105+
# A basic user can only access /alias 2 times a day
106+
for i in range(0, 2):
107+
response = self.csrf_client.get("/api/alias")
108+
self.assertEqual(response.status_code, 200)
109+
response = self.staff_csrf_client.get("/api/alias")
110+
self.assertEqual(response.status_code, 200)
111+
112+
response = self.csrf_client.get("/api/alias")
113+
# 429 - too many requests for basic user
114+
self.assertEqual(response.status_code, 429)
115+
116+
response = self.staff_csrf_client.get("/api/alias", format="json")
117+
# 200 - staff user can access API unlimited times
118+
self.assertEqual(response.status_code, 200)
119+
120+
def test_bulk_search_packages_endpoint_throttling(self):
121+
data = json.dumps({"purls": ["pkg:foo/bar"]})
122+
123+
# A basic user can only access /packages/bulk_search 5 times a day
124+
for i in range(0, 6):
125+
response = self.csrf_client.post(
126+
"/api/packages/bulk_search", data=data, content_type="application/json"
127+
)
128+
self.assertEqual(response.status_code, 200)
129+
response = self.staff_csrf_client.post(
130+
"/api/packages/bulk_search", data=data, content_type="application/json"
131+
)
132+
self.assertEqual(response.status_code, 200)
133+
134+
response = self.csrf_client.post(
135+
"/api/packages/bulk_search", data=data, content_type="application/json"
136+
)
137+
# 429 - too many requests for basic user
138+
self.assertEqual(response.status_code, 429)
139+
140+
response = self.staff_csrf_client.post(
141+
"/api/packages/bulk_search", data=data, content_type="application/json"
142+
)
143+
# 200 - staff user can access API unlimited times
144+
self.assertEqual(response.status_code, 200)
145+
146+
def test_bulk_search_cpes_endpoint_throttling(self):
147+
data = json.dumps({"cpes": ["cpe:foo/bar"]})
148+
149+
# A basic user can only access /cpes/bulk_search 5 times a day
150+
for i in range(0, 5):
151+
response = self.csrf_client.post(
152+
"/api/cpes/bulk_search", data=data, content_type="application/json"
153+
)
154+
self.assertEqual(response.status_code, 200)
155+
response = self.staff_csrf_client.post(
156+
"/api/cpes/bulk_search", data=data, content_type="application/json"
157+
)
158+
self.assertEqual(response.status_code, 200)
159+
160+
response = self.csrf_client.post(
161+
"/api/cpes/bulk_search", data=data, content_type="application/json"
162+
)
163+
# 429 - too many requests for basic user
164+
self.assertEqual(response.status_code, 429)
165+
166+
response = self.staff_csrf_client.post(
167+
"/api/cpes/bulk_search", data=data, content_type="application/json"
168+
)
169+
# 200 - staff user can access API unlimited times
170+
self.assertEqual(response.status_code, 200)

‎vulnerabilities/throttling.py‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,12 @@
88
#
99

1010
from django.contrib.auth import get_user_model
11-
from rest_framework.throttling import UserRateThrottle
11+
from rest_framework.throttling import SimpleRateThrottle
1212

1313
User = get_user_model()
1414

1515

16-
class StaffUserRateThrottle(UserRateThrottle):
16+
class StaffUserRateThrottle(SimpleRateThrottle):
1717
def allow_request(self, request, view):
1818
"""
1919
Do not apply throttling for superusers and admins.
@@ -22,3 +22,42 @@ def allow_request(self, request, view):
2222
return True
2323

2424
return super().allow_request(request, view)
25+
26+
def get_cache_key(self, request, view):
27+
"""
28+
Return the cache key to use for this request.
29+
"""
30+
if request.user.is_authenticated:
31+
ident = request.user.pk
32+
else:
33+
ident = self.get_ident(request)
34+
35+
return self.cache_format % {"scope": self.scope, "ident": ident}
36+
37+
38+
class VulnerablePackagesAPIThrottle(StaffUserRateThrottle):
39+
scope = "vulnerable_packages"
40+
41+
42+
class BulkSearchPackagesAPIThrottle(StaffUserRateThrottle):
43+
scope = "bulk_search_packages"
44+
45+
46+
class PackagesAPIThrottle(StaffUserRateThrottle):
47+
scope = "packages"
48+
49+
50+
class VulnerabilitiesAPIThrottle(StaffUserRateThrottle):
51+
scope = "vulnerabilities"
52+
53+
54+
class AliasesAPIThrottle(StaffUserRateThrottle):
55+
scope = "aliases"
56+
57+
58+
class CPEAPIThrottle(StaffUserRateThrottle):
59+
scope = "cpes"
60+
61+
62+
class BulkSearchCPEAPIThrottle(StaffUserRateThrottle):
63+
scope = "bulk_search_cpes"

‎vulnerablecode/settings.py‎

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -150,11 +150,23 @@
150150

151151
LOGIN_REDIRECT_URL = "/"
152152
LOGOUT_REDIRECT_URL = "/"
153-
THROTTLING_RATE = env.str("THROTTLING_RATE", default="1000/day")
153+
TEST_PACKAGE_THROTTLING_RATE = None
154+
TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE = None
155+
TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE = None
156+
TEST_VULNERABILITIES_THROTTLING_RATE = None
157+
TEST_CPES_THROTTLING_RATE = None
158+
TEST_BULK_SEARCH_CPES_THROTTLING_RATE = None
159+
TEST_ALIASES_THROTTLING_RATE = None
154160

155161
if IS_TESTS:
156162
VULNERABLECODEIO_REQUIRE_AUTHENTICATION = True
157-
THROTTLING_RATE = "5/day"
163+
TEST_PACKAGE_THROTTLING_RATE = "10/day"
164+
TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE = "6/day"
165+
TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE = "1/day"
166+
TEST_VULNERABILITIES_THROTTLING_RATE = "8/day"
167+
TEST_CPES_THROTTLING_RATE = "4/day"
168+
TEST_BULK_SEARCH_CPES_THROTTLING_RATE = "5/day"
169+
TEST_ALIASES_THROTTLING_RATE = "2/day"
158170

159171

160172
USE_L10N = True
@@ -190,7 +202,15 @@
190202
"DEFAULT_THROTTLE_CLASSES": [
191203
"vulnerabilities.throttling.StaffUserRateThrottle",
192204
],
193-
"DEFAULT_THROTTLE_RATES": {"user": THROTTLING_RATE},
205+
"DEFAULT_THROTTLE_RATES": {
206+
"vulnerable_packages": TEST_ALL_VULNERABLE_PACKAGE_THROTTLING_RATE or "1/hour",
207+
"bulk_search_packages": TEST_BULK_SEARCH_PACKAGE_THROTTLING_RATE or "5/hour",
208+
"packages": TEST_PACKAGE_THROTTLING_RATE or "10/minute",
209+
"vulnerabilities": TEST_VULNERABILITIES_THROTTLING_RATE or "10/minute",
210+
"aliases": TEST_ALIASES_THROTTLING_RATE or "5/minute",
211+
"cpes": TEST_CPES_THROTTLING_RATE or "5/minute",
212+
"bulk_search_cpes": TEST_BULK_SEARCH_CPES_THROTTLING_RATE or "5/hour",
213+
},
194214
"DEFAULT_PAGINATION_CLASS": "vulnerabilities.pagination.SmallResultSetPagination",
195215
# Limit the load on the Database returning a small number of records by default. https://github.com/nexB/vulnerablecode/issues/819
196216
"PAGE_SIZE": 10,

0 commit comments

Comments
 (0)