Add uWSGI support and fix background process start method

aikido_zen/background_process/__init__.py

@@ -23,22 +23,19 @@
 
 def get_process_factory():
     """
-    Return a process factory that is safe to start while an app is importing.
+    Choose how Zen's background process starts:
+    - Use the application's configured method, or Python's default if unset.
+    - Use fork instead of forkserver to avoid startup errors caused by
+      re-importing the app and calling protect() again.
+    - Always use fork under uWSGI; its executable cannot run Python subprocesses.
 
-    Python 3.14 changed the default POSIX start method from fork to forkserver.
-    Forkserver re-imports the application's main module, but Zen starts its
-    background process while that module is still importing.
-
-    Inspect the configured start method without setting multiprocessing's
-    process-wide default. If it is unset, get_all_start_methods() reports the
-    platform default as its first entry.
+    The application's own start method is unchanged.
     """
-    if sys.version_info >= (3, 14):
-        start_method = multiprocessing.get_start_method(allow_none=True)
-        if start_method is None:
-            start_method = multiprocessing.get_all_start_methods()[0]
-        if start_method == "forkserver":
-            return multiprocessing.get_context("fork").Process
+    start_method = multiprocessing.get_start_method(allow_none=True)
+    if start_method is None:
+        start_method = multiprocessing.get_all_start_methods()[0]
+    if start_method == "forkserver" or "uwsgi" in sys.modules:
+        return multiprocessing.get_context("fork").Process
     return multiprocessing.Process
 
 

aikido_zen/background_process/init_test.py

@@ -1,40 +1,37 @@
 import subprocess
 import sys
 
+import pytest
+
 import aikido_zen.background_process as background_process
 from .comms_test import reset_comms_after_test
 
 
-def test_python314_configured_forkserver_uses_fork_context(monkeypatch, mocker):
-    monkeypatch.setattr(background_process.sys, "version_info", (3, 14))
+@pytest.mark.parametrize("version", [(3, 8), (3, 13), (3, 14)])
+def test_configured_forkserver_uses_fork_context(monkeypatch, mocker, version):
+    monkeypatch.setattr(sys, "version_info", version)
     get_start_method = mocker.patch.object(
         background_process.multiprocessing,
         "get_start_method",
         return_value="forkserver",
     )
-    get_all_start_methods = mocker.patch.object(
-        background_process.multiprocessing,
-        "get_all_start_methods",
-    )
     fork_context = mocker.patch.object(
         background_process.multiprocessing,
         "get_context",
     ).return_value
 
     assert background_process.get_process_factory() == fork_context.Process
     get_start_method.assert_called_once_with(allow_none=True)
-    get_all_start_methods.assert_not_called()
     background_process.multiprocessing.get_context.assert_called_once_with("fork")
 
 
-def test_python314_unset_forkserver_default_uses_fork_context(monkeypatch, mocker):
-    monkeypatch.setattr(background_process.sys, "version_info", (3, 14))
+def test_unset_forkserver_default_uses_fork_context(mocker):
     get_start_method = mocker.patch.object(
         background_process.multiprocessing,
         "get_start_method",
         return_value=None,
     )
