Repository: elastic/apm-agent-python

Issue:
Starlette middleware not respecting "capture_body" flag
Hello :)

While investigating a bug regarding **very slow** uploads, we discovered two things about the Starlette middleware.

- It does not honor the "capture_body" flag by consuming the stream regardless the flag is turned on or not. It replaces the original stream with a mocked one, making the Starlette request.stream() feature useless because the request is already entirely in memory (using a lot of memory in case of file uploads). I think the following piece of code could check the flag and behave differently.https://github.com/elastic/apm-agent-python/blob/69beac4a5fc72249b40fc811a9ef34fd33313735/elasticapm/contrib/starlette/__init__.py#L151-L163
- The same block of code above uses byte concatenation instead of a byte array which is fine for very small buffers but **prohibitively slow** with anything bigger because of the quadratic behavior. Gigabyte-sized streams do not work at all.

I would happily send a pull request for that, but I would like some opinions first, especially after seeing this (https://github.com/elastic/apm-agent-python/pull/1528), which seems promising.

Anyway, Elastic APM is awesome. Thanks for the great work!



Discussion:
@frnkvieira Great catch! We would love to see a PR if you have time. Don't worry about waiting for #1528. :)

Keep us posted if you don't have time and I can submit a quick fix.

Resolving patch:
diff --git a/elasticapm/contrib/starlette/__init__.py b/elasticapm/contrib/starlette/__init__.py
--- a/elasticapm/contrib/starlette/__init__.py
+++ b/elasticapm/contrib/starlette/__init__.py
@@ -47,6 +47,7 @@
 from elasticapm.contrib.asyncio.traces import set_context
 from elasticapm.contrib.starlette.utils import get_body, get_data_from_request, get_data_from_response
 from elasticapm.utils.disttracing import TraceParent
+from elasticapm.utils.encoding import long_field
 from elasticapm.utils.logging import get_logger
 
 logger = get_logger("elasticapm.errors.client")
@@ -120,7 +121,7 @@ def __init__(self, app: ASGIApp, client: Optional[Client], **kwargs):
             elasticapm.instrumentation.control.instrument()
 
         # If we ever make this a general-use ASGI middleware we should use
-        # `asgiref.conpatibility.guarantee_single_callable(app)` here
+        # `asgiref.compatibility.guarantee_single_callable(app)` here
         self.app = app
 
     async def __call__(self, scope, receive, send):
@@ -131,7 +132,7 @@ async def __call__(self, scope, receive, send):
             send: send awaitable callable
         """
         # we only handle the http scope, skip anything else.
-        if scope["type"] != "http":
+        if scope["type"] != "http" or (scope["type"] == "http" and self.client.should_ignore_url(scope["path"])):
             await self.app(scope, receive, send)
             return
 
@@ -145,32 +146,36 @@ async def wrapped_send(message):
                 elasticapm.set_transaction_result(result, override=False)
             await send(message)
 
-        # When we consume the body from receive, we replace the streaming
-        # mechanism with a mocked version -- this workaround came from
-        # https://github.com/encode/starlette/issues/495#issuecomment-513138055
-        body = b""
-        while True:
-            message = await receive()
-            if not message:
-                break
-            if message["type"] == "http.request":
-                b = message.get("body", b"")
-                if b:
-                    body += b
-                if not message.get("more_body", False):
+        if self.client.config.capture_body != "off":
+
+            # When we consume the body from receive, we replace the streaming
+            # mechanism with a mocked version -- this workaround came from
+            # https://github.com/encode/starlette/issues/495#issuecomment-513138055
+            body = []
+            while True:
+                message = await receive()
+                if not message:
                     break
-            if message["type"] == "http.disconnect":
-                break
+                if message["type"] == "http.request":
+                    b = message.get("body", b"")
+                    if b:
+                        body.append(b)
+                    if not message.get("more_body", False):
+                        break
+                if message["type"] == "http.disconnect":
+                    break
+
+            async def mocked_receive() -> Message:
+                await asyncio.sleep(0)
+                return {"type": "http.request", "body": long_field(b"".join(body))}
 
-        async def _receive() -> Message:
-            await asyncio.sleep(0)
-            return {"type": "http.request", "body": body}
+            receive = mocked_receive
 
-        request = Request(scope, receive=_receive)
+        request = Request(scope, receive=receive)
         await self._request_started(request)
 
         try:
-            await self.app(scope, _receive, wrapped_send)
+            await self.app(scope, receive, wrapped_send)
             elasticapm.set_transaction_outcome(constants.OUTCOME.SUCCESS, override=False)
         except Exception:
             await self.capture_exception(
@@ -216,15 +221,12 @@ async def _request_started(self, request: Request):
         if self.client.config.capture_body != "off":
             await get_body(request)
 
-        if not self.client.should_ignore_url(request.url.path):
-            trace_parent = TraceParent.from_headers(dict(request.headers))
-            self.client.begin_transaction("request", trace_parent=trace_parent)
+        trace_parent = TraceParent.from_headers(dict(request.headers))
+        self.client.begin_transaction("request", trace_parent=trace_parent)
 
-            await set_context(
-                lambda: get_data_from_request(request, self.client.config, constants.TRANSACTION), "request"
-            )
-            transaction_name = self.get_route_name(request) or request.url.path
-            elasticapm.set_transaction_name("{} {}".format(request.method, transaction_name), override=False)
+        await set_context(lambda: get_data_from_request(request, self.client.config, constants.TRANSACTION), "request")
+        transaction_name = self.get_route_name(request) or request.url.path
+        elasticapm.set_transaction_name("{} {}".format(request.method, transaction_name), override=False)
 
     def get_route_name(self, request: Request) -> str:
         app = request.app
diff --git a/elasticapm/utils/encoding.py b/elasticapm/utils/encoding.py
--- a/elasticapm/utils/encoding.py
+++ b/elasticapm/utils/encoding.py
@@ -229,6 +229,9 @@ def long_field(data):
     If the given data, converted to string, is longer than LONG_FIELD_MAX_LENGTH,
     truncate it to LONG_FIELD_MAX_LENGTH-1, adding the "…" character at the end.
 
+    If data is bytes, truncate it to LONG_FIELD_MAX_LENGTH-3, adding b"..." to
+    the end.
+
     Returns the original data if truncation is not required.
 
     Per https://github.com/elastic/apm/blob/main/specs/agents/field-limits.md#long_field_max_length-configuration,
@@ -242,9 +245,12 @@ def long_field(data):
 
     Other fields should be truncated via `elasticapm.utils.encoding.keyword_field()`
     """
-    string = str(data) if not isinstance(data, str) else data
-    if len(string) > LONG_FIELD_MAX_LENGTH:
-        return string[: LONG_FIELD_MAX_LENGTH - 1] + "…"
+    str_or_bytes = str(data) if not isinstance(data, (str, bytes)) else data
+    if len(str_or_bytes) > LONG_FIELD_MAX_LENGTH:
+        if isinstance(str_or_bytes, bytes):
+            return str_or_bytes[: LONG_FIELD_MAX_LENGTH - 3] + b"..."
+        else:
+            return str_or_bytes[: LONG_FIELD_MAX_LENGTH - 1] + "…"
     else:
         return data
 


Test patch:
diff --git a/tests/contrib/asyncio/starlette_tests.py b/tests/contrib/asyncio/starlette_tests.py
--- a/tests/contrib/asyncio/starlette_tests.py
+++ b/tests/contrib/asyncio/starlette_tests.py
@@ -426,6 +426,31 @@ def test_static_files_only(app_static_files_only, elasticapm_client):
     assert request["socket"] == {"remote_address": "127.0.0.1"}
 
 
+def test_non_utf_8_body_in_ignored_paths_with_capture_body(app, elasticapm_client):
+    client = TestClient(app)
+    elasticapm_client.config.update(1, capture_body="all", transaction_ignore_urls="/hello")
+    response = client.post("/hello", data=b"b$\x19\xc2")
+    assert response.status_code == 200
+    assert len(elasticapm_client.events[constants.TRANSACTION]) == 0
+
+
+@pytest.mark.parametrize("elasticapm_client", [{"capture_body": "all"}], indirect=True)
+def test_long_body(app, elasticapm_client):
+    client = TestClient(app)
+
+    response = client.post(
+        "/",
+        data={"foo": "b" * 10000},
+    )
+
+    assert response.status_code == 200
+
+    assert len(elasticapm_client.events[constants.TRANSACTION]) == 1
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+    request = transaction["context"]["request"]
+    assert request["body"] == "foo=" + "b" * 9993 + "..."
+
+
 def test_static_files_only_file_notfound(app_static_files_only, elasticapm_client):
     client = TestClient(app_static_files_only)