SHOULD_RATELIMIT: pass along group & check group for rlm

Author: bitterpanda63

aikido_zen/background_process/commands/should_ratelimit.py

@@ -16,4 +16,5 @@ def process_should_ratelimit(connection_manager, data, queue=None):
         remote_address=data["remote_address"],
         user=data["user"],
         connection_manager=connection_manager,
+        group=data["group"],
     )

aikido_zen/middleware/init_test.py

@@ -145,6 +145,7 @@ def test_cache_comms_with_endpoints():
                     "url": "http://localhost:4000",
                 },
                 "user": {"id": "456"},
+                "group": None,
                 "remote_address": "::1",
             },
             receive=True,

aikido_zen/middleware/should_block_request.py

@@ -52,6 +52,7 @@ def should_block_request():
                 obj={
                     "route_metadata": route_metadata,
                     "user": context.user,
+                    "group": context.rate_limit_group,
                     "remote_address": context.remote_address,
                 },
                 receive=True,

aikido_zen/ratelimiting/__init__.py

@@ -5,47 +5,69 @@
 from .get_ratelimited_endpoint import get_ratelimited_endpoint
 
 
-def should_ratelimit_request(route_metadata, remote_address, user, connection_manager):
+def should_ratelimit_request(
+    route_metadata, remote_address, user, connection_manager, group=None
+):
     """
-    Checks if the request should be ratelimited or not
+    Checks if the request should be rate-limited or not (checks user, group id & ip)
     route_metadata object includes route, url and method
     """
     endpoints = connection_manager.conf.get_endpoints(route_metadata)
     endpoint = get_ratelimited_endpoint(endpoints, route_metadata["route"])
     if not endpoint:
         return {"block": False}
 
+    is_bypassed_ip = connection_manager.conf.is_bypassed_ip(remote_address)
+    if is_bypassed_ip:
+        return {"block": False}
+
     max_requests = int(endpoint["rateLimiting"]["maxRequests"])
     windows_size_in_ms = int(endpoint["rateLimiting"]["windowSizeInMS"])
-    is_bypassed_ip = connection_manager.conf.is_bypassed_ip(remote_address)
 
-    if is_bypassed_ip:
+    if group:
+        allowed = connection_manager.rate_limiter.is_allowed(
+            get_key_for_group(endpoint, group),
+            windows_size_in_ms,
+            max_requests,
+        )
+        if not allowed:
+            return {"block": True, "trigger": "group"}
+
+        # Do not check IP or user rate limit if group is set
         return {"block": False}
     if user:
-        uid = user["id"]
-        method = endpoint.get("method")
-        route = endpoint.get("route")
-
         allowed = connection_manager.rate_limiter.is_allowed(
-            f"{method}:{route}:user:{uid}",
+            get_key_for_user(endpoint, user),
             windows_size_in_ms,
             max_requests,
         )
         if not allowed:
             return {"block": True, "trigger": "user"}
         # Do not check IP rate limit if user is set
         return {"block": False}
-
     if remote_address:
-        method = endpoint.get("method")
-        route = endpoint.get("route")
-
         allowed = connection_manager.rate_limiter.is_allowed(
-            f"{method}:{route}:ip:{remote_address}",
+            get_key_for_ip(endpoint, remote_address),
             windows_size_in_ms,
             max_requests,
         )
         if not allowed:
             return {"block": True, "trigger": "ip"}
 
     return {"block": False}
+
+
+def get_key_for_group(endpoint, group_id):
+    method, route = endpoint.get("method"), endpoint.get("route")
+    return f"{method}:{route}:group:{group_id}"
+
+
+def get_key_for_user(endpoint, user):
+    method, route = endpoint.get("method"), endpoint.get("route")
+    user_id = user.get("id")
+    return f"{method}:{route}:user:{user_id}"
+
+
+def get_key_for_ip(endpoint, remote_address):
+    method, route = endpoint.get("method"), endpoint.get("route")
+    return f"{method}:{route}:ip:{remote_address}"