fix: Read request body after handler so receive stream is not deprived

sentry_sdk/_types.py

@@ -76,6 +76,22 @@ def removed_because_over_size_limit(cls, value=""):
             },
         )
 
+    @classmethod
+    def removed_because_body_consumed_and_not_cached(cls, value=""):
+        # type: (Any) -> AnnotatedValue
+        """The actual value was removed because the underlying stream has been consumed, without caching the value."""
+        return AnnotatedValue(
+            value=value,
+            metadata={
+                "rem": [  # Remark
+                    [
+                        "!consumed",  # Because the original stream is unavailable
+                        "x",  # The fields original value was removed
+                    ]
+                ]
+            },
+        )
+
     @classmethod
     def substituted_because_contains_sensitive_data(cls):
         # type: () -> AnnotatedValue

sentry_sdk/integrations/fastapi.py

@@ -18,6 +18,7 @@
     from sentry_sdk.integrations.starlette import (
         StarletteIntegration,
         StarletteRequestExtractor,
+        _patch_request,
     )
 except DidNotEnable:
     raise DidNotEnable("Starlette is not installed")
@@ -103,11 +104,16 @@ async def _sentry_app(*args, **kwargs):
                 return await old_app(*args, **kwargs)
 
             request = args[0]
+            _patch_request(request)
 
             _set_transaction_name_and_source(
                 sentry_sdk.get_current_scope(), integration.transaction_style, request
             )
             sentry_scope = sentry_sdk.get_isolation_scope()
+            sentry_scope._name = FastApiIntegration.identifier
+
+            response = await old_app(*args, **kwargs)
+
             extractor = StarletteRequestExtractor(request)
             info = await extractor.extract_request_info()
 
@@ -129,12 +135,11 @@ def event_processor(event, hint):
 
                 return event_processor
 
-            sentry_scope._name = FastApiIntegration.identifier
             sentry_scope.add_event_processor(
                 _make_request_event_processor(request, integration)
             )
 
-            return await old_app(*args, **kwargs)
+            return response
 
         return _sentry_app
 

sentry_sdk/integrations/graphene.py

@@ -20,7 +20,7 @@
 
 if TYPE_CHECKING:
     from collections.abc import Generator
-    from typing import Any, Dict, Union
+    from typing import Any, Dict, Union, Callable
     from graphene.language.source import Source  # type: ignore
     from graphql.execution import ExecutionResult
     from graphql.type import GraphQLSchema
@@ -48,7 +48,7 @@ def _patch_graphql():
     def _sentry_patched_graphql_sync(schema, source, *args, **kwargs):
         # type: (GraphQLSchema, Union[str, Source], Any, Any) -> ExecutionResult
         scope = sentry_sdk.get_isolation_scope()
-        scope.add_event_processor(_event_processor)
+        scope.add_event_processor(_make_event_processor(source))
 
         with graphql_span(schema, source, kwargs):
             result = old_graphql_sync(schema, source, *args, **kwargs)
@@ -75,7 +75,7 @@ async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
             return await old_graphql_async(schema, source, *args, **kwargs)
 
         scope = sentry_sdk.get_isolation_scope()
-        scope.add_event_processor(_event_processor)
+        scope.add_event_processor(_make_event_processor(source))
 
         with graphql_span(schema, source, kwargs):
             result = await old_graphql_async(schema, source, *args, **kwargs)
@@ -99,16 +99,21 @@ async def _sentry_patched_graphql_async(schema, source, *args, **kwargs):
     graphene_schema.graphql = _sentry_patched_graphql_async
 
 
-def _event_processor(event, hint):
-    # type: (Event, Dict[str, Any]) -> Event
-    if should_send_default_pii():
-        request_info = event.setdefault("request", {})
-        request_info["api_target"] = "graphql"
+def _make_event_processor(source):
+    # type: (Any) -> Callable[[Event, dict[str, Any]], Event]
+    def _event_processor(event, hint):
+        # type: (Event, Dict[str, Any]) -> Event
+        if should_send_default_pii():
+            request_info = event.setdefault("request", {})
+            request_info["api_target"] = "graphql"
+            request_info["data"] = {"query": source}
 
-    elif event.get("request", {}).get("data"):
-        del event["request"]["data"]
+        elif event.get("request", {}).get("data"):
+            del event["request"]["data"]
 
-    return event
+        return event
+
+    return _event_processor
 
 
 @contextmanager

