diff --git a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py index e7ce0bd..84f019c 100644 --- a/lite_bootstrap/bootstrappers/litestar_bootstrapper.py +++ b/lite_bootstrap/bootstrappers/litestar_bootstrapper.py @@ -25,6 +25,7 @@ capture_request_exception, disable_sentry_integrations, read_sentry_trace_headers, + should_collect_request_body, warn_about_sentry_integrations_installed_earlier, ) from lite_bootstrap.instruments.swagger_instrument import SwaggerConfig, SwaggerInstrument @@ -71,6 +72,7 @@ from sentry_sdk.data_collection import _apply_key_value_collection_filtering from sentry_sdk.integrations import logging as sentry_logging from sentry_sdk.integrations._asgi_common import _get_request_data, _RootPathInPath + from sentry_sdk.integrations._wsgi_common import request_body_within_bounds from sentry_sdk.integrations.litestar import LitestarIntegration, retrieve_user_from_scope from sentry_sdk.utils import transaction_from_function @@ -142,6 +144,8 @@ class _PendingSentryException: def _read_parsed_request_body(scope: "Scope") -> typing.Any: # noqa: ANN401 """Return the body a handler already read, without reading the stream.""" state: typing.Final = ScopeState.from_scope(scope) + if state.body is not Empty and not request_body_within_bounds(sentry_sdk.get_client(), len(state.body)): + return None if state.json is not Empty: return state.json if state.form is not Empty: @@ -154,11 +158,13 @@ def _read_parsed_request_body(scope: "Scope") -> typing.Any: # noqa: ANN401 def _build_sentry_request_data(scope: "Scope") -> dict[str, typing.Any]: request_data: typing.Final = _get_request_data(scope, root_path_in_path=_RootPathInPath.EXCLUDED) data_collection: typing.Final = sentry_sdk.get_client().options["data_collection"] - if cookies := _apply_key_value_collection_filtering( - items=ASGIConnection(scope).cookies, behaviour=data_collection["cookies"] - ): + connection = ASGIConnection(scope) + if cookies := _apply_key_value_collection_filtering(items=connection.cookies, behaviour=data_collection["cookies"]): request_data["cookies"] = cookies - if "incoming_request" in data_collection["http_bodies"] and (body := _read_parsed_request_body(scope)) is not None: + if ( + should_collect_request_body(connection.headers.get("content-length")) + and (body := _read_parsed_request_body(scope)) is not None + ): request_data["data"] = body return request_data diff --git a/lite_bootstrap/helpers/starlette_sentry.py b/lite_bootstrap/helpers/starlette_sentry.py index a82c73f..e51cafd 100644 --- a/lite_bootstrap/helpers/starlette_sentry.py +++ b/lite_bootstrap/helpers/starlette_sentry.py @@ -10,6 +10,7 @@ SENTRY_FAILED_REQUEST_STATUS_CODES, capture_request_exception, read_sentry_trace_headers, + should_collect_request_body, ) @@ -23,6 +24,7 @@ from sentry_sdk import sessions from sentry_sdk.data_collection import _apply_key_value_collection_filtering from sentry_sdk.integrations._asgi_common import _get_request_data, _RootPathInPath + from sentry_sdk.integrations._wsgi_common import request_body_within_bounds from sentry_sdk.integrations.starlette import StarletteIntegration from sentry_sdk.utils import AnnotatedValue from starlette.concurrency import run_in_threadpool @@ -40,13 +42,23 @@ class _ReceivedBody: chunks: list[bytes] = dataclasses.field(default_factory=list) complete: bool = False + size: int = 0 + exceeded_limit: bool = False def recording(self, receive: "Receive") -> "Receive": async def receive_and_record() -> "Message": message: typing.Final = await receive() - if message["type"] == "http.request": - self.chunks.append(message.get("body", b"")) - self.complete = not message.get("more_body", False) + if message["type"] == "http.request" and not self.exceeded_limit: + chunk = message.get("body", b"") + self.size += len(chunk) + if request_body_within_bounds(sentry_sdk.get_client(), self.size): + if chunk: + self.chunks.append(chunk) + self.complete = not message.get("more_body", False) + else: + self.chunks.clear() + self.complete = False + self.exceeded_limit = True return message return receive_and_record @@ -157,7 +169,7 @@ async def __call__(self, scope: "Scope", receive: "Receive", send: "Send") -> No routing_scope = dict(scope) received_body: _ReceivedBody | None = None - if "incoming_request" in sentry_sdk.get_client().options["data_collection"]["http_bodies"]: + if scope["type"] == "http" and should_collect_request_body(Headers(scope=scope).get("content-length")): received_body = _ReceivedBody() receive = received_body.recording(receive) diff --git a/lite_bootstrap/instruments/sentry_instrument.py b/lite_bootstrap/instruments/sentry_instrument.py index 714a7ea..29aeb8a 100644 --- a/lite_bootstrap/instruments/sentry_instrument.py +++ b/lite_bootstrap/instruments/sentry_instrument.py @@ -16,6 +16,7 @@ if import_checker.is_sentry_installed: import sentry_sdk + from sentry_sdk.integrations._wsgi_common import request_body_within_bounds from sentry_sdk.integrations.logging import LoggingIntegration from sentry_sdk.scrubber import EventScrubber from sentry_sdk.serializer import serialize @@ -117,6 +118,17 @@ def is_sentry_auto_enabling_integrations_on(self) -> bool: _SENTRY_TRACE_HEADERS: typing.Final = frozenset({b"sentry-trace", b"baggage"}) +def should_collect_request_body(content_length: str | None) -> bool: + client = sentry_sdk.get_client() + if "incoming_request" not in client.options["data_collection"]["http_bodies"]: + return False + try: + size = int(content_length or "0") + except ValueError: + size = 0 + return request_body_within_bounds(client, size) + + def read_sentry_trace_headers(headers: typing.Iterable[tuple[bytes, bytes]]) -> dict[str, str]: return {name.decode("latin-1"): value.decode("latin-1") for name, value in headers if name in _SENTRY_TRACE_HEADERS} diff --git a/tests/test_litestar_sentry_capture.py b/tests/test_litestar_sentry_capture.py index 64bd0d3..7af26da 100644 --- a/tests/test_litestar_sentry_capture.py +++ b/tests/test_litestar_sentry_capture.py @@ -346,6 +346,27 @@ def test_lean_capture_attaches_only_a_body_the_handler_already_read( assert event["request"].get("data") == expected_data +@pytest.mark.parametrize( + ("policy", "size", "declared_size"), + [("never", 4, "4"), ("small", 1001, "1001"), ("small", 1001, "1")], +) +def test_lean_capture_honors_body_limit( + sentry_mock: SentryTestTransport, policy: str, size: int, declared_size: str +) -> None: + params = { + "data_collection": {**DEFAULT_DATA_COLLECTION, "http_bodies": ["incoming_request"]}, + "max_request_body_size": policy, + } + application = LitestarBootstrapper( + bootstrap_config=_build_config(sentry_mock, sentry_additional_params=params) + ).bootstrap() + with TestClient(app=application, raise_server_exceptions=False) as test_client: + test_client.post("/read-raw-body-then-boom", content=b"x" * size, headers={"content-length": declared_size}) + + [event] = _error_events(sentry_mock) + assert "data" not in event["request"] + + async def test_lean_capture_continues_an_incoming_sentry_trace(sentry_mock: SentryTestTransport) -> None: # Called without TestClient: sentry-sdk's httpx integration would overwrite the outgoing sentry-trace. trace_id = "771a43a4192642f0b136d5159a501700" diff --git a/tests/test_starlette_sentry_capture.py b/tests/test_starlette_sentry_capture.py index 4ce0ad3..a19afbf 100644 --- a/tests/test_starlette_sentry_capture.py +++ b/tests/test_starlette_sentry_capture.py @@ -362,6 +362,61 @@ def test_lean_capture_attaches_only_a_body_the_handler_already_read( assert event["request"].get("data") == expected_data +@pytest.mark.parametrize(("policy", "size"), [("never", 4), ("small", 1001)]) +def test_lean_capture_honors_body_limit(sentry_mock: SentryTestTransport, policy: str, size: int) -> None: + params = { + "data_collection": {**DEFAULT_DATA_COLLECTION, "http_bodies": ["incoming_request"]}, + "max_request_body_size": policy, + } + application = _bootstrap(sentry_mock, sentry_additional_params=params) + with TestClient(app=application, raise_server_exceptions=False) as test_client: + test_client.post("/read-raw-body-then-boom", content=b"x" * size) + + [event] = _error_events(sentry_mock) + assert "data" not in event["request"] + + +@pytest.mark.parametrize("length_headers", [[], [(b"content-length", b"invalid")]]) +async def test_lean_capture_drops_streamed_body_over_limit( + sentry_mock: SentryTestTransport, length_headers: list[tuple[bytes, bytes]] +) -> None: + application = _bootstrap( + sentry_mock, + sentry_additional_params={ + "data_collection": {**DEFAULT_DATA_COLLECTION, "http_bodies": ["incoming_request"]}, + "max_request_body_size": "small", + }, + ) + path = "/read-raw-body-then-boom" + scope = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "root_path": "", + "query_string": b"", + "headers": [(b"host", b"testserver"), *length_headers], + "server": ("testserver", 80), + } + chunks = [b"a" * 600, b"b" * 600, b"c"] + receive = AsyncMock( + side_effect=[ + {"type": "http.request", "body": chunk, "more_body": index < len(chunks) - 1} + for index, chunk in enumerate(chunks) + ] + ) + + with pytest.raises(RuntimeError) as error: + await application(scope, receive, AsyncMock()) + + assert str(error.value) == b"".join(chunks).decode() + [event] = _error_events(sentry_mock) + assert "data" not in event["request"] + + async def test_lean_capture_continues_an_incoming_sentry_trace(sentry_mock: SentryTestTransport) -> None: # Called without TestClient: sentry-sdk's httpx integration would overwrite the outgoing sentry-trace. trace_id = "771a43a4192642f0b136d5159a501700"