Skip to content

Commit

Permalink
fix: error parsing response cookies in FastAPI and awsgi (#11829)
Browse files Browse the repository at this point in the history
This fix resolves an issue parsing response cookies in FastAPI and awsgi

issue: #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)
  • Loading branch information
avara1986 authored Dec 26, 2024
1 parent 998ee2c commit e29ccb0
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 9 deletions.
22 changes: 13 additions & 9 deletions ddtrace/contrib/internal/asgi/middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions releasenotes/notes/iast-fix-awsgi-368c173e1f012400.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
fixes:
- |
ASGI: This fix resolves an issue parsing response cookies in FastAPI and awsgi
48 changes: 48 additions & 0 deletions tests/contrib/asgi/test_asgi.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import asyncio
from functools import partial
import logging
import os
import random

Expand All @@ -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
Expand Down Expand Up @@ -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):
"""
Expand Down

0 comments on commit e29ccb0

Please sign in to comment.