fix: error parsing response cookies in FastAPI and awsgi [backport 2.17]

ddtrace/contrib/internal/asgi/middleware.py

@@ -91,6 +91,18 @@ async def _blocked_asgi_app(scope, receive, send):
     await send({"type": "http.response.body", "body": b""})
 
 
+def _parse_response_cookies(response_headers):
+    cookies = {}
+    try:
+        result = response_headers.get("set-cookie", "").split("=", maxsplit=1)
+        if len(result) == 2:
+            cookie_key, cookie_value = result
+            cookies[cookie_key] = cookie_value
+    except Exception:
+        log.debug("failed to extract response cookies", exc_info=True)
+    return cookies
+
+
 class TraceMiddleware:
     """
     ASGI application middleware that traces the requests.
@@ -211,7 +223,6 @@ async def __call__(self, scope, receive, send):
                 peer_ip = client[0]
             else:
                 peer_ip = None
-
             trace_utils.set_http_meta(
                 span,
                 self.integration_config,
@@ -234,15 +245,8 @@ async def wrapped_send(message):
                 except Exception:
                     log.warning("failed to extract response headers", exc_info=True)
                     response_headers = None
-
                 if span and message.get("type") == "http.response.start" and "status" in message:
-                    cookies = {}
-                    try:
-                        cookie_key, cookie_value = response_headers.get("set-cookie", "").split("=", maxsplit=1)
-                        cookies[cookie_key] = cookie_value
-                    except Exception:
-                        log.debug("failed to extract response cookies", exc_info=True)
-
+                    cookies = _parse_response_cookies(response_headers)
                     status_code = message["status"]
                     trace_utils.set_http_meta(
                         span,

releasenotes/notes/iast-fix-awsgi-368c173e1f012400.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    ASGI: This fix resolves an issue parsing response cookies in FastAPI and awsgi

tests/contrib/asgi/test_asgi.py

@@ -1,5 +1,6 @@
 import asyncio
 from functools import partial
+import logging
 import os
 import random
 
@@ -10,6 +11,7 @@
 from ddtrace.constants import ERROR_MSG
 from ddtrace.contrib.asgi import TraceMiddleware
 from ddtrace.contrib.asgi import span_from_scope
+from ddtrace.contrib.internal.asgi.middleware import _parse_response_cookies
 from ddtrace.propagation import http as http_propagation
 from tests.conftest import DEFAULT_DDTRACE_SUBPROCESS_TEST_SERVICE_NAME
 from tests.utils import DummyTracer
@@ -634,6 +636,52 @@ async def test_tasks_asgi_without_more_body(scope, tracer, test_spans):
     assert request_span.duration < 1
 
 
+@pytest.mark.asyncio
+async def test_request_parse_response_cookies(tracer, test_spans, caplog):
+    """
+    Regression test https://github.com/DataDog/dd-trace-py/issues/11818
+    """
+
+    async def tasks_cookies(scope, receive, send):
+        message = await receive()
+        if message.get("type") == "http.request":
+            await send({"type": "http.response.start", "status": 200, "headers": [[b"set-cookie", b"test_cookie"]]})
+            await send({"type": "http.response.body", "body": b"*"})
+            await asyncio.sleep(1)
+
+    with caplog.at_level(logging.DEBUG):
+        app = TraceMiddleware(tasks_cookies, tracer=tracer)
+        async with httpx.AsyncClient(app=app) as client:
+            response = await client.get("http://testserver/")
+            assert response.status_code == 200
+
+    assert "failed to extract response cookies" not in caplog.text
+
+
+@pytest.mark.parametrize(
+    "headers,expected_result",
+    [
+        ({}, {}),
+        ({"cookie": "cookie1=value1"}, {}),
+        ({"header-1": ""}, {}),
+        ({"Set-cookie": "cookie1=value1"}, {}),
+        ({"set-Cookie": "cookie1=value1"}, {}),
+        ({"SET-cookie": "cookie1=value1"}, {}),
+        ({"set-cookie": "a"}, {}),
+        ({"set-cookie": "1234"}, {}),
+        ({"set-cookie": "cookie1=value1"}, {"cookie1": "value1"}),
+        ({"set-cookie": "cookie2=value1=value2"}, {"cookie2": "value1=value2"}),
+        ({"set-cookie": "cookie3=="}, {"cookie3": "="}),
+    ],
+)
+def test__parse_response_cookies(headers, expected_result, caplog):
+    with caplog.at_level(logging.DEBUG):
+        result = _parse_response_cookies(headers)
+
+    assert "failed to extract response cookies" not in caplog.text
+    assert result == expected_result
+
+
 @pytest.mark.asyncio
 async def test_tasks_asgi_with_more_body(scope, tracer, test_spans):
     """