-    get_all_start_methods = mocker.patch.object(
+    mocker.patch.object(
         background_process.multiprocessing,
         "get_all_start_methods",
         return_value=["forkserver", "spawn", "fork"],
@@ -46,20 +43,15 @@ def test_python314_unset_forkserver_default_uses_fork_context(monkeypatch, mocke
 
     assert background_process.get_process_factory() == fork_context.Process
     get_start_method.assert_called_once_with(allow_none=True)
-    get_all_start_methods.assert_called_once_with()
     background_process.multiprocessing.get_context.assert_called_once_with("fork")
 
 
-def test_python314_configured_spawn_uses_configured_context(monkeypatch, mocker):
-    monkeypatch.setattr(background_process.sys, "version_info", (3, 14))
+@pytest.mark.parametrize("start_method", ["spawn", "fork"])
+def test_configured_method_is_preserved(mocker, start_method):
     get_start_method = mocker.patch.object(
         background_process.multiprocessing,
         "get_start_method",
-        return_value="spawn",
-    )
-    get_all_start_methods = mocker.patch.object(
-        background_process.multiprocessing,
-        "get_all_start_methods",
+        return_value=start_method,
     )
     get_context = mocker.patch.object(
         background_process.multiprocessing,
@@ -71,10 +63,69 @@ def test_python314_configured_spawn_uses_configured_context(monkeypatch, mocker)
         == background_process.multiprocessing.Process
     )
     get_start_method.assert_called_once_with(allow_none=True)
-    get_all_start_methods.assert_not_called()
     get_context.assert_not_called()
 
 
+@pytest.mark.parametrize("system", ["Darwin", "Windows"])
+@pytest.mark.parametrize("version", [(3, 8), (3, 13), (3, 14)])
+def test_unset_spawn_default_is_preserved(monkeypatch, mocker, system, version):
+    monkeypatch.setattr(sys, "version_info", version)
+    monkeypatch.setattr(background_process.platform, "system", lambda: system)
+    mocker.patch.object(
+        background_process.multiprocessing, "get_start_method", return_value=None
+    )
+    mocker.patch.object(
+        background_process.multiprocessing,
+        "get_all_start_methods",
+        return_value=(
+            ["spawn"] if system == "Windows" else ["spawn", "fork", "forkserver"]
+        ),
+    )
+    get_context = mocker.patch.object(background_process.multiprocessing, "get_context")
+
+    assert (
+        background_process.get_process_factory()
+        == background_process.multiprocessing.Process
+    )
+    get_context.assert_not_called()
+
+
+def test_unset_fork_default_is_preserved(mocker):
+    mocker.patch.object(
+        background_process.multiprocessing, "get_start_method", return_value=None
+    )
+    mocker.patch.object(
+        background_process.multiprocessing,
+        "get_all_start_methods",
+        return_value=["fork", "spawn", "forkserver"],
+    )
+    get_context = mocker.patch.object(background_process.multiprocessing, "get_context")
+
+    assert (
+        background_process.get_process_factory()
+        == background_process.multiprocessing.Process
+    )
+    get_context.assert_not_called()
+
+
+@pytest.mark.parametrize("system", ["Darwin", "Linux"])
+@pytest.mark.parametrize("start_method", [None, "spawn", "forkserver"])
+def test_uwsgi_uses_fork_context(monkeypatch, mocker, system, start_method):
+    monkeypatch.setattr(background_process.platform, "system", lambda: system)
+    monkeypatch.setitem(sys.modules, "uwsgi", mocker.Mock())
+    mocker.patch.object(
+        background_process.multiprocessing,
+        "get_start_method",
+        return_value=start_method,
+    )
+    fork_context = mocker.patch.object(
+        background_process.multiprocessing, "get_context"
+    ).return_value
+
+    assert background_process.get_process_factory() == fork_context.Process
+    background_process.multiprocessing.get_context.assert_called_once_with("fork")
+
+
 def test_get_process_factory_does_not_set_global_start_method():
     script = """
 import multiprocessing

aikido_zen/decorators/uwsgi.py

@@ -0,0 +1,22 @@
+"""
+Set up Zen in two steps:
+1. Install instrumentation when the decorator is applied.
+2. Start the background process from uWSGI's postfork hook.
+"""
+
+import aikido_zen
+
+
+def postfork(prev_func):
+    """
+    Aikido decorator for uwsgi's @postfork hook (from uwsgidecorators)
+    Function: postfork()
+    """
+
+    aikido_zen.protect(mode="daemon_disabled")
+
+    def aik_postfork():
+        aikido_zen.protect(mode="daemon_only")
+        prev_func()
+
+    return aik_postfork

aikido_zen/decorators/uwsgi_test.py

@@ -0,0 +1,38 @@
+import subprocess
+import sys
+
+
+def test_postfork_instruments_imported_functions_before_starting_daemon(monkeypatch):
+    monkeypatch.delenv("AIKIDO_DISABLE", raising=False)
+    script = """
+import threading
+from unittest.mock import patch
+
+from aikido_zen.background_process import get_comms
+from aikido_zen.decorators.uwsgi import postfork
+
+calls = []
+threads = threading.enumerate()
+
+with patch("aikido_zen.start_background_process", side_effect=lambda: calls.append("daemon")):
+    @postfork
+    def start_aikido():
+        calls.append("hook")
+
+    from os import system
+
+    assert calls == []
+    assert get_comms() is None
+    assert threading.enumerate() == threads
+
+    start_aikido()
+    assert calls == ["daemon", "hook"]
+
+    with patch("aikido_zen.vulnerabilities.run_vulnerability_scan") as scan:
+        system("exit 0")
+        scan.assert_called_once_with(
+            kind="shell_injection", op="os.system", args=("exit 0",)
+        )
+"""
+
+    subprocess.run([sys.executable, "-c", script], check=True, timeout=30)

docs/uwsgi.md

@@ -0,0 +1,28 @@
+# uWSGI
+
+## Installation/Setup
+1. Install `aikido_zen` package with pip:
+```sh
+pip install aikido_zen
+```
+
+2. uWSGI preforks worker processes from its master by default. Use the decorators below to instrument your app before its imports and start Zen's background process after each fork. Add this at the top of your app file, above any other import:
+```python
+from uwsgidecorators import postfork
+import aikido_zen.decorators.uwsgi as aik
+
+@postfork
+@aik.postfork
+def start_aikido():
+    # If you already have a postfork hook, replace pass with your own code.
+    pass
+```
+
+3. Setting your environment variables:
+Make sure to set your token in order to communicate with Aikido's servers
+```env
+AIKIDO_TOKEN="AIK_RUNTIME_YOUR_TOKEN_HERE"
+```
+
+- Enabling extra debugging (optional): ```AIKIDO_DEBUG=1```
+- Enabling blocking using an env variable (optional): ```AIKIDO_BLOCK=1```

sample-apps/flask-mysql-uwsgi/app.py

@@ -1,5 +1,13 @@
+from uwsgidecorators import postfork
 import aikido_zen # Aikido package import
-aikido_zen.protect()
+import aikido_zen.decorators.uwsgi as aik
+
+
+@postfork
+@aik.postfork
+def start_aikido():
+    pass
+
 
 import time
 from flask import Flask, render_template, request