Add excluded user ids for rateLimiting & early retrn in should ratelimit request

Author: bitterpanda63

aikido_zen/background_process/cloud_connection_manager/update_service_config.py

@@ -23,6 +23,13 @@ def update_service_config(connection_manager, res):
         received_any_stats=res.get("receivedAnyStats", True),
     )
 
+    # Handle excluded user IDs from rate limiting
+    excluded_user_ids = res.get("excludedUserIdsFromRateLimiting")
+    if isinstance(excluded_user_ids, list):
+        connection_manager.conf.update_excluded_user_ids_from_rate_limiting(
+            excluded_user_ids
+        )
+
     # Handle outbound request blocking configuration
     if "blockNewOutgoingRequests" in res:
         connection_manager.conf.set_block_new_outgoing_requests(

aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py

@@ -234,3 +234,52 @@ def test_update_service_config_block_new_outgoing_requests_only():
     assert connection_manager.conf.outbound_domains == {
         "existing.com": "allow"
     }  # Not changed
+
+
+def test_update_service_config_excluded_user_ids_from_rate_limiting():
+    """Test that update_service_config handles excludedUserIdsFromRateLimiting"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+
+    res = {
+        "success": True,
+        "excludedUserIdsFromRateLimiting": ["user1", "user2"],
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert connection_manager.conf.is_user_excluded_from_rate_limiting("user1") is True
+    assert connection_manager.conf.is_user_excluded_from_rate_limiting("user2") is True
+    assert connection_manager.conf.is_user_excluded_from_rate_limiting("user3") is False
+
+
+def test_update_service_config_excluded_user_ids_not_array():
+    """Test that update_service_config ignores non-array excludedUserIdsFromRateLimiting"""
+    connection_manager = MagicMock()
+    connection_manager.conf = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+    connection_manager.block = False
+
+    res = {
+        "success": True,
+        "excludedUserIdsFromRateLimiting": "not-an-array",
+    }
+
+    update_service_config(connection_manager, res)
+
+    assert (
+        connection_manager.conf.is_user_excluded_from_rate_limiting("not-an-array")
+        is False
+    )

aikido_zen/background_process/service_config.py

@@ -24,6 +24,7 @@ def __init__(
         )
         self.block_new_outgoing_requests = False
         self.outbound_domains = {}
+        self.excluded_user_ids_from_rate_limiting = set()
 
     def update(
         self,
@@ -75,6 +76,14 @@ def is_bypassed_ip(self, ip):
         """Checks if the IP is on the bypass list"""
         return self.bypassed_ips.has(ip)
 
+    def update_excluded_user_ids_from_rate_limiting(self, user_ids):
+        """Replaces the set of user IDs excluded from rate limiting"""
+        self.excluded_user_ids_from_rate_limiting = set(user_ids)
+
+    def is_user_excluded_from_rate_limiting(self, user_id):
+        """Checks if the user ID is excluded from rate limiting"""
+        return str(user_id) in self.excluded_user_ids_from_rate_limiting
+
     def update_outbound_domains(self, domains):
         self.outbound_domains = {
             domain["hostname"]: domain["mode"] for domain in domains

aikido_zen/background_process/service_config_test.py

@@ -319,3 +319,31 @@ def test_service_config_with_empty_allowlist():
     assert admin_endpoint["route"] == "/admin"
     assert isinstance(admin_endpoint["allowedIPAddresses"], list)
     assert len(admin_endpoint["allowedIPAddresses"]) == 0
+
+
+def test_excluded_user_ids_from_rate_limiting():
+    config = ServiceConfig(
+        endpoints=[],
+        last_updated_at=0,
+        blocked_uids=set(),
+        bypassed_ips=[],
+        received_any_stats=False,
+    )
+
+    # Initially empty
+    assert config.is_user_excluded_from_rate_limiting("user1") is False
+
+    # Update with user IDs
+    config.update_excluded_user_ids_from_rate_limiting(["user1", "user2"])
+    assert config.is_user_excluded_from_rate_limiting("user1") is True
+    assert config.is_user_excluded_from_rate_limiting("user2") is True
+    assert config.is_user_excluded_from_rate_limiting("user3") is False
+
+    # Update replaces the set
+    config.update_excluded_user_ids_from_rate_limiting(["user3"])
+    assert config.is_user_excluded_from_rate_limiting("user1") is False
+    assert config.is_user_excluded_from_rate_limiting("user3") is True
+
+    # Empty list clears all
+    config.update_excluded_user_ids_from_rate_limiting([])
+    assert config.is_user_excluded_from_rate_limiting("user3") is False

aikido_zen/ratelimiting/__init__.py

@@ -21,6 +21,9 @@ def should_ratelimit_request(
     if is_bypassed_ip:
         return {"block": False}
 
+    if user and connection_manager.conf.is_user_excluded_from_rate_limiting(user["id"]):
+        return {"block": False}
+
     max_requests = int(endpoint["rateLimiting"]["maxRequests"])
     windows_size_in_ms = int(endpoint["rateLimiting"]["windowSizeInMS"])
 

aikido_zen/ratelimiting/init_test.py

@@ -15,7 +15,9 @@ def user():
     return {"id": "user123"}
 
 
-def create_connection_manager(endpoints=[], bypassed_ips=[]):
+def create_connection_manager(
+    endpoints=[], bypassed_ips=[], excluded_user_ids_from_rate_limiting=[]
+):
     cm = MagicMock()
     cm.conf = ServiceConfig(
         endpoints=endpoints,
@@ -24,6 +26,9 @@ def create_connection_manager(endpoints=[], bypassed_ips=[]):
         bypassed_ips=bypassed_ips,
         received_any_stats=True,
     )
+    cm.conf.update_excluded_user_ids_from_rate_limiting(
+        excluded_user_ids_from_rate_limiting
+    )
     cm.rate_limiter = RateLimiter(
         max_items=5000, time_to_live_in_ms=120 * 60 * 1000  # 120 minutes
     )
@@ -511,3 +516,51 @@ def test_rate_limits_by_group_if_user_is_not_set():
         "block": True,
         "trigger": "group",
     }
+
+
+def test_does_not_rate_limit_excluded_users():
+    cm = create_connection_manager(
+        [
+            {
+                "method": "POST",
+                "route": "/login",
+                "forceProtectionOff": False,
+                "rateLimiting": {
+                    "enabled": True,
+                    "maxRequests": 3,
+                    "windowSizeInMS": 1000,
+                },
+            },
+        ],
+        excluded_user_ids_from_rate_limiting=["excluded-user-id"],
+    )
+    route_metadata = create_route_metadata()
+    excluded_user = {"id": "excluded-user-id"}
+    for _ in range(5):
+        assert should_ratelimit_request(
+            route_metadata, "1.2.3.4", excluded_user, cm
+        ) == {"block": False}
+
+
+def test_does_not_rate_limit_excluded_users_in_group():
+    cm = create_connection_manager(
+        [
+            {
+                "method": "POST",
+                "route": "/login",
+                "forceProtectionOff": False,
+                "rateLimiting": {
+                    "enabled": True,
+                    "maxRequests": 3,
+                    "windowSizeInMS": 1000,
+                },
+            },
+        ],
+        excluded_user_ids_from_rate_limiting=["excluded-user-id"],
+    )
+    route_metadata = create_route_metadata()
+    excluded_user = {"id": "excluded-user-id"}
+    for _ in range(5):
+        assert should_ratelimit_request(
+            route_metadata, "1.2.3.4", excluded_user, cm, "group1"
+        ) == {"block": False}