Send events asynchronously

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.

aikido_zen/background_process/__init__.py

@@ -3,6 +3,7 @@
 and listen for data sent by our sources and sinks
 """
 
+import hashlib
 import multiprocessing
 import os
 import platform
@@ -13,12 +14,10 @@
 from aikido_zen.helpers.hash_aikido_token import hash_aikido_token
 from aikido_zen.helpers.get_agent_session_id import get_agent_session_id
 from aikido_zen.helpers.logging import logger
+from aikido_zen.config import PKG_VERSION
 
-from aikido_zen.background_process.comms import (
-    AikidoIPCCommunications,
-    get_comms,
-    reset_comms,
-)
+from aikido_zen.background_process.ipc import IPCClient
+from aikido_zen.background_process.commands import ping
 from .aikido_background_process import AikidoBackgroundProcess
 
 
@@ -53,15 +52,15 @@ def start_background_process():
         )
         return
 
-    address = get_uds_filename()
+    address = get_socket_path()
     if platform.system() == "Windows":
         # Python does not support Windows UDS just yet, so we have to rely on INET
         address = ("127.0.0.1", 49156)
 
-    comms = AikidoIPCCommunications(address, secret_key_bytes)
+    ipc = IPCClient(address, secret_key_bytes)
 
     if isinstance(address, str) and os.path.exists(address):
-        if background_process_already_active(comms):
+        if background_process_already_active(ipc):
             return  # Don't start if the background process is already active
         try:
             os.remove(address)
@@ -72,23 +71,47 @@ def start_background_process():
     #  Daemon is set to True so that the process kills itself when the main process dies
     background_process = get_process_factory()(
         target=AikidoBackgroundProcess,
-        args=(comms.address, comms.key, get_agent_session_id()),
+        args=(ipc.address, ipc.key, get_agent_session_id()),
         name="zen-agent-process",
         daemon=True,
     )
     background_process.start()
 
 
-def get_uds_filename():
-    """Returns the address for UDS file"""
+def get_socket_path():
     temp_dir = get_temp_dir()
-    prefix = "aikido_python"
-    name = hash_aikido_token()
+    prefix = "aikido_zen"
+
+    # Include the package version to avoid reusing a background process from a previous release.
+    identity = f"{hash_aikido_token()}:{PKG_VERSION}"
+
+    name = hashlib.sha256(identity.encode()).hexdigest()
     return f"{temp_dir}/{prefix}_{name}.sock"
 
 
-def background_process_already_active(comms):
-    res = comms.send_data_to_bg_process(action="PING", obj=tuple(), receive=True)
+def check_socket_access():
+    """Returns true if access is possible, false if it's not."""
+    filename = get_socket_path()
+    directory = "/".join(filename.split("/")[:-1])
+
+    if platform.system() == "Windows":
+        # Python does not support Windows UDS just yet, so we have to rely on INET
+        return True
+
+    if not os.access(directory, os.R_OK) or not os.access(directory, os.W_OK):
+        # We need both read and write permissions in order to start the socket
+        logger.critical(
+            "Cannot start Zen: no read/write permissions for directory %s. "
+            + "You can change the path by setting "
+            + "the AIKIDO_TMP_DIR environment variable.",
+            directory,
+        )
+        return False
+    return True
+
+
+def background_process_already_active(ipc):
+    res = ipc.send_message(ping.generate())
     if res["success"] and res["data"] == "Received":
         # Ping is active, return.
         logger.debug(

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"""

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()

aikido_zen/background_process/commands/__init__.py

@@ -1,34 +1,29 @@
 from aikido_zen.helpers.logging import logger
-from aikido_zen.helpers.ipc.command_types import CommandContext
-from .put_event import PutEventCommand
-from .check_firewall_lists import process_check_firewall_lists
-from .read_property import process_read_property
-from .should_ratelimit import process_should_ratelimit
-from .ping import process_ping
-from .sync_data import process_sync_data
+from . import (
+    check_firewall_lists,
+    ping,
+    put_event,
+    read_property,
+    should_ratelimit,
+    sync_data,
+)
 
 commands_map = {
-    "SYNC_DATA": process_sync_data,
-    "READ_PROPERTY": process_read_property,
-    "SHOULD_RATELIMIT": process_should_ratelimit,
-    "PING": process_ping,
-    "CHECK_FIREWALL_LISTS": process_check_firewall_lists,
+    check_firewall_lists.IDENTIFIER: check_firewall_lists.process,
+    ping.IDENTIFIER: ping.process,
+    put_event.IDENTIFIER: put_event.process,
+    read_property.IDENTIFIER: read_property.process,
+    should_ratelimit.IDENTIFIER: should_ratelimit.process,
+    sync_data.IDENTIFIER: sync_data.process,
 }
 
-modern_commands = [PutEventCommand]
-
 
 def process_incoming_command(connection_manager, obj, conn, queue):
     inbound_identifier = obj[0]
     inbound_request = obj[1]
     if inbound_identifier in commands_map:
         func = commands_map[inbound_identifier]
-        return conn.send(func(connection_manager, inbound_request))
-
-    for cmd in modern_commands:
-        if cmd.identifier() == inbound_identifier:
-            cmd.run(CommandContext(connection_manager, queue, conn), inbound_request)
-            return None
+        return conn.send(func(connection_manager, inbound_request, queue))
 
     logger.debug("Command : `%s` not found - did not execute", inbound_identifier)
     return None

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):

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"

aikido_zen/background_process/commands/put_event.py

@@ -1,27 +1,18 @@
-from aikido_zen.helpers.ipc.command_types import Command, CommandContext, Payload
+from aikido_zen.helpers.ipc.command_types import Payload
 
+IDENTIFIER = "PUT_EVENT"
 
-class PutEventReq:
-    def __init__(self, event):
-        # Event is a dictionary containing data that is going to be reported to core
-        # The CloudConnectionManager adds the "agent" field
-        self.event = event
 
+def generate(event, timeout_in_sec=5) -> Payload:
+    return Payload(IDENTIFIER, PutEventReq(event), timeout_in_sec)
 
-class PutEventCommand(Command):
-    @classmethod
-    def identifier(cls) -> str:
-        return "put_event"
 
-    @classmethod
-    def returns_data(cls) -> bool:
-        return False
+def process(connection_manager, request, queue):
+    queue.put(request.event)
 
-    @classmethod
-    def run(cls, context: CommandContext, request: PutEventReq):
-        #  Events sent here get put in the event queue so that they are processed in the background
-        context.queue.put(request.event)
 
-    @classmethod
-    def generate(cls, request) -> Payload:
-        return Payload(cls, PutEventReq(request))
+# IPC uses pickle, so this request class must keep its import path.
+class PutEventReq:
+    def __init__(self, event):
+        # The CloudConnectionManager adds the "agent" field.
+        self.event = event