diff --git a/aikido_zen/__init__.py b/aikido_zen/__init__.py index f4939c909..d8dd2bc18 100644 --- a/aikido_zen/__init__.py +++ b/aikido_zen/__init__.py @@ -4,8 +4,6 @@ import os -from aikido_zen.background_process.test_uds_file_access import test_uds_file_access - # Re-export functions : from aikido_zen.context.users import set_user from aikido_zen.context.track import track @@ -19,7 +17,7 @@ from aikido_zen.helpers.logging import logger # Import background process -from aikido_zen.background_process import start_background_process +from aikido_zen.background_process import check_socket_access, start_background_process # Load environment variables and constants from aikido_zen.config import PKG_VERSION @@ -41,7 +39,7 @@ def protect(mode="daemon", token=""): return if gil_not_enabled(): return - if not test_uds_file_access(): + if not check_socket_access(): return # Unable to start background process if check_gevent(): return # gevent is not compatible with zen. diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index e145de8e1..d1819bd9e 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -3,6 +3,7 @@ and listen for data sent by our sources and sinks """ +import hashlib import multiprocessing import os import platform @@ -13,12 +14,10 @@ from aikido_zen.helpers.hash_aikido_token import hash_aikido_token from aikido_zen.helpers.get_agent_session_id import get_agent_session_id from aikido_zen.helpers.logging import logger +from aikido_zen.config import PKG_VERSION -from aikido_zen.background_process.comms import ( - AikidoIPCCommunications, - get_comms, - reset_comms, -) +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.commands import ping from .aikido_background_process import AikidoBackgroundProcess @@ -53,15 +52,15 @@ def start_background_process(): ) return - address = get_uds_filename() + address = get_socket_path() if platform.system() == "Windows": # Python does not support Windows UDS just yet, so we have to rely on INET address = ("127.0.0.1", 49156) - comms = AikidoIPCCommunications(address, secret_key_bytes) + ipc = IPCClient(address, secret_key_bytes) if isinstance(address, str) and os.path.exists(address): - if background_process_already_active(comms): + if background_process_already_active(ipc): return # Don't start if the background process is already active try: os.remove(address) @@ -72,23 +71,47 @@ def start_background_process(): # Daemon is set to True so that the process kills itself when the main process dies background_process = get_process_factory()( target=AikidoBackgroundProcess, - args=(comms.address, comms.key, get_agent_session_id()), + args=(ipc.address, ipc.key, get_agent_session_id()), name="zen-agent-process", daemon=True, ) background_process.start() -def get_uds_filename(): - """Returns the address for UDS file""" +def get_socket_path(): temp_dir = get_temp_dir() - prefix = "aikido_python" - name = hash_aikido_token() + prefix = "aikido_zen" + + # Include the package version to avoid reusing a background process from a previous release. + identity = f"{hash_aikido_token()}:{PKG_VERSION}" + + name = hashlib.sha256(identity.encode()).hexdigest() return f"{temp_dir}/{prefix}_{name}.sock" -def background_process_already_active(comms): - res = comms.send_data_to_bg_process(action="PING", obj=tuple(), receive=True) +def check_socket_access(): + """Returns true if access is possible, false if it's not.""" + filename = get_socket_path() + directory = "/".join(filename.split("/")[:-1]) + + if platform.system() == "Windows": + # Python does not support Windows UDS just yet, so we have to rely on INET + return True + + if not os.access(directory, os.R_OK) or not os.access(directory, os.W_OK): + # We need both read and write permissions in order to start the socket + logger.critical( + "Cannot start Zen: no read/write permissions for directory %s. " + + "You can change the path by setting " + + "the AIKIDO_TMP_DIR environment variable.", + directory, + ) + return False + return True + + +def background_process_already_active(ipc): + res = ipc.send_message(ping.generate()) if res["success"] and res["data"] == "Received": # Ping is active, return. logger.debug( diff --git a/aikido_zen/background_process/aikido_background_process.py b/aikido_zen/background_process/aikido_background_process.py index 57387a574..01ff76364 100644 --- a/aikido_zen/background_process/aikido_background_process.py +++ b/aikido_zen/background_process/aikido_background_process.py @@ -54,8 +54,7 @@ def __init__(self, address, key, session_id): add_exit_handlers() while True: - conn = listener.accept() - while True: + with listener.accept() as conn: try: data = conn.recv() # because of this no sleep needed in thread logger.debug("Incoming data : %s", data) @@ -65,12 +64,9 @@ def __init__(self, address, key, session_id): conn=conn, queue=self.queue, ) - conn.close() # Sort of EOL for Python IPC - break except Exception as e: logger.error("Exception occurred in server thread : %s", e) logger.debug("Trace \n %s", traceback.format_exc()) - break # Return back to listening for new connections def reporting_thread(self): """Reporting thread""" diff --git a/aikido_zen/background_process/aikido_background_process_test.py b/aikido_zen/background_process/aikido_background_process_test.py new file mode 100644 index 000000000..2e4fe8be2 --- /dev/null +++ b/aikido_zen/background_process/aikido_background_process_test.py @@ -0,0 +1,96 @@ +import importlib +import socket +from multiprocessing import connection +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from aikido_zen.background_process.commands import process_incoming_command +from aikido_zen.background_process.commands.put_event import PutEventReq + + +def test_event_is_acknowledged_after_enqueueing(): + operations = [] + events = Mock() + events.put.side_effect = lambda event: operations.append(("queued", event)) + connection = Mock() + connection.send.side_effect = lambda value: operations.append(("reply", value)) + event = {"name": "login.failed"} + + process_incoming_command( + None, ("PUT_EVENT", PutEventReq(event)), connection, events + ) + + assert operations == [("queued", event), ("reply", None)] + connection.send.assert_called_once_with(None) + + +def test_event_is_not_acknowledged_when_enqueueing_fails(): + events = Mock() + events.put.side_effect = RuntimeError("Queue unavailable") + connection = Mock() + + with pytest.raises(RuntimeError, match="Queue unavailable"): + process_incoming_command( + None, + ("PUT_EVENT", PutEventReq({"name": "login.failed"})), + connection, + events, + ) + + connection.send.assert_not_called() + + +@pytest.mark.parametrize("failure", ["receive", "reply", "command"]) +def test_receiver_closes_failed_connection_before_accepting_next(monkeypatch, failure): + background = importlib.import_module( + "aikido_zen.background_process.aikido_background_process" + ) + + class ListenerFinished(Exception): + pass + + failed_socket, failed_client_socket = socket.socketpair() + next_socket, next_client_socket = socket.socketpair() + failed = connection.Connection(failed_socket.detach()) + failed_client = connection.Connection(failed_client_socket.detach()) + next_peer = connection.Connection(next_socket.detach()) + next_client = connection.Connection(next_client_socket.detach()) + accepted = [] + + def accept(): + if accepted: + assert accepted[-1].closed + if len(accepted) == 2: + raise ListenerFinished + peer = (failed, next_peer)[len(accepted)] + accepted.append(peer) + return peer + + monkeypatch.setattr( + background.con, + "Listener", + lambda *args, **kwargs: SimpleNamespace(accept=accept), + ) + monkeypatch.setattr( + background, + "Thread", + lambda *args, **kwargs: SimpleNamespace(start=lambda: None), + ) + monkeypatch.setattr(background, "add_exit_handlers", lambda: None) + monkeypatch.setattr(background, "set_agent_session_id", lambda value: None) + try: + if failure != "receive": + failed_client.send(None if failure == "command" else ("PING", ())) + failed_client.close() + next_client.send(("PING", ())) + + with pytest.raises(ListenerFinished): + background.AikidoBackgroundProcess(("127.0.0.1", 0), b"fixture", "fixture") + + assert next_client.recv() == "Received" + assert failed.closed and next_peer.closed + finally: + for peer in (failed, failed_client, next_peer, next_client): + peer.close() diff --git a/aikido_zen/background_process/commands/__init__.py b/aikido_zen/background_process/commands/__init__.py index 5b36b7580..e625d6851 100644 --- a/aikido_zen/background_process/commands/__init__.py +++ b/aikido_zen/background_process/commands/__init__.py @@ -1,34 +1,29 @@ from aikido_zen.helpers.logging import logger -from aikido_zen.helpers.ipc.command_types import CommandContext -from .put_event import PutEventCommand -from .check_firewall_lists import process_check_firewall_lists -from .read_property import process_read_property -from .should_ratelimit import process_should_ratelimit -from .ping import process_ping -from .sync_data import process_sync_data +from . import ( + check_firewall_lists, + ping, + put_event, + read_property, + should_ratelimit, + sync_data, +) commands_map = { - "SYNC_DATA": process_sync_data, - "READ_PROPERTY": process_read_property, - "SHOULD_RATELIMIT": process_should_ratelimit, - "PING": process_ping, - "CHECK_FIREWALL_LISTS": process_check_firewall_lists, + check_firewall_lists.IDENTIFIER: check_firewall_lists.process, + ping.IDENTIFIER: ping.process, + put_event.IDENTIFIER: put_event.process, + read_property.IDENTIFIER: read_property.process, + should_ratelimit.IDENTIFIER: should_ratelimit.process, + sync_data.IDENTIFIER: sync_data.process, } -modern_commands = [PutEventCommand] - def process_incoming_command(connection_manager, obj, conn, queue): inbound_identifier = obj[0] inbound_request = obj[1] if inbound_identifier in commands_map: func = commands_map[inbound_identifier] - return conn.send(func(connection_manager, inbound_request)) - - for cmd in modern_commands: - if cmd.identifier() == inbound_identifier: - cmd.run(CommandContext(connection_manager, queue, conn), inbound_request) - return None + return conn.send(func(connection_manager, inbound_request, queue)) logger.debug("Command : `%s` not found - did not execute", inbound_identifier) return None diff --git a/aikido_zen/background_process/commands/check_firewall_lists.py b/aikido_zen/background_process/commands/check_firewall_lists.py index acf323043..e48dc9b2c 100644 --- a/aikido_zen/background_process/commands/check_firewall_lists.py +++ b/aikido_zen/background_process/commands/check_firewall_lists.py @@ -1,13 +1,14 @@ -"""Exports process_check_firewall_lists""" +from aikido_zen.helpers.ipc.command_types import Payload +IDENTIFIER = "CHECK_FIREWALL_LISTS" -def process_check_firewall_lists(connection_manager, data): - """ - Checks whether an IP is blocked - data: {"ip": string, "user-agent": string} - returns -> {"blocked": boolean, "type": string, "reason": string} - """ - ip = data["ip"] + +def generate(ip, user_agent, timeout_in_sec=0.01) -> Payload: + return Payload(IDENTIFIER, {"ip": ip, "user-agent": user_agent}, timeout_in_sec) + + +def process(connection_manager, request, queue): + ip = request["ip"] if ip is not None and isinstance(ip, str): # Global IP Allowlist (e.g. for geofencing) if not connection_manager.firewall_lists.is_allowed_ip(ip): @@ -22,7 +23,7 @@ def process_check_firewall_lists(connection_manager, data): "reason": reason, } - user_agent = data["user-agent"] + user_agent = request["user-agent"] if user_agent is not None and isinstance(user_agent, str): # User agent blocking (e.g. blocking AI scrapers) if connection_manager.firewall_lists.is_user_agent_blocked(user_agent): diff --git a/aikido_zen/background_process/commands/ping.py b/aikido_zen/background_process/commands/ping.py index 9363eab0c..0b24de57f 100644 --- a/aikido_zen/background_process/commands/ping.py +++ b/aikido_zen/background_process/commands/ping.py @@ -1,6 +1,11 @@ -"""exports `process_ping`""" +from aikido_zen.helpers.ipc.command_types import Payload +IDENTIFIER = "PING" -def process_ping(connection_manager, data): - """when main process quits , or during testing etc""" + +def generate(timeout_in_sec=0.1) -> Payload: + return Payload(IDENTIFIER, (), timeout_in_sec) + + +def process(connection_manager, request, queue): return "Received" diff --git a/aikido_zen/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index 790df47a3..0da133588 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -1,27 +1,18 @@ -from aikido_zen.helpers.ipc.command_types import Command, CommandContext, Payload +from aikido_zen.helpers.ipc.command_types import Payload +IDENTIFIER = "PUT_EVENT" -class PutEventReq: - def __init__(self, event): - # Event is a dictionary containing data that is going to be reported to core - # The CloudConnectionManager adds the "agent" field - self.event = event +def generate(event, timeout_in_sec=5) -> Payload: + return Payload(IDENTIFIER, PutEventReq(event), timeout_in_sec) -class PutEventCommand(Command): - @classmethod - def identifier(cls) -> str: - return "put_event" - @classmethod - def returns_data(cls) -> bool: - return False +def process(connection_manager, request, queue): + queue.put(request.event) - @classmethod - def run(cls, context: CommandContext, request: PutEventReq): - # Events sent here get put in the event queue so that they are processed in the background - context.queue.put(request.event) - @classmethod - def generate(cls, request) -> Payload: - return Payload(cls, PutEventReq(request)) +# IPC uses pickle, so this request class must keep its import path. +class PutEventReq: + def __init__(self, event): + # The CloudConnectionManager adds the "agent" field. + self.event = event diff --git a/aikido_zen/background_process/commands/read_property.py b/aikido_zen/background_process/commands/read_property.py index 1bff5e395..a3ee5928e 100644 --- a/aikido_zen/background_process/commands/read_property.py +++ b/aikido_zen/background_process/commands/read_property.py @@ -1,19 +1,20 @@ -"""Mainly exports process_read_property""" - +from aikido_zen.helpers.ipc.command_types import Payload from aikido_zen.helpers.logging import logger +IDENTIFIER = "READ_PROPERTY" + + +def generate(property_name, timeout_in_sec=0.1) -> Payload: + return Payload(IDENTIFIER, property_name, timeout_in_sec) + -def process_read_property(connection_manager, data): - """ - Takes in one arg : name of property on connection_manager, tries to read it. - Meant to get config props - """ +def process(connection_manager, request, queue): try: - return connection_manager.__dict__[data] + return connection_manager.__dict__[request] except KeyError: logger.debug( "CloudConnectionManager has no attribute %s, current connection_manager: %s", - data, + request, connection_manager, ) return None diff --git a/aikido_zen/background_process/commands/read_property_test.py b/aikido_zen/background_process/commands/read_property_test.py index a6a469e16..8f7f34799 100644 --- a/aikido_zen/background_process/commands/read_property_test.py +++ b/aikido_zen/background_process/commands/read_property_test.py @@ -1,8 +1,6 @@ import pytest from unittest.mock import MagicMock, patch -from .read_property import ( - process_read_property, -) # Replace 'your_module' with the actual module name +from . import read_property from aikido_zen.helpers.logging import logger @@ -15,17 +13,15 @@ def mock_connection_manager(): def test_process_read_property_sends_value(mock_connection_manager): - """Test that process_read_property sends the value of the property.""" data = "some_property" - process_read_property(mock_connection_manager, data) + read_property.process(mock_connection_manager, data, None) def test_process_read_property_sends_none_for_missing_property(mock_connection_manager): - """Test that process_read_property sends None for a missing property.""" data = "missing_property" - process_read_property(mock_connection_manager, data) + read_property.process(mock_connection_manager, data, None) def test_process_read_property_logs_debug_for_missing_property( @@ -35,7 +31,7 @@ def test_process_read_property_logs_debug_for_missing_property( data = "missing_property" with patch("aikido_zen.helpers.logging.logger.debug") as mock_debug: - process_read_property(mock_connection_manager, data) + read_property.process(mock_connection_manager, data, None) # Check that the debug log was called mock_debug.assert_called_once_with( diff --git a/aikido_zen/background_process/commands/should_ratelimit.py b/aikido_zen/background_process/commands/should_ratelimit.py index d51f99ce9..708a66292 100644 --- a/aikido_zen/background_process/commands/should_ratelimit.py +++ b/aikido_zen/background_process/commands/should_ratelimit.py @@ -1,20 +1,32 @@ -"""Exports `process_should_ratelimit`""" - import aikido_zen.ratelimiting as ratelimiting +from aikido_zen.helpers.ipc.command_types import Payload + +IDENTIFIER = "SHOULD_RATELIMIT" + + +def generate( + route_metadata, user, group, remote_address, timeout_in_sec=0.01 +) -> Payload: + return Payload( + IDENTIFIER, + { + "route_metadata": route_metadata, + "user": user, + "group": group, + "remote_address": remote_address, + }, + timeout_in_sec, + ) + -def process_should_ratelimit(connection_manager, data): - """ - Called to check if the context passed along as data should be rate limited - data object should be a dict including route_metadata, remote_address and user - route_metadata object includes route, url and method - """ +def process(connection_manager, request, queue): if not connection_manager: return False return ratelimiting.should_ratelimit_request( - route_metadata=data["route_metadata"], - remote_address=data["remote_address"], - user=data["user"], + route_metadata=request["route_metadata"], + remote_address=request["remote_address"], + user=request["user"], connection_manager=connection_manager, - group=data["group"], + group=request["group"], ) diff --git a/aikido_zen/background_process/commands/should_ratelimit_test.py b/aikido_zen/background_process/commands/should_ratelimit_test.py index 88a07bbb1..755c79b38 100644 --- a/aikido_zen/background_process/commands/should_ratelimit_test.py +++ b/aikido_zen/background_process/commands/should_ratelimit_test.py @@ -1,16 +1,16 @@ import pytest from unittest.mock import Mock, patch -from .should_ratelimit import process_should_ratelimit +from . import should_ratelimit @pytest.mark.parametrize( - "should_ratelimit, expected_call", + "rate_limited, expected_call", [ (True, True), (False, False), ], ) -def test_process_should_ratelimit(should_ratelimit, expected_call): +def test_process_should_ratelimit(rate_limited, expected_call): """Test rate limiting behavior based on different return values.""" # Arrange connection_manager = Mock() @@ -27,10 +27,10 @@ def test_process_should_ratelimit(should_ratelimit, expected_call): with patch( "aikido_zen.ratelimiting.should_ratelimit_request", - return_value=should_ratelimit, + return_value=rate_limited, ) as mock_should_ratelimit: # Act - process_should_ratelimit(connection_manager, data) + should_ratelimit.process(connection_manager, data, None) # Assert mock_should_ratelimit.assert_called_once_with( @@ -43,7 +43,6 @@ def test_process_should_ratelimit(should_ratelimit, expected_call): def test_process_should_ratelimit_no_connection_manager(): - """Test that process_should_ratelimit does nothing when connection_manager is not present.""" # Arrange data = { "route_metadata": { @@ -57,14 +56,13 @@ def test_process_should_ratelimit_no_connection_manager(): } # Act - result = process_should_ratelimit(None, data) # No connection_manager + result = should_ratelimit.process(None, data, None) # No connection_manager # Assert assert result is False # Ensure the function returns False def test_process_should_ratelimit_multiple_calls(): - """Test multiple calls to process_should_ratelimit.""" # Arrange connection_manager = Mock() data = { @@ -83,8 +81,8 @@ def test_process_should_ratelimit_multiple_calls(): side_effect=[True, False], ) as mock_should_ratelimit: # Act - process_should_ratelimit(connection_manager, data) # First call - process_should_ratelimit(connection_manager, data) # Second call + should_ratelimit.process(connection_manager, data, None) # First call + should_ratelimit.process(connection_manager, data, None) # Second call # Assert mock_should_ratelimit.assert_any_call( @@ -116,8 +114,8 @@ def test_process_should_ratelimit_with_different_connection_manager(): "aikido_zen.ratelimiting.should_ratelimit_request", return_value=True ) as mock_should_ratelimit: # Act - process_should_ratelimit(connection_manager1, data) - process_should_ratelimit(connection_manager2, data) + should_ratelimit.process(connection_manager1, data, None) + should_ratelimit.process(connection_manager2, data, None) # Assert assert mock_should_ratelimit.call_count == 2 diff --git a/aikido_zen/background_process/commands/sync_data.py b/aikido_zen/background_process/commands/sync_data.py index b8f22ffb6..9fa4fa46d 100644 --- a/aikido_zen/background_process/commands/sync_data.py +++ b/aikido_zen/background_process/commands/sync_data.py @@ -1,19 +1,22 @@ -"""Exports process_renew_config""" - from aikido_zen.background_process.packages import PackagesStore +from aikido_zen.helpers.ipc.command_types import Payload + +IDENTIFIER = "SYNC_DATA" + + +def generate(request, timeout_in_sec=5) -> Payload: + return Payload(IDENTIFIER, request, timeout_in_sec) -def process_sync_data(connection_manager, data): +def process(connection_manager, request, queue): """ - Synchronizes data between the thread-local cache (with a TTL of usually 1 minute) and the - background thread. Which data gets synced? - Thread -> BG Process : Routes, Hostnames, Users, Stats & middleware installed - BG Process -> Thread : Config (only if the worker hasn't received the current revision) + Application -> Background process: Routes, Hostnames, Users, Stats & middleware installed. + Background process -> Application: Config, only when the revision has changed. """ # Sync routes routes = connection_manager.routes - for route in data.get("current_routes", {}).values(): + for route in request.get("current_routes", {}).values(): route_metadata = {"method": route["method"], "route": route["path"]} if not routes.get(route_metadata): routes.initialize_route(route_metadata) @@ -27,11 +30,11 @@ def process_sync_data(connection_manager, data): routes.update_route_with_apispec(route_metadata, route["apispec"]) # Save middleware installed : - if data.get("middleware_installed", False): + if request.get("middleware_installed", False): connection_manager.middleware_installed = True # Sync hostnames - for hostnames_entry in data.get("hostnames", []): + for hostnames_entry in request.get("hostnames", []): connection_manager.hostnames.add( hostnames_entry["hostname"], hostnames_entry["port"], @@ -39,21 +42,21 @@ def process_sync_data(connection_manager, data): ) # Sync users - for user_entry in data.get("users", []): + for user_entry in request.get("users", []): connection_manager.users.add_user_from_entry(user_entry) # Sync stats - connection_manager.statistics.import_from_record(data.get("stats", {})) + connection_manager.statistics.import_from_record(request.get("stats", {})) # Sync ai stats - connection_manager.ai_stats.import_list(data.get("ai_stats", [])) + connection_manager.ai_stats.import_list(request.get("ai_stats", [])) # Sync packages - PackagesStore.import_list(data.get("packages", [])) + PackagesStore.import_list(request.get("packages", [])) if ( connection_manager.conf.last_updated_at > 0 - and data.get("config_revision") != connection_manager.conf.revision + and request.get("config_revision") != connection_manager.conf.revision ): return {"config": connection_manager.conf} return {} diff --git a/aikido_zen/background_process/commands/sync_data_test.py b/aikido_zen/background_process/commands/sync_data_test.py index 1e4778c14..da277162a 100644 --- a/aikido_zen/background_process/commands/sync_data_test.py +++ b/aikido_zen/background_process/commands/sync_data_test.py @@ -2,7 +2,7 @@ import pytest from unittest.mock import MagicMock -from .sync_data import process_sync_data +from . import sync_data from aikido_zen.background_process.routes import Routes from ..packages import PackagesStore from ...helpers.ip_matcher import IPMatcher @@ -72,7 +72,7 @@ def test_process_sync_data_initialization(setup_connection_manager): ], } - result = process_sync_data(connection_manager, data) + result = sync_data.process(connection_manager, data, None) # Check that routes were initialized correctly assert len(connection_manager.routes) == 2 @@ -147,7 +147,7 @@ def test_process_sync_data_with_last_updated_at_below_zero(setup_connection_mana "middleware_installed": True, } - result = process_sync_data(connection_manager, data) + result = sync_data.process(connection_manager, data, None) # Check that routes were initialized correctly assert len(connection_manager.routes) == 2 @@ -215,7 +215,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager } # First call to initialize the route - process_sync_data(connection_manager, data) + sync_data.process(connection_manager, data, None) # Second call to update the existing route data_update = { @@ -242,7 +242,7 @@ def test_process_sync_data_existing_route_and_hostnames(setup_connection_manager }, } - result = process_sync_data(connection_manager, data_update) + result = sync_data.process(connection_manager, data_update, None) # Check that the hit count was updated correctly assert ( @@ -277,7 +277,7 @@ def test_process_sync_data_no_routes(setup_connection_manager): connection_manager = setup_connection_manager data = {"current_routes": {}, "reqs": 0} # No requests to add - result = process_sync_data(connection_manager, data) + result = sync_data.process(connection_manager, data, None) # Check that no routes were initialized assert len(connection_manager.routes) == 0 diff --git a/aikido_zen/background_process/comms.py b/aikido_zen/background_process/comms.py deleted file mode 100644 index 48b7e402e..000000000 --- a/aikido_zen/background_process/comms.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -Holds the globally stored comms object -Exports the AikidoIPCCommunications class -""" - -import multiprocessing.connection as con -from threading import Thread -from aikido_zen.helpers.logging import logger - -# pylint: disable=invalid-name # This variable does change -comms = None - - -def get_comms(): - """ - Returns the globally stored IPC object, which you need - to communicate to our background process. - """ - return comms - - -def reset_comms(): - """This will reset communications""" - # pylint: disable=global-statement # This needs to be global - global comms - if comms: - logger.debug("Resetting communications. (comms = None)") - comms = None - - -class AikidoIPCCommunications: - """ - Facilitates Inter-Process communication - """ - - def __init__(self, address, key): - # The key needs to be in byte form - self.address = address - self.key = key - - # Set as global ipc object : - reset_comms() - # pylint: disable=global-statement # This needs to be global - global comms - comms = self - - def send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1): - """Try-catched send_data_to_bg_process""" - try: - return self._send_data_to_bg_process(action, obj, receive, timeout_in_sec) - except Exception as e: - logger.debug("Exception happened in send_data_to_bg_process : %s", e) - return {"success": False, "error": "unknown"} - - def _send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1): - """ - This creates a new client for comms to the background process - """ - if not self.key: - # If no key is set, the background process will not start - return {"success": False, "error": "invalid_key"} - - # We want to make sure that sending out this data affects the process as little as possible - # So we run it inside a separate thread with a timeout of 100ms - # If something goes wrong, it will also be encapsulated in the thread i.e. no crashes - def target(address, key, receive, data, result_obj): - try: - # Create a connection, this can get stuck : - conn = con.Client(address, authkey=None) - - # Send/Receive data : - conn.send(data) - if receive: - result_obj[1] = conn.recv() - # Close the connection : - conn.close() - result_obj[0] = True # Connection ended gracefully - except Exception as e: - logger.debug("Exception occurred in thread : %s", e) - - # Create a shared result object between the thread and this process : - result_obj = [False, None] # Needs to be an array so we can make a ref. - t = Thread( - target=target, - args=(self.address, self.key, receive, (action, obj), result_obj), - daemon=True, # This allows us to join and set a timeout after which the thread closes - ) - - # Start and join the thread for 100ms, afterwards the thread is forced to close (daemon=True) - t.start() - t.join(timeout=timeout_in_sec) - if not result_obj[0]: - logger.debug( - " Failure in communication to background process, %s(%s)", action, obj - ) - return {"success": False, "error": "timeout"} - - if receive: - return {"success": True, "data": result_obj[1]} - return {"success": True, "data": None} diff --git a/aikido_zen/background_process/comms_test.py b/aikido_zen/background_process/comms_test.py deleted file mode 100644 index e1ee26411..000000000 --- a/aikido_zen/background_process/comms_test.py +++ /dev/null @@ -1,40 +0,0 @@ -import pytest -from aikido_zen.background_process.comms import ( - AikidoIPCCommunications, - reset_comms, -) - - -@pytest.fixture(autouse=True) -def reset_comms_after_test(): - yield - reset_comms() - - -def test_comms_init(): - address = ("localhost", 9898) - key = "secret_key" - comms = AikidoIPCCommunications(address, key) - - assert comms.address == address - assert comms.key == key - - -def test_send_data_to_bg_process_exception(monkeypatch, caplog): - def mock_client(address, authkey): - raise Exception("Connection Error") - - monkeypatch.setitem(globals(), "Client", mock_client) - monkeypatch.setitem(globals(), "logger", caplog) - - comms = AikidoIPCCommunications(("localhost", 9898), "mock_key") - comms.send_data_to_bg_process("ACTION", "Test Object") - - -def test_send_data_to_bg_process_successful(monkeypatch, caplog, mocker): - comms = AikidoIPCCommunications(("localhost"), "mock_key") - mock_client = mocker.MagicMock() - monkeypatch.setattr("multiprocessing.connection.Client", mock_client) - - # Call the send_data_to_bg_process function - comms.send_data_to_bg_process("ACTION", {"key": "value"}) diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 702bc7cae..91c385a3d 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -11,7 +11,7 @@ get_agent_session_id, set_agent_session_id, ) -from .comms_test import reset_comms_after_test +from .ipc_test import reset_ipc_after_test @pytest.mark.parametrize("version", [(3, 8), (3, 13), (3, 14)]) @@ -149,6 +149,35 @@ def test_get_process_factory_does_not_set_global_start_method(): subprocess.run([sys.executable, "-c", script], check=True) +def test_socket_name_is_shared_by_independent_processes(monkeypatch): + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + script = """ +from aikido_zen.background_process import get_socket_path +print(get_socket_path()) +""" + + for _ in range(2): + address = subprocess.check_output( + [sys.executable, "-c", script], text=True, timeout=10 + ).strip() + assert address == background_process.get_socket_path() + + +def test_socket_name_separates_tokens_and_agent_versions(monkeypatch): + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + versions = ["1.2.37", "1.2.37b2", "1.2.37+custom.1"] + addresses = set() + for version in versions: + monkeypatch.setattr(background_process, "PKG_VERSION", version) + address = background_process.get_socket_path() + assert address == background_process.get_socket_path() + addresses.add(address) + + assert len(addresses) == len(versions) + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_OTHER") + assert background_process.get_socket_path() not in addresses + + def test_stale_socket_removed_by_another_worker(monkeypatch, mocker): process = mocker.patch( "aikido_zen.background_process.get_process_factory" @@ -156,11 +185,11 @@ def test_stale_socket_removed_by_another_worker(monkeypatch, mocker): monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") monkeypatch.setattr(background_process.platform, "system", lambda: "Linux") monkeypatch.setattr( - background_process, "get_uds_filename", lambda: "/tmp/aikido.sock" + background_process, "get_socket_path", lambda: "/tmp/aikido.sock" ) monkeypatch.setattr(background_process.os.path, "exists", lambda _path: True) monkeypatch.setattr( - background_process, "background_process_already_active", lambda _comms: False + background_process, "background_process_already_active", lambda _ipc: False ) monkeypatch.setattr( background_process.os, @@ -180,7 +209,7 @@ def test_started_process_receives_the_current_session_id(monkeypatch, mocker): monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") monkeypatch.setattr(background_process.platform, "system", lambda: "Linux") monkeypatch.setattr( - background_process, "get_uds_filename", lambda: "/tmp/aikido-session.sock" + background_process, "get_socket_path", lambda: "/tmp/aikido-session.sock" ) monkeypatch.setattr(background_process.os.path, "exists", lambda _path: False) diff --git a/aikido_zen/background_process/ipc.py b/aikido_zen/background_process/ipc.py new file mode 100644 index 000000000..d5581b3b6 --- /dev/null +++ b/aikido_zen/background_process/ipc.py @@ -0,0 +1,111 @@ +import os +import threading +import time +from queue import Full, Queue + +from aikido_zen.background_process.socket_connection import open_connection +from aikido_zen.helpers.ipc.command_types import Payload +from aikido_zen.helpers.logging import logger + +MAX_PENDING_EVENTS = 1000 + + +class IPCClient: + _instance = None + + # Share one worker so replacing a client does not leave another idle thread behind. + _pid = os.getpid() + _start_locks = {_pid: threading.Lock()} + _queue = Queue(maxsize=MAX_PENDING_EVENTS) + _thread = None + _warned_full = False + + def __init__(self, address, key): + self.address = address + self.key = key + + IPCClient._instance = self + + @staticmethod + def get(): + return IPCClient._instance + + @staticmethod + def reset(): + IPCClient._instance = None + + def send_message_async(self, message: Payload): + """Return local queue acceptance; delivery and its reply run on the worker.""" + if not self.key: + return {"success": False, "error": "invalid_key"} + # pylint: disable=protected-access # The alias accesses our own shared class state. + shared = IPCClient + try: + pid = os.getpid() + start_lock = shared._start_locks.get(pid) + if start_lock is None: + # A fork can copy a lock held by a thread that does not survive in the child. + start_lock = shared._start_locks.setdefault(pid, threading.Lock()) + + with start_lock: + # Pending events belong to the parent, so a forked child must start empty. + if shared._pid != pid: + shared._queue = Queue(maxsize=MAX_PENDING_EVENTS) + shared._thread = None + shared._warned_full = False + shared._pid = pid + shared._start_locks = {pid: start_lock} + if shared._thread is None: + thread = threading.Thread( + target=shared._send_events, + args=(shared._queue,), + name="aikido-event-sender", + daemon=True, + ) + thread.start() + shared._thread = thread + pending = shared._queue + + pending.put_nowait((self.send_message, message)) + return {"success": True, "queued": True} + except Full: + if not shared._warned_full: + shared._warned_full = True + logger.warning("Event delivery queue is full; dropping new events.") + return {"success": False, "error": "queue_full"} + except Exception as error: + logger.debug("Could not enqueue event: %s", error) + return {"success": False, "error": "unknown"} + + def send_message(self, message: Payload): + try: + if not self.key: + # If no key is set, the background process will not start + return {"success": False, "error": "invalid_key"} + + if message.timeout_in_sec <= 0: + return {"success": False, "error": "timeout"} + deadline = time.monotonic() + message.timeout_in_sec + + with open_connection(self.address, deadline) as conn: + conn.send((message.identifier, message.request)) + result = conn.recv() + except Exception as e: + logger.debug("Exception occurred in IPC call : %s", e) + return {"success": False, "error": "timeout"} + return {"success": True, "data": result} + + @staticmethod + def _send_events(pending): + while True: + item = pending.get() + try: + if item is None: + return + item[0](item[1]) + except Exception as error: + logger.debug("Exception occurred in event delivery: %s", error) + finally: + # The idle worker must not retain the last request it sent. + del item + pending.task_done() diff --git a/aikido_zen/background_process/ipc_test.py b/aikido_zen/background_process/ipc_test.py new file mode 100644 index 000000000..72e2d0c4b --- /dev/null +++ b/aikido_zen/background_process/ipc_test.py @@ -0,0 +1,786 @@ +import gc +import os +import socket +import struct +import threading +import time +import weakref +from contextlib import ExitStack, contextmanager +from multiprocessing import connection +from multiprocessing.connection import Listener +from queue import Queue +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from aikido_zen.background_process import ipc as module +from aikido_zen.background_process import socket_connection as transport +from aikido_zen.background_process.commands import ( + check_firewall_lists, + ping, + put_event, + read_property, +) +from aikido_zen.background_process.commands import process_incoming_command + + +@pytest.fixture(autouse=True) +def reset_ipc_after_test(): + yield + module.IPCClient.reset() + + +class Payload: + def __init__(self, size=256 * 1024): + self.body = bytearray(size) + + +@pytest.fixture +def resources(monkeypatch): + sockets = [] + connections = [] + interrupted_sends = [] + + class TrackedSocket(socket.socket): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + sockets.append(self) + + def sendall(self, *args, **kwargs): + try: + return super().sendall(*args, **kwargs) + except OSError: + interrupted_sends.append(True) + raise + + monkeypatch.setattr( + transport, + "socket", + SimpleNamespace(**{**vars(socket), "socket": TrackedSocket}), + raising=False, + ) + + class TrackedConnection(transport.SocketConnection): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + connections.append(self) + + monkeypatch.setattr(transport, "SocketConnection", TrackedConnection) + return SimpleNamespace( + sockets=sockets, + connections=connections, + interrupted_sends=interrupted_sends, + ) + + +def assert_released(resources, payloads): + gc.collect() + assert all(sock.fileno() == -1 for sock in resources.sockets) + assert all(conn.closed for conn in resources.connections) + assert all(reference() is None for reference in payloads) + + +class PeerServer: + def __init__(self, count, reply=b"", read_request=True): + self.peers = [] + self.count = count + self.reply = reply + self.read_request = read_request + self.finished = threading.Event() + self.stopped = threading.Event() + self.listener = socket.socket() + if not read_request: + self.listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096) + self.listener.bind(("127.0.0.1", 0)) + self.listener.listen(count) + self.listener.settimeout(0.1) + self.address = self.listener.getsockname() + self.thread = threading.Thread(target=self.serve, daemon=True) + self.thread.start() + + def serve(self): + try: + while len(self.peers) < self.count and not self.stopped.is_set(): + try: + peer, _ = self.listener.accept() + except socket.timeout: + continue + self.peers.append(peer) + request = None + if self.read_request: + with connection.Connection(peer.dup().detach()) as conn: + request = conn.recv() + if callable(self.reply): + self.reply(peer, request) + elif self.reply: + peer.sendall(self.reply) + except (EOFError, OSError): + pass + finally: + self.finished.set() + + def close(self): + self.stopped.set() + self.listener.close() + for peer in self.peers: + peer.close() + self.thread.join(1) + + +@pytest.fixture +def servers(): + instances = [] + + def create(*args, **kwargs): + server = PeerServer(*args, **kwargs) + instances.append(server) + return server + + yield create + for server in instances: + server.close() + + +@contextmanager +def full_accept_queue(listener): + queued = [] + with ExitStack() as stack: + for _ in range(100): + peer = stack.enter_context(socket.socket(listener.family)) + peer.settimeout(0.02) + try: + peer.connect(listener.getsockname()) + except (BlockingIOError, socket.timeout): + break + queued.append(peer) + else: + pytest.skip("Could not fill the local accept queue") + yield queued + + +@pytest.fixture +def sender(monkeypatch): + shared = module.IPCClient + monkeypatch.setattr(shared, "_pid", os.getpid()) + monkeypatch.setattr(shared, "_start_locks", {os.getpid(): threading.Lock()}) + monkeypatch.setattr( + shared, "_queue", module.Queue(maxsize=module.MAX_PENDING_EVENTS) + ) + monkeypatch.setattr(shared, "_thread", None) + monkeypatch.setattr(shared, "_warned_full", False) + yield shared + if shared._thread is not None: + shared._queue.put(None, timeout=5) + shared._thread.join(5) + assert not shared._thread.is_alive() + + +def assert_drained(sender): + with sender._queue.all_tasks_done: + assert sender._queue.all_tasks_done.wait_for( + lambda: sender._queue.unfinished_tasks == 0, timeout=5 + ) + + +def test_ipc_init(): + address = ("localhost", 9898) + key = "secret_key" + ipc = module.IPCClient(address, key) + + assert ipc.address == address + assert ipc.key == key + assert module.IPCClient.get() is ipc + module.IPCClient.reset() + assert module.IPCClient.get() is None + + +@pytest.mark.parametrize( + "message,expected", + [ + (ping.generate(timeout_in_sec=1), "Received"), + (read_property.generate("block", timeout_in_sec=1), True), + (put_event.generate({"name": "login.failed"}, timeout_in_sec=1), None), + ], +) +def test_commands_return_results_or_acknowledgments( + resources, servers, monkeypatch, message, expected +): + events = Queue() + + def respond(peer, request): + with connection.Connection(peer.dup().detach()) as conn: + process_incoming_command(SimpleNamespace(block=True), request, conn, events) + + server = servers(1, reply=respond) + + def fail_start(self): + raise AssertionError("IPC call must not start a thread") + + monkeypatch.setattr(threading.Thread, "start", fail_start) + communications = module.IPCClient(server.address, b"fixture") + assert communications.send_message(message) == { + "success": True, + "data": expected, + } + if message.identifier == "PUT_EVENT": + assert events.get_nowait() == {"name": "login.failed"} + assert_released(resources, []) + + +def test_completed_blocked_reply_is_kept_after_deadline(monkeypatch): + clock = Mock(return_value=0.0) + verdict = {"blocked": True, "type": "bot-blocking"} + + def receive(): + clock.return_value = 0.02 + return verdict + + @contextmanager + def completed_connection(address, deadline): + assert address == "fixture" + assert deadline == 0.01 + yield SimpleNamespace(send=Mock(), recv=receive) + + monkeypatch.setattr(module, "time", SimpleNamespace(monotonic=clock)) + monkeypatch.setattr(module, "open_connection", completed_connection) + ipc = module.IPCClient("fixture", b"fixture") + + assert ipc.send_message( + check_firewall_lists.generate("203.0.113.10", "browser") + ) == {"success": True, "data": verdict} + + +@pytest.mark.parametrize( + "reply", + [b"", b"\x00", struct.pack("!i", 1000) + b"\x80"], + ids=["no_reply", "partial_header", "partial_body"], +) +def test_timeout_releases_sockets_and_payloads(resources, servers, reply): + server = servers(5, reply=reply) + communications = module.IPCClient(server.address, b"fixture") + payloads = [] + for _ in range(5): + payload = Payload() + payloads.append(weakref.ref(payload)) + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.05) + ) == { + "success": False, + "error": "timeout", + } + del payload + + assert server.finished.wait(1) + assert len(server.peers) == 5 + assert_released(resources, payloads) + for peer in server.peers: + peer.settimeout(1) + try: + assert peer.recv(1) == b"" + except ConnectionResetError: + pass + + +@pytest.mark.parametrize( + "reply", + [b"", b"\x00", struct.pack("!i", 100) + b"\x80"], + ids=["before_header", "during_header", "during_body"], +) +def test_peer_disconnect_releases_sockets_and_payloads(resources, servers, reply): + def disconnect(peer, _request): + if reply: + peer.sendall(reply) + peer.shutdown(socket.SHUT_RDWR) + peer.close() + + server = servers(1, reply=disconnect) + communications = module.IPCClient(server.address, b"fixture") + payload = Payload() + reference = weakref.ref(payload) + + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) + assert not result["success"] + del payload + assert server.finished.wait(1) + assert_released(resources, [reference]) + + +def test_send_failure_releases_sockets_and_payloads(resources, servers, monkeypatch): + def fail_send(self, data): + raise BrokenPipeError("Peer closed during send") + + monkeypatch.setattr(transport.SocketConnection, "_send", fail_send) + server = servers(1) + communications = module.IPCClient(server.address, b"fixture") + payload = Payload() + reference = weakref.ref(payload) + + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) + assert not result["success"] + del payload + assert_released(resources, [reference]) + + +def test_connect_failure_closes_socket(resources): + with socket.socket() as reserved: + reserved.bind(("127.0.0.1", 0)) + communications = module.IPCClient(reserved.getsockname(), b"fixture") + result = communications.send_message(ping.generate(timeout_in_sec=0.1)) + assert not result["success"] + assert_released(resources, []) + + +def test_pickle_failure_releases_connection(resources, servers): + class Unpicklable: + def __reduce__(self): + raise ValueError("Cannot serialize payload") + + server = servers(1) + communications = module.IPCClient(server.address, b"fixture") + payload = Unpicklable() + reference = weakref.ref(payload) + + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) + assert not result["success"] + del payload + assert_released(resources, [reference]) + + +def test_invalid_reply_releases_connection(resources, servers): + server = servers(1, reply=struct.pack("!i", 4) + b"oops") + communications = module.IPCClient(server.address, b"fixture") + + result = communications.send_message(ping.generate(timeout_in_sec=1)) + assert not result["success"] + assert_released(resources, []) + + +def test_timeout_interrupts_a_blocked_send(resources, servers, monkeypatch): + create_socket = transport.socket.socket + + def small_send_buffer(*args, **kwargs): + sock = create_socket(*args, **kwargs) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4096) + return sock + + monkeypatch.setattr(transport.socket, "socket", small_send_buffer) + server = servers(1, read_request=False) + communications = module.IPCClient(server.address, b"fixture") + payload = Payload(8 * 1024 * 1024) + reference = weakref.ref(payload) + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.2) + ) == { + "success": False, + "error": "timeout", + } + del payload + + assert server.finished.wait(1) + assert_released(resources, [reference]) + if not resources.interrupted_sends: + pytest.skip("Local TCP stack buffered the request without blocking send") + for peer in server.peers: + peer.settimeout(1) + assert peer.recv(1) + + +def test_trickled_reply_does_not_restart_the_timeout(resources, servers): + def reply(peer, _request): + peer.sendall(struct.pack("!i", 100)) + for _ in range(100): + peer.sendall(b"x") + time.sleep(0.01) + + server = servers(1, reply=reply) + communications = module.IPCClient(server.address, b"fixture") + started = time.monotonic() + assert communications.send_message(ping.generate(timeout_in_sec=0.05)) == { + "success": False, + "error": "timeout", + } + assert_released(resources, []) + assert time.monotonic() - started < 0.5 + + +@pytest.mark.parametrize("timeout", [0, -1]) +def test_expired_call_does_not_open_a_socket(resources, timeout): + communications = module.IPCClient(("127.0.0.1", 1), b"fixture") + + assert communications.send_message(ping.generate(timeout_in_sec=timeout)) == { + "success": False, + "error": "timeout", + } + assert not resources.sockets + + +def test_reads_large_reply(resources, servers): + response = "x" * 1024 * 1024 + + def reply(peer, _request): + with connection.Connection(peer.dup().detach()) as conn: + conn.send(response) + + server = servers(1, reply=reply) + communications = module.IPCClient(server.address, b"fixture") + assert communications.send_message( + read_property.generate("block", timeout_in_sec=2) + ) == { + "success": True, + "data": response, + } + assert_released(resources, []) + + +@pytest.mark.parametrize("family", ["AF_INET", "AF_UNIX"]) +def test_timeout_releases_payloads_when_accept_queue_is_full( + resources, tmp_path, family +): + if family == "AF_UNIX" and not hasattr(socket, "AF_UNIX"): + pytest.skip("Unix sockets are unavailable") + with socket.socket(getattr(socket, family)) as listener: + address = ( + ("127.0.0.1", 0) if family == "AF_INET" else str(tmp_path / "ipc.sock") + ) + listener.bind(address) + listener.listen(1) + with full_accept_queue(listener): + communications = module.IPCClient(listener.getsockname(), b"fixture") + payloads = [] + for _ in range(3): + payload = Payload() + payloads.append(weakref.ref(payload)) + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.05) + ) == {"success": False, "error": "timeout"} + del payload + assert_released(resources, payloads) + + +@pytest.mark.skipif(not hasattr(socket, "AF_UNIX"), reason="Unix sockets unavailable") +def test_queued_connection_succeeds_if_backend_recovers_before_timeout( + resources, monkeypatch, tmp_path +): + with socket.socket(socket.AF_UNIX) as listener: + address = str(tmp_path / "ipc.sock") + listener.bind(address) + listener.listen(1) + listener.settimeout(2) + accept_calls = threading.Event() + server_errors = [] + + def retry_sleep(duration): + accept_calls.set() + time.sleep(duration) + + monkeypatch.setattr( + transport, "time", SimpleNamespace(**{**vars(time), "sleep": retry_sleep}) + ) + + def respond(): + try: + if not accept_calls.wait(2): + return + for _ in queued: + peer, _ = listener.accept() + peer.close() + peer, _ = listener.accept() + with peer, connection.Connection(peer.dup().detach()) as conn: + assert conn.recv() == ("PING", ()) + conn.send("Received") + except (OSError, AssertionError) as error: + server_errors.append(repr(error)) + + with full_accept_queue(listener) as queued: + server = threading.Thread(target=respond, daemon=True) + server.start() + try: + communications = module.IPCClient(address, b"fixture") + assert communications.send_message(ping.generate(timeout_in_sec=1)) == { + "success": True, + "data": "Received", + } + assert not server_errors + assert_released(resources, []) + finally: + accept_calls.set() + server.join(3) + + +def test_concurrent_calls_match_out_of_order_responses(): + listener = Listener(("127.0.0.1", 0), family="AF_INET", backlog=2, authkey=None) + first_received = threading.Event() + second_finished = threading.Event() + allow_first_reply = threading.Event() + outcomes = {} + completion_order = [] + properties = SimpleNamespace(first="first response", second="second response") + + def respond(): + with listener.accept() as first_connection: + first_request = first_connection.recv() + first_received.set() + with listener.accept() as second_connection: + second_request = second_connection.recv() + second_connection.send(getattr(properties, second_request[1])) + if allow_first_reply.wait(3): + first_connection.send(getattr(properties, first_request[1])) + + communications = module.IPCClient(listener.address, b"fixture") + + def request(name): + outcomes[name] = communications.send_message( + read_property.generate(name, timeout_in_sec=3) + ) + completion_order.append(name) + if name == "second": + second_finished.set() + + server = threading.Thread(target=respond, daemon=True) + first_caller = threading.Thread(target=request, args=("first",), daemon=True) + second_caller = threading.Thread(target=request, args=("second",), daemon=True) + server.start() + first_caller.start() + try: + assert first_received.wait(1) + second_caller.start() + assert second_finished.wait(1) + assert "first" not in outcomes + allow_first_reply.set() + first_caller.join(2) + second_caller.join(2) + assert completion_order == ["second", "first"] + assert outcomes == { + "first": {"success": True, "data": "first response"}, + "second": {"success": True, "data": "second response"}, + } + finally: + allow_first_reply.set() + first_caller.join(3) + if second_caller.ident is not None: + second_caller.join(3) + server.join(3) + listener.close() + + +def test_async_calls_return_while_delivery_is_stalled_and_share_one_worker( + sender, monkeypatch +): + entered = threading.Event() + resume = threading.Event() + worker_ids = [] + + def deliver(*args, **kwargs): + worker_ids.append(threading.get_ident()) + entered.set() + assert resume.wait(5) + + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + monkeypatch.setattr(ipc, "send_message", Mock(side_effect=deliver)) + try: + assert ipc.send_message_async(put_event.generate(0)) == { + "success": True, + "queued": True, + } + assert entered.wait(5) + first_worker = sender._thread + other = module.IPCClient(("127.0.0.1", 2), b"other") + assert module.IPCClient.get() is other + monkeypatch.setattr(other, "send_message", Mock(side_effect=deliver)) + for index in range(1, 20): + assert other.send_message_async(put_event.generate(index))["success"] + assert ipc.send_message.call_count == 1 + other.send_message.assert_not_called() + assert sender._thread is first_worker + finally: + resume.set() + assert_drained(sender) + assert worker_ids == [first_worker.ident] * 20 + calls = ipc.send_message.call_args_list + other.send_message.call_args_list + assert all(call.args[0].identifier == "PUT_EVENT" for call in calls) + assert [call.args[0].request.event for call in calls] == list(range(20)) + assert all(call.args[0].timeout_in_sec == 5 for call in calls) + + +def test_queue_full_drops_new_events_without_waiting(sender, caplog): + class Payload: + pass + + entered = threading.Event() + resume = threading.Event() + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + + def deliver(*args, **kwargs): + entered.set() + assert resume.wait(5) + + ipc.send_message = Mock(side_effect=deliver) + try: + ipc.send_message_async(put_event.generate(0)) + assert entered.wait(5) + for index in range(module.MAX_PENDING_EVENTS): + assert ipc.send_message_async(put_event.generate(index + 1))["success"] + for _ in range(3): + assert ipc.send_message_async(put_event.generate("overflow")) == { + "success": False, + "error": "queue_full", + } + dropped = Payload() + dropped_reference = weakref.ref(dropped) + assert not ipc.send_message_async(put_event.generate(dropped))["success"] + del dropped + gc.collect() + assert dropped_reference() is None + assert sender._queue.qsize() == module.MAX_PENDING_EVENTS + assert caplog.text.count("Event delivery queue is full") == 1 + finally: + resume.set() + assert_drained(sender) + assert ipc.send_message.call_count == module.MAX_PENDING_EVENTS + 1 + + +@pytest.mark.parametrize("fail", [False, True]) +def test_idle_worker_releases_payloads_after_delivery_or_failure(sender, fail): + class Payload: + pass + + delivered = [] + + class Comms(module.IPCClient): + def send_message(self, message): + delivered.append( + ( + weakref.ref(message), + weakref.ref(message.request), + weakref.ref(message.request.event), + ) + ) + if fail: + raise OSError("closed socket") + + ipc = Comms(("127.0.0.1", 1), b"fixture") + for _ in range(20): + ipc.send_message_async(put_event.generate(Payload())) + ipc_reference = weakref.ref(ipc) + module.IPCClient.reset() + del ipc + assert_drained(sender) + gc.collect() + assert len(delivered) == 20 + assert all(reference() is None for request in delivered for reference in request) + assert ipc_reference() is None + assert sender._thread.is_alive() + + +def test_invalid_key_does_not_start_the_event_sender(sender, monkeypatch): + start = Mock() + monkeypatch.setattr(module.threading.Thread, "start", start) + ipc = module.IPCClient(("127.0.0.1", 1), None) + assert ipc.send_message_async(put_event.generate(None)) == { + "success": False, + "error": "invalid_key", + } + start.assert_not_called() + assert sender._queue.empty() + + +def test_worker_start_failure_does_not_escape_into_the_application(sender, monkeypatch): + monkeypatch.setattr( + module.threading.Thread, "start", Mock(side_effect=RuntimeError("no threads")) + ) + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + assert ipc.send_message_async(put_event.generate(None)) == { + "success": False, + "error": "unknown", + } + assert sender._thread is None + + +def test_pid_change_replaces_inherited_locks_and_drops_inherited_events(sender): + old_queue = sender._queue + old_queue.put((None, put_event.generate("parent event"))) + old_lock = sender._start_locks[os.getpid()] + old_lock.acquire() + sender._pid = os.getpid() - 1 + sender._start_locks = {sender._pid: old_lock} + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + ipc.send_message = Mock() + message = put_event.generate("child event") + try: + assert ipc.send_message_async(message)["success"] + finally: + old_lock.release() + assert_drained(sender) + assert sender._queue is not old_queue + assert sender._start_locks[os.getpid()] is not old_lock + ipc.send_message.assert_called_once_with(message) + + +def test_concurrent_events_after_a_native_fork_reset_the_sender_only_once( + sender, monkeypatch +): + parent_pid = os.getpid() - 1 + sender._pid = parent_pid + sender._start_locks = {parent_pid: threading.Lock()} + reset_entered = threading.Event() + resume_reset = threading.Event() + queue_factory = module.Queue + queues = [] + + def create_queue(*args, **kwargs): + queues.append(queue_factory(*args, **kwargs)) + reset_entered.set() + assert resume_reset.wait(5) + return queues[-1] + + monkeypatch.setattr(module, "Queue", create_queue) + barrier = threading.Barrier(21) + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + ipc.send_message = Mock() + results = [] + + def enqueue(index): + barrier.wait(5) + results.append(ipc.send_message_async(put_event.generate(index))) + + producers = [threading.Thread(target=enqueue, args=(index,)) for index in range(20)] + try: + for producer in producers: + producer.start() + barrier.wait(5) + assert reset_entered.wait(5) + finally: + resume_reset.set() + for producer in producers: + producer.join(5) + assert not producer.is_alive() + assert_drained(sender) + assert len(queues) == 1 + assert len(results) == 20 + assert all(result["success"] for result in results) + assert ipc.send_message.call_count == 20 + + +def test_concurrent_producers_start_only_one_worker(sender): + barrier = threading.Barrier(21) + workers = set() + ipc = module.IPCClient(("127.0.0.1", 1), b"fixture") + ipc.send_message = Mock( + side_effect=lambda *args, **kwargs: workers.add(threading.get_ident()) + ) + + def enqueue(): + barrier.wait(5) + ipc.send_message_async(put_event.generate(None)) + + producers = [threading.Thread(target=enqueue) for _ in range(20)] + for producer in producers: + producer.start() + barrier.wait(5) + for producer in producers: + producer.join(5) + assert not producer.is_alive() + assert_drained(sender) + assert workers == {sender._thread.ident} + assert ipc.send_message.call_count == 20 diff --git a/aikido_zen/background_process/packages_test.py b/aikido_zen/background_process/packages_test.py index 66e3419a4..91cf27749 100644 --- a/aikido_zen/background_process/packages_test.py +++ b/aikido_zen/background_process/packages_test.py @@ -4,54 +4,51 @@ @pytest.fixture -def mock_get_comms(mocker): - """Fixture to mock the get_comms function.""" +def mock_get_ipc(mocker): + """Fixture to mock the get_ipc function.""" class FakeComms: def __init__(self): self.action = None self.obj = None - self.receive = False - def send_data_to_bg_process(self, action, obj, receive): - self.action = action - self.obj = obj - self.receive = receive + def send_message(self, message): + self.action = message.identifier + self.obj = message.request return {"success": True, "data": True} return FakeComms -def test_pkg_compat_check_success(mock_get_comms, mocker): +def test_pkg_compat_check_success(mock_get_ipc, mocker): """Test successful wrapping of a package.""" pkg_name = "example_package" pkg_version = "1.0.0" - mock_class = mock_get_comms() + mock_class = mock_get_ipc() # Mock the metadata.version function mocker.patch("importlib.metadata.version", return_value=pkg_version) mocker.patch( - "aikido_zen.background_process.comms.get_comms", return_value=mock_class + "aikido_zen.background_process.ipc.IPCClient.get", return_value=mock_class ) # Call the function under test is_package_compatible(pkg_name, ANY_VERSION) -def test_pkg_compat_check_retry(mock_get_comms, mocker): +def test_pkg_compat_check_retry(mock_get_ipc, mocker): """Test retry logic when sending data fails.""" pkg_name = "example_package" pkg_version = "1.0.0" - mock_class = mock_get_comms() # Create an instance of FakeComms + mock_class = mock_get_ipc() # Create an instance of FakeComms # Mock the metadata.version function mocker.patch("importlib.metadata.version", return_value=pkg_version) - # Mock the get_comms function to return the mock_class instance + # Mock the get_ipc function to return the mock_class instance mocker.patch( - "aikido_zen.background_process.comms.get_comms", return_value=mock_class + "aikido_zen.background_process.ipc.IPCClient.get", return_value=mock_class ) - # Mock the send_data_to_bg_process method to simulate failure - mock_class.send_data_to_bg_process = MagicMock( + mock_class.send_message = MagicMock( return_value={ "success": False, "data": False, @@ -62,21 +59,20 @@ def test_pkg_compat_check_retry(mock_get_comms, mocker): is_package_compatible(pkg_name, ANY_VERSION) -def test_pkg_compat_check_partial_success(mock_get_comms, mocker): +def test_pkg_compat_check_partial_success(mock_get_ipc, mocker): """Test that the function stops retrying after a successful attempt.""" pkg_name = "example_package" pkg_version = "1.0.0" - mock_class = mock_get_comms() # Create an instance of FakeComms + mock_class = mock_get_ipc() # Create an instance of FakeComms # Mock the metadata.version function mocker.patch("importlib.metadata.version", return_value=pkg_version) - # Mock the get_comms function to return the mock_class instance + # Mock the get_ipc function to return the mock_class instance mocker.patch( - "aikido_zen.background_process.comms.get_comms", return_value=mock_class + "aikido_zen.background_process.ipc.IPCClient.get", return_value=mock_class ) - # Mock the send_data_to_bg_process method to simulate a failure followed by success - mock_class.send_data_to_bg_process = MagicMock( + mock_class.send_message = MagicMock( side_effect=[ {"success": False, "data": False}, # First attempt fails {"success": True, "data": True}, # Second attempt succeeds diff --git a/aikido_zen/background_process/socket_connection.py b/aikido_zen/background_process/socket_connection.py new file mode 100644 index 000000000..79515966b --- /dev/null +++ b/aikido_zen/background_process/socket_connection.py @@ -0,0 +1,61 @@ +import io +import multiprocessing.connection as con +import socket +import time +from contextlib import contextmanager + + +class SocketConnection(con.Connection): + # Keep Connection's framing and pickling while applying deadlines to partial I/O. + # pylint: disable=arguments-differ + def __init__(self, sock, deadline): + self._socket = sock + self._deadline = deadline + super().__init__(sock.fileno()) + + def _set_timeout(self): + remaining = self._deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("IPC operation timed out") + self._socket.settimeout(remaining) + + def _send(self, buf): + self._set_timeout() + self._socket.sendall(buf) + + def _recv(self, size): + buffer = io.BytesIO() + remaining = size + while remaining > 0: + self._set_timeout() + chunk = self._socket.recv(min(remaining, 65536)) + if not chunk: + raise EOFError + buffer.write(chunk) + remaining -= len(chunk) + return buffer + + def _close(self): + self._socket.close() + + +def connect_with_deadline(sock, address, deadline): + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("IPC connection timed out") + sock.settimeout(remaining) + try: + sock.connect(address) + sock.setblocking(True) + return + except BlockingIOError: + time.sleep(min(0.001, remaining)) + + +@contextmanager +def open_connection(address, deadline): + with socket.socket(getattr(socket, con.address_type(address))) as sock: + connect_with_deadline(sock, address, deadline) + with SocketConnection(sock, deadline) as connection: + yield connection diff --git a/aikido_zen/background_process/test_uds_file_access.py b/aikido_zen/background_process/test_uds_file_access.py deleted file mode 100644 index 0c614cda9..000000000 --- a/aikido_zen/background_process/test_uds_file_access.py +++ /dev/null @@ -1,27 +0,0 @@ -import os -import platform -import sys - -from aikido_zen.background_process import get_uds_filename -from aikido_zen.helpers.logging import logger - - -def test_uds_file_access(): - """Returns true if access is possible, false if it's not.""" - filename = get_uds_filename() - directory = "/".join(filename.split("/")[:-1]) - - if platform.system() == "Windows": - # Python does not support Windows UDS just yet, so we have to rely on INET - return True - - if not os.access(directory, os.R_OK) or not os.access(directory, os.W_OK): - # We need both read and write permissions in order to start the socket - logger.critical( - "Cannot start Zen: no read/write permissions for directory %s. " - + "You can change the path by setting " - + "the AIKIDO_TMP_DIR environment variable.", - directory, - ) - return False - return True diff --git a/aikido_zen/context/track.py b/aikido_zen/context/track.py index de729718a..444da9257 100644 --- a/aikido_zen/context/track.py +++ b/aikido_zen/context/track.py @@ -2,9 +2,8 @@ from aikido_zen.helpers.create_custom_event import create_custom_event from . import get_current_context import aikido_zen.thread.thread_cache as thread_cache -import aikido_zen.background_process.comms as comms -from ..background_process.commands import PutEventCommand -from ..helpers.ipc.send_payload import send_payload +from aikido_zen.background_process.ipc import IPCClient +from ..background_process.commands import put_event MAX_EVENTS_PER_REQUEST = 25 @@ -21,7 +20,7 @@ def track(event_name): logger.info("track(...) expects a non-empty string as event name.") return - ipc = comms.get_comms() + ipc = IPCClient.get() if not ipc: return @@ -43,7 +42,7 @@ def track(event_name): if not event: return - send_payload(ipc, PutEventCommand.generate(event), (50 / 1000)) + ipc.send_message_async(put_event.generate(event)) except Exception as e: logger.debug("Exception occurred in track: %s", e) diff --git a/aikido_zen/context/track_test.py b/aikido_zen/context/track_test.py index 7134f6dba..508c0601c 100644 --- a/aikido_zen/context/track_test.py +++ b/aikido_zen/context/track_test.py @@ -40,15 +40,15 @@ def cache_with_bypassed_ips(*ips): def sent_event_with_cache(cache): set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch("aikido_zen.thread.thread_cache.get_cache", return_value=cache), patch( - "aikido_zen.background_process.comms.get_comms", return_value=comms + "aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc ): track("my-custom-event") - if not comms.send_data_to_bg_process.called: + if not ipc.send_message_async.called: return None - return comms.send_data_to_bg_process.call_args[0][1].event + return ipc.send_message_async.call_args.args[0].request.event def test_track_invalid_event_name(caplog): @@ -63,7 +63,7 @@ def test_track_empty_event_name(caplog): def test_track_without_context(caplog): with patch( - "aikido_zen.background_process.comms.get_comms", return_value=MagicMock() + "aikido_zen.background_process.ipc.IPCClient.get", return_value=MagicMock() ): track("my-event") @@ -72,7 +72,7 @@ def test_track_without_context(caplog): def test_track_without_context_only_logs_once(caplog): with patch( - "aikido_zen.background_process.comms.get_comms", return_value=MagicMock() + "aikido_zen.background_process.ipc.IPCClient.get", return_value=MagicMock() ): track("my-event") track("my-event") @@ -80,8 +80,8 @@ def test_track_without_context_only_logs_once(caplog): assert caplog.text.count("track(...) was called without a context.") == 1 -def test_track_without_comms(caplog): - with patch("aikido_zen.background_process.comms.get_comms", return_value=None): +def test_track_without_ipc(caplog): + with patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=None): track("my-event") assert "track(...) was called without a context." not in caplog.text @@ -91,23 +91,22 @@ def test_track_sends_event_over_ipc(): context = set_context_and_lifecycle() context.user = {"id": "user-1", "name": "Jane Doe"} - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), ), patch( - "aikido_zen.background_process.comms.get_comms", return_value=comms + "aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc ), patch( "aikido_zen.helpers.create_custom_event.get_unixtime_ms", return_value=111, ): track("my-custom-event") - comms.send_data_to_bg_process.assert_called_once() - identifier, request, returns_data = comms.send_data_to_bg_process.call_args[0] - assert identifier == "put_event" - assert returns_data is False - assert request.event == { + ipc.send_message_async.assert_called_once() + (message,) = ipc.send_message_async.call_args.args + assert message.identifier == "PUT_EVENT" + assert message.request.event == { "type": "custom", "time": 111, "name": "my-custom-event", @@ -165,15 +164,15 @@ def test_track_does_not_send_the_url(): def track_n_times(n, event_name="my-custom-event"): """Runs track() n times in ONE request and returns the events that were sent""" set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), - ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): + ), patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc): for _ in range(n): track(event_name) - return [c[0][1].event for c in comms.send_data_to_bg_process.call_args_list] + return [c.args[0].request.event for c in ipc.send_message_async.call_args_list] def test_track_sends_every_event_below_the_limit(): @@ -190,17 +189,17 @@ def test_track_stops_sending_events_above_the_limit(): def test_all_event_names_share_one_budget(): set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), - ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): + ), patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc): for _ in range(track_module.MAX_EVENTS_PER_REQUEST): track("payment.failed") track("order.completed") names = [ - c[0][1].event["name"] for c in comms.send_data_to_bg_process.call_args_list + c.args[0].request.event["name"] for c in ipc.send_message_async.call_args_list ] assert len(names) == track_module.MAX_EVENTS_PER_REQUEST assert "order.completed" not in names @@ -235,43 +234,43 @@ def test_the_limit_is_per_request_not_per_process(): def test_a_bypassed_ip_does_not_use_up_the_limit(): set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips("1.2.3.4"), - ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): + ), patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc): for _ in range(track_module.MAX_EVENTS_PER_REQUEST + 5): track("my-custom-event") assert get_current_context().tracked_events == 0 -def test_track_sends_the_event_with_a_short_timeout(): +def test_track_does_not_wait_for_ipc_delivery(): set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), - ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): + ), patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc): track("my-custom-event") - comms.send_data_to_bg_process.assert_called_once() - assert comms.send_data_to_bg_process.call_args[1]["timeout_in_sec"] == 0.05 + ipc.send_message_async.assert_called_once() + ipc.send_message.assert_not_called() def test_track_sends_nothing_when_the_event_could_not_be_built(): set_context_and_lifecycle() - comms = MagicMock() + ipc = MagicMock() with patch( "aikido_zen.thread.thread_cache.get_cache", return_value=cache_with_bypassed_ips(), ), patch( - "aikido_zen.background_process.comms.get_comms", return_value=comms + "aikido_zen.background_process.ipc.IPCClient.get", return_value=ipc ), patch( "aikido_zen.context.track.create_custom_event", return_value=None ): track("my-custom-event") - comms.send_data_to_bg_process.assert_not_called() + ipc.send_message_async.assert_not_called() diff --git a/aikido_zen/decorators/uwsgi_test.py b/aikido_zen/decorators/uwsgi_test.py index 9ef24c6d5..a229e15f0 100644 --- a/aikido_zen/decorators/uwsgi_test.py +++ b/aikido_zen/decorators/uwsgi_test.py @@ -8,7 +8,7 @@ def test_postfork_instruments_imported_functions_before_starting_daemon(monkeypa import threading from unittest.mock import patch -from aikido_zen.background_process import get_comms +from aikido_zen.background_process.ipc import IPCClient from aikido_zen.decorators.uwsgi import postfork calls = [] @@ -22,7 +22,7 @@ def start_aikido(): from os import system assert calls == [] - assert get_comms() is None + assert IPCClient.get() is None assert threading.enumerate() == threads start_aikido() diff --git a/aikido_zen/helpers/blocking_enabled.py b/aikido_zen/helpers/blocking_enabled.py index cb69485d9..071779513 100644 --- a/aikido_zen/helpers/blocking_enabled.py +++ b/aikido_zen/helpers/blocking_enabled.py @@ -1,6 +1,7 @@ """Helper function file, see function docstring""" -from aikido_zen.background_process import get_comms +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.commands import read_property from aikido_zen.helpers.check_env_for_blocking import check_env_for_blocking @@ -8,11 +9,9 @@ def is_blocking_enabled(): """ Checks with the background process if blocking is enabled """ - if get_comms(): - # Only check if comms are defined - should_block_res = get_comms().send_data_to_bg_process( - action="READ_PROPERTY", obj="block", receive=True - ) + ipc = IPCClient.get() + if ipc: + should_block_res = ipc.send_message(read_property.generate("block")) if should_block_res["success"]: return should_block_res["data"] return check_env_for_blocking() diff --git a/aikido_zen/helpers/ipc/command_types.py b/aikido_zen/helpers/ipc/command_types.py index dbbace9f3..bcc898303 100644 --- a/aikido_zen/helpers/ipc/command_types.py +++ b/aikido_zen/helpers/ipc/command_types.py @@ -1,40 +1,5 @@ -from abc import ABC, abstractmethod - - -class CommandContext: - def __init__(self, connection_manager, queue, connection): - self.connection_manager = connection_manager - self.queue = queue - self.connection = connection - - def send(self, response): - self.connection.send(response) - - class Payload: - def __init__(self, command, request): - self.identifier = command.identifier() - self.returns_data = command.returns_data() + def __init__(self, identifier, request, timeout_in_sec): + self.identifier = identifier self.request = request - - -class Command(ABC): - @classmethod - @abstractmethod - def identifier(cls) -> str: - pass - - @classmethod - @abstractmethod - def returns_data(cls) -> bool: - pass - - @classmethod - @abstractmethod - def run(cls, context: CommandContext, request): - pass - - @classmethod - @abstractmethod - def generate(cls, request) -> Payload: - pass + self.timeout_in_sec = timeout_in_sec diff --git a/aikido_zen/helpers/ipc/send_payload.py b/aikido_zen/helpers/ipc/send_payload.py deleted file mode 100644 index 9b5346ec3..000000000 --- a/aikido_zen/helpers/ipc/send_payload.py +++ /dev/null @@ -1,13 +0,0 @@ -from aikido_zen.background_process import AikidoIPCCommunications -from aikido_zen.helpers.ipc.command_types import Payload - - -def send_payload( - comms: AikidoIPCCommunications, payload: Payload, timeout_in_sec: float = 0.1 -): - return comms.send_data_to_bg_process( - payload.identifier, - payload.request, - payload.returns_data, - timeout_in_sec=timeout_in_sec, - ) diff --git a/aikido_zen/init_test.py b/aikido_zen/init_test.py index 15087a6d9..78f5977ca 100644 --- a/aikido_zen/init_test.py +++ b/aikido_zen/init_test.py @@ -4,7 +4,7 @@ import aikido_zen from aikido_zen import protect -from aikido_zen.background_process import get_comms, reset_comms +from aikido_zen.background_process.ipc import IPCClient from aikido_zen.helpers.token import get_token_from_env @@ -12,8 +12,8 @@ def test_protect_with_django(monkeypatch, caplog): protect() assert "starting" in caplog.text - reset_comms() - assert get_comms() == None + IPCClient.reset() + assert IPCClient.get() == None def test_protect_sets_token(): @@ -24,10 +24,10 @@ def test_protect_sets_token(): def test_protect_does_not_start_without_gil(monkeypatch): monkeypatch.setattr("aikido_zen.gil_not_enabled", lambda: True) - with patch("aikido_zen.test_uds_file_access") as test_uds_file_access, patch( + with patch("aikido_zen.check_socket_access") as check_socket_access, patch( "aikido_zen.start_background_process" ) as start_background_process: protect() - test_uds_file_access.assert_not_called() + check_socket_access.assert_not_called() start_background_process.assert_not_called() diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index 4a7fad1f9..4d78e245b 100644 --- a/aikido_zen/middleware/init_test.py +++ b/aikido_zen/middleware/init_test.py @@ -64,7 +64,7 @@ def test_bypassed_ip_skips_user_blocking(): assert should_block_request() == {"block": False} -def test_cache_comms_with_endpoints(): +def test_cache_ipc_with_endpoints(): test_utils.generate_and_set_context(user={"id": "456"}, route="/posts/:id") set_rate_limit_group("my_group") thread_cache = get_cache() @@ -85,19 +85,19 @@ def test_cache_comms_with_endpoints(): assert get_current_context().executed_middleware == False assert thread_cache.middleware_installed == False - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - mock_get_comms.return_value = None # Set the return value of get_comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + mock_get_ipc.return_value = None # Set the return value of get_ipc assert should_block_request() == {"block": False} assert get_current_context().executed_middleware == True assert thread_cache.middleware_installed == True - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms # Set the return value of get_comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Set the return value of get_ipc # No matching endpoints : assert should_block_request() == {"block": False} - mock_comms.send_data_to_bg_process.assert_not_called() + mock_ipc.send_message.assert_not_called() thread_cache.config.endpoints.append( { @@ -112,46 +112,45 @@ def test_cache_comms_with_endpoints(): } ) - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms # Set the return value of get_comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Set the return value of get_ipc # Rate-limiting disabled: assert should_block_request() == {"block": False} - mock_comms.send_data_to_bg_process.assert_not_called() + mock_ipc.send_message.assert_not_called() # Enable ratelimiting thread_cache.config.endpoints[1]["rateLimiting"]["enabled"] = True - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms # Set the return value of get_comms - mock_comms.send_data_to_bg_process.return_value = {"success": False} + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Set the return value of get_ipc + mock_ipc.send_message.return_value = {"success": False} assert should_block_request() == {"block": False} - mock_comms.send_data_to_bg_process.assert_called_with( - action="SHOULD_RATELIMIT", - obj={ - "route_metadata": { - "method": "POST", - "route": "/posts/:id", - "url": "http://localhost:8080/", - }, - "user": {"id": "456"}, - "group": "my_group", - "remote_address": "1.1.1.1", + mock_ipc.send_message.assert_called() + (message,) = mock_ipc.send_message.call_args.args + assert message.identifier == "SHOULD_RATELIMIT" + assert message.timeout_in_sec == 0.01 + assert message.request == { + "route_metadata": { + "method": "POST", + "route": "/posts/:id", + "url": "http://localhost:8080/", }, - receive=True, - timeout_in_sec=0.01, - ) + "user": {"id": "456"}, + "group": "my_group", + "remote_address": "1.1.1.1", + } - mock_comms.send_data_to_bg_process.return_value = { + mock_ipc.send_message.return_value = { "success": True, "data": {"block": False, "trigger": "my_trigger"}, } assert should_block_request() == {"block": False} - mock_comms.send_data_to_bg_process.return_value = { + mock_ipc.send_message.return_value = { "success": True, "data": {"block": True, "trigger": "my_trigger"}, } diff --git a/aikido_zen/middleware/should_block_request.py b/aikido_zen/middleware/should_block_request.py index 434f06a67..6743cca6d 100644 --- a/aikido_zen/middleware/should_block_request.py +++ b/aikido_zen/middleware/should_block_request.py @@ -3,7 +3,8 @@ from aikido_zen.helpers.logging import logger from aikido_zen.context import get_current_context from aikido_zen.thread.thread_cache import get_cache -import aikido_zen.background_process.comms as c +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.commands import should_ratelimit from aikido_zen.ratelimiting.get_ratelimited_endpoint import get_ratelimited_endpoint from aikido_zen.helpers.match_endpoints import match_endpoints @@ -42,8 +43,8 @@ def should_block_request(): route_metadata = context.get_route_metadata() endpoints = cache.get_endpoints() - comms = c.get_comms() - if not comms or not endpoints: + ipc = IPCClient.get() + if not ipc or not endpoints: return {"block": False} matched_endpoints = match_endpoints(route_metadata, endpoints) # Ratelimiting : @@ -51,16 +52,13 @@ def should_block_request(): matched_endpoints, context.route ): # As an optimization check if the route is rate limited before sending over IPC - ratelimit_res = comms.send_data_to_bg_process( - action="SHOULD_RATELIMIT", - obj={ - "route_metadata": route_metadata, - "user": context.user, - "group": context.rate_limit_group, - "remote_address": context.remote_address, - }, - receive=True, - timeout_in_sec=(10 / 1000), # 10ms + ratelimit_res = ipc.send_message( + should_ratelimit.generate( + route_metadata=route_metadata, + user=context.user, + group=context.rate_limit_group, + remote_address=context.remote_address, + ), ) if ratelimit_res["success"] and ratelimit_res["data"]["block"]: cache.stats.on_rate_limit() diff --git a/aikido_zen/sinks/tests/asyncpg_test.py b/aikido_zen/sinks/tests/asyncpg_test.py index 706a665ac..c6a429946 100644 --- a/aikido_zen/sinks/tests/asyncpg_test.py +++ b/aikido_zen/sinks/tests/asyncpg_test.py @@ -2,7 +2,7 @@ from unittest.mock import patch import aikido_zen.sinks.asyncpg import aikido_zen.sinks.os -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient kind = "sql_injection" op = "MySQLdb.connections.query" @@ -20,7 +20,7 @@ async def database_conn(): @pytest.mark.asyncio async def test_conn_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -41,7 +41,7 @@ async def test_conn_execute(database_conn): @pytest.mark.asyncio async def test_conn_fetchrow(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -58,7 +58,7 @@ async def test_conn_fetchrow(database_conn): @pytest.mark.asyncio async def test_conn_fetch(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -75,7 +75,7 @@ async def test_conn_fetch(database_conn): @pytest.mark.asyncio async def test_conn_fetchval(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -92,7 +92,7 @@ async def test_conn_fetchval(database_conn): @pytest.mark.asyncio async def test_conn_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -113,7 +113,7 @@ async def test_conn_execute_parameterized(database_conn): @pytest.mark.asyncio async def test_conn_transaction(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -138,7 +138,7 @@ async def test_conn_transaction(database_conn): @pytest.mark.asyncio async def test_conn_cursor(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/clickhouse_driver_test.py b/aikido_zen/sinks/tests/clickhouse_driver_test.py index 45794daa0..b31a41870 100644 --- a/aikido_zen/sinks/tests/clickhouse_driver_test.py +++ b/aikido_zen/sinks/tests/clickhouse_driver_test.py @@ -1,13 +1,13 @@ import aikido_zen.sinks.clickhouse_driver import pytest -from aikido_zen.background_process import reset_comms +from aikido_zen.background_process.ipc import IPCClient from aikido_zen.errors import AikidoSQLInjection import aikido_zen.test_utils as test_utils @pytest.fixture(autouse=True) def setup(monkeypatch): - reset_comms() + IPCClient.reset() monkeypatch.setenv("AIKIDO_BLOCK", "1") diff --git a/aikido_zen/sinks/tests/motor_test.py b/aikido_zen/sinks/tests/motor_test.py index a6a880b59..268a4f1a0 100644 --- a/aikido_zen/sinks/tests/motor_test.py +++ b/aikido_zen/sinks/tests/motor_test.py @@ -2,7 +2,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.pymongo -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient import pymongo.errors as mongo_errs @@ -16,7 +16,7 @@ def db(): @pytest.mark.asyncio async def test_replace_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -34,7 +34,7 @@ async def test_replace_one(db): @pytest.mark.asyncio async def test_update_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -50,7 +50,7 @@ async def test_update_one(db): @pytest.mark.asyncio async def test_update_many(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -66,7 +66,7 @@ async def test_update_many(db): @pytest.mark.asyncio async def test_delete_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -82,7 +82,7 @@ async def test_delete_one(db): @pytest.mark.asyncio async def test_delete_many(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -98,7 +98,7 @@ async def test_delete_many(db): @pytest.mark.asyncio async def test_find_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -115,7 +115,7 @@ async def test_find_one(db): @pytest.mark.asyncio async def test_count_documents(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -132,7 +132,7 @@ async def test_count_documents(db): @pytest.mark.asyncio async def test_find_one_and_delete(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -149,7 +149,7 @@ async def test_find_one_and_delete(db): @pytest.mark.asyncio async def test_find_one_and_replace(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -168,7 +168,7 @@ async def test_find_one_and_replace(db): @pytest.mark.asyncio async def test_find_one_and_update(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -185,7 +185,7 @@ async def test_find_one_and_update(db): @pytest.mark.asyncio async def test_find_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -196,7 +196,7 @@ async def test_find_empty(db): @pytest.mark.asyncio async def test_find_not_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -213,7 +213,7 @@ async def test_find_not_empty(db): @pytest.mark.asyncio async def test_find_raw_batches(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -230,7 +230,7 @@ async def test_find_raw_batches(db): @pytest.mark.asyncio async def test_distinct(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -247,7 +247,7 @@ async def test_distinct(db): @pytest.mark.asyncio async def test_distinct_kwargs(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -264,7 +264,7 @@ async def test_distinct_kwargs(db): @pytest.mark.asyncio async def test_distinct_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/mysqlclient_test.py b/aikido_zen/sinks/tests/mysqlclient_test.py index 74ded55e0..c5d7769fc 100644 --- a/aikido_zen/sinks/tests/mysqlclient_test.py +++ b/aikido_zen/sinks/tests/mysqlclient_test.py @@ -1,7 +1,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.mysqlclient -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient kind = "sql_injection" op = "MySQLdb.connections.query" @@ -15,7 +15,7 @@ def database_conn(): def test_cursor_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -36,7 +36,7 @@ def test_cursor_execute(database_conn): def test_cursor_execute_no_args(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -62,7 +62,7 @@ def test_cursor_execute_no_args(database_conn): def test_cursor_executemany(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -89,7 +89,7 @@ def test_cursor_executemany(database_conn): def test_cursor_execute_with_fstring(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/psycopg2_test.py b/aikido_zen/sinks/tests/psycopg2_test.py index f53844876..e8c50a69d 100644 --- a/aikido_zen/sinks/tests/psycopg2_test.py +++ b/aikido_zen/sinks/tests/psycopg2_test.py @@ -3,7 +3,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.psycopg2 -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient kind = "sql_injection" op = "pymysql.connections.query" @@ -54,7 +54,7 @@ def database_conn_empty_cursor(): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -75,7 +75,7 @@ def test_cursor_execute(database_conn): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_execute2(database_conn_empty_cursor): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -95,7 +95,7 @@ def test_cursor_execute2(database_conn_empty_cursor): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_execute3(database_conn_immutable_cursor): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -115,7 +115,7 @@ def test_cursor_execute3(database_conn_immutable_cursor): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -132,7 +132,7 @@ def test_cursor_execute_parameterized(database_conn): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_execute_parameterized2(database_conn_empty_cursor): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -152,7 +152,7 @@ def test_cursor_execute_parameterized2(database_conn_empty_cursor): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_executemany(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -177,7 +177,7 @@ def test_cursor_executemany(database_conn): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_executemany2(database_conn_empty_cursor): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -204,7 +204,7 @@ def test_cursor_executemany2(database_conn_empty_cursor): @pytest.mark.skipif(skip_python_3_13, reason="Skipping test on Python 3.13") def test_cursor_executemany3(database_conn_immutable_cursor): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/psycopg_test.py b/aikido_zen/sinks/tests/psycopg_test.py index e9a8a89f3..ae5847430 100644 --- a/aikido_zen/sinks/tests/psycopg_test.py +++ b/aikido_zen/sinks/tests/psycopg_test.py @@ -4,7 +4,7 @@ import pytest_asyncio import aikido_zen.sinks.psycopg -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient @pytest.fixture @@ -17,7 +17,7 @@ def database_conn(): def test_cursor_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -38,7 +38,7 @@ def test_cursor_execute(database_conn): def test_cursor_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() cursor = database_conn.cursor() query = "SELECT * FROM dogs WHERE dog_name = %s" params = ("Fido",) @@ -59,7 +59,7 @@ def test_cursor_execute_parameterized(database_conn): def test_cursor_executemany(database_conn): - reset_comms() + IPCClient.reset() cursor = database_conn.cursor() query = "INSERT INTO dogs (dog_name, isadmin) VALUES (%s, %s)" params = [("doggo1", False), ("Rex", False), ("Buddy", True)] @@ -85,7 +85,7 @@ def test_cursor_executemany(database_conn): def test_cursor_copy(database_conn): - reset_comms() + IPCClient.reset() cursor = database_conn.cursor() query = "COPY dogs FROM STDIN" @@ -108,7 +108,7 @@ def test_cursor_copy(database_conn): async def test_async_cursor_execute(): import psycopg - reset_comms() + IPCClient.reset() async with await psycopg.AsyncConnection.connect( host="127.0.0.1", user="user", password="password", dbname="db" ) as conn: @@ -134,7 +134,7 @@ async def test_async_cursor_execute(): async def test_async_cursor_execute_parameterized(): import psycopg - reset_comms() + IPCClient.reset() query = "SELECT * FROM dogs WHERE dog_name = %s" params = ("Fido",) @@ -163,7 +163,7 @@ async def test_async_cursor_execute_parameterized(): async def test_async_cursor_executemany(): import psycopg - reset_comms() + IPCClient.reset() query = "INSERT INTO dogs (dog_name, isadmin) VALUES (%s, %s)" params = [("doggo1", False), ("Rex", False), ("Buddy", True)] @@ -194,7 +194,7 @@ async def test_async_cursor_executemany(): async def test_async_cursor_copy(): import psycopg - reset_comms() + IPCClient.reset() query = "COPY dogs FROM STDIN" async with await psycopg.AsyncConnection.connect( diff --git a/aikido_zen/sinks/tests/pymongo_test.py b/aikido_zen/sinks/tests/pymongo_test.py index 1756a1133..2fa6ae8df 100644 --- a/aikido_zen/sinks/tests/pymongo_test.py +++ b/aikido_zen/sinks/tests/pymongo_test.py @@ -1,7 +1,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.pymongo -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient import pymongo.errors as mongo_errs from collections import defaultdict, OrderedDict, ChainMap @@ -15,7 +15,7 @@ def db(): def test_replace_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -30,7 +30,7 @@ def test_replace_one(db): def test_replace_one_with_chainmap(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -46,7 +46,7 @@ def test_replace_one_with_chainmap(db): def test_update_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -61,7 +61,7 @@ def test_update_one(db): def test_update_many(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -76,7 +76,7 @@ def test_update_many(db): def test_delete_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -91,7 +91,7 @@ def test_delete_one(db): def test_delete_many(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -106,7 +106,7 @@ def test_delete_many(db): def test_find_one(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -122,7 +122,7 @@ def test_find_one(db): def test_count_documents(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -138,7 +138,7 @@ def test_count_documents(db): def test_find_one_and_delete(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -154,7 +154,7 @@ def test_find_one_and_delete(db): def test_find_one_and_replace(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -170,7 +170,7 @@ def test_find_one_and_replace(db): def test_find_one_and_update(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -186,7 +186,7 @@ def test_find_one_and_update(db): def test_find_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -196,7 +196,7 @@ def test_find_empty(db): def test_find_not_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -212,7 +212,7 @@ def test_find_not_empty(db): def test_find_raw_batches(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -228,7 +228,7 @@ def test_find_raw_batches(db): def test_distinct(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -244,7 +244,7 @@ def test_distinct(db): def test_distinct_kwargs(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -260,7 +260,7 @@ def test_distinct_kwargs(db): def test_distinct_empty(db): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/pymysql_test.py b/aikido_zen/sinks/tests/pymysql_test.py index 6c8381883..b9aa89a13 100644 --- a/aikido_zen/sinks/tests/pymysql_test.py +++ b/aikido_zen/sinks/tests/pymysql_test.py @@ -1,7 +1,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.pymysql -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient kind = "sql_injection" op = "pymysql.connections.query" @@ -15,7 +15,7 @@ def database_conn(): def test_cursor_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -34,7 +34,7 @@ def test_cursor_execute(database_conn): def test_cursor_execute_context(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -53,7 +53,7 @@ def test_cursor_execute_context(database_conn): def test_cursor_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -75,7 +75,7 @@ def test_cursor_execute_parameterized(database_conn): def test_cursor_executemany(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -102,7 +102,7 @@ def test_cursor_executemany(database_conn): def test_cursor_execute_with_fstring(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -128,7 +128,7 @@ def test_cursor_execute_with_fstring(database_conn): def test_cursor_execute_no_args(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sinks/tests/requests_and_urllib3_test.py b/aikido_zen/sinks/tests/requests_and_urllib3_test.py index 566633b3f..8bf5fe47c 100644 --- a/aikido_zen/sinks/tests/requests_and_urllib3_test.py +++ b/aikido_zen/sinks/tests/requests_and_urllib3_test.py @@ -4,7 +4,7 @@ from aikido_zen.context import Context, current_context from aikido_zen.thread.thread_cache import ThreadCache, get_cache from aikido_zen.errors import AikidoSSRF -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient import aikido_zen.sinks.socket import aikido_zen.sinks.http_client import requests @@ -50,7 +50,7 @@ def set_context_and_lifecycle(url, host=None): def ssrf_check(monkeypatch, url, requests_only=False): - reset_comms() + IPCClient.reset() set_context_and_lifecycle(url) monkeypatch.setenv("AIKIDO_BLOCK", "1") with pytest.raises(AikidoSSRF): @@ -172,7 +172,7 @@ def test_different_capitalization_raises_ssrf(monkeypatch): def test_srrf_with_request_to_itself_urllib3(monkeypatch): http = urllib3.PoolManager() - reset_comms() + IPCClient.reset() set_context_and_lifecycle("http://localhost:5000/test/1", host="localhost:5000") monkeypatch.setenv("AIKIDO_BLOCK", "1") @@ -211,7 +211,7 @@ def test_ssrf_imds_ip(monkeypatch): def test_zero_padded_ip(monkeypatch): monkeypatch.setenv("AIKIDO_BLOCK", "1") - reset_comms() + IPCClient.reset() url = "http://0127.0.0.01:5000" set_context_and_lifecycle(url) diff --git a/aikido_zen/sinks/tests/sqlite3_test.py b/aikido_zen/sinks/tests/sqlite3_test.py index 9ce689efa..ab4643778 100644 --- a/aikido_zen/sinks/tests/sqlite3_test.py +++ b/aikido_zen/sinks/tests/sqlite3_test.py @@ -1,7 +1,7 @@ import pytest from unittest.mock import patch import aikido_zen.sinks.sqlite3 -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient kind = "sql_injection" @@ -19,7 +19,7 @@ def database_conn(): def test_cursor_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -41,7 +41,7 @@ def test_cursor_execute(database_conn): def test_cursor_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -61,7 +61,7 @@ def test_cursor_execute_parameterized(database_conn): def test_cursor_execute_no_args(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -85,7 +85,7 @@ def test_cursor_execute_no_args(database_conn): def test_cursor_execute_with_fstring(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -111,7 +111,7 @@ def test_cursor_execute_with_fstring(database_conn): def test_cursor_executemany(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -137,7 +137,7 @@ def test_cursor_executemany(database_conn): def test_cursor_executescript(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -159,7 +159,7 @@ def test_cursor_executescript(database_conn): def test_connection_execute(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -176,7 +176,7 @@ def test_connection_execute(database_conn): def test_connection_execute_parameterized(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -194,7 +194,7 @@ def test_connection_execute_parameterized(database_conn): def test_connection_executemany(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -220,7 +220,7 @@ def test_connection_executemany(database_conn): def test_connection_executescript(database_conn): - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: @@ -390,7 +390,7 @@ def test_cursor_with_custom_factory(database_conn): class CustomCursor(sqlite3.Cursor): pass - reset_comms() + IPCClient.reset() with patch( "aikido_zen.vulnerabilities.run_vulnerability_scan" ) as mock_run_vulnerability_scan: diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py index 9d25f765a..1aa58a4b0 100644 --- a/aikido_zen/sources/flask/flask_test.py +++ b/aikido_zen/sources/flask/flask_test.py @@ -2,7 +2,7 @@ import signal from unittest.mock import patch import aikido_zen.sources.flask -from aikido_zen.background_process.comms import reset_comms +from aikido_zen.background_process.ipc import IPCClient from aikido_zen.context import current_context, get_current_context sample_environ = { @@ -80,7 +80,7 @@ def test_flask_all_3_func_with_view_args_and_invalid_json_body(): with patch( "aikido_zen.sources.functions.request_handler.request_handler" ) as mock_request_handler: - reset_comms() + IPCClient.reset() current_context.set(None) mock_request_handler.return_value = None @@ -125,7 +125,7 @@ def test_flask_all_3_func_with_malformed_cookie(): with patch( "aikido_zen.sources.functions.request_handler.request_handler" ) as mock_request_handler: - reset_comms() + IPCClient.reset() current_context.set(None) mock_request_handler.return_value = None @@ -156,7 +156,7 @@ def test_flask_all_3_func_with_invalid_body(): with patch( "aikido_zen.sources.functions.request_handler.request_handler" ) as mock_request_handler: - reset_comms() + IPCClient.reset() current_context.set(None) mock_request_handler.return_value = None @@ -192,7 +192,7 @@ def test_flask_all_3_func(): with patch( "aikido_zen.sources.functions.request_handler.request_handler" ) as mock_request_handler: - reset_comms() + IPCClient.reset() current_context.set(None) mock_request_handler.return_value = None @@ -222,7 +222,7 @@ def test_flask_all_3_func(): def test_startup_flask(): - reset_comms() + IPCClient.reset() from flask import Flask app = Flask(__name__) diff --git a/aikido_zen/sources/functions/request_handler.py b/aikido_zen/sources/functions/request_handler.py index 4c56c9cf9..19b6f4b0d 100644 --- a/aikido_zen/sources/functions/request_handler.py +++ b/aikido_zen/sources/functions/request_handler.py @@ -7,9 +7,8 @@ from aikido_zen.helpers.create_attack_wave_event import create_attack_wave_event from aikido_zen.thread.thread_cache import get_cache from .ip_allowed_to_access_route import ip_allowed_to_access_route -import aikido_zen.background_process.comms as c -from ...background_process.commands import PutEventCommand -from ...helpers.ipc.send_payload import send_payload +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.commands import check_firewall_lists, put_event from ...helpers.serialize_to_json import serialize_to_json from ...storage.attack_wave_detector_store import attack_wave_detector_store @@ -64,15 +63,11 @@ def pre_response(): # Do a check on firewall lists, this happens in background because of the heavy data. # For the timeout we notice the request during heavy loads usually takes 2ms - 2.5ms, we set timeout at 10ms. # That way we have a very small timeout with very little risk of not blocking ips. - comms = c.get_comms() - check_fw_lists_res = comms.send_data_to_bg_process( - action="CHECK_FIREWALL_LISTS", - obj={ - "ip": context.remote_address, - "user-agent": context.get_user_agent(), - }, - receive=True, - timeout_in_sec=(10 / 1000), + ipc = IPCClient.get() + check_fw_lists_res = ipc.send_message( + check_firewall_lists.generate( + ip=context.remote_address, user_agent=context.get_user_agent() + ), ) if not check_fw_lists_res["success"] or not check_fw_lists_res["data"]["blocked"]: return @@ -115,9 +110,9 @@ def post_response(status_code): event = create_attack_wave_event(context) logger.debug("Attack wave: %s", serialize_to_json(event)[:5000]) - # Report in background to core (send event over IPC) - if c.get_comms() and event: - send_payload(c.get_comms(), PutEventCommand.generate(event)) + ipc = IPCClient.get() + if ipc and event: + ipc.send_message_async(put_event.generate(event)) # Check if the current route is useful for API discovery is_curr_route_useful = is_useful_route( diff --git a/aikido_zen/sources/functions/request_handler_test.py b/aikido_zen/sources/functions/request_handler_test.py index 45e9239d9..d4e1d77e6 100644 --- a/aikido_zen/sources/functions/request_handler_test.py +++ b/aikido_zen/sources/functions/request_handler_test.py @@ -4,7 +4,7 @@ from unittest.mock import patch, MagicMock from aikido_zen.thread.thread_cache import get_cache, ThreadCache from .request_handler import request_handler -from ...background_process.commands import process_check_firewall_lists +from ...background_process.commands import check_firewall_lists from ...background_process.service_config import ServiceConfig from ...context import Context, current_context from ...helpers.headers import Headers @@ -46,10 +46,10 @@ def __init__(self): self.conn_manager.attack_wave_detector = AttackWaveDetector() self.attacks = [] - def send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1): - if action != "CHECK_FIREWALL_LISTS": + def send_message(self, message): + if message.identifier != "CHECK_FIREWALL_LISTS": return {"success": False} - res = process_check_firewall_lists(self.conn_manager, obj) + res = check_firewall_lists.process(self.conn_manager, message.request, None) return { "success": True, "data": res, @@ -98,11 +98,11 @@ def test_post_response_useful_route(mock_context): } -@patch("aikido_zen.background_process.get_comms") -def test_post_response_not_useful_route(mock_get_comms, mock_context): +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_post_response_not_useful_route(mock_get_ipc, mock_context): """Test post_response when the route is not useful.""" - comms = MagicMock() - mock_get_comms.return_value = comms + ipc = MagicMock() + mock_get_ipc.return_value = ipc cache = ThreadCache() # Creates a new cache assert cache.routes.routes == {} @@ -111,14 +111,14 @@ def test_post_response_not_useful_route(mock_get_comms, mock_context): request_handler("post_response", status_code=500) assert cache.routes.routes == {} - comms.send_data_to_bg_process.assert_not_called() + ipc.send_message.assert_not_called() -@patch("aikido_zen.background_process.get_comms") -def test_post_response_no_context(mock_get_comms): +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_post_response_no_context(mock_get_ipc): """Test post_response when there is no context.""" - comms = MagicMock() - mock_get_comms.return_value = comms + ipc = MagicMock() + mock_get_ipc.return_value = ipc cache = ThreadCache() # Creates a new cache assert cache.routes.routes == {} @@ -128,7 +128,7 @@ def test_post_response_no_context(mock_get_comms): result = request_handler("post_response", status_code=200) assert cache.routes.routes == {} - comms.send_data_to_bg_process.assert_not_called() + ipc.send_message.assert_not_called() def test_bypassed_ip_no_stats_in_init(): @@ -211,15 +211,15 @@ def create_service_config(): def patch_firewall_lists(func): def wrapper(*args, **kwargs): - with patch("aikido_zen.background_process.comms.get_comms") as mock_comms: - comms = MyMockComms() - mock_comms.return_value = comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_ipc: + ipc = MyMockComms() + mock_ipc.return_value = ipc sig = inspect.signature(func) if "attacks" in sig.parameters: - kwargs["attacks"] = comms.attacks + kwargs["attacks"] = ipc.attacks if "firewall_lists" in sig.parameters: - kwargs["firewall_lists"] = comms.firewall_lists + kwargs["firewall_lists"] = ipc.firewall_lists return func(*args, **kwargs) diff --git a/aikido_zen/thread/process_worker_loader.py b/aikido_zen/thread/process_worker_loader.py index 7917e6992..50c4c1d9e 100644 --- a/aikido_zen/thread/process_worker_loader.py +++ b/aikido_zen/thread/process_worker_loader.py @@ -1,3 +1,4 @@ +import contextvars import multiprocessing import threading @@ -10,13 +11,10 @@ def load_worker(): - """ - Loads in a new process worker if one does not already exist for the current process - """ + """Start the cache synchronization thread on the first request in each process.""" if get_current_context() is None: return # don't start a worker if it's not related to a request. - # The name is aikido-process-worker- + the current PID thread_name = "aikido-process-worker-" + str(multiprocessing.current_process().pid) with _load_worker_lock: @@ -24,16 +22,19 @@ def load_worker(): if any(thread.name == thread_name for thread in threading.enumerate()): return - # The first request can arrive before the cache has config, so sync now - # instead of waiting for the periodic sync. Below the thread check so it - # runs once per process rather than on every request. - if not thread_cache.is_config_loaded(): - try: - thread_cache.renew() - except Exception as e: - logger.warning("An error occurred during data synchronization: %s", e) - - # Create a new daemon thread that will handle communication to and from background agent - thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) - thread.daemon = True - thread.start() + # Run setup without a request context to prevent self-deadlock when hooks + # triggered by imports or error logging call load_worker() again. + contextvars.Context().run(_start_worker, thread_name) + + +def _start_worker(thread_name): + # Try to load config for the first request instead of waiting for periodic sync. + if not thread_cache.is_config_loaded(): + try: + thread_cache.renew() + except Exception as e: + logger.warning("An error occurred during data synchronization: %s", e) + + thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) + thread.daemon = True + thread.start() diff --git a/aikido_zen/thread/process_worker_loader_test.py b/aikido_zen/thread/process_worker_loader_test.py index 55378f272..7877f4601 100644 --- a/aikido_zen/thread/process_worker_loader_test.py +++ b/aikido_zen/thread/process_worker_loader_test.py @@ -1,3 +1,4 @@ +import threading from unittest.mock import MagicMock, patch import pytest @@ -54,7 +55,9 @@ def fake_thread(target=None, name=None, **_kwargs): def test_startup_sync_runs_on_the_first_request(): - assert load_worker_times(1) == 1 + renew = MagicMock() + assert load_worker_times(1, renew=renew) == 1 + renew.assert_called_once_with() def test_startup_sync_is_attempted_only_once_per_process(): @@ -86,3 +89,54 @@ def test_a_failed_sync_is_not_retried_on_the_request_path(): renew = MagicMock(side_effect=Exception("IPC down")) assert load_worker_times(20, renew=renew) == 1 + + +def test_load_worker_called_again_during_the_startup_sync_does_not_deadlock(): + renew = MagicMock(side_effect=loader.load_worker) + + def first_request(): + set_context() + loader.load_worker() + + request = threading.Thread(target=first_request, daemon=True) + with patch.object(loader, "_load_worker_lock", threading.Lock()), patch( + "aikido_zen.thread.thread_cache.renew", renew + ), patch("threading.enumerate", return_value=[]), patch("threading.Thread"): + request.start() + request.join(timeout=5) + + assert not request.is_alive() + assert renew.call_count == 1 + + +def test_failed_sync_logging_does_not_reenter_worker_startup(): + renew = MagicMock(side_effect=Exception("IPC down")) + + with patch.object(loader, "_load_worker_lock") as lock, patch.object( + loader.logger, "warning", side_effect=lambda *_args: loader.load_worker() + ) as warning: + lock.__enter__.side_effect = [ + None, + AssertionError("recursive lock acquisition"), + ] + + assert load_worker_times(1, renew=renew) == 1 + + lock.__enter__.assert_called_once() + warning.assert_called_once() + + +def test_load_worker_can_retry_after_thread_start_fails(): + set_context() + thread = MagicMock() + thread.start.side_effect = [RuntimeError("thread start failed"), None] + + with patch.object(thread_cache, "renew"), patch.object( + loader.threading, "Thread", return_value=thread + ), patch.object(loader.threading, "enumerate", return_value=[]): + with pytest.raises(RuntimeError, match="thread start failed"): + loader.load_worker() + + loader.load_worker() + + assert thread.start.call_count == 2 diff --git a/aikido_zen/thread/process_worker_test.py b/aikido_zen/thread/process_worker_test.py index 40bb05f0c..c54aeec19 100644 --- a/aikido_zen/thread/process_worker_test.py +++ b/aikido_zen/thread/process_worker_test.py @@ -1,5 +1,17 @@ -from unittest.mock import patch +import socket +import threading +from multiprocessing.connection import Connection +from unittest.mock import MagicMock, patch +import pytest + +from aikido_zen.api_discovery.get_data_schema import get_data_schema +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.cloud_connection_manager import ( + CloudConnectionManager, +) +from aikido_zen.background_process.commands import sync_data +from aikido_zen.background_process.routes import Routes import aikido_zen.thread.process_worker as process_worker import aikido_zen.thread.thread_cache as thread_cache @@ -50,3 +62,60 @@ def test_stops_retrying_quickly_when_the_config_never_arrives(): ) assert sleeps == [process_worker.RENEW_CACHE_EVERY_X_SEC] * 2 + + +def test_periodic_sync_delivers_large_batch_after_receiver_stall(monkeypatch): + cache = thread_cache.ThreadCache() + for index in range(1000): + metadata = {"method": "POST", "route": f"/api/resource/{index}/update"} + cache.routes.increment_route(metadata) + cache.stats.increment_total_hits() + body = {f"property_{field:03d}": "value" for field in range(50)} + cache.routes.update_route_with_apispec( + metadata, {"body": {"type": "json", "schema": get_data_schema(body)}} + ) + manager = CloudConnectionManager(False, MagicMock(), "test-token", None) + manager.routes = Routes(1000) + manager.conf.last_updated_at = 1 + manager.conf.revision = 0 + errors = [] + monkeypatch.setattr(thread_cache, "global_thread_cache", cache) + monkeypatch.setattr(IPCClient, "_instance", None) + + with socket.socket() as listener: + listener.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536) + listener.bind(("127.0.0.1", 0)) + listener.listen(1) + listener.settimeout(2) + IPCClient(listener.getsockname(), b"fixture") + + def receive(): + try: + peer, _ = listener.accept() + with peer, Connection(peer.dup().detach()) as connection: + threading.Event().wait(0.3) + action, payload = connection.recv() + assert action == "SYNC_DATA" + connection.send(sync_data.process(manager, payload, None)) + except (OSError, EOFError, AssertionError) as error: + errors.append(error) + + server = threading.Thread(target=receive, daemon=True) + server.start() + try: + with patch.object( + process_worker.time, "sleep", side_effect=KeyboardInterrupt + ): + with pytest.raises(KeyboardInterrupt): + process_worker.aikido_process_worker_thread() + assert cache.config.revision == 0 + assert manager.statistics.get_record()["requests"]["total"] == 1000 + assert sum(route["hits"] for route in manager.routes) == 1000 + assert all( + len(route["apispec"]["body"]["schema"]["properties"]) == 50 + for route in manager.routes + ) + finally: + server.join(2) + assert not server.is_alive() + assert not errors diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 2aa829523..ee75888d8 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -1,6 +1,7 @@ """Exports class ThreadConfig""" -import aikido_zen.background_process.comms as comms +from aikido_zen.background_process.ipc import IPCClient +from aikido_zen.background_process.commands import sync_data from aikido_zen.background_process.packages import PackagesStore from aikido_zen.background_process.routes import Routes from aikido_zen.background_process.service_config import ServiceConfig @@ -57,7 +58,8 @@ def _clear_synced_deltas(self): PackagesStore.clear() def renew(self): - if not comms.get_comms(): + ipc = IPCClient.get() + if not ipc: return # Clear deltas before the IPC, not after. Clearing post-response would @@ -83,11 +85,7 @@ def renew(self): "packages": PackagesStore.export(), } self._clear_synced_deltas() - res = comms.get_comms().send_data_to_bg_process( - action="SYNC_DATA", - obj=payload, - receive=True, - ) + res = ipc.send_message(sync_data.generate(payload)) if not res["success"] or not res["data"]: return diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index b56024b6b..6a8e2eec9 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -10,7 +10,7 @@ CloudConnectionManager, ) from aikido_zen.background_process.cloud_connection_manager.on_start import on_start -from aikido_zen.background_process.commands.sync_data import process_sync_data +from aikido_zen.background_process.commands import sync_data from aikido_zen.background_process.routes import Routes from aikido_zen.sources.functions.request_handler import post_response from . import process_worker_loader @@ -48,20 +48,20 @@ def sync(monkeypatch): context = Context(body={"name": "example"}) context.route, context.method, context.remote_address = "/test", "POST", "1.2.3.4" token = current_context.set(context) - comms = MagicMock() - comms.send_data_to_bg_process.side_effect = lambda action, obj, receive: { + ipc = MagicMock() + ipc.send_message.side_effect = lambda message: { "success": True, - "data": deepcopy(process_sync_data(manager, deepcopy(obj))), + "data": deepcopy(sync_data.process(manager, deepcopy(message.request), None)), } generate_schema = MagicMock(wraps=get_api_info) - monkeypatch.setattr("aikido_zen.background_process.comms.get_comms", lambda: comms) + monkeypatch.setattr("aikido_zen.background_process.ipc.IPCClient.get", lambda: ipc) monkeypatch.setattr( "aikido_zen.sources.functions.request_handler.get_cache", lambda: cache ) monkeypatch.setattr( "aikido_zen.sources.functions.request_handler.get_api_info", generate_schema ) - yield cache, manager, comms, generate_schema, context + yield cache, manager, ipc, generate_schema, context current_context.reset(token) @@ -130,8 +130,8 @@ def keys(self): thread_cache.routes.routes["GET:/first"] ) outgoing = thread_cache.routes.get_routes_with_hits() - with patch("aikido_zen.background_process.comms.get_comms") as get_comms: - get_comms.return_value.send_data_to_bg_process.return_value = { + with patch("aikido_zen.background_process.ipc.IPCClient.get") as get_ipc: + get_ipc.return_value.send_message.return_value = { "success": True, "data": {}, } @@ -154,9 +154,9 @@ def test_increment_total_hits(thread_cache): assert thread_cache.stats.get_record()["requests"]["total"] == 2 -def test_renew_with_no_comms(thread_cache: ThreadCache): +def test_renew_with_no_ipc(thread_cache: ThreadCache): """Test that renew does not proceed if there are no communications available.""" - with patch("aikido_zen.background_process.comms.get_comms", return_value=None): + with patch("aikido_zen.background_process.ipc.IPCClient.get", return_value=None): thread_cache.renew() assert isinstance(thread_cache.config.bypassed_ips, IPMatcher) assert thread_cache.get_endpoints() == [] @@ -272,11 +272,11 @@ def test_load_worker_starts_thread_when_cache_renewal_fails( ) -@patch("aikido_zen.background_process.comms.get_comms") -def test_renew_with_invalid_response(mock_get_comms, thread_cache: ThreadCache): +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_renew_with_invalid_response(mock_get_ipc, thread_cache: ThreadCache): """Test that renew handles an invalid response gracefully.""" - mock_get_comms.return_value = MagicMock() - mock_get_comms.return_value.send_data_to_bg_process.return_value = { + mock_get_ipc.return_value = MagicMock() + mock_get_ipc.return_value.send_message.return_value = { "success": True, "data": { "config": "not_a_service_config", # Invalid type @@ -315,10 +315,10 @@ def increment_in_thread(): ) # 10 threads incrementing 100 times -@patch("aikido_zen.background_process.comms.get_comms") -def test_renew_updates_config(mock_get_comms, thread_cache: ThreadCache): - mock_get_comms.return_value = MagicMock() - mock_get_comms.return_value.send_data_to_bg_process.return_value = { +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_renew_updates_config(mock_get_ipc, thread_cache: ThreadCache): + mock_get_ipc.return_value = MagicMock() + mock_get_ipc.return_value.send_message.return_value = { "success": True, "data": { "config": ServiceConfig( @@ -358,11 +358,10 @@ def test_renew_updates_config(mock_get_comms, thread_cache: ThreadCache): assert thread_cache.is_user_blocked("user123") -@patch("aikido_zen.background_process.comms.get_comms") -def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCache): - """Test that renew calls send_data_to_bg_process with correct arguments.""" - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_renew_called_with_correct_args(mock_get_ipc, thread_cache: ThreadCache): + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Setup initial state thread_cache.stats.increment_total_hits() @@ -395,84 +394,82 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach } assert PackagesStore.export() == [] - # Assert that send_data_to_bg_process was called with the correct arguments - mock_comms.send_data_to_bg_process.assert_called_once_with( - action="SYNC_DATA", - obj={ - "config_revision": -1, - "current_routes": { - "GET:/test": { - "method": "GET", - "path": "/test", - "hits": 1, - "hits_delta_since_sync": 1, - "apispec": {}, - } + mock_ipc.send_message.assert_called_once() + (message,) = mock_ipc.send_message.call_args.args + assert message.identifier == "SYNC_DATA" + assert message.timeout_in_sec == 5 + assert message.request == { + "config_revision": -1, + "current_routes": { + "GET:/test": { + "method": "GET", + "path": "/test", + "hits": 1, + "hits_delta_since_sync": 1, + "apispec": {}, + } + }, + "stats": { + "startedAt": -1, + "endedAt": -1, + "requests": { + "total": 2, + "rateLimited": 0, + "aborted": 0, + "attacksDetected": {"blocked": 1, "total": 3}, + "attackWaves": {"total": 0, "blocked": 0}, }, - "stats": { - "startedAt": -1, - "endedAt": -1, - "requests": { - "total": 2, - "rateLimited": 0, - "aborted": 0, - "attacksDetected": {"blocked": 1, "total": 3}, - "attackWaves": {"total": 0, "blocked": 0}, + "operations": { + "op1": { + "attacksDetected": {"blocked": 1, "total": 2}, + "kind": "sql_op", + "total": 1, }, - "operations": { - "op1": { - "attacksDetected": {"blocked": 1, "total": 2}, - "kind": "sql_op", - "total": 1, - }, - "op2": { - "attacksDetected": {"blocked": 0, "total": 1}, - "kind": "sql_op", - "total": 1, - }, + "op2": { + "attacksDetected": {"blocked": 0, "total": 1}, + "kind": "sql_op", + "total": 1, }, }, - "ai_stats": [ - { - "provider": "openai", - "model": "gpt-4o", - "calls": 1, - "tokens": {"input": 6427, "output": 200, "total": 6627}, - }, - { - "provider": "openai", - "model": "gpt-4o2", - "calls": 2, - "tokens": {"input": 656, "output": 1167, "total": 1823}, - }, - { - "provider": "openai", - "model": "gpt-4o3", - "calls": 1, - "tokens": {"input": 8223, "output": 173, "total": 8396}, - }, - ], - "middleware_installed": False, - "hostnames": [], - "users": [], - "packages": [ - { - "name": "test-package-1", - "version": "4.3.0", - "requiredAt": -1, - "cleared": False, - } - ], }, - receive=True, - ) + "ai_stats": [ + { + "provider": "openai", + "model": "gpt-4o", + "calls": 1, + "tokens": {"input": 6427, "output": 200, "total": 6627}, + }, + { + "provider": "openai", + "model": "gpt-4o2", + "calls": 2, + "tokens": {"input": 656, "output": 1167, "total": 1823}, + }, + { + "provider": "openai", + "model": "gpt-4o3", + "calls": 1, + "tokens": {"input": 8223, "output": 173, "total": 8396}, + }, + ], + "middleware_installed": False, + "hostnames": [], + "users": [], + "packages": [ + { + "name": "test-package-1", + "version": "4.3.0", + "requiredAt": -1, + "cleared": False, + } + ], + } -@patch("aikido_zen.background_process.comms.get_comms") -def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): - """Test that renew calls send_data_to_bg_process with correct arguments.""" - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_sync_data_for_users(mock_get_ipc, thread_cache: ThreadCache): + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc test_utils.generate_and_set_context(ip="5.6.7.8") # Setup initial state @@ -485,84 +482,81 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): with test_utils.patch_time(time_ms=-1): thread_cache.renew() - # Assert that send_data_to_bg_process was called with the correct arguments - mock_comms.send_data_to_bg_process.assert_called_once_with( - action="SYNC_DATA", - obj={ - "config_revision": -1, - "current_routes": {}, - "stats": { - "startedAt": -1, - "endedAt": -1, - "requests": { - "total": 1, - "rateLimited": 0, - "aborted": 0, - "attacksDetected": {"total": 0, "blocked": 0}, - "attackWaves": {"total": 0, "blocked": 0}, - }, - "operations": {}, + mock_ipc.send_message.assert_called_once() + (message,) = mock_ipc.send_message.call_args.args + assert message.identifier == "SYNC_DATA" + assert message.timeout_in_sec == 5 + assert message.request == { + "config_revision": -1, + "current_routes": {}, + "stats": { + "startedAt": -1, + "endedAt": -1, + "requests": { + "total": 1, + "rateLimited": 0, + "aborted": 0, + "attacksDetected": {"total": 0, "blocked": 0}, + "attackWaves": {"total": 0, "blocked": 0}, }, - "middleware_installed": False, - "hostnames": [], - "ai_stats": [], - "packages": [], - "users": [ - { - "id": "123", - "name": "test", - "lastIpAddress": "5.6.7.8", - "firstSeenAt": 1, - "lastSeenAt": 1, - }, - { - "id": "567", - "name": "test", - "lastIpAddress": "5.6.7.8", - "firstSeenAt": 1, - "lastSeenAt": 1, - }, - ], + "operations": {}, }, - receive=True, - ) + "middleware_installed": False, + "hostnames": [], + "ai_stats": [], + "packages": [], + "users": [ + { + "id": "123", + "name": "test", + "lastIpAddress": "5.6.7.8", + "firstSeenAt": 1, + "lastSeenAt": 1, + }, + { + "id": "567", + "name": "test", + "lastIpAddress": "5.6.7.8", + "firstSeenAt": 1, + "lastSeenAt": 1, + }, + ], + } -@patch("aikido_zen.background_process.comms.get_comms") -def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCache): - """Test that renew calls send_data_to_bg_process with empty routes.""" - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_renew_called_with_empty_routes(mock_get_ipc, thread_cache: ThreadCache): + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc with test_utils.patch_time(time_ms=-1): thread_cache.renew() - # Assert that send_data_to_bg_process was called with the correct arguments - mock_comms.send_data_to_bg_process.assert_called_once_with( - action="SYNC_DATA", - obj={ - "config_revision": -1, - "current_routes": {}, - "stats": { - "startedAt": -1, - "endedAt": -1, - "requests": { - "total": 0, - "rateLimited": 0, - "aborted": 0, - "attacksDetected": {"total": 0, "blocked": 0}, - "attackWaves": {"total": 0, "blocked": 0}, - }, - "operations": {}, + mock_ipc.send_message.assert_called_once() + (message,) = mock_ipc.send_message.call_args.args + assert message.identifier == "SYNC_DATA" + assert message.timeout_in_sec == 5 + assert message.request == { + "config_revision": -1, + "current_routes": {}, + "stats": { + "startedAt": -1, + "endedAt": -1, + "requests": { + "total": 0, + "rateLimited": 0, + "aborted": 0, + "attacksDetected": {"total": 0, "blocked": 0}, + "attackWaves": {"total": 0, "blocked": 0}, }, - "middleware_installed": False, - "hostnames": [], - "users": [], - "ai_stats": [], - "packages": [], + "operations": {}, }, - receive=True, - ) + "middleware_installed": False, + "hostnames": [], + "users": [], + "ai_stats": [], + "packages": [], + } def test_sync_includes_updates_to_routes_already_in_payload(sync): @@ -636,25 +630,23 @@ def request_after_collection(): def test_sync_counts_each_request_once( sync, config_loaded, first_sync_fails, concurrent_path ): - cache, manager, comms, _, context = sync + cache, manager, ipc, _, context = sync manager.conf.last_updated_at = 1 if config_loaded else -1 - deliver = comms.send_data_to_bg_process.side_effect + deliver = ipc.send_message.side_effect cache.stats.increment_total_hits() post_response(200) - def first_sync(action, obj, receive): - response = ( - {"success": False} if first_sync_fails else deliver(action, obj, receive) - ) + def first_sync(message): + response = {"success": False} if first_sync_fails else deliver(message) if concurrent_path: context.route, context.body = concurrent_path, {"during": "sync"} cache.stats.increment_total_hits() post_response(200) return response - comms.send_data_to_bg_process.side_effect = first_sync + ipc.send_message.side_effect = first_sync cache.renew() - comms.send_data_to_bg_process.side_effect = deliver + ipc.send_message.side_effect = deliver cache.renew() cache.renew() @@ -689,34 +681,34 @@ def test_sync_merges_optional_fields_without_reusing_previous_heartbeat_schema(s def test_overlapping_syncs_do_not_replay_hits(sync): - cache, manager, comms, _, _ = sync - deliver = comms.send_data_to_bg_process.side_effect + cache, manager, ipc, _, _ = sync + deliver = ipc.send_message.side_effect post_response(200) - def first_sync(action, obj, receive): - comms.send_data_to_bg_process.side_effect = deliver + def first_sync(message): + ipc.send_message.side_effect = deliver post_response(200) cache.renew() - return deliver(action, obj, receive) + return deliver(message) - comms.send_data_to_bg_process.side_effect = first_sync + ipc.send_message.side_effect = first_sync cache.renew() assert manager.routes.get({"method": "POST", "route": "/test"})["hits"] == 2 assert cache.routes.get_routes_with_hits() == {} def test_config_is_sent_once_per_worker_per_cloud_response(sync): - cache, manager, comms, _, _ = sync + cache, manager, ipc, _, _ = sync workers = (cache, ThreadCache()) - deliver = comms.send_data_to_bg_process.side_effect + deliver = ipc.send_message.side_effect replies = [] - def record_reply(action, obj, receive): - response = deliver(action, obj, receive) + def record_reply(message): + response = deliver(message) replies.append(response["data"]) return response - comms.send_data_to_bg_process.side_effect = record_reply + ipc.send_message.side_effect = record_reply for _ in range(2): manager.update_service_config({"success": True, "configUpdatedAt": 1}) for worker in workers: @@ -740,28 +732,31 @@ def test_initial_config_is_sent_after_failed_cloud_startup(sync): cache.renew() assert cache.config.revision == 0 assert cache.config.last_updated_at > 0 - assert process_sync_data(manager, {"config_revision": cache.config.revision}) == {} + assert ( + sync_data.process(manager, {"config_revision": cache.config.revision}, None) + == {} + ) def test_config_is_retried_after_lost_reply(sync): - cache, manager, comms, _, _ = sync - deliver = comms.send_data_to_bg_process.side_effect + cache, manager, ipc, _, _ = sync + deliver = ipc.send_message.side_effect - def lose_reply(action, obj, receive): - assert "config" in deliver(action, obj, receive)["data"] + def lose_reply(message): + assert "config" in deliver(message)["data"] return {"success": False} - comms.send_data_to_bg_process.side_effect = lose_reply + ipc.send_message.side_effect = lose_reply cache.renew() assert cache.config.last_updated_at == -1 - comms.send_data_to_bg_process.side_effect = deliver + ipc.send_message.side_effect = deliver cache.renew() assert cache.config.last_updated_at == manager.conf.last_updated_at - request = comms.send_data_to_bg_process.call_args.kwargs["obj"] - assert "config" in process_sync_data(manager, request) + request = ipc.send_message.call_args.args[0].request + assert "config" in sync_data.process(manager, request, None) cache.renew() - request = comms.send_data_to_bg_process.call_args.kwargs["obj"] - assert process_sync_data(manager, request) == {} + request = ipc.send_message.call_args.args[0].request + assert sync_data.process(manager, request, None) == {} @pytest.mark.parametrize("config_loaded", [False, True]) @@ -770,7 +765,7 @@ def lose_reply(action, obj, receive): def test_sampling_limit_resets_after_fresh_config( sync, config_loaded, report_succeeds, first_sync_fails ): - cache, manager, comms, generate_schema, _ = sync + cache, manager, ipc, generate_schema, _ = sync manager.conf.last_updated_at = 1 if config_loaded else -1 manager.report_api_event.return_value = { "success": report_succeeds, @@ -785,19 +780,19 @@ def test_sampling_limit_resets_after_fresh_config( expected_samples = 40 if config_loaded else 20 assert generate_schema.call_count == expected_samples - deliver = comms.send_data_to_bg_process.side_effect + deliver = ipc.send_message.side_effect for heartbeat_at in (31_000, 131_000): with test_utils.patch_time(time_ms=heartbeat_at): manager.send_heartbeat() post_response(200) assert generate_schema.call_count == expected_samples if first_sync_fails: - comms.send_data_to_bg_process.side_effect = None - comms.send_data_to_bg_process.return_value = {"success": False} + ipc.send_message.side_effect = None + ipc.send_message.return_value = {"success": False} cache.renew() post_response(200) assert generate_schema.call_count == expected_samples - comms.send_data_to_bg_process.side_effect = deliver + ipc.send_message.side_effect = deliver cache.renew() for _ in range(21): post_response(200) @@ -811,23 +806,23 @@ def test_sampling_limit_resets_after_fresh_config( def test_sampling_reset_preserves_requests_collected_during_sync( sync, concurrent_path, requests_during_sync ): - cache, manager, comms, generate_schema, context = sync + cache, manager, ipc, generate_schema, context = sync manager.report_api_event.return_value = {"success": True, "configUpdatedAt": 1} for _ in range(3): post_response(200) cache.renew() with test_utils.patch_time(time_ms=31_000): manager.send_heartbeat() - deliver = comms.send_data_to_bg_process.side_effect + deliver = ipc.send_message.side_effect - def request_during_sync(action, obj, receive): - response = deliver(action, obj, receive) + def request_during_sync(message): + response = deliver(message) context.route, context.body = concurrent_path, {"during": "sync"} for _ in range(requests_during_sync): post_response(200) return response - comms.send_data_to_bg_process.side_effect = request_during_sync + ipc.send_message.side_effect = request_during_sync cache.renew() metadata = {"method": "POST", "route": concurrent_path} pending = cache.routes.get(metadata) @@ -836,7 +831,7 @@ def request_during_sync(action, obj, receive): assert pending["apispec"]["body"]["schema"]["properties"] == { "during": {"type": "string"} } - comms.send_data_to_bg_process.side_effect = deliver + ipc.send_message.side_effect = deliver cache.renew() reported = manager.routes.get(metadata) assert reported["hits"] == requests_during_sync @@ -847,14 +842,14 @@ def request_during_sync(action, obj, receive): @pytest.mark.parametrize("response_data", [{}, {"config": None}]) -@patch("aikido_zen.background_process.comms.get_comms") +@patch("aikido_zen.background_process.ipc.IPCClient.get") def test_renew_preserves_increments_during_ipc( - mock_get_comms, thread_cache: ThreadCache, response_data + mock_get_ipc, thread_cache: ThreadCache, response_data ): """Increments arriving during the IPC call survive on the response path - the snapshot is sent, but the live counter keeps the concurrent increment.""" - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc thread_cache.stats.increment_total_hits() thread_cache.stats.increment_total_hits() @@ -863,23 +858,23 @@ def simulate_concurrent_increment(*args, **kwargs): thread_cache.stats.increment_total_hits() return {"success": True, "data": response_data} - mock_comms.send_data_to_bg_process.side_effect = simulate_concurrent_increment + mock_ipc.send_message.side_effect = simulate_concurrent_increment thread_cache.renew() - sent_total = mock_comms.send_data_to_bg_process.call_args.kwargs["obj"]["stats"][ - "requests" - ]["total"] + sent_total = mock_ipc.send_message.call_args.args[0].request["stats"]["requests"][ + "total" + ] assert sent_total == 2 assert thread_cache.stats.get_record()["requests"]["total"] == 1 -@patch("aikido_zen.background_process.comms.get_comms") +@patch("aikido_zen.background_process.ipc.IPCClient.get") def test_renew_drops_failed_batch_and_preserves_new_increments( - mock_get_comms, thread_cache: ThreadCache + mock_get_ipc, thread_cache: ThreadCache ): - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc thread_cache.stats.increment_total_hits() thread_cache.stats.increment_total_hits() @@ -892,7 +887,7 @@ def fail_after_concurrent_increment(*args, **kwargs): thread_cache.stats.on_detected_attack_wave(blocked=False) return {"success": False} - mock_comms.send_data_to_bg_process.side_effect = fail_after_concurrent_increment + mock_ipc.send_message.side_effect = fail_after_concurrent_increment thread_cache.renew() @@ -905,11 +900,10 @@ def fail_after_concurrent_increment(*args, **kwargs): assert thread_cache.ai_stats.get_stats() == [] -@patch("aikido_zen.background_process.comms.get_comms") -def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache): - """Test that renew calls send_data_to_bg_process with zero requests.""" - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms +@patch("aikido_zen.background_process.ipc.IPCClient.get") +def test_renew_called_with_no_requests(mock_get_ipc, thread_cache: ThreadCache): + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Setup initial state with a route but no requests thread_cache.routes.initialize_route({"method": "GET", "route": "/test"}) @@ -917,29 +911,28 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache with test_utils.patch_time(time_ms=-1): thread_cache.renew() - # Assert that send_data_to_bg_process was called with the correct arguments - mock_comms.send_data_to_bg_process.assert_called_once_with( - action="SYNC_DATA", - obj={ - "config_revision": -1, - "current_routes": {}, - "stats": { - "startedAt": -1, - "endedAt": -1, - "requests": { - "total": 0, - "rateLimited": 0, - "aborted": 0, - "attacksDetected": {"total": 0, "blocked": 0}, - "attackWaves": {"total": 0, "blocked": 0}, - }, - "operations": {}, + mock_ipc.send_message.assert_called_once() + (message,) = mock_ipc.send_message.call_args.args + assert message.identifier == "SYNC_DATA" + assert message.timeout_in_sec == 5 + assert message.request == { + "config_revision": -1, + "current_routes": {}, + "stats": { + "startedAt": -1, + "endedAt": -1, + "requests": { + "total": 0, + "rateLimited": 0, + "aborted": 0, + "attacksDetected": {"total": 0, "blocked": 0}, + "attackWaves": {"total": 0, "blocked": 0}, }, - "middleware_installed": False, - "hostnames": [], - "users": [], - "ai_stats": [], - "packages": [], + "operations": {}, }, - receive=True, - ) + "middleware_installed": False, + "hostnames": [], + "users": [], + "ai_stats": [], + "packages": [], + } diff --git a/aikido_zen/vulnerabilities/__init__.py b/aikido_zen/vulnerabilities/__init__.py index 436e63094..72c198e26 100644 --- a/aikido_zen/vulnerabilities/__init__.py +++ b/aikido_zen/vulnerabilities/__init__.py @@ -12,7 +12,7 @@ AikidoPathTraversal, AikidoSSRF, ) -import aikido_zen.background_process.comms as comm +from aikido_zen.background_process.ipc import IPCClient from aikido_zen.helpers.logging import logger from aikido_zen.helpers.get_clean_stacktrace import get_clean_stacktrace from aikido_zen.helpers.blocking_enabled import is_blocking_enabled @@ -29,9 +29,8 @@ from .path_traversal.check_context_for_path_traversal import ( check_context_for_path_traversal, ) -from ..background_process.commands import PutEventCommand +from ..background_process.commands import put_event from ..helpers.create_detected_attack_api_event import create_detected_attack_api_event -from ..helpers.ipc.send_payload import send_payload def run_vulnerability_scan(kind, op, args): @@ -44,7 +43,7 @@ def run_vulnerability_scan(kind, op, args): if is_protection_forced_off_cached(context): return - comms = comm.get_comms() + ipc = IPCClient.get() thread_cache = get_cache() if not context and kind != "ssrf": # Make a special exception for SSRF, which checks itself if context is set. @@ -107,7 +106,7 @@ def run_vulnerability_scan(kind, op, args): injection_results, context, blocked, stack ) logger.debug("Attack: %s", serialize_to_json(event)[:5000]) - if comms and event: - send_payload(comms, PutEventCommand.generate(event)) + if ipc and event: + ipc.send_message_async(put_event.generate(event)) if blocked: raise error_type(*error_args) diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py index a24c5a4d4..eed2c0370 100644 --- a/aikido_zen/vulnerabilities/init_test.py +++ b/aikido_zen/vulnerabilities/init_test.py @@ -18,9 +18,9 @@ def run_around_tests(): @pytest.fixture -def mock_comms(): - """Fixture to mock the get_comms function.""" - with patch("aikido_zen.background_process.comms.get_comms") as mock: +def mock_ipc(): + """Fixture to mock the get_ipc function.""" + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock: yield mock @@ -129,24 +129,25 @@ def test_sql_injection_with_route_params(caplog, get_context, monkeypatch): ) -def test_sql_injection_with_comms(caplog, get_context, monkeypatch): +def test_sql_injection_with_ipc(caplog, get_context, monkeypatch): get_cache().config.last_updated_at = 1 get_context.set_as_current_context() monkeypatch.setenv("AIKIDO_BLOCK", "1") - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - # Create a mock comms object - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms # Set the return value of get_comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + # Create a mock ipc object + mock_ipc = MagicMock() + mock_ipc.send_message.return_value = {"success": True, "data": True} + mock_get_ipc.return_value = mock_ipc # Set the return value of get_ipc with pytest.raises(AikidoSQLInjection): run_vulnerability_scan( kind="sql_injection", op="test_op", args=("INSERT * INTO VALUES ('doggoss2', TRUE);", "mysql"), ) - mock_comms.send_data_to_bg_process.assert_called_once() - call_args = mock_comms.send_data_to_bg_process.call_args[0] - assert call_args[0] == "put_event" - assert call_args[1].event["request"] == { + mock_ipc.send_message_async.assert_called_once() + (message,) = mock_ipc.send_message_async.call_args.args + assert message.identifier == "PUT_EVENT" + assert message.request.event["request"] == { "ipAddress": "198.51.100.23", "method": "GET", "route": "/hello", @@ -154,8 +155,8 @@ def test_sql_injection_with_comms(caplog, get_context, monkeypatch): "url": "http://localhost:8080/hello", "userAgent": None, } - del call_args[1].event["attack"]["stack"] # Hard to test - assert call_args[1].event["attack"] == { + del message.request.event["attack"]["stack"] # Hard to test + assert message.request.event["attack"] == { "blocked": True, "kind": "sql_injection", "metadata": { @@ -215,10 +216,10 @@ def test_ssrf_vulnerability_scan_protection_gets_forced_off(get_context): def test_sql_injection_with_protection_forced_off(caplog, get_context, monkeypatch): get_context.set_as_current_context() monkeypatch.setenv("AIKIDO_BLOCK", "1") - with patch("aikido_zen.background_process.comms.get_comms") as mock_get_comms: - # Create a mock comms object - mock_comms = MagicMock() - mock_get_comms.return_value = mock_comms # Set the return value of get_comms + with patch("aikido_zen.background_process.ipc.IPCClient.get") as mock_get_ipc: + # Create a mock ipc object + mock_ipc = MagicMock() + mock_get_ipc.return_value = mock_ipc # Set the return value of get_ipc with pytest.raises(AikidoSQLInjection): run_vulnerability_scan( kind="sql_injection", diff --git a/docs/background_process.md b/docs/background_process.md index b7e2a5668..44dfb4315 100644 --- a/docs/background_process.md +++ b/docs/background_process.md @@ -14,4 +14,4 @@ Many Python applications utilize multiple threads, and some even employ multiple Upon starting the application, a background process is initiated that listens on a Unix domain socket for incoming data. This process also sends periodic heartbeats with statistics to Aikido's servers and immediately relays information regarding any detected attacks. Additionally, the background process contains configuration settings that can be fetched in real-time. If a thread or subprocess needs to access this configuration, it retrieves it from this background process at the beginning of its cycle and caches it for the duration of the request, thereby improving performance. ## What is AIKIDO_TMP_DIR? -The `AIKIDO_TMP_DIR` environment variable specifies the directory used for temporary files related to the Aikido application. By default, this variable is set to /tmp. However, users can customize it to point to a directory of their choice. The Unix domain socket file, aikido.sock, is stored in this directory to facilitate Inter-Process Communication (IPC) between processes. +The `AIKIDO_TMP_DIR` environment variable specifies the directory used for temporary files related to the Aikido application. By default, this variable is set to /tmp. However, users can customize it to point to a directory of their choice. The Unix domain socket file, `aikido_zen_.sock`, is stored in this directory to facilitate Inter-Process Communication (IPC) between processes. diff --git a/end2end/e2e.sh b/end2end/e2e.sh index 3e535dbe1..0ed1b6f4b 100755 --- a/end2end/e2e.sh +++ b/end2end/e2e.sh @@ -150,9 +150,8 @@ APP_PORT_DISABLED="${APP_PORT_DISABLED:-8081}" # 1. Kill any lingering processes from previous runs AIKIDO_SOCK=$(AIKIDO_TOKEN="AIK_secret_token" poetry run python3 -c " import os; os.environ['AIKIDO_TOKEN']='AIK_secret_token' -from aikido_zen.helpers.hash_aikido_token import hash_aikido_token -from aikido_zen.helpers.get_temp_dir import get_temp_dir -print(f'{get_temp_dir()}/aikido_python_{hash_aikido_token()}.sock') +from aikido_zen.background_process import get_socket_path +print(get_socket_path()) " 2>/dev/null) if [[ -n "$AIKIDO_SOCK" ]]; then rm -f "$AIKIDO_SOCK"