Capture Flask request bodies only when needed

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)

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

aikido_zen/helpers/lazy_body_stream.py

@@ -0,0 +1,47 @@
+# Proxy members are supplied by wrapt's C extension.
+# pylint: disable=abstract-method,no-member
+
+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__)

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

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,18 +26,73 @@ 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)
     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 None or not (
+        req.content_length or req.environ.get("wsgi.input_terminated")
+    ):
+        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
 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")
@@ -61,3 +127,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)