ratelimiter: Allow to configure different limits for authenticated users

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>

diff --git a/src/dbl/api/__init__.py b/src/dbl/api/__init__.py
index b3aa397739ec336e0e7f00cee00559d8f153dde9..2b29f318374de46ff2d0cd1b85bba7c398dca0ad 100644 (file)
--- a/src/dbl/api/__init__.py
+++ b/src/dbl/api/__init__.py
@@ -97,14 +97,15 @@ async def require_current_user(
 
        return user
 
-def limit(*, minutes, limit, key=None):
+def limit(*, minutes, limit, auth_limit=None, key=None):
        """
                This decorator takes a limit of how many requests per minute can be processed.
        """
        limiter = {
-               "minutes" : minutes,
-               "limit"   : limit,
-               "key"     : key,
+               "minutes"    : minutes,
+               "limit"      : limit,
+               "auth_limit" : auth_limit,
+               "key"        : key,
        }
 
        # Store the configuration with the handler

diff --git a/src/dbl/ratelimiter.py b/src/dbl/ratelimiter.py
index 0ff51c963787425d7158ece390faed257aa424f2..920156132f8b8ff8056664fffb8725b4c354bf14 100644 (file)
--- a/src/dbl/ratelimiter.py
+++ b/src/dbl/ratelimiter.py
@@ -82,7 +82,8 @@ class RateLimiter(object):
 
 
 class RateLimiterRequest(object):
-       def __init__(self, backend, request: fastapi.Request, *, minutes, limit, key=None):
+       def __init__(self, backend, request: fastapi.Request,
+                       *, minutes, limit, auth_limit=None, key=None):
                self.backend = backend
 
                # Store the request
@@ -93,7 +94,8 @@ class RateLimiterRequest(object):
 
                # Save the limits
                self.minutes = minutes
-               self.limit   = limit
+               self.limit = limit
+               self.auth_limit = auth_limit
 
                # Create a default key if none given
                if key is None:
@@ -130,21 +132,32 @@ class RateLimiterRequest(object):
                """
                        Returns True if the request is prohibited by the rate limiter
                """
+               limit = self.limit
+
                # Fetch the API key
                if self.api_key:
                        # Fetch the API key
                        key = await self.backend.auth(api_key)
 
-                       # If the API key was found and is not ratelimited, we skip further checks
-                       if key and not key.is_ratelimited:
-                               return False
+                       # Check if the user is actually properly authenticated
+                       if key:
+                               # If the API keyis not ratelimited, we skip further checks
+                               if not key.is_ratelimited:
+                                       return False
+
+                               # If authenticated users have a different limit set,
+                               # we change the checked limit to that
+                               if self.auth_limit:
+                                       limit = self.auth_limit
 
                # Fetch the number of past requests
                self.requests = await self.get_requests()
 
-               # The client is rate-limited when more requests have been
-               # received than allowed.
-               if self.requests >= self.limit:
+               # Assert that some limit has been set
+               assert limit
+
+               # The client is rate-limited when more requests have been received than allowed
+               if self.requests >= limit:
                        return True
 
                # Increment the request counter