diff --git a/aikido_zen/background_process/cloud_connection_manager/update_service_config.py b/aikido_zen/background_process/cloud_connection_manager/update_service_config.py index e09f2c3d1..4184d42b6 100644 --- a/aikido_zen/background_process/cloud_connection_manager/update_service_config.py +++ b/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( diff --git a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py b/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py index be6ddfed7..6080a3515 100644 --- a/aikido_zen/background_process/cloud_connection_manager/update_service_config_test.py +++ b/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 + ) diff --git a/aikido_zen/background_process/service_config.py b/aikido_zen/background_process/service_config.py index a72f6dbaf..fc659cb7d 100644 --- a/aikido_zen/background_process/service_config.py +++ b/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 diff --git a/aikido_zen/background_process/service_config_test.py b/aikido_zen/background_process/service_config_test.py index d34611693..d4b7d38e0 100644 --- a/aikido_zen/background_process/service_config_test.py +++ b/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 diff --git a/aikido_zen/ratelimiting/__init__.py b/aikido_zen/ratelimiting/__init__.py index dcf0927ac..4d60bdd01 100644 --- a/aikido_zen/ratelimiting/__init__.py +++ b/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"]) diff --git a/aikido_zen/ratelimiting/init_test.py b/aikido_zen/ratelimiting/init_test.py index 966d8c7ac..dfd14acf6 100644 --- a/aikido_zen/ratelimiting/init_test.py +++ b/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}