sentry_sdk/integrations/starlette.py

@@ -422,6 +422,41 @@ def _is_async_callable(obj):
     )
 
 
+def _patch_request(request):
+    # type: (Request) -> None
+    _original_body = request.body
+    _original_json = request.json
+    _original_form = request.form
+
+    def restore_original_methods():
+        # type: () -> None
+        request.body = _original_body
+        request.json = _original_json
+        request.form = _original_form
+
+    async def sentry_body():
+        # type: () -> bytes
+        request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
+        restore_original_methods()
+        return await _original_body()
+
+    async def sentry_json():
+        # type: () -> Any
+        request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
+        restore_original_methods()
+        return await _original_json()
+
+    async def sentry_form():
+        # type: () -> Any
+        request.scope.setdefault("state", {})["sentry_sdk.is_body_cached"] = True
+        restore_original_methods()
+        return await _original_form()
+
+    request.body = sentry_body  # type: ignore
+    request.json = sentry_json  # type: ignore
+    request.form = sentry_form  # type: ignore
+
+
 def patch_request_response():
     # type: () -> None
     old_request_response = starlette.routing.request_response
@@ -442,6 +477,7 @@ async def _sentry_async_func(*args, **kwargs):
                     return await old_func(*args, **kwargs)
 
                 request = args[0]
+                _patch_request(request)
 
                 _set_transaction_name_and_source(
                     sentry_sdk.get_current_scope(),
@@ -450,6 +486,10 @@ async def _sentry_async_func(*args, **kwargs):
                 )
 
                 sentry_scope = sentry_sdk.get_isolation_scope()
+                sentry_scope._name = StarletteIntegration.identifier
+
+                response = await old_func(*args, **kwargs)
+
                 extractor = StarletteRequestExtractor(request)
                 info = await extractor.extract_request_info()
 
@@ -471,12 +511,11 @@ def event_processor(event, hint):
 
                     return event_processor
 
-                sentry_scope._name = StarletteIntegration.identifier
                 sentry_scope.add_event_processor(
                     _make_request_event_processor(request, integration)
                 )
 
-                return await old_func(*args, **kwargs)
+                return response
 
             func = _sentry_async_func
 
@@ -623,10 +662,23 @@ async def extract_request_info(self):
                 request_info["data"] = AnnotatedValue.removed_because_over_size_limit()
                 return request_info
 
+            # Avoid hangs by not parsing body when ASGI stream is consumed
+            is_body_cached = (
+                "state" in self.request.scope
+                and "sentry_sdk.is_body_cached" in self.request.scope["state"]
+                and self.request.scope["state"]["sentry_sdk.is_body_cached"]
+            )
+            if await self.request.is_disconnected() and not is_body_cached:
+                request_info["data"] = (
+                    AnnotatedValue.removed_because_body_consumed_and_not_cached()
+                )
+                return request_info
+
             # Add JSON body, if it is a JSON request
             json = await self.json()
             if json:
                 request_info["data"] = json
+                print("in fastapi json", request_info)
                 return request_info
 
             # Add form as key/value pairs, if request has form data

tests/integrations/fastapi/test_fastapi.py

@@ -24,6 +24,8 @@
 from tests.integrations.conftest import parametrize_test_configurable_status_codes
 from tests.integrations.starlette import test_starlette
 
+BODY_JSON = {"some": "json", "for": "testing", "nested": {"numbers": 123}}
+
 
 def fastapi_app_factory():
     app = FastAPI()
@@ -72,6 +74,29 @@ async def _thread_ids_async():
     return app
 
 
+def test_stream_available_in_handler(sentry_init):
+    sentry_init(
+        integrations=[StarletteIntegration(), FastApiIntegration()],
+    )
+
+    app = FastAPI()
+
+    @app.post("/consume")
+    async def _consume_stream_body(request):
+        # Avoid cache by constructing new request
+        wrapped_request = Request(request.scope, request.receive)
+
+        assert await wrapped_request.json() == BODY_JSON
+
+        return {"status": "ok"}
+
+    client = TestClient(app)
+    client.post(
+        "/consume",
+        json=BODY_JSON,
+    )
+
+
 @pytest.mark.asyncio
 async def test_response(sentry_init, capture_events):
     # FastAPI is heavily based on Starlette so we also need
