From 07f9a98d5d536546906467157a58988adcfd11de Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Thu, 29 Oct 2026 17:01:46 +0100
Subject: [PATCH 1/6] Release captured Flask request bodies when responses
 close

---
 aikido_zen/sources/flask/__init__.py   |  68 +++-
 aikido_zen/sources/flask/flask_test.py | 437 ++++++++++++++++++++++++-
 2 files changed, 499 insertions(+), 6 deletions(-)

diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index bc73df788..9903cb3b6 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -1,4 +1,8 @@
-from aikido_zen.context import Context
+from sys import exc_info
+from traceback import clear_frames
+from weakref import ref
+
+from aikido_zen.context import Context, current_context, get_current_context
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
 from .extract_cookies import extract_cookies_from_flask_request_and_save_data
@@ -21,10 +25,63 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 
 
 @after
-def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
+def _request_after(func, instance, args, kwargs, return_value):
+    from flask.globals import request_ctx
+
     if not hasattr(return_value, "status_code"):
         return
-    funcs.request_handler(stage="post_response", status_code=return_value.status_code)
+    error = get_argument(args, kwargs, 0, "e")
+    if error is None:
+        funcs.request_handler(
+            stage="post_response", status_code=return_value.status_code
+        )
+    if error is None and request_ctx.request.routing_exception is None:
+        return
+    request_ref = ref(request_ctx.request)
+
+    # Finished traceback frames retain cached bodies until cyclic GC runs.
+    def release_request():
+        if exc_info()[0] is not None:
+            return
+        pending = [error] if error is not None else []
+        seen = set()
+        while pending:
+            exception = pending.pop()
+            if exception is None or id(exception) in seen:
+                continue
+            seen.add(id(exception))
+            pending.extend((exception.__cause__, exception.__context__))
+            clear_frames(exception.__traceback__)
+        req = request_ref()
+        if req is not None and req.routing_exception is not None:
+            routing_error = req.routing_exception
+            routing_error.__traceback__ = None
+            routing_error.__context__ = None
+            routing_error.__cause__ = None
+
+    return_value.call_on_close(release_request)
+
+
+def _call_with_context_cleanup(func, instance, args, kwargs):
+    from werkzeug.wsgi import ClosingIterator
+
+    if get_current_context() is None:
+        return func(*args, **kwargs)
+    context_ref = ref(get_current_context())
+
+    def release_context():
+        if get_current_context() is context_ref():
+            current_context.set(None)
+
+    try:
+        response = func(*args, **kwargs)
+        # Wrapping file responses prevents the WSGI server's sendfile optimization.
+        if context_ref() is None or context_ref().body is None:
+            return response
+        return ClosingIterator(response, release_context)
+    except BaseException:
+        release_context()
+        raise
 
 
 @before_modify_return
