Add PerRequestCacheThrottlingValidator

Author: paveldedik

example/apps/test_security/tests/__init__.py

@@ -4,4 +4,5 @@
 from .input_request_log import InputRequestLogTestCase
 from .output_request_log import OutputRequestLogTestCase
 from .partitioned_log import PartitionedLogTestCase
+from .throttling import PerRequestCacheThrottlingValidatorTestCase
 from .utils import UtilsTestCase

example/apps/test_security/tests/throttling.py

@@ -0,0 +1,40 @@
+from django.core.cache import cache
+from django.test import RequestFactory, TestCase, override_settings
+
+from germanium.tools import assert_raises
+
+from security.throttling.exception import ThrottlingException
+from security.throttling.validators import PerRequestCacheThrottlingValidator
+
+
+class PerRequestCacheThrottlingValidatorTestCase(TestCase):
+
+    def setUp(self):
+        cache.clear()
+        self.factory = RequestFactory()
+
+    def _make_request(self, path='/test/', method='GET', ip='127.0.0.1'):
+        request = self.factory.generic(method, path)
+        request.META['REMOTE_ADDR'] = ip
+        return request
+
+    @override_settings(SECURITY_THROTTLING_ENABLED=True)
+    def test_validate_raises_throttling_exception_when_over_limit(self):
+        validator = PerRequestCacheThrottlingValidator(60, 2)
+        request = self._make_request()
+        validator.validate(request)
+        validator.validate(request)
+        with assert_raises(ThrottlingException):
+            validator.validate(request)
+
+    @override_settings(SECURITY_THROTTLING_ENABLED=True)
+    def test_different_paths_are_counted_independently(self):
+        validator = PerRequestCacheThrottlingValidator(60, 2)
+        request_a = self._make_request(path='/path-a/')
+        request_b = self._make_request(path='/path-b/')
+        validator.validate(request_a)
+        validator.validate(request_a)
+        validator.validate(request_b)
+        validator.validate(request_b)
+        with assert_raises(ThrottlingException):
+            validator.validate(request_a)

security/throttling/validators.py

@@ -1,5 +1,8 @@
+import hashlib
+import time
 from datetime import timedelta
 
+from django.core.cache import cache
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from django.urls import resolve, Resolver404
@@ -55,6 +58,29 @@ def _validate(self, request):
         ) < self.throttle_at
 
 
+class PerRequestCacheThrottlingValidator(PerRequestThrottlingValidator):
+    """
+    Throttling validator that counts requests using Django's cache backend.
+    Unlike PerRequestThrottlingValidator, this doesn't rely on async input request logs,
+    making throttling effective immediately.
+    """
+
+    def _get_cache_key(self, request):
+        ip = get_client_ip(request)[0] or ''
+        window = int(time.time() / self.timeframe)
+        raw = '{}:{}:{}:{}'.format(ip, request.method.upper(), request.path, window)
+        return 'throttle:req:{}'.format(hashlib.md5(raw.encode()).hexdigest())
+
+    def _validate(self, request):
+        key = self._get_cache_key(request)
+        try:
+            count = cache.incr(key)
+        except ValueError:
+            cache.set(key, 1, timeout=self.timeframe)
+            count = 1
+        return count <= self.throttle_at
+
+
 class LoginThrottlingValidator(ThrottlingValidator):
 
     slug = None