@@ -223,6 +248,7 @@ def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, en
 @pytest.mark.asyncio
 async def test_original_request_not_scrubbed(sentry_init, capture_events):
     sentry_init(
+        default_integrations=False,
         integrations=[StarletteIntegration(), FastApiIntegration()],
         traces_sample_rate=1.0,
     )

tests/integrations/starlette/test_starlette.py

@@ -26,6 +26,7 @@
     AuthenticationError,
     SimpleUser,
 )
+from starlette.requests import Request
 from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
@@ -435,6 +436,7 @@ async def test_starletterequestextractor_extract_request_info(sentry_init):
     side_effect = [_mock_receive(msg) for msg in JSON_RECEIVE_MESSAGES]
     starlette_request._receive = mock.Mock(side_effect=side_effect)
 
+    starlette_request.scope["state"] = {"sentry_sdk.is_body_cached": True}
     extractor = StarletteRequestExtractor(starlette_request)
 
     request_info = await extractor.extract_request_info()
@@ -447,6 +449,37 @@ async def test_starletterequestextractor_extract_request_info(sentry_init):
     assert request_info["data"] == BODY_JSON
 
 
+@pytest.mark.asyncio
+async def test_starletterequestextractor_extract_request_info_not_cached(sentry_init):
+    sentry_init(
+        send_default_pii=True,
+        integrations=[StarletteIntegration()],
+    )
+    scope = SCOPE.copy()
+    scope["headers"] = [
+        [b"content-type", b"application/json"],
+        [b"content-length", str(len(json.dumps(BODY_JSON))).encode()],
+        [b"cookie", b"yummy_cookie=choco; tasty_cookie=strawberry"],
+    ]
+
+    starlette_request = starlette.requests.Request(scope)
+
+    # Mocking async `_receive()` that works in Python 3.7+
+    side_effect = [_mock_receive(msg) for msg in JSON_RECEIVE_MESSAGES]
+    starlette_request._receive = mock.Mock(side_effect=side_effect)
+
+    extractor = StarletteRequestExtractor(starlette_request)
+
+    request_info = await extractor.extract_request_info()
+
+    assert request_info
+    assert request_info["cookies"] == {
+        "tasty_cookie": "strawberry",
+        "yummy_cookie": "choco",
+    }
+    assert request_info["data"].metadata == {"rem": [["!consumed", "x"]]}
+
+
 @pytest.mark.asyncio
 async def test_starletterequestextractor_extract_request_info_no_pii(sentry_init):
     sentry_init(
@@ -466,6 +499,7 @@ async def test_starletterequestextractor_extract_request_info_no_pii(sentry_init
     side_effect = [_mock_receive(msg) for msg in JSON_RECEIVE_MESSAGES]
     starlette_request._receive = mock.Mock(side_effect=side_effect)
 
+    starlette_request.scope["state"] = {"sentry_sdk.is_body_cached": True}
     extractor = StarletteRequestExtractor(starlette_request)
 
     request_info = await extractor.extract_request_info()
@@ -475,6 +509,32 @@ async def test_starletterequestextractor_extract_request_info_no_pii(sentry_init
     assert request_info["data"] == BODY_JSON
 
 
+def test_stream_available_in_handler(sentry_init):
+    sentry_init(
+        integrations=[StarletteIntegration()],
+    )
+
+    async def _consume_stream_body(request):
+        # Avoid cache by constructing new request
+        wrapped_request = Request(request.scope, request.receive)
+
+        assert await wrapped_request.json() == BODY_JSON
+
+        return starlette.responses.JSONResponse({"status": "ok"})
+
+    app = starlette.applications.Starlette(
+        routes=[
+            starlette.routing.Route("/consume", _consume_stream_body, methods=["POST"]),
+        ],
+    )
+
+    client = TestClient(app)
+    client.post(
+        "/consume",
+        json=BODY_JSON,
+    )
+
+
 @pytest.mark.parametrize(
     "url,transaction_style,expected_transaction,expected_source",
     [
@@ -942,7 +1002,11 @@ def test_active_thread_id(sentry_init, capture_envelopes, teardown_profiling, en
 
 
 def test_original_request_not_scrubbed(sentry_init, capture_events):
-    sentry_init(integrations=[StarletteIntegration()])
+    sentry_init(
+        default_integrations=False,
+        integrations=[StarletteIntegration()],
+        traces_sample_rate=1.0,
+    )
 
     events = capture_events()