ratelimiter: Use the API key as bucket for authenticated users

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

diff --git a/src/dbl/ratelimiter.py b/src/dbl/ratelimiter.py
index 0b56d488748a97a723eb5d4c7c14a65fdb215862..f006c47421cd9f7af0757855c896c686e2577b77 100644 (file)
--- a/src/dbl/ratelimiter.py
+++ b/src/dbl/ratelimiter.py
@@ -105,8 +105,21 @@ class RateLimiterRequest(object):
 
                # Store the key and address
                self.key = key
+
+               # Fetch the client IP address
                self.address, port = self.request.client
 
+               # Check if the user is authenticated
+               self.api_key = self.request.headers.get("X-Api-Key")
+
+               # Make the bucket
+               # Use the API key for authenticated users,
+               # or otherwise use the client IP address
+               if self.api_key:
+                       self.bucket = self.api_key
+               else:
+                       self.bucket = "%s" % self.address
+
                # What is the current time?
                self.now = datetime.datetime.utcnow()
 
@@ -150,7 +163,7 @@ class RateLimiterRequest(object):
                        .where(
                                ratelimiter.c.key        == self.key,
                                ratelimiter.c.timestamp  >= since,
-                               ratelimiter.c.bucket     == "%s" % self.address,
+                               ratelimiter.c.bucket     == self.bucket,
                        )
                )
 
@@ -196,7 +209,7 @@ class RateLimiterRequest(object):
                        .values({
                                "key"        : self.key,
                                "timestamp"  : now,
-                               "bucket"     : "%s" % self.address,
+                               "bucket"     : self.bucket,
                                "requests"   : 1,
                                "expires_at" : expires_at,
                        })