@@ -58,6 +115,9 @@ def patch(m):
     available, so we can extract data from it.
     (github src: https://github.com/pallets/flask/blob/bc143499cf1137a271a7cf75bdd3e16e43ede2f0/src/flask/app.py#L1529)
     """
+    # The response wrapper needs the context created by the outer hook.
+    patch_function(m, "Flask.__call__", _call_with_context_cleanup)
     patch_function(m, "Flask.__call__", _call)
     patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_before)
-    patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_after)
+    patch_function(m, "Flask.full_dispatch_request", _request_after)
+    patch_function(m, "Flask.handle_exception", _request_after)
diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py
index 1aa58a4b0..9d5fe225f 100644
--- a/aikido_zen/sources/flask/flask_test.py
+++ b/aikido_zen/sources/flask/flask_test.py
@@ -1,9 +1,29 @@
-import pytest
+import gc
+from contextvars import copy_context
 import signal
+import traceback
+import weakref
+from io import BytesIO
 from unittest.mock import patch
-import aikido_zen.sources.flask
+
+import pytest
+from flask import (
+    Flask,
+    Response,
+    copy_current_request_context,
+    request,
+    stream_with_context,
+)
+from werkzeug.exceptions import HTTPException, NotFound
+from werkzeug.test import EnvironBuilder
+from werkzeug.wsgi import ClosingIterator, FileWrapper
+
+import aikido_zen.sources.flask as source
 from aikido_zen.background_process.ipc import IPCClient
 from aikido_zen.context import current_context, get_current_context
+from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
+    context_contains_sql_injection,
+)
 
 sample_environ = {
     "REQUEST_METHOD": "POST",
@@ -231,3 +251,416 @@ def test_startup_flask():
         app.run()
     except TimeoutException:
         pass
+
+
+@pytest.fixture
+def app(monkeypatch):
+    monkeypatch.setattr(source, "should_block", lambda: None)
+    monkeypatch.setattr(source.funcs, "request_handler", lambda **kwargs: None)
+    current_context.set(None)
+    yield Flask(__name__)
+    current_context.set(None)
+
+
+@pytest.fixture
+def without_gc():
+    enabled = gc.isenabled()
+    gc.collect()
+    gc.disable()
+    try:
+        yield
+    finally:
+        if enabled:
+            gc.enable()
+        gc.collect()
+
+
+@pytest.mark.parametrize("status", [404, 405, 308])
+@pytest.mark.parametrize("custom_handler", [False, True])
+def test_routing_requests_release_cached_bodies_without_gc(
+    app, without_gc, status, custom_handler
+):
+    refs = []
+    app.add_url_rule("/known", endpoint="known", view_func=lambda: "ok")
+    app.add_url_rule(
+        "/slash/", endpoint="slash", view_func=lambda: "ok", methods=["POST"]
+    )
+    if custom_handler:
+        app.register_error_handler(HTTPException, lambda error: ("handled", error.code))
+
+    @app.before_request
+    def capture():
+        refs.append(weakref.ref(request._get_current_object()))
+        assert get_current_context().body == "x" * 1024 * 1024
+
+    path = {404: "/missing", 405: "/known", 308: "/slash"}[status]
+    client = app.test_client()
+    for _ in range(3):
+        response = client.post(
+            path, data=b"x" * 1024 * 1024, content_type="application/octet-stream"
+        )
+        assert response.status_code == status
+        response.close()
+        del response
+
+    assert all(ref() is None for ref in refs)
+
+
+@pytest.mark.parametrize("status", [404, 405])
+def test_custom_handlers_and_teardown_keep_body_inspection(app, without_gc, status):
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+    payload = {"name": attack}
+    refs = []
+    stages = []
+    app.add_url_rule("/known", view_func=lambda: "ok")
+
+    def inspect_body():
+        assert get_current_context().body == payload
+        assert request.get_json() == payload
+        result = context_contains_sql_injection(
+            sql=f"INSERT INTO pets(name) VALUES ('{attack}')",
+            operation="psycopg2.Cursor.execute",
+            context=get_current_context(),
+            dialect="postgres",
+        )
+        assert result["source"] == "body"
+        assert result["payload"] == attack
+        assert request.routing_exception.__traceback__ is not None
+
+    @app.errorhandler(HTTPException)
+    def handler(error):
+        refs.append(weakref.ref(request._get_current_object()))
+        inspect_body()
+        stages.append("handler")
+        return "handled", error.code
+
+    @app.teardown_request
+    def teardown(error):
+        inspect_body()
+        stages.append("teardown")
+
+    response = app.test_client().post(
+        "/missing" if status == 404 else "/known", json=payload
+    )
+    assert response.status_code == status
+    assert response.data == b"handled"
+    response.close()
+    del response
+
+    assert stages == ["handler", "teardown"]
+    assert refs[0]() is None
+
+
+def test_custom_flask_error_override_keeps_body_and_releases_request(app, without_gc):
+    refs = []
+    payload = {"name": "value"}
+
+    class CustomFlask(Flask):
+        def handle_http_exception(self, error):
+            refs.append(weakref.ref(request._get_current_object()))
+            assert get_current_context().body == payload
+            assert request.get_json() == payload
+            return self.make_response(("handled", error.code))
+
+    response = CustomFlask(__name__).test_client().post("/missing", json=payload)
+    assert response.status_code == 404
+    response.close()
+    del response
+    assert refs[0]() is None
+
+
+@pytest.mark.parametrize("finish_stream", [False, True])
+def test_streamed_error_response_keeps_body_until_stream_closes(
+    app, without_gc, finish_stream
+):
+    refs = []
+    stages = []
+    payload = {"name": "value"}
+
+    @app.errorhandler(404)
+    def handler(error):
+        refs.append(weakref.ref(request._get_current_object()))
+
+        @stream_with_context
+        def generate():
+            try:
+                for part in (b"first", b"second"):
+                    assert get_current_context().body == payload
+                    assert request.get_json() == payload
+                    assert error.__traceback__ is not None
+                    stages.append(part)
+                    yield part
+            finally:
+                assert get_current_context().body == payload
+                assert request.get_json() == payload
+
+        return Response(generate(), status=404)
+
+    response = app.test_client().post("/missing", json=payload, buffered=False)
+    assert response.status_code == 404
+    assert refs[0]() is not None
+    if finish_stream:
+        assert b"".join(response.response) == b"firstsecond"
+    response.close()
+    del response
+
+    assert stages == ([b"first", b"second"] if finish_stream else [b"first"])
+    assert refs[0]() is None
+
+
+def test_routing_exception_raised_by_stream_keeps_its_diagnostic_chain(app):
+    @app.errorhandler(404)
+    def handler(error):
+        def generate():
+            yield b"first"
+            raise error
+
+        return Response(generate(), status=404)
+
+    response = app.test_client().post("/missing", json={"name": "value"})
+    with pytest.raises(NotFound) as raised:
+        try:
+            b"".join(response.response)
+        finally:
+            response.close()
+    assert raised.value.__traceback__ is not None
+    assert raised.value.__context__ is not None
+
+
+def test_copied_request_context_can_read_body_after_original_request_closes(app):
+    callbacks = []
+    payload = {"name": "value"}
+
+    @app.errorhandler(404)
+    def handler(error):
+        @copy_current_request_context
+        def later():
+            return request.get_json()
+
+        callbacks.append(later)
+        return "handled", 404
+
+    response = app.test_client().post("/missing", json=payload)
+    response.close()
+    assert callbacks.pop()() == payload
+
+
+def test_propagating_routing_exception_keeps_its_diagnostic_chain(app):
+    app.config.update(TESTING=True, TRAP_HTTP_EXCEPTIONS=True)
+
+    with pytest.raises(NotFound) as raised:
+        app.test_client().post("/missing", json={"name": "value"})
+
+    assert raised.value.__context__ is not None
+    functions = []
+    trace = raised.value.__traceback__
+    while trace is not None:
+        functions.append(trace.tb_frame.f_code.co_name)
+        trace = trace.tb_next
+    assert "dispatch_request" in functions
+    assert "match_request" in functions
+
+
+def test_request_close_still_closes_uploaded_files(app):
+    files = []
+
+    @app.errorhandler(404)
+    def handler(error):
+        assert get_current_context().body.getlist("name") == ["upload"]
+        files.append(request.files["file"])
+        assert files[0].read() == b"x" * 1024 * 1024
+        return "handled", 404
+
+    response = app.test_client().post(
+        "/missing",
+        data={"name": "upload", "file": (BytesIO(b"x" * 1024 * 1024), "file.bin")},
+    )
+    assert response.status_code == 404
+    response.close()
+    assert files[0].closed
+
+
+@pytest.mark.parametrize("missing_route", [False, True])
+def test_handled_server_errors_release_bodies_and_preserve_tracebacks(
+    app, without_gc, missing_route
+):
+    refs = []
+    errors = []
+    formatted = []
+    payload = {"name": "value"}
+    app.logger.disabled = True
+
+    @app.before_request
+    def capture():
+        refs.append(weakref.ref(request._get_current_object()))
+        assert get_current_context().body == payload
+        if missing_route:
+            raise ValueError("application error")
+
+    @app.post("/error")
+    def view():
+        assert request.get_json() == payload
+        try:
+            fail()
+        except LookupError as cause:
+            raise ValueError("application error") from cause
+
+    def fail():
+        local_request = request._get_current_object()
+        assert local_request.get_json() == payload
+        raise LookupError("inner error")
+
+    @app.errorhandler(500)
+    def handler(error):
+        original = error.original_exception
+        errors.append(original)
+        formatted.append(
+            traceback.format_exception(type(original), original, original.__traceback__)
+        )
+        assert get_current_context().body == payload
+        return "handled", 500
+
+    response = app.test_client().post(
+        "/missing" if missing_route else "/error", json=payload
+    )
+    assert response.status_code == 500
+    response.close()
+    del response
+
+    assert refs[0]() is None
+    error = errors[0]
+    assert error.__traceback__ is not None
+    assert (
+        traceback.format_exception(type(error), error, error.__traceback__)
+        == formatted[0]
+    )
+
+
+@pytest.mark.parametrize("status", [200, 404, 500])
+def test_response_close_releases_sdk_context_without_another_request(app, status):
+    refs = []
+    payload = {"name": "value"}
+    app.logger.disabled = True
+
+    @app.before_request
+    def capture():
+        refs.append(weakref.ref(get_current_context()))
+        assert get_current_context().body == payload
+
+    @app.post("/body")
+    def view():
+        if status == 500:
+            raise ValueError("application error")
+        return "ok"
+
+    response = app.test_client().post(
+        "/missing" if status == 404 else "/body", json=payload
+    )
+    assert response.status_code == status
+    assert get_current_context().body == payload
+    response.close()
+    assert get_current_context() is None
+    assert refs[0]() is None
+
+
+def test_late_stream_close_callbacks_keep_body_inspection_until_finished(app):
+    payload = {"name": "value"}
+    closed = []
+
+    @app.post("/body")
+    def view():
+        def on_close():
+            assert get_current_context().body == payload
+            closed.append(True)
+
+        def generate():
+            response.call_on_close(on_close)
+            assert get_current_context().body == payload
+            yield b"ok"
+
+        response = Response(generate())
+        return response
+
+    response = app.test_client().post("/body", json=payload, buffered=False)
+    assert response.data == b"ok"
+    response.close()
+    assert closed == [True]
+    assert get_current_context() is None
+
+
+def test_response_close_preserves_another_context(app):
+    @app.post("/body")
+    def view():
+        return "ok"
+
+    response = app.test_client().post("/body", json={"name": "value"})
+    other = source.Context(body={"name": "other"})
+    other.set_as_current_context()
+    response.close()
+    assert get_current_context() is other
+
+
+def test_response_close_preserves_application_copies_of_context(app):
+    copies = []
+    payload = {"name": "value"}
+
+    @app.post("/body")
+    def view():
+        copies.append(copy_context())
+        return "ok"
+
+    response = app.test_client().post("/body", json=payload)
+    response.close()
+    assert get_current_context() is None
+    assert copies[0].run(get_current_context).body == payload
+
+
+def test_propagating_application_error_releases_sdk_context(app):
+    app.config["TESTING"] = True
+
+    @app.post("/body")
+    def view():
+        raise ValueError("application error")
+
+    with pytest.raises(ValueError):
+        app.test_client().post("/body", json={"name": "value"})
+    assert get_current_context() is None
+
+
+def test_bodyless_file_response_keeps_wsgi_file_wrapper(app):
+    @app.get("/file")
+    def view():
+        return Response(FileWrapper(BytesIO(b"file data")), direct_passthrough=True)
+
+    builder = EnvironBuilder(path="/file")
+    try:
+        response = app(builder.get_environ(), lambda status, headers: None)
+        assert isinstance(response, FileWrapper)
+        assert b"".join(response) == b"file data"
+        response.close()
+    finally:
+        builder.close()
+
+
+def test_wsgi_middleware_close_callbacks_keep_body_inspection(app):
+    payload = {"name": "value"}
+    closed = []
+    original = app.wsgi_app
+
+    def middleware(environ, start_response):
+        def on_close():
+            assert get_current_context().body == payload
+            closed.append(True)
+
+        return ClosingIterator(original(environ, start_response), on_close)
+
+    app.wsgi_app = middleware
+
+    @app.post("/body")
+    def view():
+        return "ok"
+
+    response = app.test_client().post("/body", json=payload)
+    response.close()
+    assert closed == [True]
+    assert get_current_context() is None

From 3675385871dc46ea4f6003a82f476f4e044e7111 Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Thu, 29 Oct 2026 18:11:47 +0100
Subject: [PATCH 2/6] Limit request body cleanup to Flask routing errors

---
 aikido_zen/sources/flask/__init__.py   |  70 +++------
 aikido_zen/sources/flask/flask_test.py | 190 -------------------------
 2 files changed, 16 insertions(+), 244 deletions(-)

diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index 9903cb3b6..b9a7c30e3 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -1,8 +1,7 @@
 from sys import exc_info
-from traceback import clear_frames
 from weakref import ref
 
-from aikido_zen.context import Context, current_context, get_current_context
+from aikido_zen.context import Context
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
 from .extract_cookies import extract_cookies_from_flask_request_and_save_data
@@ -25,63 +24,29 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 
 
 @after
-def _request_after(func, instance, args, kwargs, return_value):
+def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
     from flask.globals import request_ctx
 
     if not hasattr(return_value, "status_code"):
         return
-    error = get_argument(args, kwargs, 0, "e")
-    if error is None:
-        funcs.request_handler(
-            stage="post_response", status_code=return_value.status_code
-        )
-    if error is None and request_ctx.request.routing_exception is None:
+    funcs.request_handler(stage="post_response", status_code=return_value.status_code)
+
+    if request_ctx.request.routing_exception is None:
         return
     request_ref = ref(request_ctx.request)
 
-    # Finished traceback frames retain cached bodies until cyclic GC runs.
-    def release_request():
-        if exc_info()[0] is not None:
-            return
-        pending = [error] if error is not None else []
-        seen = set()
-        while pending:
-            exception = pending.pop()
-            if exception is None or id(exception) in seen:
-                continue
-            seen.add(id(exception))
-            pending.extend((exception.__cause__, exception.__context__))
-            clear_frames(exception.__traceback__)
+    def release_routing_exception():
         req = request_ref()
-        if req is not None and req.routing_exception is not None:
-            routing_error = req.routing_exception
-            routing_error.__traceback__ = None
-            routing_error.__context__ = None
-            routing_error.__cause__ = None
-
-    return_value.call_on_close(release_request)
-
-
-def _call_with_context_cleanup(func, instance, args, kwargs):
-    from werkzeug.wsgi import ClosingIterator
-
-    if get_current_context() is None:
-        return func(*args, **kwargs)
-    context_ref = ref(get_current_context())
-
-    def release_context():
-        if get_current_context() is context_ref():
-            current_context.set(None)
+        if req is None or exc_info()[0] is not None:
+            return
+        error = req.routing_exception
+        if error is not None:
+            # The request stores the routing error, whose tracebacks retain the request.
+            error.__traceback__ = None
+            error.__context__ = None
+            error.__cause__ = None
 
-    try:
-        response = func(*args, **kwargs)
-        # Wrapping file responses prevents the WSGI server's sendfile optimization.
-        if context_ref() is None or context_ref().body is None:
-            return response
-        return ClosingIterator(response, release_context)
-    except BaseException:
-        release_context()
-        raise
+    return_value.call_on_close(release_routing_exception)
 
 
 @before_modify_return
@@ -115,9 +80,6 @@ def patch(m):
     available, so we can extract data from it.
     (github src: https://github.com/pallets/flask/blob/bc143499cf1137a271a7cf75bdd3e16e43ede2f0/src/flask/app.py#L1529)
     """
-    # The response wrapper needs the context created by the outer hook.
-    patch_function(m, "Flask.__call__", _call_with_context_cleanup)
     patch_function(m, "Flask.__call__", _call)
     patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_before)
-    patch_function(m, "Flask.full_dispatch_request", _request_after)
-    patch_function(m, "Flask.handle_exception", _request_after)
+    patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_after)
diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py
index 9d5fe225f..010e710bb 100644
--- a/aikido_zen/sources/flask/flask_test.py
+++ b/aikido_zen/sources/flask/flask_test.py
@@ -1,7 +1,5 @@
 import gc
-from contextvars import copy_context
 import signal
-import traceback
 import weakref
 from io import BytesIO
 from unittest.mock import patch
@@ -15,8 +13,6 @@
     stream_with_context,
 )
 from werkzeug.exceptions import HTTPException, NotFound
-from werkzeug.test import EnvironBuilder
-from werkzeug.wsgi import ClosingIterator, FileWrapper
 
 import aikido_zen.sources.flask as source
 from aikido_zen.background_process.ipc import IPCClient
@@ -478,189 +474,3 @@ def handler(error):
     assert response.status_code == 404
     response.close()
     assert files[0].closed
