diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 3b6459164..23d1d69b3 100644 --- a/aikido_zen/background_process/__init__.py +++ b/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 diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 3032e097d..7483efe06 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -1,21 +1,20 @@ 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", @@ -23,18 +22,16 @@ 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_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 diff --git a/aikido_zen/decorators/uwsgi.py b/aikido_zen/decorators/uwsgi.py new file mode 100644 index 000000000..000649297 --- /dev/null +++ b/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 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 new file mode 100644 index 000000000..422e5eee5 --- /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. 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``` 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