From 2aa0cbcbe7f98317265a7e18bf9d59aa0a1f9cb3 Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 01:09:28 +0200 Subject: [PATCH 1/8] Send events asynchronously --- aikido_zen/background_process/__init__.py | 2 +- .../aikido_background_process.py | 6 +- .../aikido_background_process_test.py | 96 ++++ .../background_process/commands/__init__.py | 4 +- .../background_process/commands/put_event.py | 4 - aikido_zen/background_process/comms.py | 118 ++-- aikido_zen/background_process/comms_test.py | 502 +++++++++++++++++- .../background_process/packages_test.py | 4 +- aikido_zen/context/track.py | 2 +- aikido_zen/context/track_test.py | 23 +- aikido_zen/helpers/blocking_enabled.py | 2 +- aikido_zen/helpers/ipc/command_types.py | 12 +- aikido_zen/helpers/ipc/event_sender.py | 73 +++ aikido_zen/helpers/ipc/event_sender_test.py | 236 ++++++++ aikido_zen/helpers/ipc/send_payload.py | 8 +- aikido_zen/middleware/init_test.py | 1 - aikido_zen/middleware/should_block_request.py | 1 - .../sources/functions/request_handler.py | 1 - .../sources/functions/request_handler_test.py | 2 +- aikido_zen/thread/thread_cache.py | 1 - aikido_zen/thread/thread_cache_test.py | 28 +- aikido_zen/vulnerabilities/init_test.py | 4 +- 22 files changed, 1000 insertions(+), 130 deletions(-) create mode 100644 aikido_zen/background_process/aikido_background_process_test.py create mode 100644 aikido_zen/helpers/ipc/event_sender.py create mode 100644 aikido_zen/helpers/ipc/event_sender_test.py diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index e145de8e1..24dfac283 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -88,7 +88,7 @@ def get_uds_filename(): def background_process_already_active(comms): - res = comms.send_data_to_bg_process(action="PING", obj=tuple(), receive=True) + res = comms.send_data_to_bg_process(action="PING", obj=tuple()) 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..e3f4efd80 --- /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..3a43943d2 100644 --- a/aikido_zen/background_process/commands/__init__.py +++ b/aikido_zen/background_process/commands/__init__.py @@ -27,8 +27,8 @@ def process_incoming_command(connection_manager, obj, conn, queue): for cmd in modern_commands: if cmd.identifier() == inbound_identifier: - cmd.run(CommandContext(connection_manager, queue, conn), inbound_request) - return None + result = cmd.run(CommandContext(connection_manager, queue), inbound_request) + return conn.send(result) logger.debug("Command : `%s` not found - did not execute", inbound_identifier) return None diff --git a/aikido_zen/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index 790df47a3..bf4f3aa40 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -13,10 +13,6 @@ class PutEventCommand(Command): def identifier(cls) -> str: return "put_event" - @classmethod - def returns_data(cls) -> bool: - return False - @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 diff --git a/aikido_zen/background_process/comms.py b/aikido_zen/background_process/comms.py index 48b7e402e..ecc22de94 100644 --- a/aikido_zen/background_process/comms.py +++ b/aikido_zen/background_process/comms.py @@ -3,14 +3,65 @@ Exports the AikidoIPCCommunications class """ +import io import multiprocessing.connection as con -from threading import Thread +import socket +import time from aikido_zen.helpers.logging import logger +from aikido_zen.helpers.ipc.event_sender import event_sender # pylint: disable=invalid-name # This variable does change comms = None +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)) + + def get_comms(): """ Returns the globally stored IPC object, which you need @@ -44,15 +95,24 @@ def __init__(self, address, key): global comms comms = self - def send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0.1): + def enqueue_event(self, action, obj): + if not self.key: + return {"success": False, "error": "invalid_key"} + try: + return event_sender.enqueue(self, action, obj) + except Exception as error: + logger.debug("Could not enqueue event: %s", error) + return {"success": False, "error": "unknown"} + + def send_data_to_bg_process(self, action, obj, 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) + return self._send_data_to_bg_process(action, obj, 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): + def _send_data_to_bg_process(self, action, obj, timeout_in_sec=0.1): """ This creates a new client for comms to the background process """ @@ -60,41 +120,19 @@ def _send_data_to_bg_process(self, action, obj, receive=False, timeout_in_sec=0. # 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 - ) + if timeout_in_sec <= 0: return {"success": False, "error": "timeout"} + deadline = time.monotonic() + timeout_in_sec - if receive: - return {"success": True, "data": result_obj[1]} - return {"success": True, "data": None} + with socket.socket(getattr(socket, con.address_type(self.address))) as sock: + try: + connect_with_deadline(sock, self.address, deadline) + with SocketConnection(sock, deadline) as conn: + conn.send((action, obj)) + result = conn.recv() + if time.monotonic() >= deadline: + return {"success": False, "error": "timeout"} + except Exception as e: + logger.debug("Exception occurred in IPC call : %s", e) + return {"success": False, "error": "timeout"} + return {"success": True, "data": result} diff --git a/aikido_zen/background_process/comms_test.py b/aikido_zen/background_process/comms_test.py index e1ee26411..42f361581 100644 --- a/aikido_zen/background_process/comms_test.py +++ b/aikido_zen/background_process/comms_test.py @@ -1,40 +1,504 @@ +import gc +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 + import pytest -from aikido_zen.background_process.comms import ( - AikidoIPCCommunications, - reset_comms, -) + +from aikido_zen.background_process import comms as module +from aikido_zen.background_process.commands import process_incoming_command +from aikido_zen.background_process.commands.put_event import PutEventReq @pytest.fixture(autouse=True) def reset_comms_after_test(): yield - reset_comms() + module.reset_comms() + + +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( + module, + "socket", + SimpleNamespace(**{**vars(socket), "socket": TrackedSocket}), + raising=False, + ) + + class TrackedConnection(module.SocketConnection): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + connections.append(self) + + monkeypatch.setattr(module, "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 def test_comms_init(): address = ("localhost", 9898) key = "secret_key" - comms = AikidoIPCCommunications(address, key) + comms = module.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") +@pytest.mark.parametrize( + "action,payload,expected", + [ + ("PING", (), "Received"), + ("READ_PROPERTY", "block", True), + ("put_event", PutEventReq({"name": "login.failed"}), None), + ], +) +def test_commands_return_results_or_acknowledgments( + resources, servers, monkeypatch, action, payload, 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.AikidoIPCCommunications(server.address, b"fixture") + assert communications.send_data_to_bg_process( + action, payload, timeout_in_sec=1 + ) == {"success": True, "data": expected} + if action == "put_event": + assert events.get_nowait() == {"name": "login.failed"} + assert_released(resources, []) + + +@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.AikidoIPCCommunications(server.address, b"fixture") + payloads = [] + for _ in range(5): + payload = Payload() + payloads.append(weakref.ref(payload)) + assert communications.send_data_to_bg_process( + "put_event", 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.AikidoIPCCommunications(server.address, b"fixture") + payload = Payload() + reference = weakref.ref(payload) + + result = communications.send_data_to_bg_process( + "put_event", 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(module.SocketConnection, "_send", fail_send) + server = servers(1) + communications = module.AikidoIPCCommunications(server.address, b"fixture") + payload = Payload() + reference = weakref.ref(payload) + + result = communications.send_data_to_bg_process( + "put_event", 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.AikidoIPCCommunications( + reserved.getsockname(), b"fixture" + ) + result = communications.send_data_to_bg_process("PING", (), 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.AikidoIPCCommunications(server.address, b"fixture") + payload = Unpicklable() + reference = weakref.ref(payload) + + result = communications.send_data_to_bg_process( + "put_event", 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.AikidoIPCCommunications(server.address, b"fixture") + + result = communications.send_data_to_bg_process("PING", (), timeout_in_sec=1) + assert not result["success"] + assert_released(resources, []) + + +def test_timeout_interrupts_a_blocked_send(resources, servers, monkeypatch): + create_socket = module.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(module.socket, "socket", small_send_buffer) + server = servers(1, read_request=False) + communications = module.AikidoIPCCommunications(server.address, b"fixture") + payload = Payload(8 * 1024 * 1024) + reference = weakref.ref(payload) + assert communications.send_data_to_bg_process( + "put_event", 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.AikidoIPCCommunications(server.address, b"fixture") + started = time.monotonic() + assert communications.send_data_to_bg_process("PING", (), 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.AikidoIPCCommunications(("127.0.0.1", 1), b"fixture") + + assert communications.send_data_to_bg_process( + "PING", (), 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.AikidoIPCCommunications(server.address, b"fixture") + assert communications.send_data_to_bg_process( + "READ_PROPERTY", "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.AikidoIPCCommunications( + listener.getsockname(), b"fixture" + ) + payloads = [] + for _ in range(3): + payload = Payload() + payloads.append(weakref.ref(payload)) + assert communications.send_data_to_bg_process( + "put_event", 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( + module, "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.AikidoIPCCommunications(address, b"fixture") + assert communications.send_data_to_bg_process( + "PING", (), timeout_in_sec=1 + ) == {"success": True, "data": "Received"} + assert not server_errors + assert_released(resources, []) + finally: + accept_calls.set() + server.join(3) + - monkeypatch.setitem(globals(), "Client", mock_client) - monkeypatch.setitem(globals(), "logger", caplog) +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") - comms = AikidoIPCCommunications(("localhost", 9898), "mock_key") - comms.send_data_to_bg_process("ACTION", "Test Object") + 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.AikidoIPCCommunications(listener.address, b"fixture") -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) + def request(name): + outcomes[name] = communications.send_data_to_bg_process( + "READ_PROPERTY", name, timeout_in_sec=3 + ) + completion_order.append(name) + if name == "second": + second_finished.set() - # Call the send_data_to_bg_process function - comms.send_data_to_bg_process("ACTION", {"key": "value"}) + 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() diff --git a/aikido_zen/background_process/packages_test.py b/aikido_zen/background_process/packages_test.py index 66e3419a4..f613af77b 100644 --- a/aikido_zen/background_process/packages_test.py +++ b/aikido_zen/background_process/packages_test.py @@ -11,12 +11,10 @@ class FakeComms: def __init__(self): self.action = None self.obj = None - self.receive = False - def send_data_to_bg_process(self, action, obj, receive): + def send_data_to_bg_process(self, action, obj, timeout_in_sec=0.1): self.action = action self.obj = obj - self.receive = receive return {"success": True, "data": True} return FakeComms diff --git a/aikido_zen/context/track.py b/aikido_zen/context/track.py index de729718a..2cce0545b 100644 --- a/aikido_zen/context/track.py +++ b/aikido_zen/context/track.py @@ -43,7 +43,7 @@ def track(event_name): if not event: return - send_payload(ipc, PutEventCommand.generate(event), (50 / 1000)) + send_payload(ipc, PutEventCommand.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..b313ca7d9 100644 --- a/aikido_zen/context/track_test.py +++ b/aikido_zen/context/track_test.py @@ -46,9 +46,9 @@ def sent_event_with_cache(cache): ): track("my-custom-event") - if not comms.send_data_to_bg_process.called: + if not comms.enqueue_event.called: return None - return comms.send_data_to_bg_process.call_args[0][1].event + return comms.enqueue_event.call_args[0][1].event def test_track_invalid_event_name(caplog): @@ -103,10 +103,9 @@ def test_track_sends_event_over_ipc(): ): 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] + comms.enqueue_event.assert_called_once() + identifier, request = comms.enqueue_event.call_args[0] assert identifier == "put_event" - assert returns_data is False assert request.event == { "type": "custom", "time": 111, @@ -173,7 +172,7 @@ def track_n_times(n, event_name="my-custom-event"): 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[0][1].event for c in comms.enqueue_event.call_args_list] def test_track_sends_every_event_below_the_limit(): @@ -199,9 +198,7 @@ def test_all_event_names_share_one_budget(): track("payment.failed") track("order.completed") - names = [ - c[0][1].event["name"] for c in comms.send_data_to_bg_process.call_args_list - ] + names = [c[0][1].event["name"] for c in comms.enqueue_event.call_args_list] assert len(names) == track_module.MAX_EVENTS_PER_REQUEST assert "order.completed" not in names @@ -246,7 +243,7 @@ def test_a_bypassed_ip_does_not_use_up_the_limit(): 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() @@ -256,8 +253,8 @@ def test_track_sends_the_event_with_a_short_timeout(): ), patch("aikido_zen.background_process.comms.get_comms", return_value=comms): 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 + comms.enqueue_event.assert_called_once() + comms.send_data_to_bg_process.assert_not_called() def test_track_sends_nothing_when_the_event_could_not_be_built(): @@ -274,4 +271,4 @@ def test_track_sends_nothing_when_the_event_could_not_be_built(): ): track("my-custom-event") - comms.send_data_to_bg_process.assert_not_called() + comms.enqueue_event.assert_not_called() diff --git a/aikido_zen/helpers/blocking_enabled.py b/aikido_zen/helpers/blocking_enabled.py index cb69485d9..46dce63ef 100644 --- a/aikido_zen/helpers/blocking_enabled.py +++ b/aikido_zen/helpers/blocking_enabled.py @@ -11,7 +11,7 @@ def is_blocking_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 + action="READ_PROPERTY", obj="block" ) if should_block_res["success"]: return should_block_res["data"] diff --git a/aikido_zen/helpers/ipc/command_types.py b/aikido_zen/helpers/ipc/command_types.py index dbbace9f3..364954bbf 100644 --- a/aikido_zen/helpers/ipc/command_types.py +++ b/aikido_zen/helpers/ipc/command_types.py @@ -2,19 +2,14 @@ class CommandContext: - def __init__(self, connection_manager, queue, connection): + def __init__(self, connection_manager, queue): 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() self.request = request @@ -24,11 +19,6 @@ class Command(ABC): def identifier(cls) -> str: pass - @classmethod - @abstractmethod - def returns_data(cls) -> bool: - pass - @classmethod @abstractmethod def run(cls, context: CommandContext, request): diff --git a/aikido_zen/helpers/ipc/event_sender.py b/aikido_zen/helpers/ipc/event_sender.py new file mode 100644 index 000000000..b2c3dad3a --- /dev/null +++ b/aikido_zen/helpers/ipc/event_sender.py @@ -0,0 +1,73 @@ +import os +import threading +from queue import Full, Queue + +from aikido_zen.helpers.logging import logger + +MAX_PENDING_EVENTS = 100 +EVENT_DELIVERY_TIMEOUT_IN_SEC = 5 + + +class EventSender: + def __init__(self): + pid = os.getpid() + self._start_locks = {pid: threading.Lock()} + self._reset_queue(pid) + + def _reset_queue(self, pid): + self._queue = Queue(maxsize=MAX_PENDING_EVENTS) + self._thread = None + self._warned_full = False + self._pid = pid + + def enqueue(self, comms, action, request): + pid = os.getpid() + start_lock = self._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 = self._start_locks.setdefault(pid, threading.Lock()) + + with start_lock: + # Pending events belong to the parent, so a forked child must start empty. + if self._pid != pid: + self._reset_queue(pid) + self._start_locks = {pid: start_lock} + if self._thread is None: + thread = threading.Thread( + target=self._send_events, + args=(self._queue,), + name="aikido-event-sender", + daemon=True, + ) + thread.start() + self._thread = thread + pending = self._queue + + try: + pending.put_nowait((comms, action, request)) + return {"success": True, "queued": True} + except Full: + if not self._warned_full: + self._warned_full = True + logger.warning("Event delivery queue is full; dropping new events.") + return {"success": False, "error": "queue_full"} + + @staticmethod + def _send_events(pending): + while True: + item = pending.get() + try: + if item is None: + return + item[0].send_data_to_bg_process( + item[1], item[2], timeout_in_sec=EVENT_DELIVERY_TIMEOUT_IN_SEC + ) + 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() + + +event_sender = EventSender() diff --git a/aikido_zen/helpers/ipc/event_sender_test.py b/aikido_zen/helpers/ipc/event_sender_test.py new file mode 100644 index 000000000..d43d0f011 --- /dev/null +++ b/aikido_zen/helpers/ipc/event_sender_test.py @@ -0,0 +1,236 @@ +import gc +import os +import threading +import weakref +from unittest.mock import Mock + +import pytest + +from aikido_zen.background_process.comms import AikidoIPCCommunications +from aikido_zen.background_process import comms as communications +from aikido_zen.helpers.ipc import event_sender as module + + +@pytest.fixture +def sender(): + instance = module.EventSender() + yield instance + if instance._thread is not None: + instance._queue.put(None, timeout=5) + instance._thread.join(5) + assert not instance._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_events_return_while_delivery_is_stalled_and_use_one_worker(sender): + entered = threading.Event() + resume = threading.Event() + worker_ids = [] + + def deliver(*args, **kwargs): + worker_ids.append(threading.get_ident()) + entered.set() + assert resume.wait(5) + + comms = Mock() + comms.send_data_to_bg_process.side_effect = deliver + try: + assert sender.enqueue(comms, "put_event", 0) == { + "success": True, + "queued": True, + } + assert entered.wait(5) + first_worker = sender._thread + for index in range(1, 20): + assert sender.enqueue(comms, "put_event", index)["success"] + assert comms.send_data_to_bg_process.call_count == 1 + assert sender._thread is first_worker + finally: + resume.set() + assert_drained(sender) + assert worker_ids == [first_worker.ident] * 20 + assert [ + call.args[1] for call in comms.send_data_to_bg_process.call_args_list + ] == list(range(20)) + assert all( + call.kwargs["timeout_in_sec"] == 5 + for call in comms.send_data_to_bg_process.call_args_list + ) + + +def test_queue_full_drops_new_events_without_waiting(sender, caplog): + class Payload: + pass + + entered = threading.Event() + resume = threading.Event() + comms = Mock() + + def deliver(*args, **kwargs): + entered.set() + assert resume.wait(5) + + comms.send_data_to_bg_process.side_effect = deliver + try: + sender.enqueue(comms, "put_event", 0) + assert entered.wait(5) + for index in range(module.MAX_PENDING_EVENTS): + assert sender.enqueue(comms, "put_event", index + 1)["success"] + for _ in range(3): + assert sender.enqueue(comms, "put_event", "overflow") == { + "success": False, + "error": "queue_full", + } + dropped = Payload() + dropped_reference = weakref.ref(dropped) + assert not sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.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: + def send_data_to_bg_process(self, action, obj, timeout_in_sec): + delivered.append(weakref.ref(obj)) + if fail: + raise OSError("closed socket") + + comms = Comms() + for _ in range(20): + sender.enqueue(comms, "put_event", Payload()) + assert_drained(sender) + gc.collect() + assert len(delivered) == 20 + assert all(reference() is None for reference in delivered) + assert sender._thread.is_alive() + + +def test_invalid_key_does_not_start_the_event_sender(monkeypatch): + enqueue = Mock() + monkeypatch.setattr(module.event_sender, "enqueue", enqueue) + comms = AikidoIPCCommunications(("127.0.0.1", 1), None) + assert comms.enqueue_event("put_event", None) == { + "success": False, + "error": "invalid_key", + } + enqueue.assert_not_called() + + +def test_worker_start_failure_does_not_escape_into_the_application(sender, monkeypatch): + monkeypatch.setattr(communications, "event_sender", sender) + monkeypatch.setattr( + module.threading.Thread, "start", Mock(side_effect=RuntimeError("no threads")) + ) + comms = AikidoIPCCommunications(("127.0.0.1", 1), b"fixture") + assert comms.enqueue_event("put_event", 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", "parent event")) + old_lock = sender._start_locks[os.getpid()] + old_lock.acquire() + sender._pid = os.getpid() - 1 + sender._start_locks = {sender._pid: old_lock} + comms = Mock() + try: + assert sender.enqueue(comms, "put_event", "child event")["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 + comms.send_data_to_bg_process.assert_called_once_with( + "put_event", "child event", timeout_in_sec=5 + ) + + +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) + comms = Mock() + results = [] + + def enqueue(index): + barrier.wait(5) + results.append(sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.call_count == 20 + + +def test_concurrent_producers_start_only_one_worker(sender): + barrier = threading.Barrier(21) + workers = set() + comms = Mock() + comms.send_data_to_bg_process.side_effect = lambda *args, **kwargs: workers.add( + threading.get_ident() + ) + + def enqueue(): + barrier.wait(5) + sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.call_count == 20 diff --git a/aikido_zen/helpers/ipc/send_payload.py b/aikido_zen/helpers/ipc/send_payload.py index 9b5346ec3..230045df9 100644 --- a/aikido_zen/helpers/ipc/send_payload.py +++ b/aikido_zen/helpers/ipc/send_payload.py @@ -2,12 +2,8 @@ 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( +def send_payload(comms: AikidoIPCCommunications, payload: Payload): + return comms.enqueue_event( payload.identifier, payload.request, - payload.returns_data, - timeout_in_sec=timeout_in_sec, ) diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index 4a7fad1f9..d680a3110 100644 --- a/aikido_zen/middleware/init_test.py +++ b/aikido_zen/middleware/init_test.py @@ -141,7 +141,6 @@ def test_cache_comms_with_endpoints(): "group": "my_group", "remote_address": "1.1.1.1", }, - receive=True, timeout_in_sec=0.01, ) diff --git a/aikido_zen/middleware/should_block_request.py b/aikido_zen/middleware/should_block_request.py index 434f06a67..f1cb52a40 100644 --- a/aikido_zen/middleware/should_block_request.py +++ b/aikido_zen/middleware/should_block_request.py @@ -59,7 +59,6 @@ def should_block_request(): "group": context.rate_limit_group, "remote_address": context.remote_address, }, - receive=True, timeout_in_sec=(10 / 1000), # 10ms ) if ratelimit_res["success"] and ratelimit_res["data"]["block"]: diff --git a/aikido_zen/sources/functions/request_handler.py b/aikido_zen/sources/functions/request_handler.py index 4c56c9cf9..8ac339021 100644 --- a/aikido_zen/sources/functions/request_handler.py +++ b/aikido_zen/sources/functions/request_handler.py @@ -71,7 +71,6 @@ def pre_response(): "ip": context.remote_address, "user-agent": context.get_user_agent(), }, - receive=True, timeout_in_sec=(10 / 1000), ) if not check_fw_lists_res["success"] or not check_fw_lists_res["data"]["blocked"]: diff --git a/aikido_zen/sources/functions/request_handler_test.py b/aikido_zen/sources/functions/request_handler_test.py index 45e9239d9..6a6715b41 100644 --- a/aikido_zen/sources/functions/request_handler_test.py +++ b/aikido_zen/sources/functions/request_handler_test.py @@ -46,7 +46,7 @@ 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): + def send_data_to_bg_process(self, action, obj, timeout_in_sec=0.1): if action != "CHECK_FIREWALL_LISTS": return {"success": False} res = process_check_firewall_lists(self.conn_manager, obj) diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 2aa829523..9349e807f 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -86,7 +86,6 @@ def renew(self): res = comms.get_comms().send_data_to_bg_process( action="SYNC_DATA", obj=payload, - receive=True, ) 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..fb95b0d25 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -49,7 +49,7 @@ def sync(monkeypatch): 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: { + comms.send_data_to_bg_process.side_effect = lambda action, obj: { "success": True, "data": deepcopy(process_sync_data(manager, deepcopy(obj))), } @@ -464,7 +464,6 @@ def test_renew_called_with_correct_args(mock_get_comms, thread_cache: ThreadCach } ], }, - receive=True, ) @@ -524,7 +523,6 @@ def test_sync_data_for_users(mock_get_comms, thread_cache: ThreadCache): }, ], }, - receive=True, ) @@ -561,7 +559,6 @@ def test_renew_called_with_empty_routes(mock_get_comms, thread_cache: ThreadCach "ai_stats": [], "packages": [], }, - receive=True, ) @@ -642,10 +639,8 @@ def test_sync_counts_each_request_once( 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(action, obj): + response = {"success": False} if first_sync_fails else deliver(action, obj) if concurrent_path: context.route, context.body = concurrent_path, {"during": "sync"} cache.stats.increment_total_hits() @@ -693,11 +688,11 @@ def test_overlapping_syncs_do_not_replay_hits(sync): deliver = comms.send_data_to_bg_process.side_effect post_response(200) - def first_sync(action, obj, receive): + def first_sync(action, obj): comms.send_data_to_bg_process.side_effect = deliver post_response(200) cache.renew() - return deliver(action, obj, receive) + return deliver(action, obj) comms.send_data_to_bg_process.side_effect = first_sync cache.renew() @@ -711,8 +706,8 @@ def test_config_is_sent_once_per_worker_per_cloud_response(sync): deliver = comms.send_data_to_bg_process.side_effect replies = [] - def record_reply(action, obj, receive): - response = deliver(action, obj, receive) + def record_reply(action, obj): + response = deliver(action, obj) replies.append(response["data"]) return response @@ -747,8 +742,8 @@ def test_config_is_retried_after_lost_reply(sync): cache, manager, comms, _, _ = sync deliver = comms.send_data_to_bg_process.side_effect - def lose_reply(action, obj, receive): - assert "config" in deliver(action, obj, receive)["data"] + def lose_reply(action, obj): + assert "config" in deliver(action, obj)["data"] return {"success": False} comms.send_data_to_bg_process.side_effect = lose_reply @@ -820,8 +815,8 @@ def test_sampling_reset_preserves_requests_collected_during_sync( manager.send_heartbeat() deliver = comms.send_data_to_bg_process.side_effect - def request_during_sync(action, obj, receive): - response = deliver(action, obj, receive) + def request_during_sync(action, obj): + response = deliver(action, obj) context.route, context.body = concurrent_path, {"during": "sync"} for _ in range(requests_during_sync): post_response(200) @@ -941,5 +936,4 @@ def test_renew_called_with_no_requests(mock_get_comms, thread_cache: ThreadCache "ai_stats": [], "packages": [], }, - receive=True, ) diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py index a24c5a4d4..fb314cf16 100644 --- a/aikido_zen/vulnerabilities/init_test.py +++ b/aikido_zen/vulnerabilities/init_test.py @@ -143,8 +143,8 @@ def test_sql_injection_with_comms(caplog, get_context, monkeypatch): 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] + mock_comms.enqueue_event.assert_called_once() + call_args = mock_comms.enqueue_event.call_args[0] assert call_args[0] == "put_event" assert call_args[1].event["request"] == { "ipAddress": "198.51.100.23", From 741671ea17a042d5a427416c4d2ded79d29290f1 Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 12:58:19 +0200 Subject: [PATCH 2/8] Unify IPC messaging and command definitions --- aikido_zen/background_process/__init__.py | 17 +- .../background_process/commands/__init__.py | 35 +- .../commands/check_firewall_lists.py | 19 +- .../background_process/commands/ping.py | 11 +- .../background_process/commands/put_event.py | 29 +- .../commands/read_property.py | 19 +- .../commands/read_property_test.py | 12 +- .../commands/should_ratelimit.py | 36 +- .../commands/should_ratelimit_test.py | 22 +- .../background_process/commands/sync_data.py | 33 +- .../commands/sync_data_test.py | 12 +- aikido_zen/background_process/comms.py | 138 ----- aikido_zen/background_process/init_test.py | 4 +- aikido_zen/background_process/ipc.py | 113 +++++ .../{comms_test.py => ipc_test.py} | 400 ++++++++++++--- .../background_process/packages_test.py | 38 +- .../background_process/socket_connection.py | 61 +++ aikido_zen/context/track.py | 9 +- aikido_zen/context/track_test.py | 60 +-- aikido_zen/decorators/uwsgi_test.py | 4 +- aikido_zen/helpers/blocking_enabled.py | 11 +- aikido_zen/helpers/ipc/command_types.py | 31 +- aikido_zen/helpers/ipc/event_sender.py | 73 --- aikido_zen/helpers/ipc/event_sender_test.py | 236 --------- aikido_zen/helpers/ipc/send_payload.py | 9 - aikido_zen/init_test.py | 6 +- aikido_zen/middleware/init_test.py | 60 +-- aikido_zen/middleware/should_block_request.py | 23 +- aikido_zen/sinks/tests/asyncpg_test.py | 16 +- .../sinks/tests/clickhouse_driver_test.py | 4 +- aikido_zen/sinks/tests/motor_test.py | 34 +- aikido_zen/sinks/tests/mysqlclient_test.py | 10 +- aikido_zen/sinks/tests/psycopg2_test.py | 18 +- aikido_zen/sinks/tests/psycopg_test.py | 18 +- aikido_zen/sinks/tests/pymongo_test.py | 36 +- aikido_zen/sinks/tests/pymysql_test.py | 14 +- .../sinks/tests/requests_and_urllib3_test.py | 8 +- aikido_zen/sinks/tests/sqlite3_test.py | 24 +- aikido_zen/sources/flask/flask_test.py | 12 +- .../sources/functions/request_handler.py | 24 +- .../sources/functions/request_handler_test.py | 38 +- .../thread/process_worker_loader_test.py | 4 +- aikido_zen/thread/process_worker_test.py | 71 ++- aikido_zen/thread/thread_cache.py | 11 +- aikido_zen/thread/thread_cache_test.py | 477 +++++++++--------- aikido_zen/vulnerabilities/__init__.py | 11 +- aikido_zen/vulnerabilities/init_test.py | 37 +- 47 files changed, 1200 insertions(+), 1188 deletions(-) delete mode 100644 aikido_zen/background_process/comms.py create mode 100644 aikido_zen/background_process/ipc.py rename aikido_zen/background_process/{comms_test.py => ipc_test.py} (51%) create mode 100644 aikido_zen/background_process/socket_connection.py delete mode 100644 aikido_zen/helpers/ipc/event_sender.py delete mode 100644 aikido_zen/helpers/ipc/event_sender_test.py delete mode 100644 aikido_zen/helpers/ipc/send_payload.py diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 24dfac283..160749a4e 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -14,11 +14,8 @@ from aikido_zen.helpers.get_agent_session_id import get_agent_session_id from aikido_zen.helpers.logging import logger -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 @@ -58,10 +55,10 @@ def start_background_process(): # 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,7 +69,7 @@ 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, ) @@ -87,8 +84,8 @@ def get_uds_filename(): return f"{temp_dir}/{prefix}_{name}.sock" -def background_process_already_active(comms): - res = comms.send_data_to_bg_process(action="PING", obj=tuple()) +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/commands/__init__.py b/aikido_zen/background_process/commands/__init__.py index 3a43943d2..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: - result = cmd.run(CommandContext(connection_manager, queue), inbound_request) - return conn.send(result) + 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 bf4f3aa40..f187069c9 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -1,23 +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" +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 ecc22de94..000000000 --- a/aikido_zen/background_process/comms.py +++ /dev/null @@ -1,138 +0,0 @@ -""" -Holds the globally stored comms object -Exports the AikidoIPCCommunications class -""" - -import io -import multiprocessing.connection as con -import socket -import time -from aikido_zen.helpers.logging import logger -from aikido_zen.helpers.ipc.event_sender import event_sender - -# pylint: disable=invalid-name # This variable does change -comms = None - - -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)) - - -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 enqueue_event(self, action, obj): - if not self.key: - return {"success": False, "error": "invalid_key"} - try: - return event_sender.enqueue(self, action, obj) - except Exception as error: - logger.debug("Could not enqueue event: %s", error) - return {"success": False, "error": "unknown"} - - def send_data_to_bg_process(self, action, obj, timeout_in_sec=0.1): - """Try-catched send_data_to_bg_process""" - try: - return self._send_data_to_bg_process(action, obj, 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, 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"} - - if timeout_in_sec <= 0: - return {"success": False, "error": "timeout"} - deadline = time.monotonic() + timeout_in_sec - - with socket.socket(getattr(socket, con.address_type(self.address))) as sock: - try: - connect_with_deadline(sock, self.address, deadline) - with SocketConnection(sock, deadline) as conn: - conn.send((action, obj)) - result = conn.recv() - if time.monotonic() >= deadline: - return {"success": False, "error": "timeout"} - except Exception as e: - logger.debug("Exception occurred in IPC call : %s", e) - return {"success": False, "error": "timeout"} - return {"success": True, "data": result} diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 702bc7cae..ab52057a9 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)]) @@ -160,7 +160,7 @@ def test_stale_socket_removed_by_another_worker(monkeypatch, mocker): ) 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, diff --git a/aikido_zen/background_process/ipc.py b/aikido_zen/background_process/ipc.py new file mode 100644 index 000000000..63245701b --- /dev/null +++ b/aikido_zen/background_process/ipc.py @@ -0,0 +1,113 @@ +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() + if time.monotonic() >= deadline: + return {"success": False, "error": "timeout"} + 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/comms_test.py b/aikido_zen/background_process/ipc_test.py similarity index 51% rename from aikido_zen/background_process/comms_test.py rename to aikido_zen/background_process/ipc_test.py index 42f361581..4af5a0b88 100644 --- a/aikido_zen/background_process/comms_test.py +++ b/aikido_zen/background_process/ipc_test.py @@ -1,4 +1,5 @@ import gc +import os import socket import struct import threading @@ -9,18 +10,20 @@ 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 comms as module +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 ping, put_event, read_property from aikido_zen.background_process.commands import process_incoming_command -from aikido_zen.background_process.commands.put_event import PutEventReq @pytest.fixture(autouse=True) -def reset_comms_after_test(): +def reset_ipc_after_test(): yield - module.reset_comms() + module.IPCClient.reset() class Payload: @@ -47,18 +50,18 @@ def sendall(self, *args, **kwargs): raise monkeypatch.setattr( - module, + transport, "socket", SimpleNamespace(**{**vars(socket), "socket": TrackedSocket}), raising=False, ) - class TrackedConnection(module.SocketConnection): + class TrackedConnection(transport.SocketConnection): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) connections.append(self) - monkeypatch.setattr(module, "SocketConnection", TrackedConnection) + monkeypatch.setattr(transport, "SocketConnection", TrackedConnection) return SimpleNamespace( sockets=sockets, connections=connections, @@ -151,25 +154,52 @@ def full_accept_queue(listener): yield queued -def test_comms_init(): +@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" - comms = module.AikidoIPCCommunications(address, key) + ipc = module.IPCClient(address, key) - assert comms.address == address - assert comms.key == 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( - "action,payload,expected", + "message,expected", [ - ("PING", (), "Received"), - ("READ_PROPERTY", "block", True), - ("put_event", PutEventReq({"name": "login.failed"}), None), + (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, action, payload, expected + resources, servers, monkeypatch, message, expected ): events = Queue() @@ -183,11 +213,12 @@ def fail_start(self): raise AssertionError("IPC call must not start a thread") monkeypatch.setattr(threading.Thread, "start", fail_start) - communications = module.AikidoIPCCommunications(server.address, b"fixture") - assert communications.send_data_to_bg_process( - action, payload, timeout_in_sec=1 - ) == {"success": True, "data": expected} - if action == "put_event": + 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, []) @@ -199,14 +230,17 @@ def fail_start(self): ) def test_timeout_releases_sockets_and_payloads(resources, servers, reply): server = servers(5, reply=reply) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") payloads = [] for _ in range(5): payload = Payload() payloads.append(weakref.ref(payload)) - assert communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=0.05 - ) == {"success": False, "error": "timeout"} + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.05) + ) == { + "success": False, + "error": "timeout", + } del payload assert server.finished.wait(1) @@ -233,13 +267,11 @@ def disconnect(peer, _request): peer.close() server = servers(1, reply=disconnect) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") payload = Payload() reference = weakref.ref(payload) - result = communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=1 - ) + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) assert not result["success"] del payload assert server.finished.wait(1) @@ -250,15 +282,13 @@ def test_send_failure_releases_sockets_and_payloads(resources, servers, monkeypa def fail_send(self, data): raise BrokenPipeError("Peer closed during send") - monkeypatch.setattr(module.SocketConnection, "_send", fail_send) + monkeypatch.setattr(transport.SocketConnection, "_send", fail_send) server = servers(1) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") payload = Payload() reference = weakref.ref(payload) - result = communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=1 - ) + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) assert not result["success"] del payload assert_released(resources, [reference]) @@ -267,10 +297,8 @@ def fail_send(self, data): def test_connect_failure_closes_socket(resources): with socket.socket() as reserved: reserved.bind(("127.0.0.1", 0)) - communications = module.AikidoIPCCommunications( - reserved.getsockname(), b"fixture" - ) - result = communications.send_data_to_bg_process("PING", (), timeout_in_sec=0.1) + 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, []) @@ -281,13 +309,11 @@ def __reduce__(self): raise ValueError("Cannot serialize payload") server = servers(1) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") payload = Unpicklable() reference = weakref.ref(payload) - result = communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=1 - ) + result = communications.send_message(put_event.generate(payload, timeout_in_sec=1)) assert not result["success"] del payload assert_released(resources, [reference]) @@ -295,29 +321,32 @@ def __reduce__(self): def test_invalid_reply_releases_connection(resources, servers): server = servers(1, reply=struct.pack("!i", 4) + b"oops") - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") - result = communications.send_data_to_bg_process("PING", (), timeout_in_sec=1) + 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 = module.socket.socket + 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(module.socket, "socket", small_send_buffer) + monkeypatch.setattr(transport.socket, "socket", small_send_buffer) server = servers(1, read_request=False) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") payload = Payload(8 * 1024 * 1024) reference = weakref.ref(payload) - assert communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=0.2 - ) == {"success": False, "error": "timeout"} + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.2) + ) == { + "success": False, + "error": "timeout", + } del payload assert server.finished.wait(1) @@ -337,9 +366,9 @@ def reply(peer, _request): time.sleep(0.01) server = servers(1, reply=reply) - communications = module.AikidoIPCCommunications(server.address, b"fixture") + communications = module.IPCClient(server.address, b"fixture") started = time.monotonic() - assert communications.send_data_to_bg_process("PING", (), timeout_in_sec=0.05) == { + assert communications.send_message(ping.generate(timeout_in_sec=0.05)) == { "success": False, "error": "timeout", } @@ -349,11 +378,12 @@ def reply(peer, _request): @pytest.mark.parametrize("timeout", [0, -1]) def test_expired_call_does_not_open_a_socket(resources, timeout): - communications = module.AikidoIPCCommunications(("127.0.0.1", 1), b"fixture") + communications = module.IPCClient(("127.0.0.1", 1), b"fixture") - assert communications.send_data_to_bg_process( - "PING", (), timeout_in_sec=timeout - ) == {"success": False, "error": "timeout"} + assert communications.send_message(ping.generate(timeout_in_sec=timeout)) == { + "success": False, + "error": "timeout", + } assert not resources.sockets @@ -365,10 +395,13 @@ def reply(peer, _request): conn.send(response) server = servers(1, reply=reply) - communications = module.AikidoIPCCommunications(server.address, b"fixture") - assert communications.send_data_to_bg_process( - "READ_PROPERTY", "block", timeout_in_sec=2 - ) == {"success": True, "data": response} + 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, []) @@ -385,15 +418,13 @@ def test_timeout_releases_payloads_when_accept_queue_is_full( listener.bind(address) listener.listen(1) with full_accept_queue(listener): - communications = module.AikidoIPCCommunications( - listener.getsockname(), b"fixture" - ) + communications = module.IPCClient(listener.getsockname(), b"fixture") payloads = [] for _ in range(3): payload = Payload() payloads.append(weakref.ref(payload)) - assert communications.send_data_to_bg_process( - "put_event", payload, timeout_in_sec=0.05 + assert communications.send_message( + put_event.generate(payload, timeout_in_sec=0.05) ) == {"success": False, "error": "timeout"} del payload assert_released(resources, payloads) @@ -416,7 +447,7 @@ def retry_sleep(duration): time.sleep(duration) monkeypatch.setattr( - module, "time", SimpleNamespace(**{**vars(time), "sleep": retry_sleep}) + transport, "time", SimpleNamespace(**{**vars(time), "sleep": retry_sleep}) ) def respond(): @@ -437,10 +468,11 @@ def respond(): server = threading.Thread(target=respond, daemon=True) server.start() try: - communications = module.AikidoIPCCommunications(address, b"fixture") - assert communications.send_data_to_bg_process( - "PING", (), timeout_in_sec=1 - ) == {"success": True, "data": "Received"} + 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: @@ -467,11 +499,11 @@ def respond(): if allow_first_reply.wait(3): first_connection.send(getattr(properties, first_request[1])) - communications = module.AikidoIPCCommunications(listener.address, b"fixture") + communications = module.IPCClient(listener.address, b"fixture") def request(name): - outcomes[name] = communications.send_data_to_bg_process( - "READ_PROPERTY", name, timeout_in_sec=3 + outcomes[name] = communications.send_message( + read_property.generate(name, timeout_in_sec=3) ) completion_order.append(name) if name == "second": @@ -502,3 +534,225 @@ def request(name): 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 f613af77b..91cf27749 100644 --- a/aikido_zen/background_process/packages_test.py +++ b/aikido_zen/background_process/packages_test.py @@ -4,52 +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 - def send_data_to_bg_process(self, action, obj, timeout_in_sec=0.1): - self.action = action - self.obj = obj + 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, @@ -60,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/context/track.py b/aikido_zen/context/track.py index 2cce0545b..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)) + 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 b313ca7d9..48a68a3d6 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.enqueue_event.called: + if not ipc.send_message_async.called: return None - return comms.enqueue_event.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,22 +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.enqueue_event.assert_called_once() - identifier, request = comms.enqueue_event.call_args[0] - assert identifier == "put_event" - 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", @@ -164,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.enqueue_event.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(): @@ -189,16 +189,18 @@ 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.enqueue_event.call_args_list] + names = [ + 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 @@ -232,11 +234,11 @@ 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") @@ -246,29 +248,29 @@ def test_a_bypassed_ip_does_not_use_up_the_limit(): 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.enqueue_event.assert_called_once() - comms.send_data_to_bg_process.assert_not_called() + 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.enqueue_event.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 46dce63ef..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" - ) + 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 364954bbf..bcc898303 100644 --- a/aikido_zen/helpers/ipc/command_types.py +++ b/aikido_zen/helpers/ipc/command_types.py @@ -1,30 +1,5 @@ -from abc import ABC, abstractmethod - - -class CommandContext: - def __init__(self, connection_manager, queue): - self.connection_manager = connection_manager - self.queue = queue - - class Payload: - def __init__(self, command, request): - self.identifier = command.identifier() + 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 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/event_sender.py b/aikido_zen/helpers/ipc/event_sender.py deleted file mode 100644 index b2c3dad3a..000000000 --- a/aikido_zen/helpers/ipc/event_sender.py +++ /dev/null @@ -1,73 +0,0 @@ -import os -import threading -from queue import Full, Queue - -from aikido_zen.helpers.logging import logger - -MAX_PENDING_EVENTS = 100 -EVENT_DELIVERY_TIMEOUT_IN_SEC = 5 - - -class EventSender: - def __init__(self): - pid = os.getpid() - self._start_locks = {pid: threading.Lock()} - self._reset_queue(pid) - - def _reset_queue(self, pid): - self._queue = Queue(maxsize=MAX_PENDING_EVENTS) - self._thread = None - self._warned_full = False - self._pid = pid - - def enqueue(self, comms, action, request): - pid = os.getpid() - start_lock = self._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 = self._start_locks.setdefault(pid, threading.Lock()) - - with start_lock: - # Pending events belong to the parent, so a forked child must start empty. - if self._pid != pid: - self._reset_queue(pid) - self._start_locks = {pid: start_lock} - if self._thread is None: - thread = threading.Thread( - target=self._send_events, - args=(self._queue,), - name="aikido-event-sender", - daemon=True, - ) - thread.start() - self._thread = thread - pending = self._queue - - try: - pending.put_nowait((comms, action, request)) - return {"success": True, "queued": True} - except Full: - if not self._warned_full: - self._warned_full = True - logger.warning("Event delivery queue is full; dropping new events.") - return {"success": False, "error": "queue_full"} - - @staticmethod - def _send_events(pending): - while True: - item = pending.get() - try: - if item is None: - return - item[0].send_data_to_bg_process( - item[1], item[2], timeout_in_sec=EVENT_DELIVERY_TIMEOUT_IN_SEC - ) - 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() - - -event_sender = EventSender() diff --git a/aikido_zen/helpers/ipc/event_sender_test.py b/aikido_zen/helpers/ipc/event_sender_test.py deleted file mode 100644 index d43d0f011..000000000 --- a/aikido_zen/helpers/ipc/event_sender_test.py +++ /dev/null @@ -1,236 +0,0 @@ -import gc -import os -import threading -import weakref -from unittest.mock import Mock - -import pytest - -from aikido_zen.background_process.comms import AikidoIPCCommunications -from aikido_zen.background_process import comms as communications -from aikido_zen.helpers.ipc import event_sender as module - - -@pytest.fixture -def sender(): - instance = module.EventSender() - yield instance - if instance._thread is not None: - instance._queue.put(None, timeout=5) - instance._thread.join(5) - assert not instance._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_events_return_while_delivery_is_stalled_and_use_one_worker(sender): - entered = threading.Event() - resume = threading.Event() - worker_ids = [] - - def deliver(*args, **kwargs): - worker_ids.append(threading.get_ident()) - entered.set() - assert resume.wait(5) - - comms = Mock() - comms.send_data_to_bg_process.side_effect = deliver - try: - assert sender.enqueue(comms, "put_event", 0) == { - "success": True, - "queued": True, - } - assert entered.wait(5) - first_worker = sender._thread - for index in range(1, 20): - assert sender.enqueue(comms, "put_event", index)["success"] - assert comms.send_data_to_bg_process.call_count == 1 - assert sender._thread is first_worker - finally: - resume.set() - assert_drained(sender) - assert worker_ids == [first_worker.ident] * 20 - assert [ - call.args[1] for call in comms.send_data_to_bg_process.call_args_list - ] == list(range(20)) - assert all( - call.kwargs["timeout_in_sec"] == 5 - for call in comms.send_data_to_bg_process.call_args_list - ) - - -def test_queue_full_drops_new_events_without_waiting(sender, caplog): - class Payload: - pass - - entered = threading.Event() - resume = threading.Event() - comms = Mock() - - def deliver(*args, **kwargs): - entered.set() - assert resume.wait(5) - - comms.send_data_to_bg_process.side_effect = deliver - try: - sender.enqueue(comms, "put_event", 0) - assert entered.wait(5) - for index in range(module.MAX_PENDING_EVENTS): - assert sender.enqueue(comms, "put_event", index + 1)["success"] - for _ in range(3): - assert sender.enqueue(comms, "put_event", "overflow") == { - "success": False, - "error": "queue_full", - } - dropped = Payload() - dropped_reference = weakref.ref(dropped) - assert not sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.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: - def send_data_to_bg_process(self, action, obj, timeout_in_sec): - delivered.append(weakref.ref(obj)) - if fail: - raise OSError("closed socket") - - comms = Comms() - for _ in range(20): - sender.enqueue(comms, "put_event", Payload()) - assert_drained(sender) - gc.collect() - assert len(delivered) == 20 - assert all(reference() is None for reference in delivered) - assert sender._thread.is_alive() - - -def test_invalid_key_does_not_start_the_event_sender(monkeypatch): - enqueue = Mock() - monkeypatch.setattr(module.event_sender, "enqueue", enqueue) - comms = AikidoIPCCommunications(("127.0.0.1", 1), None) - assert comms.enqueue_event("put_event", None) == { - "success": False, - "error": "invalid_key", - } - enqueue.assert_not_called() - - -def test_worker_start_failure_does_not_escape_into_the_application(sender, monkeypatch): - monkeypatch.setattr(communications, "event_sender", sender) - monkeypatch.setattr( - module.threading.Thread, "start", Mock(side_effect=RuntimeError("no threads")) - ) - comms = AikidoIPCCommunications(("127.0.0.1", 1), b"fixture") - assert comms.enqueue_event("put_event", 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", "parent event")) - old_lock = sender._start_locks[os.getpid()] - old_lock.acquire() - sender._pid = os.getpid() - 1 - sender._start_locks = {sender._pid: old_lock} - comms = Mock() - try: - assert sender.enqueue(comms, "put_event", "child event")["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 - comms.send_data_to_bg_process.assert_called_once_with( - "put_event", "child event", timeout_in_sec=5 - ) - - -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) - comms = Mock() - results = [] - - def enqueue(index): - barrier.wait(5) - results.append(sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.call_count == 20 - - -def test_concurrent_producers_start_only_one_worker(sender): - barrier = threading.Barrier(21) - workers = set() - comms = Mock() - comms.send_data_to_bg_process.side_effect = lambda *args, **kwargs: workers.add( - threading.get_ident() - ) - - def enqueue(): - barrier.wait(5) - sender.enqueue(comms, "put_event", 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 comms.send_data_to_bg_process.call_count == 20 diff --git a/aikido_zen/helpers/ipc/send_payload.py b/aikido_zen/helpers/ipc/send_payload.py deleted file mode 100644 index 230045df9..000000000 --- a/aikido_zen/helpers/ipc/send_payload.py +++ /dev/null @@ -1,9 +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): - return comms.enqueue_event( - payload.identifier, - payload.request, - ) diff --git a/aikido_zen/init_test.py b/aikido_zen/init_test.py index 15087a6d9..18cb259de 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(): diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index d680a3110..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,45 +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/", }, - 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 f1cb52a40..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,15 +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, - }, - 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 8ac339021..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,14 +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(), - }, - 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 @@ -114,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 6a6715b41..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, 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_test.py b/aikido_zen/thread/process_worker_loader_test.py index 55378f272..ecc70e17f 100644 --- a/aikido_zen/thread/process_worker_loader_test.py +++ b/aikido_zen/thread/process_worker_loader_test.py @@ -54,7 +54,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(): 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 9349e807f..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,10 +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, - ) + 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 fb95b0d25..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: { + 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,83 +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, - } - ], }, - ) + "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 @@ -484,82 +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": {}, }, - ) + "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": {}, }, - ) + "middleware_installed": False, + "hostnames": [], + "users": [], + "ai_stats": [], + "packages": [], + } def test_sync_includes_updates_to_routes_already_in_payload(sync): @@ -633,23 +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): - response = {"success": False} if first_sync_fails else deliver(action, obj) + 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() @@ -684,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): - 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) + 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): - response = deliver(action, obj) + 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: @@ -735,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): - assert "config" in deliver(action, obj)["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]) @@ -765,7 +765,7 @@ def lose_reply(action, obj): 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, @@ -780,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) @@ -806,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): - response = deliver(action, obj) + 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) @@ -831,7 +831,7 @@ def request_during_sync(action, obj): 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 @@ -842,14 +842,14 @@ def request_during_sync(action, obj): @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() @@ -858,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() @@ -887,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() @@ -900,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"}) @@ -912,28 +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": {}, }, - ) + "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 fb314cf16..f4ea9d78d 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.enqueue_event.assert_called_once() - call_args = mock_comms.enqueue_event.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", From 8236158f0c6dd3977843b5ad2fc2485e4a070660 Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 13:49:11 +0200 Subject: [PATCH 3/8] Align event IPC identifier with other commands --- .../background_process/aikido_background_process_test.py | 4 ++-- aikido_zen/background_process/commands/put_event.py | 2 +- aikido_zen/background_process/ipc_test.py | 4 ++-- aikido_zen/context/track_test.py | 2 +- aikido_zen/vulnerabilities/init_test.py | 2 +- 5 files changed, 7 insertions(+), 7 deletions(-) diff --git a/aikido_zen/background_process/aikido_background_process_test.py b/aikido_zen/background_process/aikido_background_process_test.py index e3f4efd80..2e4fe8be2 100644 --- a/aikido_zen/background_process/aikido_background_process_test.py +++ b/aikido_zen/background_process/aikido_background_process_test.py @@ -19,7 +19,7 @@ def test_event_is_acknowledged_after_enqueueing(): event = {"name": "login.failed"} process_incoming_command( - None, ("put_event", PutEventReq(event)), connection, events + None, ("PUT_EVENT", PutEventReq(event)), connection, events ) assert operations == [("queued", event), ("reply", None)] @@ -34,7 +34,7 @@ def test_event_is_not_acknowledged_when_enqueueing_fails(): with pytest.raises(RuntimeError, match="Queue unavailable"): process_incoming_command( None, - ("put_event", PutEventReq({"name": "login.failed"})), + ("PUT_EVENT", PutEventReq({"name": "login.failed"})), connection, events, ) diff --git a/aikido_zen/background_process/commands/put_event.py b/aikido_zen/background_process/commands/put_event.py index f187069c9..0da133588 100644 --- a/aikido_zen/background_process/commands/put_event.py +++ b/aikido_zen/background_process/commands/put_event.py @@ -1,6 +1,6 @@ from aikido_zen.helpers.ipc.command_types import Payload -IDENTIFIER = "put_event" +IDENTIFIER = "PUT_EVENT" def generate(event, timeout_in_sec=5) -> Payload: diff --git a/aikido_zen/background_process/ipc_test.py b/aikido_zen/background_process/ipc_test.py index 4af5a0b88..9af487846 100644 --- a/aikido_zen/background_process/ipc_test.py +++ b/aikido_zen/background_process/ipc_test.py @@ -218,7 +218,7 @@ def fail_start(self): "success": True, "data": expected, } - if message.identifier == "put_event": + if message.identifier == "PUT_EVENT": assert events.get_nowait() == {"name": "login.failed"} assert_released(resources, []) @@ -570,7 +570,7 @@ def deliver(*args, **kwargs): 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 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) diff --git a/aikido_zen/context/track_test.py b/aikido_zen/context/track_test.py index 48a68a3d6..508c0601c 100644 --- a/aikido_zen/context/track_test.py +++ b/aikido_zen/context/track_test.py @@ -105,7 +105,7 @@ def test_track_sends_event_over_ipc(): ipc.send_message_async.assert_called_once() (message,) = ipc.send_message_async.call_args.args - assert message.identifier == "put_event" + assert message.identifier == "PUT_EVENT" assert message.request.event == { "type": "custom", "time": 111, diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py index f4ea9d78d..eed2c0370 100644 --- a/aikido_zen/vulnerabilities/init_test.py +++ b/aikido_zen/vulnerabilities/init_test.py @@ -146,7 +146,7 @@ def test_sql_injection_with_ipc(caplog, get_context, monkeypatch): ) mock_ipc.send_message_async.assert_called_once() (message,) = mock_ipc.send_message_async.call_args.args - assert message.identifier == "put_event" + assert message.identifier == "PUT_EVENT" assert message.request.event["request"] == { "ipAddress": "198.51.100.23", "method": "GET", From a4e51de65fec9bce72daf79f05a49d3286b80ee5 Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 15:41:27 +0200 Subject: [PATCH 4/8] Isolate IPC sockets and preserve completed replies --- aikido_zen/background_process/__init__.py | 5 +- aikido_zen/background_process/init_test.py | 82 ++++++++++++++++++++++ aikido_zen/background_process/ipc.py | 2 - aikido_zen/background_process/ipc_test.py | 30 +++++++- 4 files changed, 115 insertions(+), 4 deletions(-) diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 160749a4e..041d3f7aa 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 @@ -80,7 +81,9 @@ def get_uds_filename(): """Returns the address for UDS file""" temp_dir = get_temp_dir() prefix = "aikido_python" - name = hash_aikido_token() + # This key is shared by workers and changes between application runs. + identity = hash_aikido_token().encode() + multiprocessing.current_process().authkey + name = hashlib.sha256(identity).hexdigest() return f"{temp_dir}/{prefix}_{name}.sock" diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index ab52057a9..9d38e3c7b 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -1,3 +1,5 @@ +import json +import multiprocessing import subprocess import sys @@ -149,6 +151,86 @@ def test_get_process_factory_does_not_set_global_start_method(): subprocess.run([sys.executable, "-c", script], check=True) +def test_socket_name_is_stable_and_separates_tokens(monkeypatch): + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + current = background_process.get_uds_filename() + + assert current == background_process.get_uds_filename() + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_OTHER") + assert current != background_process.get_uds_filename() + + +@pytest.mark.parametrize("method", multiprocessing.get_all_start_methods()) +def test_socket_name_is_shared_by_workers_and_unique_to_each_run( + monkeypatch, tmp_path, method +): + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + script = tmp_path / "worker_socket_names.py" + script.write_text( + """ +import json +import multiprocessing +import sys + +def worker(pipe): + from aikido_zen.background_process import get_uds_filename + pipe.send(get_uds_filename()) + pipe.close() + +if __name__ == "__main__": + context = multiprocessing.get_context(sys.argv[1]) + readers, processes = [], [] + for _ in range(2): + reader, writer = context.Pipe(duplex=False) + process = context.Process(target=worker, args=(writer,)) + process.start() + writer.close() + readers.append(reader) + processes.append(process) + for process in processes: + process.join(10) + assert process.exitcode == 0 + from aikido_zen.background_process import get_uds_filename + names = [get_uds_filename()] + [reader.recv() for reader in readers] + print(json.dumps(names)) +""", + encoding="utf-8", + ) + runs = [ + json.loads( + subprocess.check_output( + [sys.executable, str(script), method], text=True, timeout=25 + ) + ) + for _ in range(2) + ] + + assert all(len(set(names)) == 1 for names in runs) + assert runs[0][0] != runs[1][0] + + +def test_startup_does_not_reuse_a_legacy_socket(monkeypatch, mocker): + factory = mocker.patch.object( + background_process, "get_process_factory" + ).return_value + active = mocker.patch.object( + background_process, "background_process_already_active", return_value=True + ) + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + monkeypatch.setenv("AIKIDO_TMP_DIR", "/tmp") + monkeypatch.setattr(background_process.platform, "system", lambda: "Linux") + legacy = f"/tmp/aikido_python_{background_process.hash_aikido_token()}.sock" + monkeypatch.setattr( + background_process.os.path, "exists", lambda path: path == legacy + ) + + background_process.start_background_process() + + active.assert_not_called() + assert factory.call_args.kwargs["args"][0] != legacy + factory.return_value.start.assert_called_once_with() + + def test_stale_socket_removed_by_another_worker(monkeypatch, mocker): process = mocker.patch( "aikido_zen.background_process.get_process_factory" diff --git a/aikido_zen/background_process/ipc.py b/aikido_zen/background_process/ipc.py index 63245701b..d5581b3b6 100644 --- a/aikido_zen/background_process/ipc.py +++ b/aikido_zen/background_process/ipc.py @@ -90,8 +90,6 @@ def send_message(self, message: Payload): with open_connection(self.address, deadline) as conn: conn.send((message.identifier, message.request)) result = conn.recv() - if time.monotonic() >= deadline: - return {"success": False, "error": "timeout"} except Exception as e: logger.debug("Exception occurred in IPC call : %s", e) return {"success": False, "error": "timeout"} diff --git a/aikido_zen/background_process/ipc_test.py b/aikido_zen/background_process/ipc_test.py index 9af487846..72e2d0c4b 100644 --- a/aikido_zen/background_process/ipc_test.py +++ b/aikido_zen/background_process/ipc_test.py @@ -16,7 +16,12 @@ 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 ping, put_event, read_property +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 @@ -223,6 +228,29 @@ def fail_start(self): 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"], From 654c6fb6ca53570b30a7ef02aa7e454c58ddba9d Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 16:22:13 +0200 Subject: [PATCH 5/8] Use versioned IPC socket paths --- aikido_zen/background_process/__init__.py | 16 +-- aikido_zen/background_process/init_test.py | 101 +++++------------- .../test_uds_file_access.py | 4 +- docs/background_process.md | 2 +- end2end/e2e.sh | 5 +- 5 files changed, 38 insertions(+), 90 deletions(-) diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 041d3f7aa..b9d4bdd4c 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -14,6 +14,7 @@ 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.ipc import IPCClient from aikido_zen.background_process.commands import ping @@ -51,7 +52,7 @@ 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) @@ -77,13 +78,14 @@ def start_background_process(): 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" - # This key is shared by workers and changes between application runs. - identity = hash_aikido_token().encode() + multiprocessing.current_process().authkey - name = hashlib.sha256(identity).hexdigest() + 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" diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 9d38e3c7b..91c385a3d 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -1,5 +1,3 @@ -import json -import multiprocessing import subprocess import sys @@ -151,84 +149,33 @@ def test_get_process_factory_does_not_set_global_start_method(): subprocess.run([sys.executable, "-c", script], check=True) -def test_socket_name_is_stable_and_separates_tokens(monkeypatch): - monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") - current = background_process.get_uds_filename() - - assert current == background_process.get_uds_filename() - monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_OTHER") - assert current != background_process.get_uds_filename() - - -@pytest.mark.parametrize("method", multiprocessing.get_all_start_methods()) -def test_socket_name_is_shared_by_workers_and_unique_to_each_run( - monkeypatch, tmp_path, method -): +def test_socket_name_is_shared_by_independent_processes(monkeypatch): monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") - script = tmp_path / "worker_socket_names.py" - script.write_text( - """ -import json -import multiprocessing -import sys - -def worker(pipe): - from aikido_zen.background_process import get_uds_filename - pipe.send(get_uds_filename()) - pipe.close() + script = """ +from aikido_zen.background_process import get_socket_path +print(get_socket_path()) +""" -if __name__ == "__main__": - context = multiprocessing.get_context(sys.argv[1]) - readers, processes = [], [] for _ in range(2): - reader, writer = context.Pipe(duplex=False) - process = context.Process(target=worker, args=(writer,)) - process.start() - writer.close() - readers.append(reader) - processes.append(process) - for process in processes: - process.join(10) - assert process.exitcode == 0 - from aikido_zen.background_process import get_uds_filename - names = [get_uds_filename()] + [reader.recv() for reader in readers] - print(json.dumps(names)) -""", - encoding="utf-8", - ) - runs = [ - json.loads( - subprocess.check_output( - [sys.executable, str(script), method], text=True, timeout=25 - ) - ) - for _ in range(2) - ] - - assert all(len(set(names)) == 1 for names in runs) - assert runs[0][0] != runs[1][0] - - -def test_startup_does_not_reuse_a_legacy_socket(monkeypatch, mocker): - factory = mocker.patch.object( - background_process, "get_process_factory" - ).return_value - active = mocker.patch.object( - background_process, "background_process_already_active", return_value=True - ) - monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") - monkeypatch.setenv("AIKIDO_TMP_DIR", "/tmp") - monkeypatch.setattr(background_process.platform, "system", lambda: "Linux") - legacy = f"/tmp/aikido_python_{background_process.hash_aikido_token()}.sock" - monkeypatch.setattr( - background_process.os.path, "exists", lambda path: path == legacy - ) + address = subprocess.check_output( + [sys.executable, "-c", script], text=True, timeout=10 + ).strip() + assert address == background_process.get_socket_path() - background_process.start_background_process() - active.assert_not_called() - assert factory.call_args.kwargs["args"][0] != legacy - factory.return_value.start.assert_called_once_with() +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): @@ -238,7 +185,7 @@ 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( @@ -262,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/test_uds_file_access.py b/aikido_zen/background_process/test_uds_file_access.py index 0c614cda9..61a605502 100644 --- a/aikido_zen/background_process/test_uds_file_access.py +++ b/aikido_zen/background_process/test_uds_file_access.py @@ -2,13 +2,13 @@ import platform import sys -from aikido_zen.background_process import get_uds_filename +from aikido_zen.background_process import get_socket_path 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() + filename = get_socket_path() directory = "/".join(filename.split("/")[:-1]) if platform.system() == "Windows": 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" From 5ec7f5c4d0686df8eba8a836e4db6ab73673c65e Mon Sep 17 00:00:00 2001 From: teta2k Date: Tue, 20 Oct 2026 16:35:10 +0200 Subject: [PATCH 6/8] Move socket access check into startup module --- aikido_zen/__init__.py | 6 ++--- aikido_zen/background_process/__init__.py | 21 +++++++++++++++ .../test_uds_file_access.py | 27 ------------------- aikido_zen/init_test.py | 4 +-- 4 files changed, 25 insertions(+), 33 deletions(-) delete mode 100644 aikido_zen/background_process/test_uds_file_access.py 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 b9d4bdd4c..d1819bd9e 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -89,6 +89,27 @@ def get_socket_path(): return f"{temp_dir}/{prefix}_{name}.sock" +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": 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 61a605502..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_socket_path -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_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 diff --git a/aikido_zen/init_test.py b/aikido_zen/init_test.py index 18cb259de..78f5977ca 100644 --- a/aikido_zen/init_test.py +++ b/aikido_zen/init_test.py @@ -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() From 468c9e3e8a880147d821324137c8f4f2b407718a Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Wed, 21 Oct 2026 16:44:40 +0200 Subject: [PATCH 7/8] Run the startup sync without the request context --- aikido_zen/thread/process_worker_loader.py | 5 ++++- .../thread/process_worker_loader_test.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/aikido_zen/thread/process_worker_loader.py b/aikido_zen/thread/process_worker_loader.py index 7917e6992..c71fe5c6e 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 @@ -29,7 +30,9 @@ def load_worker(): # runs once per process rather than on every request. if not thread_cache.is_config_loaded(): try: - thread_cache.renew() + # Unpickling the reply runs Zen hooks that call load_worker() again. + # With no request context, it returns instead of waiting on this lock. + contextvars.Context().run(thread_cache.renew) except Exception as e: logger.warning("An error occurred during data synchronization: %s", e) diff --git a/aikido_zen/thread/process_worker_loader_test.py b/aikido_zen/thread/process_worker_loader_test.py index ecc70e17f..b67fe1e21 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 @@ -88,3 +89,21 @@ 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 From 802960a78399231efe7aadc95d183c74b2abebdd Mon Sep 17 00:00:00 2001 From: teta2k Date: Mon, 26 Oct 2026 23:31:15 +0100 Subject: [PATCH 8/8] Prevent worker startup deadlocks during error logging --- aikido_zen/thread/process_worker_loader.py | 36 +++++++++---------- .../thread/process_worker_loader_test.py | 33 +++++++++++++++++ 2 files changed, 50 insertions(+), 19 deletions(-) diff --git a/aikido_zen/thread/process_worker_loader.py b/aikido_zen/thread/process_worker_loader.py index c71fe5c6e..50c4c1d9e 100644 --- a/aikido_zen/thread/process_worker_loader.py +++ b/aikido_zen/thread/process_worker_loader.py @@ -11,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: @@ -25,18 +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: - # Unpickling the reply runs Zen hooks that call load_worker() again. - # With no request context, it returns instead of waiting on this lock. - contextvars.Context().run(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 b67fe1e21..7877f4601 100644 --- a/aikido_zen/thread/process_worker_loader_test.py +++ b/aikido_zen/thread/process_worker_loader_test.py @@ -107,3 +107,36 @@ def first_request(): 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