-
-
-@pytest.mark.parametrize("missing_route", [False, True])
-def test_handled_server_errors_release_bodies_and_preserve_tracebacks(
-    app, without_gc, missing_route
-):
-    refs = []
-    errors = []
-    formatted = []
-    payload = {"name": "value"}
-    app.logger.disabled = True
-
-    @app.before_request
-    def capture():
-        refs.append(weakref.ref(request._get_current_object()))
-        assert get_current_context().body == payload
-        if missing_route:
-            raise ValueError("application error")
-
-    @app.post("/error")
-    def view():
-        assert request.get_json() == payload
-        try:
-            fail()
-        except LookupError as cause:
-            raise ValueError("application error") from cause
-
-    def fail():
-        local_request = request._get_current_object()
-        assert local_request.get_json() == payload
-        raise LookupError("inner error")
-
-    @app.errorhandler(500)
-    def handler(error):
-        original = error.original_exception
-        errors.append(original)
-        formatted.append(
-            traceback.format_exception(type(original), original, original.__traceback__)
-        )
-        assert get_current_context().body == payload
-        return "handled", 500
-
-    response = app.test_client().post(
-        "/missing" if missing_route else "/error", json=payload
-    )
-    assert response.status_code == 500
-    response.close()
-    del response
-
-    assert refs[0]() is None
-    error = errors[0]
-    assert error.__traceback__ is not None
-    assert (
-        traceback.format_exception(type(error), error, error.__traceback__)
-        == formatted[0]
-    )
-
-
-@pytest.mark.parametrize("status", [200, 404, 500])
-def test_response_close_releases_sdk_context_without_another_request(app, status):
-    refs = []
-    payload = {"name": "value"}
-    app.logger.disabled = True
-
-    @app.before_request
-    def capture():
-        refs.append(weakref.ref(get_current_context()))
-        assert get_current_context().body == payload
-
-    @app.post("/body")
-    def view():
-        if status == 500:
-            raise ValueError("application error")
-        return "ok"
-
-    response = app.test_client().post(
-        "/missing" if status == 404 else "/body", json=payload
-    )
-    assert response.status_code == status
-    assert get_current_context().body == payload
-    response.close()
-    assert get_current_context() is None
-    assert refs[0]() is None
-
-
-def test_late_stream_close_callbacks_keep_body_inspection_until_finished(app):
-    payload = {"name": "value"}
-    closed = []
-
-    @app.post("/body")
-    def view():
-        def on_close():
-            assert get_current_context().body == payload
-            closed.append(True)
-
-        def generate():
-            response.call_on_close(on_close)
-            assert get_current_context().body == payload
-            yield b"ok"
-
-        response = Response(generate())
-        return response
-
-    response = app.test_client().post("/body", json=payload, buffered=False)
-    assert response.data == b"ok"
-    response.close()
-    assert closed == [True]
-    assert get_current_context() is None
-
-
-def test_response_close_preserves_another_context(app):
-    @app.post("/body")
-    def view():
-        return "ok"
-
-    response = app.test_client().post("/body", json={"name": "value"})
-    other = source.Context(body={"name": "other"})
-    other.set_as_current_context()
-    response.close()
-    assert get_current_context() is other
-
-
-def test_response_close_preserves_application_copies_of_context(app):
-    copies = []
-    payload = {"name": "value"}
-
-    @app.post("/body")
-    def view():
-        copies.append(copy_context())
-        return "ok"
-
-    response = app.test_client().post("/body", json=payload)
-    response.close()
-    assert get_current_context() is None
-    assert copies[0].run(get_current_context).body == payload
-
-
-def test_propagating_application_error_releases_sdk_context(app):
-    app.config["TESTING"] = True
-
-    @app.post("/body")
-    def view():
-        raise ValueError("application error")
-
-    with pytest.raises(ValueError):
-        app.test_client().post("/body", json={"name": "value"})
-    assert get_current_context() is None
-
-
-def test_bodyless_file_response_keeps_wsgi_file_wrapper(app):
-    @app.get("/file")
-    def view():
-        return Response(FileWrapper(BytesIO(b"file data")), direct_passthrough=True)
-
-    builder = EnvironBuilder(path="/file")
-    try:
-        response = app(builder.get_environ(), lambda status, headers: None)
-        assert isinstance(response, FileWrapper)
-        assert b"".join(response) == b"file data"
-        response.close()
-    finally:
-        builder.close()
-
-
-def test_wsgi_middleware_close_callbacks_keep_body_inspection(app):
-    payload = {"name": "value"}
-    closed = []
-    original = app.wsgi_app
-
-    def middleware(environ, start_response):
-        def on_close():
-            assert get_current_context().body == payload
-            closed.append(True)
-
-        return ClosingIterator(original(environ, start_response), on_close)
-
-    app.wsgi_app = middleware
-
-    @app.post("/body")
-    def view():
-        return "ok"
-
-    response = app.test_client().post("/body", json=payload)
-    response.close()
-    assert closed == [True]
-    assert get_current_context() is None

From 1b2ad715e165839e9ac2585dd6ac58a65f3c9a45 Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Thu, 29 Oct 2026 18:30:07 +0100
Subject: [PATCH 3/6] Clean up Flask routing errors for passthrough responses

---
 aikido_zen/sources/flask/__init__.py   |  8 +++++++-
 aikido_zen/sources/flask/flask_test.py | 27 ++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index b9a7c30e3..6bcc74f92 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -26,6 +26,7 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 @after
 def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
     from flask.globals import request_ctx
+    from werkzeug.wsgi import ClosingIterator
 
     if not hasattr(return_value, "status_code"):
         return
@@ -46,7 +47,12 @@ def release_routing_exception():
             error.__context__ = None
             error.__cause__ = None
 
-    return_value.call_on_close(release_routing_exception)
+    if return_value.direct_passthrough:
+        return_value.response = ClosingIterator(
+            return_value.response, release_routing_exception
+        )
+    else:
+        return_value.call_on_close(release_routing_exception)
 
 
 @before_modify_return
diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py
index 010e710bb..43a043999 100644
--- a/aikido_zen/sources/flask/flask_test.py
+++ b/aikido_zen/sources/flask/flask_test.py
@@ -10,6 +10,7 @@
     Response,
     copy_current_request_context,
     request,
+    send_file,
     stream_with_context,
 )
 from werkzeug.exceptions import HTTPException, NotFound
@@ -366,8 +367,9 @@ def handle_http_exception(self, error):
 
 
 @pytest.mark.parametrize("finish_stream", [False, True])
+@pytest.mark.parametrize("direct_passthrough", [False, True])
 def test_streamed_error_response_keeps_body_until_stream_closes(
-    app, without_gc, finish_stream
+    app, without_gc, finish_stream, direct_passthrough
 ):
     refs = []
     stages = []
@@ -390,7 +392,7 @@ def generate():
                 assert get_current_context().body == payload
                 assert request.get_json() == payload
 
-        return Response(generate(), status=404)
+        return Response(generate(), status=404, direct_passthrough=direct_passthrough)
 
     response = app.test_client().post("/missing", json=payload, buffered=False)
     assert response.status_code == 404
@@ -404,6 +406,27 @@ def generate():
     assert refs[0]() is None
 
 
+def test_file_error_response_releases_request_and_closes_file(app, without_gc):
+    refs = []
+    file = BytesIO(b"not found")
+    payload = {"name": "value"}
+
+    @app.errorhandler(404)
+    def handler(error):
+        refs.append(weakref.ref(request._get_current_object()))
+        assert request.get_json() == payload
+        return send_file(file, mimetype="text/plain"), 404
+
+    response = app.test_client().post("/missing", json=payload)
+    assert response.status_code == 404
+    assert response.data == b"not found"
+    response.close()
+    del response
+
+    assert file.closed
+    assert refs[0]() is None
+
+
 def test_routing_exception_raised_by_stream_keeps_its_diagnostic_chain(app):
     @app.errorhandler(404)
     def handler(error):

From 89f8f7cecc6c3f10594eb52ca3ff42c2ea6d42ab Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Thu, 29 Oct 2026 20:30:01 +0100
Subject: [PATCH 4/6] Skip unused bodies for default Flask routing errors

---
 aikido_zen/sources/flask/__init__.py   |  71 +++--
 aikido_zen/sources/flask/flask_test.py | 377 ++++++++++++-------------
 2 files changed, 225 insertions(+), 223 deletions(-)

diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index 6bcc74f92..323131bbd 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -1,6 +1,3 @@
-from sys import exc_info
-from weakref import ref
-
 from aikido_zen.context import Context
 from aikido_zen.helpers.get_argument import get_argument
 from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
@@ -18,42 +15,58 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 
     req = request_ctx.request
     extract_cookies_from_flask_request_and_save_data(req)
-    extract_form_data_from_flask_request_and_save_data(req)
+    if _should_capture_body(instance, req):
+        extract_form_data_from_flask_request_and_save_data(req)
     extract_view_args_from_flask_request_and_save_data(req)
     extract_query_from_flask_request_and_save_data(req)
 
 
+def _should_capture_body(app, request):
+    error = request.routing_exception
+    if (
+        error is None
+        or error.code not in (404, 405)
+        or request.blueprints
+        or app.trap_http_exception(error)
+    ):
+        return True
+
+    # Routing errors can still reach application hooks and error handlers.
+    if any(
+        (
+            app.before_request_funcs,
+            app.url_value_preprocessors,
+            app.after_request_funcs,
+            app.teardown_request_funcs,
+            app.teardown_appcontext_funcs,
+        )
+    ):
+        return True
+
+    from flask import signals
+
+    if getattr(signals, "signals_available", True) and any(
+        signal.has_receivers_for(app)
+        for signal in (
+            signals.request_started,
+            signals.request_finished,
+            signals.request_tearing_down,
+            signals.appcontext_tearing_down,
+            signals.appcontext_popped,
+        )
+    ):
+        return True
+
+    handlers = app.error_handler_spec.get(None, {})
+    return isinstance(error, (*handlers.get(error.code, {}), *handlers.get(None, {})))
+
+
 @after
 def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
-    from flask.globals import request_ctx
-    from werkzeug.wsgi import ClosingIterator
-
     if not hasattr(return_value, "status_code"):
         return
     funcs.request_handler(stage="post_response", status_code=return_value.status_code)
 
-    if request_ctx.request.routing_exception is None:
-        return
-    request_ref = ref(request_ctx.request)
-
-    def release_routing_exception():
-        req = request_ref()
-        if req is None or exc_info()[0] is not None:
-            return
-        error = req.routing_exception
-        if error is not None:
-            # The request stores the routing error, whose tracebacks retain the request.
-            error.__traceback__ = None
-            error.__context__ = None
-            error.__cause__ = None
-
-    if return_value.direct_passthrough:
-        return_value.response = ClosingIterator(
-            return_value.response, release_routing_exception
-        )
-    else:
-        return_value.call_on_close(release_routing_exception)
-
 
 @before_modify_return
 def _call(func, instance, args, kwargs):
diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py
index 43a043999..8ff8fe67e 100644
--- a/aikido_zen/sources/flask/flask_test.py
+++ b/aikido_zen/sources/flask/flask_test.py
@@ -1,20 +1,10 @@
-import gc
+import pytest
 import signal
-import weakref
 from io import BytesIO
-from unittest.mock import patch
-
-import pytest
-from flask import (
-    Flask,
-    Response,
-    copy_current_request_context,
-    request,
-    send_file,
-    stream_with_context,
-)
-from werkzeug.exceptions import HTTPException, NotFound
 
+from flask import Flask, request, signals
+from werkzeug.exceptions import HTTPException
+from unittest.mock import patch
 import aikido_zen.sources.flask as source
 from aikido_zen.background_process.ipc import IPCClient
 from aikido_zen.context import current_context, get_current_context
@@ -250,6 +240,18 @@ def test_startup_flask():
         pass
 
 
+class CountingInput(BytesIO):
+    reads = 0
+
+    def read(self, *args, **kwargs):
+        self.reads += 1
+        return super().read(*args, **kwargs)
+
+    def readinto(self, *args, **kwargs):
+        self.reads += 1
+        return super().readinto(*args, **kwargs)
+
+
 @pytest.fixture
 def app(monkeypatch):
     monkeypatch.setattr(source, "should_block", lambda: None)
@@ -259,61 +261,101 @@ def app(monkeypatch):
     current_context.set(None)
 
 
