Fix error with missing content-type header in Starlette (#763)

* Run black (how was this missed?)

* Don't assume "content-type" header is present

* Better error handling around body capture

* Add test for POST with body

* Fix isort

Author: basepi

elasticapm/contrib/starlette/utils.py

@@ -30,12 +30,11 @@
 
 import json
 
-from starlette.requests import Request
-from starlette.types import Message
-from starlette.responses import Response
-
 from elasticapm.conf import constants
 from elasticapm.utils import compat, get_url_dict
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.types import Message
 
 
 async def get_data_from_request(request: Request, capture_body=False, capture_headers=True) -> dict:
@@ -51,21 +50,22 @@ async def get_data_from_request(request: Request, capture_body=False, capture_he
     """
     result = {
         "method": request.method,
-        "socket": {
-            "remote_address": _get_client_ip(request),
-            "encrypted": request.url.is_secure
-        },
+        "socket": {"remote_address": _get_client_ip(request), "encrypted": request.url.is_secure},
         "cookies": request.cookies,
     }
     if capture_headers:
         result["headers"] = dict(request.headers)
 
     if request.method in constants.HTTP_WITH_BODY:
-        body = await get_body(request)
-        if request.headers['content-type'] == "application/x-www-form-urlencoded":
-            body = await query_params_to_dict(body)
-        else:
-            body = json.loads(body)
+        body = None
+        try:
+            body = await get_body(request)
+            if request.headers.get("content-type") == "application/x-www-form-urlencoded":
+                body = await query_params_to_dict(body)
+            else:
+                body = json.loads(body)
+        except Exception:
+            pass
         if body is not None:
             result["body"] = body if capture_body else "[REDACTED]"
 
@@ -107,6 +107,7 @@ async def set_body(request: Request, body: bytes):
         request (Request)
         body (bytes)
     """
+
     async def receive() -> Message:
         return {"type": "http.request", "body": body}
 
@@ -155,9 +156,9 @@ async def query_params_to_dict(query_params: str) -> dict:
 
 
 def _get_client_ip(request: Request):
-    x_forwarded_for = request.headers.get('HTTP_X_FORWARDED_FOR')
+    x_forwarded_for = request.headers.get("HTTP_X_FORWARDED_FOR")
     if x_forwarded_for:
-        ip = x_forwarded_for.split(',')[0]
+        ip = x_forwarded_for.split(",")[0]
     else:
-        ip = request.headers.get('REMOTE_ADDR')
+        ip = request.headers.get("REMOTE_ADDR")
     return ip

tests/contrib/asyncio/starlette_tests.py

@@ -33,15 +33,14 @@
 starlette = pytest.importorskip("starlette")  # isort:skip
 
 import mock
-
 from starlette.applications import Starlette
 from starlette.responses import PlainTextResponse
 from starlette.testclient import TestClient
-from elasticapm.utils.disttracing import TraceParent
-from elasticapm import async_capture_span
 
+from elasticapm import async_capture_span
 from elasticapm.conf import constants
 from elasticapm.contrib.starlette import ElasticAPM
+from elasticapm.utils.disttracing import TraceParent
 
 pytestmark = [pytest.mark.starlette]
 
@@ -50,7 +49,7 @@
 def app(elasticapm_client):
     app = Starlette()
 
-    @app.route("/")
+    @app.route("/", methods=["GET", "POST"])
     async def hi(request):
         with async_capture_span("test"):
             pass
@@ -68,11 +67,14 @@ async def raise_exception(request):
 def test_get(app, elasticapm_client):
     client = TestClient(app)
 
-    response = client.get('/', headers={
-        constants.TRACEPARENT_HEADER_NAME: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
-        constants.TRACESTATE_HEADER_NAME: "foo=bar,bar=baz",
-        "REMOTE_ADDR": "127.0.0.1",
-    })
+    response = client.get(
+        "/",
+        headers={
+            constants.TRACEPARENT_HEADER_NAME: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
+            constants.TRACESTATE_HEADER_NAME: "foo=bar,bar=baz",
+            "REMOTE_ADDR": "127.0.0.1",
+        },
+    )
 
     assert response.status_code == 200
 
@@ -93,15 +95,52 @@ def test_get(app, elasticapm_client):
     assert span["name"] == "test"
 
 
-def test_exception(app, elasticapm_client):
+@pytest.mark.parametrize("elasticapm_client", [{"capture_body": "all"}], indirect=True)
+def test_post(app, elasticapm_client):
     client = TestClient(app)
 
-    with pytest.raises(ValueError):
-        client.get('/raise-exception', headers={
+    response = client.post(
+        "/",
+        headers={
             constants.TRACEPARENT_HEADER_NAME: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
             constants.TRACESTATE_HEADER_NAME: "foo=bar,bar=baz",
             "REMOTE_ADDR": "127.0.0.1",
-        })
+        },
+        data={"foo": "bar"},
+    )
+
+    assert response.status_code == 200
+
+    assert len(elasticapm_client.events[constants.TRANSACTION]) == 1
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+
+    assert transaction["name"] == "POST /"
+    assert transaction["result"] == "HTTP 2xx"
+    assert transaction["type"] == "request"
+    assert transaction["span_count"]["started"] == 1
+    request = transaction["context"]["request"]
+    request["method"] == "GET"
+    request["socket"] == {"remote_address": "127.0.0.1", "encrypted": False}
+    assert request["body"]["foo"] == "bar"
+
+    assert span["name"] == "test"
+
+
+def test_exception(app, elasticapm_client):
+    client = TestClient(app)
+
+    with pytest.raises(ValueError):
+        client.get(
+            "/raise-exception",
+            headers={
+                constants.TRACEPARENT_HEADER_NAME: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
+                constants.TRACESTATE_HEADER_NAME: "foo=bar,bar=baz",
+                "REMOTE_ADDR": "127.0.0.1",
+            },
+        )
 
     assert len(elasticapm_client.events[constants.TRANSACTION]) == 1
     transaction = elasticapm_client.events[constants.TRANSACTION][0]
@@ -129,10 +168,13 @@ def test_traceparent_handling(app, elasticapm_client, header_name):
     with mock.patch(
         "elasticapm.contrib.flask.TraceParent.from_string", wraps=TraceParent.from_string
     ) as wrapped_from_string:
-        response = client.get('/', headers={
-            header_name: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
-            constants.TRACESTATE_HEADER_NAME: "foo=bar,baz=bazzinga",
-        })
+        response = client.get(
+            "/",
+            headers={
+                header_name: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-03",
+                constants.TRACESTATE_HEADER_NAME: "foo=bar,baz=bazzinga",
+            },
+        )
 
     assert response.status_code == 200