Merge pull request #2230 from strikaco/throttle

Refactors throttle to use Django caching (partial fix for #2223).

evennia/accounts/accounts.py

@@ -52,10 +52,10 @@ _MUDINFO_CHANNEL = None
 
 # Create throttles for too many account-creations and login attempts
 CREATION_THROTTLE = Throttle(
-    limit=settings.CREATION_THROTTLE_LIMIT, timeout=settings.CREATION_THROTTLE_TIMEOUT
+    name='creation', limit=settings.CREATION_THROTTLE_LIMIT, timeout=settings.CREATION_THROTTLE_TIMEOUT
 )
 LOGIN_THROTTLE = Throttle(
-    limit=settings.LOGIN_THROTTLE_LIMIT, timeout=settings.LOGIN_THROTTLE_TIMEOUT
+    name='login', limit=settings.LOGIN_THROTTLE_LIMIT, timeout=settings.LOGIN_THROTTLE_TIMEOUT
 )
 
 

evennia/server/tests/test_misc.py

@@ -90,8 +90,8 @@ class ThrottleTest(EvenniaTest):
     """
 
     def test_throttle(self):
-        ips = ("94.100.176.153", "45.56.148.77", "5.196.1.129")
+        ips = ("256.256.256.257", "257.257.257.257", "258.258.258.258")
-        kwargs = {"limit": 5, "timeout": 15 * 60}
+        kwargs = {"name": "testing", "limit": 5, "timeout": 15 * 60}
 
         throttle = Throttle(**kwargs)
 
@@ -124,3 +124,14 @@ class ThrottleTest(EvenniaTest):
 
         # There should only be (cache_size * num_ips) total in the Throttle cache
         self.assertEqual(sum([len(cache[x]) for x in cache.keys()]), throttle.cache_size * len(ips))
+        
+        # Make sure the cache is populated
+        self.assertTrue(throttle.get())
+        
+        # Remove the test IPs from the throttle cache 
+        # (in case persistent storage was configured by the user)
+        for ip in ips:
+            self.assertTrue(throttle.remove(ip))
+            
+        # Make sure the cache is empty
+        self.assertFalse(throttle.get())

evennia/server/throttle.py

@@ -1,4 +1,5 @@
-from collections import defaultdict, deque
+from django.core.cache import caches
+from collections import deque
 from evennia.utils import logger
 import time
 
@@ -12,8 +13,8 @@ class Throttle(object):
 
     This version of the throttle is usable by both the terminal server as well
     as the web server, imposes limits on memory consumption by using deques
-    with length limits instead of open-ended lists, and removes sparse keys when
+    with length limits instead of open-ended lists, and uses native Django 
-    no recent failures have been recorded.
+    caches for automatic key eviction and persistence configurability.
     """
 
     error_msg = "Too many failed attempts; you must wait a few minutes before trying again."
@@ -23,6 +24,7 @@ class Throttle(object):
         Allows setting of throttle parameters.
 
         Keyword Args:
+            name (str): Name of this throttle.
             limit (int): Max number of failures before imposing limiter
             timeout (int): number of timeout seconds after
                 max number of tries has been reached.
@@ -30,9 +32,37 @@ class Throttle(object):
                 rolling window; this is NOT the same as the limit after which
                 the throttle is imposed!
         """
-        self.storage = defaultdict(deque)
+        try:
-        self.cache_size = self.limit = kwargs.get("limit", 5)
+            self.storage = caches['throttle']
+        except Exception as e:
+            logger.log_trace("Throttle: Errors encountered; using default cache.")
+            self.storage = caches['default']
+            
+        self.name = kwargs.get('name', 'undefined-throttle')
+        self.limit = kwargs.get("limit", 5)
+        self.cache_size = kwargs.get('cache_size', self.limit)
         self.timeout = kwargs.get("timeout", 5 * 60)
+        
+    def get_cache_key(self, *args, **kwargs):
+        """
+        Creates a 'prefixed' key containing arbitrary terms to prevent key
+        collisions in the same namespace.
+        
+        """
+        return '-'.join((self.name, *args))
+        
+    def touch(self, key, *args, **kwargs):
+        """
+        Refreshes the timeout on a given key and ensures it is recorded in the
+        key register.
+        
+        Args:
+            key(str): Key of entry to renew.
+            
+        """
+        cache_key = self.get_cache_key(key)
+        if self.storage.touch(cache_key, self.timeout):
+            self.record_key(key)
 
     def get(self, ip=None):
         """
@@ -50,9 +80,18 @@ class Throttle(object):
 
         """
         if ip:
-            return self.storage.get(ip, deque(maxlen=self.cache_size))
+            cache_key = self.get_cache_key(str(ip))
+            return self.storage.get(cache_key, deque(maxlen=self.cache_size))
         else:
-            return self.storage
+            keys_key = self.get_cache_key('keys')
+            keys = self.storage.get_or_set(keys_key, set(), self.timeout)
+            data = self.storage.get_many((self.get_cache_key(x) for x in keys))
+            
+            found_keys = set(data.keys())
+            if len(keys) != len(found_keys):
+                self.storage.set(keys_key, found_keys, self.timeout)
+                
+            return data
 
     def update(self, ip, failmsg="Exceeded threshold."):
         """
@@ -67,24 +106,77 @@ class Throttle(object):
             None
 
         """
+        cache_key = self.get_cache_key(ip)
+        
         # Get current status
         previously_throttled = self.check(ip)
 
-        # Enforce length limits
+        # Get previous failures, if any
-        if not self.storage[ip].maxlen:
+        entries = self.storage.get(cache_key, [])
-            self.storage[ip] = deque(maxlen=self.cache_size)
+        entries.append(time.time())
-
+        
-        self.storage[ip].append(time.time())
+        # Store updated record
+        self.storage.set(cache_key, deque(entries, maxlen=self.cache_size), self.timeout)
 
         # See if this update caused a change in status
         currently_throttled = self.check(ip)
 
         # If this makes it engage, log a single activation event
         if not previously_throttled and currently_throttled:
-            logger.log_sec(
+            logger.log_sec(f"Throttle Activated: {failmsg} (IP: {ip}, {self.limit} hits in {self.timeout} seconds.)")
-                "Throttle Activated: %s (IP: %s, %i hits in %i seconds.)"
+            
-                % (failmsg, ip, self.limit, self.timeout)
+        self.record_ip(ip)
-            )
+        
+    def remove(self, ip, *args, **kwargs):
+        """
+        Clears data stored for an IP from the throttle.
+        
+        Args:
+            ip(str): IP to clear.
+        
+        """
+        exists = self.get(ip)
+        if not exists: return False
+        
+        cache_key = self.get_cache_key(ip)
+        self.storage.delete(cache_key)
+        self.unrecord_ip(ip)
+        
+        # Return True if NOT exists
+        return ~bool(self.get(ip))
+        
+    def record_ip(self, ip, *args, **kwargs):
+        """
+        Tracks keys as they are added to the cache (since there is no way to 
+        get a list of keys after-the-fact).
+        
+        Args:
+            ip(str): IP being added to cache. This should be the original
+                IP, not the cache-prefixed key.
+            
+        """
+        keys_key = self.get_cache_key('keys')
+        keys = self.storage.get(keys_key, set())
+        keys.add(ip)
+        self.storage.set(keys_key, keys, self.timeout)
+        return True
+        
+    def unrecord_ip(self, ip, *args, **kwargs):
+        """
+        Forces removal of a key from the key registry.
+        
+        Args:
+            ip(str): IP to remove from list of keys.
+            
+        """
+        keys_key = self.get_cache_key('keys')
+        keys = self.storage.get(keys_key, set())
+        try: 
+            keys.remove(ip)
+            self.storage.set(keys_key, keys, self.timeout)
+            return True
+        except KeyError:
+            return False
 
     def check(self, ip):
         """
@@ -102,17 +194,20 @@ class Throttle(object):
         """
         now = time.time()
         ip = str(ip)
+        
+        cache_key = self.get_cache_key(ip)
 
         # checking mode
-        latest_fails = self.storage[ip]
+        latest_fails = self.storage.get(cache_key)
         if latest_fails and len(latest_fails) >= self.limit:
             # too many fails recently
             if now - latest_fails[-1] < self.timeout:
                 # too soon - timeout in play
+                self.touch(cache_key)
                 return True
             else:
                 # timeout has passed. clear faillist
-                del self.storage[ip]
+                self.remove(ip)
                 return False
         else:
-            return False
+            return False

evennia/settings_default.py

@@ -864,6 +864,21 @@ WEBCLIENT_OPTIONS = {
     # messages
 }
 
+# Django cache settings
+# https://docs.djangoproject.com/en/dev/topics/cache/#setting-up-the-cache
+CACHES = {
+    'default': {
+        'BACKEND': 'django.core.cache.backends.locmem.LocMemCache',
+    },
+    'throttle': {
+        'BACKEND': 'django.core.cache.backends.locmem.LocMemCache',
+        'TIMEOUT': 60 * 5,
+        'OPTIONS': {
+            'MAX_ENTRIES': 2000
+        }
+    }
+}
+
 # We setup the location of the website template as well as the admin site.
 TEMPLATES = [
     {