From e29ccb0a68bc73bf3d7f071e66e1c8a20eb0c801 Mon Sep 17 00:00:00 2001 From: Alberto Vara Date: Thu, 26 Dec 2024 02:08:11 +0100 Subject: [PATCH] fix: error parsing response cookies in FastAPI and awsgi (#11829) This fix resolves an issue parsing response cookies in FastAPI and awsgi issue: https://github.com/DataDog/dd-trace-py/issues/11818 ## Checklist - [x] PR author has checked that all the criteria below are met - The PR description includes an overview of the change - The PR description articulates the motivation for the change - The change includes tests OR the PR description describes a testing strategy - The PR description notes risks associated with the change, if any - Newly-added code is easy to change - The change follows the [library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) - The change includes or references documentation updates if necessary - Backport labels are set (if [applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)) ## Reviewer Checklist - [x] Reviewer has checked that all the criteria below are met - Title is accurate - All changes are related to the pull request's stated goal - Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - Testing strategy adequately addresses listed risks - Newly-added code is easy to change - Release note makes sense to a user of the library - If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment - Backport labels are set in a manner that is consistent with the [release branch maintenance policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting) --- ddtrace/contrib/internal/asgi/middleware.py | 22 +++++---- .../iast-fix-awsgi-368c173e1f012400.yaml | 4 ++ tests/contrib/asgi/test_asgi.py | 48 +++++++++++++++++++ 3 files changed, 65 insertions(+), 9 deletions(-) create mode 100644 releasenotes/notes/iast-fix-awsgi-368c173e1f012400.yaml diff --git a/ddtrace/contrib/internal/asgi/middleware.py b/ddtrace/contrib/internal/asgi/middleware.py index e32d2994a3d..98d352cf75f 100644 --- a/ddtrace/contrib/internal/asgi/middleware.py +++ b/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, diff --git a/releasenotes/notes/iast-fix-awsgi-368c173e1f012400.yaml b/releasenotes/notes/iast-fix-awsgi-368c173e1f012400.yaml new file mode 100644 index 00000000000..4d40945744a --- /dev/null +++ b/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 diff --git a/tests/contrib/asgi/test_asgi.py b/tests/contrib/asgi/test_asgi.py index 40935990fc2..c1e1f1c7328 100644 --- a/tests/contrib/asgi/test_asgi.py +++ b/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): """