Initialize worker cache before first request

.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 }}

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
 

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,

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

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

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

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: