From e0f1cc13cfc391e1ea436a4d89baa0a3a26f64d7 Mon Sep 17 00:00:00 2001 From: teta2k Date: Thu, 23 Jul 2026 22:29:49 +0200 Subject: [PATCH 1/2] Initialize worker cache before first request --- .github/workflows/end2end.yml | 7 +- aikido_zen/middleware/init_test.py | 1 + .../sources/functions/request_handler_test.py | 4 +- aikido_zen/thread/process_worker_loader.py | 28 ++++- aikido_zen/thread/thread_cache.py | 4 + aikido_zen/thread/thread_cache_test.py | 103 ++++++++++++++++++ aikido_zen/vulnerabilities/init_test.py | 1 + 7 files changed, 138 insertions(+), 10 deletions(-) diff --git a/.github/workflows/end2end.yml b/.github/workflows/end2end.yml index 541eec5a8..121c296e4 100644 --- a/.github/workflows/end2end.yml +++ b/.github/workflows/end2end.yml @@ -73,7 +73,10 @@ jobs: - name: Start application working-directory: ./sample-apps/${{ matrix.app.name }} run: | - nohup make run > output.log & tail -f output.log & sleep 20 - nohup make runZenDisabled & sleep 20 + nohup make runZenDisabled & + sleep 20 + nohup make run > output.log & + tail -f output.log & + sleep 20 - name: Run end2end tests for application run: tail -f ./sample-apps/${{ matrix.app.name }}/output.log & poetry run pytest ./${{ matrix.app.testfile }} diff --git a/aikido_zen/middleware/init_test.py b/aikido_zen/middleware/init_test.py index d7f4bf1e8..4a7fad1f9 100644 --- a/aikido_zen/middleware/init_test.py +++ b/aikido_zen/middleware/init_test.py @@ -81,6 +81,7 @@ def test_cache_comms_with_endpoints(): }, } ] + thread_cache.config.last_updated_at = 1 assert get_current_context().executed_middleware == False assert thread_cache.middleware_installed == False diff --git a/aikido_zen/sources/functions/request_handler_test.py b/aikido_zen/sources/functions/request_handler_test.py index 901a3a9e0..e3d0781c3 100644 --- a/aikido_zen/sources/functions/request_handler_test.py +++ b/aikido_zen/sources/functions/request_handler_test.py @@ -200,7 +200,7 @@ def create_service_config(): "allowedIPAddresses": ["1.1.1.1", "2.2.2.2", "3.3.3.3"], } ], - last_updated_at=None, + last_updated_at=1, blocked_uids=set(), bypassed_ips=[], received_any_stats=False, @@ -238,7 +238,7 @@ def test_bypassed_ip_skips_all_checks(firewall_lists): "allowedIPAddresses": ["1.1.1.1"], # 192.168.1.1 not in this list } ], - last_updated_at=None, + last_updated_at=1, blocked_uids=set(), bypassed_ips=["192.168.1.1"], received_any_stats=False, diff --git a/aikido_zen/thread/process_worker_loader.py b/aikido_zen/thread/process_worker_loader.py index 87edbfc16..cca4d090d 100644 --- a/aikido_zen/thread/process_worker_loader.py +++ b/aikido_zen/thread/process_worker_loader.py @@ -2,8 +2,12 @@ import threading from aikido_zen.context import get_current_context +from aikido_zen.helpers.logging import logger +from aikido_zen.thread import thread_cache from aikido_zen.thread.process_worker import aikido_process_worker_thread +_load_worker_lock = threading.Lock() + def load_worker(): """ @@ -14,10 +18,22 @@ def load_worker(): # The name is aikido-process-worker- + the current PID thread_name = "aikido-process-worker-" + str(multiprocessing.current_process().pid) - if any(thread.name == thread_name for thread in threading.enumerate()): - return # The thread already exists, returning. - # Create a new daemon thread tht will handle communication to and from background agent - thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) - thread.daemon = True - thread.start() + with _load_worker_lock: + # The first HTTP request in a worker process may arrive before its local cache + # has received config. Synchronize with the background process immediately + # instead of waiting for the 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) + + # Each worker process should have only one periodic synchronization thread. + if any(thread.name == thread_name for thread in threading.enumerate()): + return + + # Create a new daemon thread tht will handle communication to and from background agent + thread = threading.Thread(target=aikido_process_worker_thread, name=thread_name) + thread.daemon = True + thread.start() diff --git a/aikido_zen/thread/thread_cache.py b/aikido_zen/thread/thread_cache.py index 6eb60bddf..a7129ebd0 100644 --- a/aikido_zen/thread/thread_cache.py +++ b/aikido_zen/thread/thread_cache.py @@ -123,5 +123,9 @@ def get_cache(): return global_thread_cache +def is_config_loaded(): + return global_thread_cache.config.last_updated_at > 0 + + def renew(): global_thread_cache.renew() diff --git a/aikido_zen/thread/thread_cache_test.py b/aikido_zen/thread/thread_cache_test.py index 21e294025..ed6f7e68e 100644 --- a/aikido_zen/thread/thread_cache_test.py +++ b/aikido_zen/thread/thread_cache_test.py @@ -3,6 +3,7 @@ import aikido_zen.test_utils as test_utils from aikido_zen.background_process.routes import Routes +from . import process_worker_loader from .thread_cache import ThreadCache, get_cache from .. import set_user from ..background_process.packages import PackagesStore @@ -105,6 +106,108 @@ def test_renew_with_no_comms(thread_cache: ThreadCache): } +@patch.object(process_worker_loader.thread_cache, "renew") +@patch.object( + process_worker_loader.thread_cache, "is_config_loaded", return_value=False +) +@patch.object(process_worker_loader.threading, "Thread") +@patch.object(process_worker_loader.threading, "enumerate", return_value=[]) +@patch.object(process_worker_loader, "get_current_context", return_value=object()) +def test_load_worker_renews_cache_before_starting_thread( + _mock_context, + _mock_enumerate, + mock_thread_type, + _mock_is_config_loaded, + mock_renew, +): + call_order = [] + thread = mock_thread_type.return_value + thread.start.side_effect = lambda: call_order.append("start") + mock_renew.side_effect = lambda: call_order.append("renew") + + process_worker_loader.load_worker() + + assert call_order == ["renew", "start"] + assert thread.daemon is True + + +@patch.object(process_worker_loader.thread_cache, "renew") +@patch.object(process_worker_loader.thread_cache, "is_config_loaded", return_value=True) +@patch.object(process_worker_loader.threading, "Thread") +@patch.object(process_worker_loader.threading, "enumerate") +@patch.object(process_worker_loader, "get_current_context", return_value=object()) +def test_load_worker_does_not_initialize_when_worker_is_running( + _mock_context, + mock_enumerate, + mock_thread_type, + _mock_is_config_loaded, + mock_renew, +): + worker = MagicMock() + worker.name = "aikido-process-worker-" + str( + process_worker_loader.multiprocessing.current_process().pid + ) + mock_enumerate.return_value = [worker] + + process_worker_loader.load_worker() + + mock_renew.assert_not_called() + mock_thread_type.assert_not_called() + + +@patch.object(process_worker_loader.thread_cache, "renew") +@patch.object( + process_worker_loader.thread_cache, "is_config_loaded", return_value=False +) +@patch.object(process_worker_loader.threading, "Thread") +@patch.object(process_worker_loader.threading, "enumerate") +@patch.object(process_worker_loader, "get_current_context", return_value=object()) +def test_load_worker_retries_cache_initialization_when_worker_is_running( + _mock_context, + mock_enumerate, + mock_thread_type, + _mock_is_config_loaded, + mock_renew, +): + worker = MagicMock() + worker.name = "aikido-process-worker-" + str( + process_worker_loader.multiprocessing.current_process().pid + ) + mock_enumerate.return_value = [worker] + + process_worker_loader.load_worker() + + mock_renew.assert_called_once_with() + mock_thread_type.assert_not_called() + + +@patch.object(process_worker_loader.logger, "warning") +@patch.object(process_worker_loader.thread_cache, "renew") +@patch.object( + process_worker_loader.thread_cache, "is_config_loaded", return_value=False +) +@patch.object(process_worker_loader.threading, "Thread") +@patch.object(process_worker_loader.threading, "enumerate", return_value=[]) +@patch.object(process_worker_loader, "get_current_context", return_value=object()) +def test_load_worker_starts_thread_when_cache_renewal_fails( + _mock_context, + _mock_enumerate, + mock_thread_type, + _mock_is_config_loaded, + mock_renew, + mock_warning, +): + error = RuntimeError("sync failed") + mock_renew.side_effect = error + + process_worker_loader.load_worker() + + mock_thread_type.return_value.start.assert_called_once_with() + mock_warning.assert_called_once_with( + "An error occurred during data synchronization: %s", error + ) + + @patch("aikido_zen.background_process.comms.get_comms") def test_renew_with_invalid_response(mock_get_comms, thread_cache: ThreadCache): """Test that renew handles an invalid response gracefully.""" diff --git a/aikido_zen/vulnerabilities/init_test.py b/aikido_zen/vulnerabilities/init_test.py index 54e10e89f..a24c5a4d4 100644 --- a/aikido_zen/vulnerabilities/init_test.py +++ b/aikido_zen/vulnerabilities/init_test.py @@ -130,6 +130,7 @@ def test_sql_injection_with_route_params(caplog, get_context, monkeypatch): def test_sql_injection_with_comms(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: From c2bd93df3431ff28450f259fb2cf8f44bda68653 Mon Sep 17 00:00:00 2001 From: teta2k Date: Fri, 24 Jul 2026 10:21:21 +0200 Subject: [PATCH 2/2] Update aikido_zen/thread/process_worker_loader.py Co-authored-by: Hans Ott --- aikido_zen/thread/process_worker_loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aikido_zen/thread/process_worker_loader.py b/aikido_zen/thread/process_worker_loader.py index cca4d090d..355a8159d 100644 --- a/aikido_zen/thread/process_worker_loader.py +++ b/aikido_zen/thread/process_worker_loader.py @@ -33,7 +33,7 @@ def load_worker(): if any(thread.name == thread_name for thread in threading.enumerate()): return - # Create a new daemon thread tht will handle communication to and from background agent + # 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()