From 72962f145309111f3cc852f5f19543760b08e19d Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Wed, 7 Oct 2026 18:56:50 +0300 Subject: [PATCH 1/3] Add uWSGI support and fix background process start method --- aikido_zen/background_process/__init__.py | 23 +++++-------- aikido_zen/background_process/init_test.py | 39 ++++++++++------------ aikido_zen/decorators/uwsgi.py | 18 ++++++++++ docs/uwsgi.md | 28 ++++++++++++++++ sample-apps/flask-mysql-uwsgi/app.py | 10 +++++- 5 files changed, 81 insertions(+), 37 deletions(-) create mode 100644 aikido_zen/decorators/uwsgi.py create mode 100644 docs/uwsgi.md diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 3b6459164..887f45fc3 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -6,7 +6,6 @@ import multiprocessing import os import platform -import sys from aikido_zen.helpers.token import get_token_from_env from aikido_zen.helpers.get_temp_dir import get_temp_dir @@ -25,20 +24,16 @@ def get_process_factory(): """ Return a process factory that is safe to start while an app is importing. - 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. + "spawn" (macOS's default) and "forkserver" (Linux's default on Python + 3.14+) both re-exec sys.executable to boot a fresh interpreter, which + breaks inside hosts that embed their own Python interpreter (e.g. uWSGI). + "fork" doesn't re-exec, so force it unless the app explicitly configured + its own method. Windows has no "fork" context, so leave it alone there. """ - 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 + configured = multiprocessing.get_start_method(allow_none=True) + needs_fork = configured is None or configured == "forkserver" + if needs_fork and platform.system() != "Windows": + return multiprocessing.get_context("fork").Process return multiprocessing.Process diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 3032e097d..3badfcafb 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -5,17 +5,12 @@ 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)) +def test_configured_forkserver_uses_fork_context(mocker): 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", @@ -23,22 +18,15 @@ def test_python314_configured_forkserver_uses_fork_context(monkeypatch, mocker): 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_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( - background_process.multiprocessing, - "get_all_start_methods", - return_value=["forkserver", "spawn", "fork"], - ) fork_context = mocker.patch.object( background_process.multiprocessing, "get_context", @@ -46,21 +34,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)) +def test_configured_spawn_uses_configured_context(mocker): 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", - ) get_context = mocker.patch.object( background_process.multiprocessing, "get_context", @@ -71,7 +53,20 @@ 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() + + +def test_windows_unset_default_uses_configured_context(monkeypatch, mocker): + monkeypatch.setattr(background_process.platform, "system", lambda: "Windows") + mocker.patch.object( + background_process.multiprocessing, "get_start_method", return_value=None + ) + 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() diff --git a/aikido_zen/decorators/uwsgi.py b/aikido_zen/decorators/uwsgi.py new file mode 100644 index 000000000..5c3f53f1c --- /dev/null +++ b/aikido_zen/decorators/uwsgi.py @@ -0,0 +1,18 @@ +""" +Includes all the wrappers for uwsgi apps +""" + +import aikido_zen + + +def postfork(prev_func): + """ + Aikido decorator for uwsgi's @postfork hook (from uwsgidecorators) + Function: postfork() + """ + + def aik_postfork(): + aikido_zen.protect() + prev_func() + + return aik_postfork diff --git a/docs/uwsgi.md b/docs/uwsgi.md new file mode 100644 index 000000000..217762c3f --- /dev/null +++ b/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, so Zen must be started with uWSGI's own `postfork` decorator instead of a plain `aikido_zen.protect()` call. 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``` diff --git a/sample-apps/flask-mysql-uwsgi/app.py b/sample-apps/flask-mysql-uwsgi/app.py index 8e9f4979a..3483569e9 100644 --- a/sample-apps/flask-mysql-uwsgi/app.py +++ b/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 From 3c2fcdbe0bafd8e17d8c1fb8bf6811e48fc2f81a Mon Sep 17 00:00:00 2001 From: teta2k Date: Thu, 8 Oct 2026 09:32:01 +0200 Subject: [PATCH 2/3] Fix background process selection and uWSGI instrumentation --- aikido_zen/background_process/__init__.py | 24 +++++--- aikido_zen/background_process/init_test.py | 68 ++++++++++++++++++++-- aikido_zen/decorators/uwsgi.py | 4 +- aikido_zen/decorators/uwsgi_test.py | 38 ++++++++++++ docs/uwsgi.md | 2 +- 5 files changed, 119 insertions(+), 17 deletions(-) create mode 100644 aikido_zen/decorators/uwsgi_test.py diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 887f45fc3..c5d592957 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -6,6 +6,7 @@ import multiprocessing import os import platform +import sys from aikido_zen.helpers.token import get_token_from_env from aikido_zen.helpers.get_temp_dir import get_temp_dir @@ -22,17 +23,22 @@ 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, without changing the start + method used by the application's own processes. - "spawn" (macOS's default) and "forkserver" (Linux's default on Python - 3.14+) both re-exec sys.executable to boot a fresh interpreter, which - breaks inside hosts that embed their own Python interpreter (e.g. uWSGI). - "fork" doesn't re-exec, so force it unless the app explicitly configured - its own method. Windows has no "fork" context, so leave it alone there. + Use the application's configured method, or Python's default if unset: + - fork or spawn: use that method. + - forkserver: use fork instead. Forkserver can rerun app startup and call + protect() again, which tries to start another Zen process before the new + process finishes initializing. Python rejects this with RuntimeError. + + Under uWSGI, always use fork: its executable cannot run the Python + commands needed by spawn or forkserver. """ - configured = multiprocessing.get_start_method(allow_none=True) - needs_fork = configured is None or configured == "forkserver" - if needs_fork and platform.system() != "Windows": + 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 diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 3badfcafb..7483efe06 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -1,11 +1,15 @@ import subprocess import sys +import pytest + import aikido_zen.background_process as background_process from .comms_test import reset_comms_after_test -def test_configured_forkserver_uses_fork_context(mocker): +@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", @@ -21,12 +25,17 @@ def test_configured_forkserver_uses_fork_context(mocker): background_process.multiprocessing.get_context.assert_called_once_with("fork") -def test_unset_default_uses_fork_context(mocker): +def test_unset_forkserver_default_uses_fork_context(mocker): get_start_method = mocker.patch.object( background_process.multiprocessing, "get_start_method", return_value=None, ) + mocker.patch.object( + background_process.multiprocessing, + "get_all_start_methods", + return_value=["forkserver", "spawn", "fork"], + ) fork_context = mocker.patch.object( background_process.multiprocessing, "get_context", @@ -37,11 +46,12 @@ def test_unset_default_uses_fork_context(mocker): background_process.multiprocessing.get_context.assert_called_once_with("fork") -def test_configured_spawn_uses_configured_context(mocker): +@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", + return_value=start_method, ) get_context = mocker.patch.object( background_process.multiprocessing, @@ -56,11 +66,21 @@ def test_configured_spawn_uses_configured_context(mocker): get_context.assert_not_called() -def test_windows_unset_default_uses_configured_context(monkeypatch, mocker): - monkeypatch.setattr(background_process.platform, "system", lambda: "Windows") +@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 ( @@ -70,6 +90,42 @@ def test_windows_unset_default_uses_configured_context(monkeypatch, mocker): 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 diff --git a/aikido_zen/decorators/uwsgi.py b/aikido_zen/decorators/uwsgi.py index 5c3f53f1c..1d358b56c 100644 --- a/aikido_zen/decorators/uwsgi.py +++ b/aikido_zen/decorators/uwsgi.py @@ -11,8 +11,10 @@ def postfork(prev_func): Function: postfork() """ + aikido_zen.protect(mode="daemon_disabled") + def aik_postfork(): - aikido_zen.protect() + aikido_zen.protect(mode="daemon_only") prev_func() return aik_postfork diff --git a/aikido_zen/decorators/uwsgi_test.py b/aikido_zen/decorators/uwsgi_test.py new file mode 100644 index 000000000..9ef24c6d5 --- /dev/null +++ b/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) diff --git a/docs/uwsgi.md b/docs/uwsgi.md index 217762c3f..41452bd22 100644 --- a/docs/uwsgi.md +++ b/docs/uwsgi.md @@ -6,7 +6,7 @@ pip install aikido_zen ``` -2. uWSGI preforks worker processes from its master by default, so Zen must be started with uWSGI's own `postfork` decorator instead of a plain `aikido_zen.protect()` call. Add this at the top of your app file, above any other import : +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 From 8d26fbe1bdd9b9ca1237b3290dce4776dd691a47 Mon Sep 17 00:00:00 2001 From: teta2k Date: Thu, 8 Oct 2026 09:51:37 +0200 Subject: [PATCH 3/3] Simplify process startup comments and uWSGI docs --- aikido_zen/background_process/__init__.py | 16 ++++++---------- aikido_zen/decorators/uwsgi.py | 4 +++- docs/uwsgi.md | 4 ++-- 3 files changed, 11 insertions(+), 13 deletions(-) diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index c5d592957..23d1d69b3 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -23,17 +23,13 @@ def get_process_factory(): """ - Choose how Zen's background process starts, without changing the start - method used by the application's own processes. + 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. - Use the application's configured method, or Python's default if unset: - - fork or spawn: use that method. - - forkserver: use fork instead. Forkserver can rerun app startup and call - protect() again, which tries to start another Zen process before the new - process finishes initializing. Python rejects this with RuntimeError. - - Under uWSGI, always use fork: its executable cannot run the Python - commands needed by spawn or forkserver. + The application's own start method is unchanged. """ start_method = multiprocessing.get_start_method(allow_none=True) if start_method is None: diff --git a/aikido_zen/decorators/uwsgi.py b/aikido_zen/decorators/uwsgi.py index 1d358b56c..000649297 100644 --- a/aikido_zen/decorators/uwsgi.py +++ b/aikido_zen/decorators/uwsgi.py @@ -1,5 +1,7 @@ """ -Includes all the wrappers for uwsgi apps +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 diff --git a/docs/uwsgi.md b/docs/uwsgi.md index 41452bd22..422e5eee5 100644 --- a/docs/uwsgi.md +++ b/docs/uwsgi.md @@ -1,7 +1,7 @@ # uWSGI ## Installation/Setup -1. Install `aikido_zen` package with pip : +1. Install `aikido_zen` package with pip: ```sh pip install aikido_zen ``` @@ -18,7 +18,7 @@ def start_aikido(): pass ``` -3. Setting your environment variables : +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"