-@pytest.fixture
-def without_gc():
-    enabled = gc.isenabled()
-    gc.collect()
-    gc.disable()
-    try:
-        yield
-    finally:
-        if enabled:
-            gc.enable()
-        gc.collect()
-
-
-@pytest.mark.parametrize("status", [404, 405, 308])
-@pytest.mark.parametrize("custom_handler", [False, True])
-def test_routing_requests_release_cached_bodies_without_gc(
-    app, without_gc, status, custom_handler
-):
-    refs = []
-    app.add_url_rule("/known", endpoint="known", view_func=lambda: "ok")
-    app.add_url_rule(
-        "/slash/", endpoint="slash", view_func=lambda: "ok", methods=["POST"]
+@pytest.mark.parametrize("status", [404, 405])
+def test_default_routing_rejection_does_not_read_body(app, status, monkeypatch):
+    stages = []
+    monkeypatch.setattr(
+        source.funcs, "request_handler", lambda **kwargs: stages.append(kwargs)
+    )
+    app.add_url_rule("/known", view_func=lambda: "ok")
+    stream = CountingInput(b"x" * 24 * 1024 * 1024)
+    path = "/missing" if status == 404 else "/known"
+    response = app.test_client(use_cookies=False).post(
+        path + "?name=value",
+        input_stream=stream,
+        content_type="application/octet-stream",
+        headers={"Cookie": "name=value"},
     )
-    if custom_handler:
-        app.register_error_handler(HTTPException, lambda error: ("handled", error.code))
+    assert response.status_code == status
+    assert stream.reads == 0
+    context = get_current_context()
+    assert context.body is None
+    assert context.query == {"name": ["value"]}
+    assert context.cookies == {"name": "value"}
+    assert context.method == "POST"
+    assert stages == [
+        {"stage": "init"},
+        {"stage": "post_response", "status_code": status},
+    ]
+    response.close()
 
-    @app.before_request
-    def capture():
-        refs.append(weakref.ref(request._get_current_object()))
-        assert get_current_context().body == "x" * 1024 * 1024
 
-    path = {404: "/missing", 405: "/known", 308: "/slash"}[status]
+def test_matched_route_still_captures_body_before_application_reads(app):
+    @app.post("/body")
+    def view():
+        assert get_current_context().body == {"name": "value"}
+        assert request.get_json() == {"name": "value"}
+        return "ok"
+
+    response = app.test_client().post("/body", json={"name": "value"})
+    assert response.status_code == 200
+    response.close()
+
+
+def test_skipping_unused_body_keeps_attack_wave_detection(app, monkeypatch):
+    from aikido_zen.vulnerabilities.attack_wave_detection.attack_wave_detector import (
+        AttackWaveDetector,
+    )
+
+    detector = AttackWaveDetector(attack_wave_threshold=2)
+    results = []
+
+    def handle(stage, status_code=0):
+        if stage == "post_response":
+            results.append(detector.is_attack_wave(get_current_context(), status_code))
+
+    monkeypatch.setattr(source.funcs, "request_handler", handle)
     client = app.test_client()
-    for _ in range(3):
-        response = client.post(
-            path, data=b"x" * 1024 * 1024, content_type="application/octet-stream"
-        )
-        assert response.status_code == status
+    for _ in range(2):
+        response = client.post("/.env", data=b"body")
+        assert response.status_code == 404
+        assert get_current_context().body is None
         response.close()
-        del response
+    assert results == [False, True]
+
+
+@pytest.mark.parametrize(
+    "hook",
+    [
+        "before_request",
+        "url_value_preprocessor",
+        "after_request",
+        "teardown_request",
+        "teardown_appcontext",
+    ],
+)
+def test_routing_rejections_keep_body_available_to_application_hooks(app, hook):
+    captured = []
 
-    assert all(ref() is None for ref in refs)
+    def callback(*args):
+        captured.append(get_current_context().body)
+        if hook == "after_request":
+            return args[0]
+
+    getattr(app, hook)(callback)
+    response = app.test_client().post("/missing", json={"name": "value"})
+    assert response.status_code == 404
+    response.close()
+    assert captured == [{"name": "value"}]
 
 
 @pytest.mark.parametrize("status", [404, 405])
-def test_custom_handlers_and_teardown_keep_body_inspection(app, without_gc, status):
+@pytest.mark.parametrize("handler_type", ["code", "http_exception", "exception"])
+def test_custom_routing_error_handler_keeps_body_inspection(app, status, handler_type):
     attack = "Malicious Pet', 'Gru from the Minions') --"
-    payload = {"name": attack}
-    refs = []
-    stages = []
     app.add_url_rule("/known", view_func=lambda: "ok")
 
-    def inspect_body():
-        assert get_current_context().body == payload
-        assert request.get_json() == payload
+    def handler(error):
         result = context_contains_sql_injection(
             sql=f"INSERT INTO pets(name) VALUES ('{attack}')",
             operation="psycopg2.Cursor.execute",
@@ -322,178 +364,125 @@ def inspect_body():
         )
         assert result["source"] == "body"
         assert result["payload"] == attack
-        assert request.routing_exception.__traceback__ is not None
-
-    @app.errorhandler(HTTPException)
-    def handler(error):
-        refs.append(weakref.ref(request._get_current_object()))
-        inspect_body()
-        stages.append("handler")
+        assert request.get_json() == {"name": attack}
         return "handled", error.code
 
-    @app.teardown_request
-    def teardown(error):
-        inspect_body()
-        stages.append("teardown")
-
-    response = app.test_client().post(
-        "/missing" if status == 404 else "/known", json=payload
-    )
+    error = {"code": status, "http_exception": HTTPException, "exception": Exception}[
+        handler_type
+    ]
+    app.register_error_handler(error, handler)
+    path = "/missing" if status == 404 else "/known"
+    response = app.test_client().post(path, json={"name": attack})
     assert response.status_code == status
     assert response.data == b"handled"
     response.close()
-    del response
-
-    assert stages == ["handler", "teardown"]
-    assert refs[0]() is None
-
-
-def test_custom_flask_error_override_keeps_body_and_releases_request(app, without_gc):
-    refs = []
-    payload = {"name": "value"}
-
-    class CustomFlask(Flask):
-        def handle_http_exception(self, error):
-            refs.append(weakref.ref(request._get_current_object()))
-            assert get_current_context().body == payload
-            assert request.get_json() == payload
-            return self.make_response(("handled", error.code))
-
-    response = CustomFlask(__name__).test_client().post("/missing", json=payload)
-    assert response.status_code == 404
-    response.close()
-    del response
-    assert refs[0]() is None
 
 
-@pytest.mark.parametrize("finish_stream", [False, True])
-@pytest.mark.parametrize("direct_passthrough", [False, True])
-def test_streamed_error_response_keeps_body_until_stream_closes(
-    app, without_gc, finish_stream, direct_passthrough
-):
-    refs = []
-    stages = []
-    payload = {"name": "value"}
-
-    @app.errorhandler(404)
-    def handler(error):
-        refs.append(weakref.ref(request._get_current_object()))
-
-        @stream_with_context
-        def generate():
-            try:
-                for part in (b"first", b"second"):
-                    assert get_current_context().body == payload
-                    assert request.get_json() == payload
-                    assert error.__traceback__ is not None
-                    stages.append(part)
-                    yield part
-            finally:
-                assert get_current_context().body == payload
-                assert request.get_json() == payload
-
-        return Response(generate(), status=404, direct_passthrough=direct_passthrough)
-
-    response = app.test_client().post("/missing", json=payload, buffered=False)
+def test_unrelated_error_handlers_do_not_force_body_capture(app):
+    app.register_error_handler(500, lambda error: ("server error", 500))
+    app.register_error_handler(ValueError, lambda error: ("value error", 400))
+    stream = CountingInput(b"body")
+    response = app.test_client().post("/missing", input_stream=stream)
     assert response.status_code == 404
-    assert refs[0]() is not None
-    if finish_stream:
-        assert b"".join(response.response) == b"firstsecond"
+    assert stream.reads == 0
+    assert get_current_context().body is None
     response.close()
-    del response
-
-    assert stages == ([b"first", b"second"] if finish_stream else [b"first"])
-    assert refs[0]() is None
 
 
-def test_file_error_response_releases_request_and_closes_file(app, without_gc):
-    refs = []
-    file = BytesIO(b"not found")
-    payload = {"name": "value"}
+def test_trapped_routing_error_keeps_body_for_server_error_handler(app):
+    app.config["TRAP_HTTP_EXCEPTIONS"] = True
+    captured = []
 
-    @app.errorhandler(404)
     def handler(error):
-        refs.append(weakref.ref(request._get_current_object()))
-        assert request.get_json() == payload
-        return send_file(file, mimetype="text/plain"), 404
+        captured.append(get_current_context().body)
+        return "handled", 500
 
-    response = app.test_client().post("/missing", json=payload)
-    assert response.status_code == 404
-    assert response.data == b"not found"
+    app.register_error_handler(500, handler)
+    response = app.test_client().post("/missing", json={"name": "value"})
+    assert response.status_code == 500
+    assert response.data == b"handled"
+    assert captured == [{"name": "value"}]
     response.close()
-    del response
-
-    assert file.closed
-    assert refs[0]() is None
 
 
-def test_routing_exception_raised_by_stream_keeps_its_diagnostic_chain(app):
-    @app.errorhandler(404)
-    def handler(error):
-        def generate():
-            yield b"first"
-            raise error
-
-        return Response(generate(), status=404)
+def test_before_request_can_handle_unmatched_request_with_body(app):
+    @app.before_request
+    def handle():
+        assert get_current_context().body == {"name": "value"}
+        assert request.get_json() == {"name": "value"}
+        return "handled"
 
     response = app.test_client().post("/missing", json={"name": "value"})
-    with pytest.raises(NotFound) as raised:
-        try:
-            b"".join(response.response)
-        finally:
-            response.close()
-    assert raised.value.__traceback__ is not None
-    assert raised.value.__context__ is not None
-
-
-def test_copied_request_context_can_read_body_after_original_request_closes(app):
-    callbacks = []
-    payload = {"name": "value"}
-
-    @app.errorhandler(404)
-    def handler(error):
-        @copy_current_request_context
-        def later():
-            return request.get_json()
+    assert response.status_code == 200
+    assert response.data == b"handled"
+    response.close()
 
-        callbacks.append(later)
-        return "handled", 404
 
-    response = app.test_client().post("/missing", json=payload)
-    response.close()
-    assert callbacks.pop()() == payload
+@pytest.mark.parametrize(
+    "name",
+    [
+        "request_started",
+        "request_finished",
+        "request_tearing_down",
+        "appcontext_tearing_down",
+        "appcontext_popped",
+    ],
+)
+def test_routing_rejections_keep_body_available_to_signal_receivers(app, name):
+    pytest.importorskip("blinker")
+    captured = []
 
+    def receiver(sender, **kwargs):
+        captured.append(get_current_context().body)
 
-def test_propagating_routing_exception_keeps_its_diagnostic_chain(app):
-    app.config.update(TESTING=True, TRAP_HTTP_EXCEPTIONS=True)
+    signal = getattr(signals, name)
+    with signal.connected_to(receiver, app):
+        response = app.test_client().post("/missing", json={"name": "value"})
+        assert response.status_code == 404
+        response.close()
+    assert captured == [{"name": "value"}]
 
-    with pytest.raises(NotFound) as raised:
-        app.test_client().post("/missing", json={"name": "value"})
 
-    assert raised.value.__context__ is not None
-    functions = []
-    trace = raised.value.__traceback__
-    while trace is not None:
-        functions.append(trace.tb_frame.f_code.co_name)
-        trace = trace.tb_next
-    assert "dispatch_request" in functions
-    assert "match_request" in functions
+def test_matched_route_returning_404_keeps_existing_body_capture(app):
+    @app.post("/body")
+    def view():
+        assert get_current_context().body == {"name": "value"}
+        return "missing", 404
 
+    response = app.test_client().post("/body", json={"name": "value"})
+    assert response.status_code == 404
+    response.close()
 
-def test_request_close_still_closes_uploaded_files(app):
-    files = []
 
+def test_custom_error_handler_keeps_multipart_files_and_fields(app):
     @app.errorhandler(404)
     def handler(error):
         assert get_current_context().body.getlist("name") == ["upload"]
-        files.append(request.files["file"])
-        assert files[0].read() == b"x" * 1024 * 1024
+        assert len(request.files["file"].read()) == 10 * 1024 * 1024
         return "handled", 404
 
     response = app.test_client().post(
         "/missing",
-        data={"name": "upload", "file": (BytesIO(b"x" * 1024 * 1024), "file.bin")},
+        data={"name": "upload", "file": (BytesIO(b"x" * 10 * 1024 * 1024), "file.bin")},
     )
     assert response.status_code == 404
+    assert response.data == b"handled"
+    response.close()
+
+
+def test_routing_exception_traceback_is_preserved_after_response_closes(app):
+    errors = []
+
+    @app.errorhandler(404)
+    def missing(error):
+        errors.append(error)
+        return "missing", 404
+
+    response = app.test_client().post("/missing", json={"name": "value"})
+    error = errors[0]
+    traceback, context, cause = error.__traceback__, error.__context__, error.__cause__
+    assert traceback is not None
     response.close()
-    assert files[0].closed
+    assert error.__traceback__ is traceback
+    assert error.__context__ is context
+    assert error.__cause__ is cause

From bed72aed52e3a5ecb6363dd6aef50dde53993b0c Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Thu, 29 Oct 2026 23:47:53 +0100
Subject: [PATCH 5/6] Capture Flask request bodies on demand

---
 aikido_zen/context/__init__.py                |  31 ++
 aikido_zen/context/lazy_body_test.py          | 102 ++++
 aikido_zen/helpers/lazy_body_stream.py        |  44 ++
 aikido_zen/helpers/lazy_body_stream_test.py   |  90 ++++
 aikido_zen/sources/flask/__init__.py          | 116 +++--
 aikido_zen/sources/flask/body_test.py         | 473 ++++++++++++++++++
 aikido_zen/sources/flask/extract_form_data.py |  15 +-
 aikido_zen/sources/flask/flask_test.py        | 257 +---------
 8 files changed, 819 insertions(+), 309 deletions(-)
 create mode 100644 aikido_zen/context/lazy_body_test.py
 create mode 100644 aikido_zen/helpers/lazy_body_stream.py
 create mode 100644 aikido_zen/helpers/lazy_body_stream_test.py
 create mode 100644 aikido_zen/sources/flask/body_test.py

diff --git a/aikido_zen/context/__init__.py b/aikido_zen/context/__init__.py
index 82de90b61..de3930878 100644
--- a/aikido_zen/context/__init__.py
+++ b/aikido_zen/context/__init__.py
@@ -6,6 +6,7 @@
 import json
 from json import JSONDecodeError
 from time import sleep
+from threading import RLock
 from typing import Dict
 from urllib.parse import parse_qs
 
@@ -108,6 +109,36 @@ def set_as_current_context(self):
     def set_cookies(self, cookies):
         self.cookies = cookies
 
+    @property
+    def body(self):
+        self.load_body()
+        return self.__dict__.get("body")
+
+    @body.setter
+    def body(self, value):
+        self.__dict__.pop("_body_loader", None)
+        self.__dict__["body"] = value
+
+    def set_body_loader(self, loader):
+        self.__dict__["body"] = None
+        self.__dict__["_body_loader"] = loader
+        self.__dict__["_body_lock"] = RLock()
+
+    def load_body(self):
+        lock = self.__dict__.get("_body_lock")
+        if lock is None:
+            return
+        with lock:
+            # The loader can read through hooks that access this property again.
+            loader = self.__dict__.pop("_body_loader", None)
+            if loader is not None:
+                try:
+                    self.set_body(loader())
+                except Exception as e:
+                    logger.debug("Exception occurred whilst loading body: %s", e)
+                finally:
+                    self.__dict__.pop("_body_lock", None)
+
     def set_body(self, body):
         try:
             self.set_body_internal(body)
diff --git a/aikido_zen/context/lazy_body_test.py b/aikido_zen/context/lazy_body_test.py
new file mode 100644
index 000000000..04172a6f4
--- /dev/null
+++ b/aikido_zen/context/lazy_body_test.py
@@ -0,0 +1,102 @@
+import pickle
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeout
+from threading import Event
+from unittest.mock import Mock
+
+import pytest
+
+from . import Context
+
+
+def test_body_is_loaded_once_on_first_access():
+    context = Context()
+    loader = Mock(return_value=b'{"name":"input"}')
+    context.set_body_loader(loader)
+    loader.assert_not_called()
+    assert context.body == {"name": "input"}
+    assert context.body == {"name": "input"}
+    loader.assert_called_once_with()
+
+
+def test_body_loader_is_not_reentered():
+    context = Context()
+    loader = Mock(side_effect=lambda: context.body or b"body")
+    context.set_body_loader(loader)
+    assert context.body == "body"
+    loader.assert_called_once_with()
+
+
+def test_concurrent_body_access_waits_for_the_complete_body():
+    context = Context()
+    started = Event()
+    second_started = Event()
+    release = Event()
+
+    def load():
+        started.set()
+        assert release.wait(5)
+        return b'{"name":"input"}'
+
+    loader = Mock(side_effect=load)
+    context.set_body_loader(loader)
+
+    def read_second():
+        second_started.set()
+        return context.body
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        first = pool.submit(lambda: context.body)
+        assert started.wait(5)
+        second = pool.submit(read_second)
+        try:
+            assert second_started.wait(5)
+            with pytest.raises(FutureTimeout):
+                second.result(0.1)
+        finally:
+            release.set()
+        assert first.result(5) == {"name": "input"}
+        assert second.result(5) == {"name": "input"}
+    loader.assert_called_once_with()
+
+
+def test_explicit_body_replaces_pending_loader():
+    context = Context()
+    loader = Mock(return_value=b"old")
+    context.set_body_loader(loader)
+    context.set_body(b"replacement")
+    assert context.body == "replacement"
+    loader.assert_not_called()
+
+
+def test_failed_loader_is_not_retried_on_every_inspection():
+    context = Context()
+    loader = Mock(side_effect=ValueError("body unavailable"))
+    context.set_body_loader(loader)
+    assert context.body is None
+    assert context.body is None
+    loader.assert_called_once_with()
+
+
+def test_pickle_loads_body_without_serializing_loader():
+    context = Context()
+    loader = Mock(return_value=b'{"name":"input"}')
+    context.set_body_loader(loader)
+    restored = pickle.loads(pickle.dumps(context))
+    assert restored.body == {"name": "input"}
+    assert "_body_loader" not in restored.__dict__
+    loader.assert_called_once_with()
+
+
+def test_body_loader_is_used_by_schema_discovery_without_application_read():
+    from aikido_zen.api_discovery.get_api_info import get_body_info
+    from aikido_zen.helpers.headers import Headers
+
+    context = Context()
+    context.headers = Headers({"CONTENT_TYPE": ["application/json"]})
+    loader = Mock(return_value=b'{"name":"input"}')
+    context.set_body_loader(loader)
+    assert get_body_info(context) == {
+        "type": "json",
+        "schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+    }
+    loader.assert_called_once_with()
diff --git a/aikido_zen/helpers/lazy_body_stream.py b/aikido_zen/helpers/lazy_body_stream.py
new file mode 100644
index 000000000..888064011
--- /dev/null
+++ b/aikido_zen/helpers/lazy_body_stream.py
@@ -0,0 +1,44 @@
+from wrapt import ObjectProxy
+
+
+class LazyBodyStream(ObjectProxy):
+    _READ_METHODS = (
+        "read",
+        "readall",
+        "read1",
+        "readinto",
+        "readinto1",
+        "readline",
+        "readlines",
+        "peek",
+    )
+
+    def __init__(self, stream, before_read):
+        super().__init__(stream)
+        self._self_before_read = before_read
+
+    def _prepare(self):
+        callback = self._self_before_read
+        # Body capture reads this stream too, so it must not trigger itself.
+        self._self_before_read = None
+        if callback is not None:
+            callback()
+
+    def __getattr__(self, name):
+        value = getattr(self.__wrapped__, name)
+        if name not in self._READ_METHODS or self._self_before_read is None:
+            return value
+
+        def read(*args, **kwargs):
+            self._prepare()
+            return getattr(self.__wrapped__, name)(*args, **kwargs)
+
+        return read
+
+    def __iter__(self):
+        self._prepare()
+        return iter(self.__wrapped__)
+
+    def __next__(self):
+        self._prepare()
+        return next(self.__wrapped__)
diff --git a/aikido_zen/helpers/lazy_body_stream_test.py b/aikido_zen/helpers/lazy_body_stream_test.py
new file mode 100644
index 000000000..d9b956d00
--- /dev/null
+++ b/aikido_zen/helpers/lazy_body_stream_test.py
@@ -0,0 +1,90 @@
+from io import BytesIO
+from unittest.mock import Mock
+
+import pytest
+
+from .lazy_body_stream import LazyBodyStream
+
+
+@pytest.mark.parametrize(
+    "method,args", [("read", (2,)), ("readline", ()), ("readlines", ())]
+)
+def test_capture_runs_before_reading_and_only_once(method, args):
+    original = BytesIO(b"first\nsecond\n")
+    callback = Mock(side_effect=lambda: original.seek(0))
+    stream = LazyBodyStream(original, callback)
+    assert original.tell() == 0
+    expected = getattr(BytesIO(b"first\nsecond\n"), method)(*args)
+    assert getattr(stream, method)(*args) == expected
+    stream.read()
+    callback.assert_called_once_with()
+
+
+def test_readinto_is_captured():
+    callback = Mock()
+    stream = LazyBodyStream(BytesIO(b"body"), callback)
+    buffer = bytearray(4)
+    assert stream.readinto(buffer) == 4
+    assert buffer == b"body"
+    callback.assert_called_once_with()
+
+
+def test_iteration_is_captured():
+    callback = Mock()
+    stream = LazyBodyStream(BytesIO(b"first\nsecond\n"), callback)
+    assert list(stream) == [b"first\n", b"second\n"]
+    callback.assert_called_once_with()
+
+
+def test_next_is_captured():
+    callback = Mock()
+    stream = LazyBodyStream(BytesIO(b"first\nsecond\n"), callback)
+    assert next(stream) == b"first\n"
+    assert next(stream) == b"second\n"
+    callback.assert_called_once_with()
+
+
+def test_loader_can_read_the_same_stream_without_recursion():
+    captured = []
+    stream = LazyBodyStream(BytesIO(b"body"), lambda: captured.append(stream.read()))
+    assert stream.read() == b""
+    assert captured == [b"body"]
+
+
+def test_read_uses_a_stream_replaced_by_the_loader():
+    def replace():
+        stream.__wrapped__ = BytesIO(b"replacement")
+
+    stream = LazyBodyStream(BytesIO(b"original"), replace)
+    assert stream.read(3) == b"rep"
+    assert stream.read() == b"lacement"
+
+
+def test_does_not_add_unsupported_read_methods():
+    class Input:
+        def read(self, size=-1):
+            return b"body"
+
+    stream = LazyBodyStream(Input(), Mock())
+    assert not hasattr(stream, "readinto")
+    assert stream.read() == b"body"
+
+
+def test_non_read_operations_do_not_load_body():
+    callback = Mock()
+    stream = LazyBodyStream(BytesIO(b"body"), callback)
+    assert stream.tell() == 0
+    assert stream.readable()
+    stream.close()
+    assert stream.closed
+    callback.assert_not_called()
+
+
+def test_original_read_error_is_propagated():
+    original = BytesIO(b"body")
+    original.close()
+    callback = Mock()
+    stream = LazyBodyStream(original, callback)
+    with pytest.raises(ValueError):
+        stream.read()
+    callback.assert_called_once_with()
diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index 323131bbd..4a06b995e 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -1,8 +1,19 @@
-from aikido_zen.context import Context
+from functools import partial
+from contextvars import copy_context
+from weakref import proxy, ref
+
+from aikido_zen.context import Context, get_current_context
 from aikido_zen.helpers.get_argument import get_argument
-from aikido_zen.sinks import on_import, patch_function, before_modify_return, after
+from aikido_zen.helpers.lazy_body_stream import LazyBodyStream
+from aikido_zen.sinks import (
+    on_import,
+    patch_function,
+    before_modify_return,
+    before,
+    after,
+)
 from .extract_cookies import extract_cookies_from_flask_request_and_save_data
-from .extract_form_data import extract_form_data_from_flask_request_and_save_data
+from .extract_form_data import extract_body_from_flask_request
 from .extract_query import extract_query_from_flask_request_and_save_data
 from .extract_view_args import extract_view_args_from_flask_request_and_save_data
 import aikido_zen.sources.functions.request_handler as funcs
@@ -15,59 +26,72 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
 
     req = request_ctx.request
     extract_cookies_from_flask_request_and_save_data(req)
-    if _should_capture_body(instance, req):
-        extract_form_data_from_flask_request_and_save_data(req)
     extract_view_args_from_flask_request_and_save_data(req)
     extract_query_from_flask_request_and_save_data(req)
-
-
-def _should_capture_body(app, request):
-    error = request.routing_exception
-    if (
-        error is None
-        or error.code not in (404, 405)
-        or request.blueprints
-        or app.trap_http_exception(error)
+    context = get_current_context()
+    if context is not None and (
+        req.content_length or req.environ.get("wsgi.input_terminated")
     ):
-        return True
-
-    # Routing errors can still reach application hooks and error handlers.
-    if any(
-        (
-            app.before_request_funcs,
-            app.url_value_preprocessors,
-            app.after_request_funcs,
-            app.teardown_request_funcs,
-            app.teardown_appcontext_funcs,
+        context.set_body_loader(partial(extract_body_from_flask_request, proxy(req)))
+        req.environ["aikido_zen.body_context"] = ref(context)
+        req.environ["wsgi.input"] = LazyBodyStream(
+            req.environ["wsgi.input"], context.load_body
         )
-    ):
-        return True
-
-    from flask import signals
-
-    if getattr(signals, "signals_available", True) and any(
-        signal.has_receivers_for(app)
-        for signal in (
-            signals.request_started,
-            signals.request_finished,
-            signals.request_tearing_down,
-            signals.appcontext_tearing_down,
-            signals.appcontext_popped,
-        )
-    ):
-        return True
-
-    handlers = app.error_handler_spec.get(None, {})
-    return isinstance(error, (*handlers.get(error.code, {}), *handlers.get(None, {})))
+        req.__dict__["stream"] = LazyBodyStream(req.stream, context.load_body)
 
 
 @after
 def _full_dispatch_request_after(func, instance, args, kwargs, return_value):
     if not hasattr(return_value, "status_code"):
         return
+    context = get_current_context()
+    if context is not None and "_body_loader" in context.__dict__:
+        from flask import has_request_context
+        from flask.globals import request_ctx
+        from werkzeug.wsgi import ClosingIterator
+
+        # Response generators can need these bindings after Flask pops them.
+        saved_context = copy_context()
+        request_ref = ref(request_ctx.request)
+        loader = context.__dict__["_body_loader"]
+
+        def load_body():
+            try:
+                if saved_context is not None:
+                    return saved_context.run(loader)
+                return loader()
+            finally:
+                # Deferred parsing can open uploads after Flask's request cleanup.
+                if (
+                    not has_request_context()
+                    or request_ctx.request is not request_ref()
+                ):
+                    req = request_ref()
+                    if req is not None:
+                        req.close()
+
+        def release_request_context():
+            nonlocal saved_context
+            saved_context = None
+
+        context.set_body_loader(load_body)
+        if return_value.direct_passthrough:
+            return_value.response = ClosingIterator(
+                return_value.response, release_request_context
+            )
+        else:
+            return_value.call_on_close(release_request_context)
     funcs.request_handler(stage="post_response", status_code=return_value.status_code)
 
 
+@before
+def _read_body_before(func, instance, args, kwargs):
+    context_ref = instance.environ.get("aikido_zen.body_context")
+    context = context_ref() if context_ref is not None else None
+    if context is not None:
+        context.load_body()
+
+
 @before_modify_return
 def _call(func, instance, args, kwargs):
     environ = get_argument(args, kwargs, 0, "environ")
@@ -102,3 +126,9 @@ def patch(m):
     patch_function(m, "Flask.__call__", _call)
     patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_before)
     patch_function(m, "Flask.full_dispatch_request", _full_dispatch_request_after)
+
+
+@on_import("flask.wrappers", "flask", version_requirement="2.2.4")
+def patch_request(m):
+    patch_function(m, "Request.get_data", _read_body_before)
+    patch_function(m, "Request._load_form_data", _read_body_before)
diff --git a/aikido_zen/sources/flask/body_test.py b/aikido_zen/sources/flask/body_test.py
new file mode 100644
index 000000000..182f907a9
--- /dev/null
+++ b/aikido_zen/sources/flask/body_test.py
@@ -0,0 +1,473 @@
+import gc
+import json
+from io import BytesIO
+from weakref import ref
+
+import pytest
+from flask import Flask, Response, copy_current_request_context, request, send_file
+from werkzeug.datastructures import MultiDict
+
+import aikido_zen.sources.flask as source
+from aikido_zen.api_discovery.get_api_info import get_api_info
+from aikido_zen.context import current_context, get_current_context
+from aikido_zen.vulnerabilities.nosql_injection import detect_nosql_injection
+from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
+    context_contains_sql_injection,
+)
+
+
+class CountingInput(BytesIO):
+    reads = 0
+
+    def read(self, *args, **kwargs):
+        self.reads += 1
+        return super().read(*args, **kwargs)
+
+    def readinto(self, *args, **kwargs):
+        self.reads += 1
+        return super().readinto(*args, **kwargs)
+
+
+@pytest.fixture
+def app(monkeypatch):
+    monkeypatch.setattr(source, "should_block", lambda: None)
+    monkeypatch.setattr(source.funcs, "request_handler", lambda **kwargs: None)
+    current_context.set(None)
+    yield Flask(__name__)
+    current_context.set(None)
+
+
+@pytest.mark.parametrize("status", [404, 405, 308])
+@pytest.mark.parametrize("hooks", [False, True])
+def test_unread_error_body_is_not_captured(app, status, hooks):
+    app.add_url_rule("/known", view_func=lambda: "ok")
+    app.add_url_rule(
+        "/redirect/", endpoint="redirect", view_func=lambda: "ok", methods=["POST"]
+    )
+    if hooks:
+        app.before_request(lambda: None)
+        app.after_request(lambda response: response)
+        app.register_error_handler(404, lambda error: ("missing", 404))
+    stream = CountingInput(b"x" * 24 * 1024 * 1024)
+    path = {404: "/missing", 405: "/known", 308: "/redirect"}[status]
+    response = app.test_client().post(
+        path, input_stream=stream, content_type="application/octet-stream"
+    )
+    assert response.status_code == status
+    assert stream.reads == 0
+    assert get_current_context().__dict__["body"] is None
+    response.close()
+
+
+@pytest.mark.parametrize(
+    "reader",
+    [
+        "data",
+        "get_data",
+        "get_data_text",
+        "get_data_no_cache",
+        "json",
+        "get_json",
+        "get_json_force",
+    ],
+)
+def test_json_capture_before_body_reader_returns(app, reader):
+    @app.post("/body")
+    def view():
+        if reader == "data":
+            data = request.data
+        elif reader == "get_data_text":
+            data = request.get_data(as_text=True)
+        elif reader == "get_data_no_cache":
+            data = request.get_data(cache=False)
+        elif reader == "json":
+            data = request.json
+        elif reader == "get_json_force":
+            data = request.get_json(force=True)
+        else:
+            data = getattr(request, reader)()
+        assert data
+        assert get_current_context().body == {"name": "value"}
+        return "ok"
+
+    content_type = "text/plain" if reader == "get_json_force" else "application/json"
+    response = app.test_client().post(
+        "/body", data=b'{"name":"value"}', content_type=content_type
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+@pytest.mark.parametrize(
+    "reader",
+    [
+        "read",
+        "readline",
+        "readlines",
+        "readinto",
+        "iter",
+        "next",
+        "input_stream",
+        "wsgi_input",
+    ],
+)
+def test_raw_stream_reads_keep_body_available_for_inspection(app, reader):
+    @app.post("/body")
+    def view():
+        if reader == "wsgi_input":
+            request.environ["wsgi.input"].read()
+        elif reader == "input_stream":
+            request.input_stream.read()
+        elif reader == "readinto":
+            request.stream.readinto(bytearray(3))
+        elif reader == "iter":
+            list(request.stream)
+        elif reader == "next":
+            next(request.stream, None)
+        else:
+            getattr(request.stream, reader)()
+        assert get_current_context().body == {"name": "value"}
+        return "ok"
+
+    response = app.test_client().post(
+        "/body", data=b'{"name":"value"}', content_type="application/json"
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+def test_inspection_can_load_body_without_application_read(app):
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+
+    @app.post("/body")
+    def view():
+        result = context_contains_sql_injection(
+            sql=f"INSERT INTO pets(name) VALUES ('{attack}')",
+            operation="psycopg2.Cursor.execute",
+            context=get_current_context(),
+            dialect="postgres",
+        )
+        assert result["source"] == "body"
+        assert result["payload"] == attack
+        return "ok"
+
+    response = app.test_client().post("/body", json={"name": attack})
+    assert response.status_code == 200
+    response.close()
+
+
+def test_nosql_inspection_can_load_body_without_application_read(app):
+    @app.post("/body")
+    def view():
+        result = detect_nosql_injection(get_current_context(), {"title": {"$ne": None}})
+        assert result["source"] == "body"
+        assert result["injection"]
+        return "ok"
+
+    response = app.test_client().post("/body", json={"title": {"$ne": None}})
+    assert response.status_code == 200
+    response.close()
+
+
+def test_schema_discovery_can_load_unused_body(app, monkeypatch):
+    schemas = []
+
+    def discover(stage, status_code=0):
+        if stage == "post_response":
+            schemas.append(get_api_info(get_current_context()))
+
+    monkeypatch.setattr(source.funcs, "request_handler", discover)
+    app.add_url_rule("/body", methods=["POST"], view_func=lambda: "ok")
+    response = app.test_client().post("/body", json={"name": "value"})
+    assert response.status_code == 200
+    assert schemas[0]["body"] == {
+        "type": "json",
+        "schema": {"type": "object", "properties": {"name": {"type": "string"}}},
+    }
+    response.close()
+
+
+def test_form_fields_and_duplicates_are_preserved(app):
+    @app.post("/body")
+    def view():
+        assert request.form.getlist("name") == ["first", "second"]
+        assert get_current_context().body.getlist("name") == ["first", "second"]
+        return "ok"
+
+    response = app.test_client().post(
+        "/body", data=MultiDict([("name", "first"), ("name", "second")])
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+def test_large_file_is_available_and_not_added_to_inspected_form(app):
+    @app.post("/body")
+    def view():
+        assert len(request.files["file"].read()) == 10 * 1024 * 1024
+        assert get_current_context().body.getlist("name") == ["upload"]
+        assert "file" not in get_current_context().body
+        assert request.form["name"] == "upload"
+        return "ok"
+
+    response = app.test_client().post(
+        "/body",
+        data={
+            "name": "upload",
+            "file": (BytesIO(b"x" * 10 * 1024 * 1024), "large.bin"),
+        },
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+def test_large_json_remains_supported(app):
+    @app.post("/body")
+    def view():
+        data = request.get_json()
+        assert len(data["padding"]) == 10 * 1024 * 1024
+        assert get_current_context().body == data
+        return "ok"
+
+    response = app.test_client().post("/body", json={"padding": "x" * 10 * 1024 * 1024})
+    assert response.status_code == 200
+    response.close()
+
+
+def test_chunked_body_without_content_length_is_captured(app):
+    @app.post("/body")
+    def view():
+        assert request.get_json() == {"name": "value"}
+        assert get_current_context().body == {"name": "value"}
+        return "ok"
+
+    response = app.test_client().post(
+        "/body",
+        data=b'{"name":"value"}',
+        content_type="application/json",
+        environ_overrides={"CONTENT_LENGTH": "", "wsgi.input_terminated": True},
+    )
+    assert response.status_code == 200
+    response.close()
+
+
+def test_application_json_mutations_do_not_change_captured_input(app):
+    @app.post("/body")
+    def view():
+        request.get_json()["name"] = "changed"
+        assert get_current_context().body == {"name": "original"}
+        return "ok"
+
+    response = app.test_client().post("/body", json={"name": "original"})
+    assert response.status_code == 200
+    response.close()
+
+
+def test_direct_passthrough_response_keeps_body_available(app):
+    @app.post("/body")
+    def view():
+        def generate():
+            gc.collect()
+            yield json.dumps(get_current_context().body).encode()
+
+        return Response(generate(), direct_passthrough=True)
+
+    response = app.test_client().post("/body", json={"name": "value"}, buffered=True)
+    assert response.status_code == 200
+    assert json.loads(response.data) == {"name": "value"}
+    response.close()
+
+
+def test_file_error_leaves_unused_body_unread(app, tmp_path):
+    path = tmp_path / "missing.txt"
+    path.write_bytes(b"missing")
+
+    @app.errorhandler(404)
+    def missing(error):
+        return send_file(path), 404
+
+    stream = CountingInput(b"x" * 1024 * 1024)
+    response = app.test_client().post("/missing", input_stream=stream)
+    assert response.status_code == 404
+    assert response.data == b"missing"
+    response.close()
+    assert stream.reads == 0
+
+
+def test_routing_exception_traceback_is_preserved_after_response_close(app):
+    errors = []
+
+    @app.errorhandler(404)
+    def missing(error):
+        errors.append(error)
+        return "missing", 404
+
+    response = app.test_client().post("/missing", json={"name": "value"})
+    error = errors[0]
+    traceback, context, cause = error.__traceback__, error.__context__, error.__cause__
+    assert traceback is not None
+    response.close()
+    assert error.__traceback__ is traceback
+    assert error.__context__ is context
+    assert error.__cause__ is cause
+
+
+@pytest.mark.parametrize("inspect_first", [False, True])
+@pytest.mark.parametrize("intervening_request", [False, True])
+@pytest.mark.parametrize("encoding", ["json", "form", "multipart"])
+def test_copied_request_context_can_capture_body_after_response_closes(
+    app, inspect_first, intervening_request, encoding
+):
+    callbacks = []
+    uploads = []
+    attack = "Malicious Pet', 'Gru from the Minions') --"
+    app.add_url_rule("/other", view_func=lambda: "ok", methods=["POST"])
+
+    @app.errorhandler(404)
+    def missing(error):
+        context = get_current_context()
+
+        @copy_current_request_context
+        def later():
+            if inspect_first:
+                result = context_contains_sql_injection(
+                    sql=f"INSERT INTO pets(name) VALUES ('{attack}')",
+                    operation="psycopg2.Cursor.execute",
+                    context=context,
+                    dialect="postgres",
+                )
+                assert result["source"] == "body"
+            if encoding == "json":
+                assert request.get_json() == {"name": attack}
+            else:
+                assert request.form["name"] == attack
+            assert context.body["name"] == attack
+            if encoding == "multipart":
+                assert request.files["file"].read() == b"upload"
+                uploads.append(request.files["file"].stream)
+
+        callbacks.append(later)
+        return "missing", 404
+
+    client = app.test_client()
+    if encoding == "json":
+        response = client.post("/missing", json={"name": attack})
+    else:
+        data = {"name": attack}
+        if encoding == "multipart":
+            data["file"] = (BytesIO(b"upload"), "file.txt")
+        response = client.post("/missing", data=data)
+    assert response.status_code == 404
+    response.close()
+    assert get_current_context().__dict__["body"] is None
+    if intervening_request:
+        other = client.post("/other", json={"other": "request"})
+        other.close()
+    callbacks.pop()()
+    if intervening_request:
+        assert get_current_context().__dict__["body"] is None
+    assert all(upload.closed for upload in uploads)
+
+
+def test_body_can_be_captured_by_application_response_close_callback(app):
+    captured = []
+
+    @app.post("/body")
+    def view():
+        req = request._get_current_object()
+        context = get_current_context()
+        response = Response("ok")
+
+        def on_close():
+            captured.append((req.get_json(), context.body))
+
+        response.call_on_close(on_close)
+        return response
+
+    response = app.test_client().post("/body", json={"name": "value"})
+    response.close()
+    assert captured == [({"name": "value"}, {"name": "value"})]
+
+
+def test_streamed_response_can_inspect_body_after_request_teardown(app):
+    @app.post("/body")
+    def view():
+        def generate():
+            gc.collect()
+            yield json.dumps(get_current_context().body)
+
+        return generate()
+
+    response = app.test_client().post("/body", json={"name": "value"}, buffered=True)
+    assert response.status_code == 200
+    assert json.loads(response.data) == {"name": "value"}
+    response.close()
+
+
+def test_uploads_opened_during_response_streaming_are_closed(app):
+    streams = []
+
+    @app.post("/body")
+    def view():
+        req = request._get_current_object()
+
+        def generate():
+            assert get_current_context().body.getlist("name") == ["upload"]
+            streams.append(req.files["file"].stream)
+            yield b"ok"
+
+        return generate()
+
+    response = app.test_client().post(
+        "/body",
+        data={"name": "upload", "file": (BytesIO(b"file"), "file.bin")},
+        buffered=True,
+    )
+    assert response.status_code == 200
+    response.close()
+    assert streams[0].closed
+
+
+def test_unread_request_is_not_retained_by_body_loader(app):
+    requests = []
+
+    @app.before_request
+    def remember():
+        requests.append(ref(request._get_current_object()))
+
+    response = app.test_client().post("/missing", data=b"body")
+    response.close()
+    del response
+    gc.collect()
+    assert requests[0]() is None
+    assert get_current_context().__dict__["body"] is None
+
+
+@pytest.mark.parametrize(
+    "data,status", [(b'{"invalid":', 400), (b'{"name":"\xff"}', 400)]
+)
+def test_invalid_json_errors_are_preserved(app, data, status):
+    @app.post("/body")
+    def view():
+        request.get_json()
+        return "ok"
+
+    response = app.test_client().post(
+        "/body", data=data, content_type="application/json"
+    )
+    assert response.status_code == status
+    response.close()
+
+
+def test_application_body_limit_is_preserved(app):
+    app.config["MAX_CONTENT_LENGTH"] = 8
+
+    @app.post("/body")
+    def view():
+        assert get_current_context().query == {"search": ["value"]}
+        request.get_data()
+        return "ok"
+
+    response = app.test_client().post("/body?search=value", data=b"x" * 32)
+    assert response.status_code == 413
+    response.close()
diff --git a/aikido_zen/sources/flask/extract_form_data.py b/aikido_zen/sources/flask/extract_form_data.py
index 991ed16cf..5ed4c4571 100644
--- a/aikido_zen/sources/flask/extract_form_data.py
+++ b/aikido_zen/sources/flask/extract_form_data.py
@@ -1,16 +1,11 @@
-from aikido_zen.context import get_current_context
 from aikido_zen.helpers.logging import logger
 
 
-def extract_form_data_from_flask_request_and_save_data(req):
-    """Extract form data from flask request"""
-    context = get_current_context()
+def extract_body_from_flask_request(req):
     try:
-        if context:
-            if req.form:
-                context.set_body(req.form)
-            else:
-                context.set_body(req.data)
-            context.set_as_current_context()
+        if req.form:
+            return req.form
+        return req.data
     except Exception as e:
         logger.debug("Exception occurred whilst extracting flask body data: %s", e)
+    return None
diff --git a/aikido_zen/sources/flask/flask_test.py b/aikido_zen/sources/flask/flask_test.py
index 8ff8fe67e..1aa58a4b0 100644
--- a/aikido_zen/sources/flask/flask_test.py
+++ b/aikido_zen/sources/flask/flask_test.py
@@ -1,16 +1,9 @@
 import pytest
 import signal
-from io import BytesIO
-
-from flask import Flask, request, signals
-from werkzeug.exceptions import HTTPException
 from unittest.mock import patch
-import aikido_zen.sources.flask as source
+import aikido_zen.sources.flask
 from aikido_zen.background_process.ipc import IPCClient
 from aikido_zen.context import current_context, get_current_context
-from aikido_zen.vulnerabilities.sql_injection.context_contains_sql_injection import (
-    context_contains_sql_injection,
-)
 
 sample_environ = {
     "REQUEST_METHOD": "POST",
@@ -238,251 +231,3 @@ def test_startup_flask():
         app.run()
     except TimeoutException:
         pass
-
-
-class CountingInput(BytesIO):
-    reads = 0
-
-    def read(self, *args, **kwargs):
-        self.reads += 1
-        return super().read(*args, **kwargs)
-
-    def readinto(self, *args, **kwargs):
-        self.reads += 1
-        return super().readinto(*args, **kwargs)
-
-
-@pytest.fixture
-def app(monkeypatch):
-    monkeypatch.setattr(source, "should_block", lambda: None)
-    monkeypatch.setattr(source.funcs, "request_handler", lambda **kwargs: None)
-    current_context.set(None)
-    yield Flask(__name__)
-    current_context.set(None)
-
-
-@pytest.mark.parametrize("status", [404, 405])
-def test_default_routing_rejection_does_not_read_body(app, status, monkeypatch):
-    stages = []
-    monkeypatch.setattr(
-        source.funcs, "request_handler", lambda **kwargs: stages.append(kwargs)
-    )
-    app.add_url_rule("/known", view_func=lambda: "ok")
-    stream = CountingInput(b"x" * 24 * 1024 * 1024)
-    path = "/missing" if status == 404 else "/known"
-    response = app.test_client(use_cookies=False).post(
-        path + "?name=value",
-        input_stream=stream,
-        content_type="application/octet-stream",
-        headers={"Cookie": "name=value"},
-    )
-    assert response.status_code == status
-    assert stream.reads == 0
-    context = get_current_context()
-    assert context.body is None
-    assert context.query == {"name": ["value"]}
-    assert context.cookies == {"name": "value"}
-    assert context.method == "POST"
-    assert stages == [
-        {"stage": "init"},
-        {"stage": "post_response", "status_code": status},
-    ]
-    response.close()
-
-
-def test_matched_route_still_captures_body_before_application_reads(app):
-    @app.post("/body")
-    def view():
-        assert get_current_context().body == {"name": "value"}
-        assert request.get_json() == {"name": "value"}
-        return "ok"
-
-    response = app.test_client().post("/body", json={"name": "value"})
-    assert response.status_code == 200
-    response.close()
-
-
-def test_skipping_unused_body_keeps_attack_wave_detection(app, monkeypatch):
-    from aikido_zen.vulnerabilities.attack_wave_detection.attack_wave_detector import (
-        AttackWaveDetector,
-    )
-
-    detector = AttackWaveDetector(attack_wave_threshold=2)
-    results = []
-
-    def handle(stage, status_code=0):
-        if stage == "post_response":
-            results.append(detector.is_attack_wave(get_current_context(), status_code))
-
-    monkeypatch.setattr(source.funcs, "request_handler", handle)
-    client = app.test_client()
-    for _ in range(2):
-        response = client.post("/.env", data=b"body")
-        assert response.status_code == 404
-        assert get_current_context().body is None
-        response.close()
-    assert results == [False, True]
-
-
-@pytest.mark.parametrize(
-    "hook",
-    [
-        "before_request",
-        "url_value_preprocessor",
-        "after_request",
-        "teardown_request",
-        "teardown_appcontext",
-    ],
-)
-def test_routing_rejections_keep_body_available_to_application_hooks(app, hook):
-    captured = []
-
-    def callback(*args):
-        captured.append(get_current_context().body)
-        if hook == "after_request":
-            return args[0]
-
-    getattr(app, hook)(callback)
-    response = app.test_client().post("/missing", json={"name": "value"})
-    assert response.status_code == 404
-    response.close()
-    assert captured == [{"name": "value"}]
-
-
-@pytest.mark.parametrize("status", [404, 405])
-@pytest.mark.parametrize("handler_type", ["code", "http_exception", "exception"])
-def test_custom_routing_error_handler_keeps_body_inspection(app, status, handler_type):
-    attack = "Malicious Pet', 'Gru from the Minions') --"
-    app.add_url_rule("/known", view_func=lambda: "ok")
-
-    def handler(error):
-        result = context_contains_sql_injection(
-            sql=f"INSERT INTO pets(name) VALUES ('{attack}')",
-            operation="psycopg2.Cursor.execute",
-            context=get_current_context(),
-            dialect="postgres",
-        )
-        assert result["source"] == "body"
-        assert result["payload"] == attack
-        assert request.get_json() == {"name": attack}
-        return "handled", error.code
-
-    error = {"code": status, "http_exception": HTTPException, "exception": Exception}[
-        handler_type
-    ]
-    app.register_error_handler(error, handler)
-    path = "/missing" if status == 404 else "/known"
-    response = app.test_client().post(path, json={"name": attack})
-    assert response.status_code == status
-    assert response.data == b"handled"
-    response.close()
-
-
-def test_unrelated_error_handlers_do_not_force_body_capture(app):
-    app.register_error_handler(500, lambda error: ("server error", 500))
-    app.register_error_handler(ValueError, lambda error: ("value error", 400))
-    stream = CountingInput(b"body")
-    response = app.test_client().post("/missing", input_stream=stream)
-    assert response.status_code == 404
-    assert stream.reads == 0
-    assert get_current_context().body is None
-    response.close()
-
-
-def test_trapped_routing_error_keeps_body_for_server_error_handler(app):
-    app.config["TRAP_HTTP_EXCEPTIONS"] = True
-    captured = []
-
-    def handler(error):
-        captured.append(get_current_context().body)
-        return "handled", 500
-
-    app.register_error_handler(500, handler)
-    response = app.test_client().post("/missing", json={"name": "value"})
-    assert response.status_code == 500
-    assert response.data == b"handled"
-    assert captured == [{"name": "value"}]
-    response.close()
-
-
-def test_before_request_can_handle_unmatched_request_with_body(app):
-    @app.before_request
-    def handle():
-        assert get_current_context().body == {"name": "value"}
-        assert request.get_json() == {"name": "value"}
-        return "handled"
-
-    response = app.test_client().post("/missing", json={"name": "value"})
-    assert response.status_code == 200
-    assert response.data == b"handled"
-    response.close()
-
-
-@pytest.mark.parametrize(
-    "name",
-    [
-        "request_started",
-        "request_finished",
-        "request_tearing_down",
-        "appcontext_tearing_down",
-        "appcontext_popped",
-    ],
-)
-def test_routing_rejections_keep_body_available_to_signal_receivers(app, name):
-    pytest.importorskip("blinker")
-    captured = []
-
-    def receiver(sender, **kwargs):
-        captured.append(get_current_context().body)
-
-    signal = getattr(signals, name)
-    with signal.connected_to(receiver, app):
-        response = app.test_client().post("/missing", json={"name": "value"})
-        assert response.status_code == 404
-        response.close()
-    assert captured == [{"name": "value"}]
-
-
-def test_matched_route_returning_404_keeps_existing_body_capture(app):
-    @app.post("/body")
-    def view():
-        assert get_current_context().body == {"name": "value"}
-        return "missing", 404
-
-    response = app.test_client().post("/body", json={"name": "value"})
-    assert response.status_code == 404
-    response.close()
-
-
-def test_custom_error_handler_keeps_multipart_files_and_fields(app):
-    @app.errorhandler(404)
-    def handler(error):
-        assert get_current_context().body.getlist("name") == ["upload"]
-        assert len(request.files["file"].read()) == 10 * 1024 * 1024
-        return "handled", 404
-
-    response = app.test_client().post(
-        "/missing",
-        data={"name": "upload", "file": (BytesIO(b"x" * 10 * 1024 * 1024), "file.bin")},
-    )
-    assert response.status_code == 404
-    assert response.data == b"handled"
-    response.close()
-
-
-def test_routing_exception_traceback_is_preserved_after_response_closes(app):
-    errors = []
-
-    @app.errorhandler(404)
-    def missing(error):
-        errors.append(error)
-        return "missing", 404
-
-    response = app.test_client().post("/missing", json={"name": "value"})
-    error = errors[0]
-    traceback, context, cause = error.__traceback__, error.__context__, error.__cause__
-    assert traceback is not None
-    response.close()
-    assert error.__traceback__ is traceback
-    assert error.__context__ is context
-    assert error.__cause__ is cause

From 9c6c9d09d3ff3740779e460fbd9096148f30e359 Mon Sep 17 00:00:00 2001
From: teta2k <teta2k@yahoo.com>
Date: Fri, 30 Oct 2026 11:56:23 +0100
Subject: [PATCH 6/6] Fix Flask body capture test isolation and proxy lint

---
 aikido_zen/helpers/lazy_body_stream.py |  3 +++
 aikido_zen/sources/flask/__init__.py   | 15 ++++++++-------
 aikido_zen/sources/flask/body_test.py  |  5 ++++-
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/aikido_zen/helpers/lazy_body_stream.py b/aikido_zen/helpers/lazy_body_stream.py
index 888064011..d67e396ea 100644
--- a/aikido_zen/helpers/lazy_body_stream.py
+++ b/aikido_zen/helpers/lazy_body_stream.py
@@ -1,3 +1,6 @@
+# Proxy members are supplied by wrapt's C extension.
+# pylint: disable=abstract-method,no-member
+
 from wrapt import ObjectProxy
 
 
diff --git a/aikido_zen/sources/flask/__init__.py b/aikido_zen/sources/flask/__init__.py
index 4a06b995e..9e34b34b1 100644
--- a/aikido_zen/sources/flask/__init__.py
+++ b/aikido_zen/sources/flask/__init__.py
@@ -29,15 +29,16 @@ def _full_dispatch_request_before(func, instance, args, kwargs):
     extract_view_args_from_flask_request_and_save_data(req)
     extract_query_from_flask_request_and_save_data(req)
     context = get_current_context()
-    if context is not None and (
+    if context is None or not (
         req.content_length or req.environ.get("wsgi.input_terminated")
     ):
-        context.set_body_loader(partial(extract_body_from_flask_request, proxy(req)))
-        req.environ["aikido_zen.body_context"] = ref(context)
-        req.environ["wsgi.input"] = LazyBodyStream(
-            req.environ["wsgi.input"], context.load_body
-        )
-        req.__dict__["stream"] = LazyBodyStream(req.stream, context.load_body)
+        return
+    context.set_body_loader(partial(extract_body_from_flask_request, proxy(req)))
+    req.environ["aikido_zen.body_context"] = ref(context)
+    req.environ["wsgi.input"] = LazyBodyStream(
+        req.environ["wsgi.input"], context.load_body
+    )
+    req.__dict__["stream"] = LazyBodyStream(req.stream, context.load_body)
 
 
 @after
diff --git a/aikido_zen/sources/flask/body_test.py b/aikido_zen/sources/flask/body_test.py
index 182f907a9..5c46dbb0f 100644
--- a/aikido_zen/sources/flask/body_test.py
+++ b/aikido_zen/sources/flask/body_test.py
@@ -281,10 +281,12 @@ def generate():
 def test_file_error_leaves_unused_body_unread(app, tmp_path):
     path = tmp_path / "missing.txt"
     path.write_bytes(b"missing")
+    # File I/O during the request legitimately triggers body inspection.
+    file = path.open("rb")
 
     @app.errorhandler(404)
     def missing(error):
-        return send_file(path), 404
+        return send_file(file, download_name=path.name, mimetype="text/plain"), 404
 
     stream = CountingInput(b"x" * 1024 * 1024)
     response = app.test_client().post("/missing", input_stream=stream)
@@ -292,6 +294,7 @@ def missing(error):
     assert response.data == b"missing"
     response.close()
     assert stream.reads == 0
+    assert file.closed
 
 
 def test_routing_exception_traceback_is_preserved_after_response_close(app):