fix(wsgi): handle django wsgi app being an iterator [backport 1.20] (#7077)

Backport 229378a from #6858 to 1.20.

This change fixes #5325 by
having the WSGI middleware handle cases when the wrapped iterable object
is an iterator. This condition can arise when a Django app attempts to
send its "request finished" signal, in which case it may cause
connection leaks. Standard methods of distinguishing an iterable from
its iterator, such as checking for the presence of __iter__ and __next__
methods, don't work in this case for unknown reasons. Instead of adding
brittle special-case detection logic to the middleware, this new
argument allows users to indicate when this is the case.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] 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)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

Co-authored-by: Emmett Butler <[email protected]>

Author: github-actions[bot]

ddtrace/contrib/wsgi/wsgi.py

@@ -58,14 +58,16 @@ class _DDWSGIMiddlewareBase(object):
     :param tracer: Tracer instance to use the middleware with. Defaults to the global tracer.
     :param int_config: Integration specific configuration object.
     :param pin: Set tracing metadata on a particular traced connection
+    :param app_is_iterator: Boolean indicating whether the wrapped app is a Python iterator
     """
 
-    def __init__(self, application, tracer, int_config, pin):
-        # type: (Iterable, Tracer, Config, Pin) -> None
+    def __init__(self, application, tracer, int_config, pin, app_is_iterator=False):
+        # type: (Iterable, Tracer, Config, Pin, bool) -> None
         self.app = application
         self.tracer = tracer
         self._config = int_config
         self._pin = pin
+        self.app_is_iterator = app_is_iterator
 
     @property
     def _request_span_name(self):
@@ -88,7 +90,7 @@ def _response_span_name(self):
     def __call__(self, environ, start_response):
         # type: (Iterable, Callable) -> wrapt.ObjectProxy
         headers = get_request_headers(environ)
-        closing_iterator = ()
+        closing_iterable = ()
         not_blocked = True
         with core.context_with_data(
             "wsgi.__call__",
@@ -101,7 +103,7 @@ def __call__(self, environ, start_response):
             if core.get_item(HTTP_REQUEST_BLOCKED):
                 status, headers, content = core.dispatch("wsgi.block.started", ctx, construct_url)[0][0]
                 start_response(str(status), headers)
-                closing_iterator = [content]
+                closing_iterable = [content]
                 not_blocked = False
 
             def blocked_view():
@@ -113,17 +115,17 @@ def blocked_view():
             if not_blocked:
                 core.dispatch("wsgi.request.prepare", ctx, start_response)
                 try:
-                    closing_iterator = self.app(environ, ctx.get_item("intercept_start_response"))
+                    closing_iterable = self.app(environ, ctx.get_item("intercept_start_response"))
                 except BaseException:
                     core.dispatch("wsgi.app.exception", ctx)
                     raise
                 else:
-                    core.dispatch("wsgi.app.success", ctx, closing_iterator)
+                    core.dispatch("wsgi.app.success", ctx, closing_iterable)
                 if core.get_item(HTTP_REQUEST_BLOCKED):
                     _, _, content = core.dispatch("wsgi.block.started", ctx, construct_url)[0][0]
-                    closing_iterator = [content]
+                    closing_iterable = [content]
 
-            return core.dispatch("wsgi.request.complete", ctx, closing_iterator)[0][0]
+            return core.dispatch("wsgi.request.complete", ctx, closing_iterable, self.app_is_iterator)[0][0]
 
     def _traced_start_response(self, start_response, request_span, app_span, status, environ, exc_info=None):
         # type: (Callable, Span, Span, str, Dict, Any) -> None
@@ -203,15 +205,18 @@ class DDWSGIMiddleware(_DDWSGIMiddlewareBase):
     :param tracer: Tracer instance to use the middleware with. Defaults to the global tracer.
     :param span_modifier: Span modifier that can add tags to the root span.
                             Defaults to using the request method and url in the resource.
+    :param app_is_iterator: Boolean indicating whether the wrapped WSGI app is a Python iterator
     """
 
     _request_span_name = schematize_url_operation("wsgi.request", protocol="http", direction=SpanDirection.INBOUND)
     _application_span_name = "wsgi.application"
     _response_span_name = "wsgi.response"
 
-    def __init__(self, application, tracer=None, span_modifier=default_wsgi_span_modifier):
-        # type: (Iterable, Optional[Tracer], Callable[[Span, Dict[str, str]], None]) -> None
-        super(DDWSGIMiddleware, self).__init__(application, tracer or ddtrace.tracer, config.wsgi, None)
+    def __init__(self, application, tracer=None, span_modifier=default_wsgi_span_modifier, app_is_iterator=False):
+        # type: (Iterable, Optional[Tracer], Callable[[Span, Dict[str, str]], None], bool) -> None
+        super(DDWSGIMiddleware, self).__init__(
+            application, tracer or ddtrace.tracer, config.wsgi, None, app_is_iterator=app_is_iterator
+        )
         self.span_modifier = span_modifier
 
     def _traced_start_response(self, start_response, request_span, app_span, status, environ, exc_info=None):

ddtrace/tracing/trace_handlers.py

@@ -28,8 +28,13 @@
 
 
 class _TracedIterable(wrapt.ObjectProxy):
-    def __init__(self, wrapped, span, parent_span):
-        super(_TracedIterable, self).__init__(wrapped)
+    def __init__(self, wrapped, span, parent_span, wrapped_is_iterator=False):
+        self._self_wrapped_is_iterator = wrapped_is_iterator
+        if self._self_wrapped_is_iterator:
+            super(_TracedIterable, self).__init__(wrapped)
+            self._wrapped_iterator = iter(wrapped)
+        else:
+            super(_TracedIterable, self).__init__(iter(wrapped))
         self._self_span = span
         self._self_parent_span = parent_span
         self._self_span_finished = False
@@ -39,7 +44,10 @@ def __iter__(self):
 
     def __next__(self):
         try:
-            return next(self.__wrapped__)
+            if self._self_wrapped_is_iterator:
+                return next(self._wrapped_iterator)
+            else:
+                return next(self.__wrapped__)
         except StopIteration:
             self._finish_spans()
             raise
@@ -158,15 +166,15 @@ def _on_request_prepare(ctx, start_response):
     ctx.set_item("intercept_start_response", intercept_start_response)
 
 
-def _on_app_success(ctx, closing_iterator):
+def _on_app_success(ctx, closing_iterable):
     app_span = ctx.get_item("app_span")
     middleware = ctx.get_item("middleware")
     modifier = (
         middleware._application_call_modifier
         if hasattr(middleware, "_application_call_modifier")
         else middleware._application_span_modifier
     )
-    modifier(app_span, ctx.get_item("environ"), closing_iterator)
+    modifier(app_span, ctx.get_item("environ"), closing_iterable)
     app_span.finish()
 
 
@@ -179,7 +187,7 @@ def _on_app_exception(ctx):
     req_span.finish()
 
 
-def _on_request_complete(ctx, closing_iterator):
+def _on_request_complete(ctx, closing_iterable, app_is_iterator):
     middleware = ctx.get_item("middleware")
     req_span = ctx.get_item("req_span")
     # start flask.response span. This span will be finished after iter(result) is closed.
@@ -199,9 +207,9 @@ def _on_request_complete(ctx, closing_iterator):
         if hasattr(middleware, "_response_call_modifier")
         else middleware._response_span_modifier
     )
-    modifier(resp_span, closing_iterator)
+    modifier(resp_span, closing_iterable)
 
-    return _TracedIterable(iter(closing_iterator), resp_span, req_span)
+    return _TracedIterable(closing_iterable, resp_span, req_span, wrapped_is_iterator=app_is_iterator)
 
 
 def _on_response_context_started(ctx):

docs/spelling_wordlist.txt

@@ -77,6 +77,7 @@ doctest
 dogpile
 dogpile.cache
 dogstatsd
+dunder
 dsn
 elasticsearch
 elasticsearch1

releasenotes/notes/wsgi-iter-cca70a8d4429e925.yaml

@@ -0,0 +1,10 @@
+---
+features:
+  - |
+    wsgi: This change introduces the keyword argument app_is_iterator to the DDWSGIMiddleware constructor.
+    It's provided as a workaround for an issue where the Datadog WSGI middleware would fail to handle WSGI
+    apps that are not their own iterators. This condition can arise when a Django app attempts to send its
+    "request finished" signal, in which case it may cause connection leaks. Standard methods of distinguishing
+    an iterable from its iterator, such as checking for the presence of iter and next dunder methods, don't
+    work in this case for unknown reasons. Instead of adding brittle special-case detection logic to the
+    middleware, this new argument allows users to indicate when this is the case.

tests/.suitespec.json

@@ -632,6 +632,7 @@
             "@contrib",
             "@appsec",
             "@asgi",
+            "@wsgi",
             "@django",
             "@dbapi",
             "@pg",
@@ -1078,6 +1079,7 @@
             "@tracing",
             "@wsgi",
             "tests/contrib/wsgi/*",
+            "tests/contrib/django/test_django_wsgi.py",
             "tests/contrib/uwsgi/__init__.py",
             "tests/snapshots/tests.contrib.wsgi.*"
         ],

tests/contrib/django/test_django_wsgi.py

@@ -0,0 +1,77 @@
+import logging
+import os
+import subprocess
+
+import django
+from django.core.signals import request_finished
+from django.core.wsgi import get_wsgi_application
+from django.dispatch import receiver
+from django.http import HttpResponse
+import pytest
+
+from ddtrace.contrib.wsgi import DDWSGIMiddleware
+from ddtrace.internal.compat import PY2
+from ddtrace.internal.compat import PY3
+from tests.webclient import Client
+
+
+filepath, extension = os.path.splitext(__file__)
+ROOT_URLCONF = os.path.basename(filepath)
+WSGI_APPLICATION = os.path.basename(filepath) + ".app"
+DEBUG = True
+SERVER_PORT = 8000
+SENTINEL_LOG = "request finished signal received"
+
+log = logging.getLogger(__name__)
+
+
+@receiver(request_finished)
+def log_request_finished(*_, **__):
+    log.warning(SENTINEL_LOG)
+
+
+def handler(_):
+    return HttpResponse("Hello!")
+
+
+if PY3:
+    from django.urls import path
+
+    urlpatterns = [path("", handler)]
+    # it would be better to check for app_is_iterator programmatically, but Django WSGI apps behave like
+    # iterators for the purpose of DDWSGIMiddleware despite not having both "__next__" and "__iter__" methods
+    app = DDWSGIMiddleware(get_wsgi_application(), app_is_iterator=True)
+
+
+@pytest.mark.skipif(
+    django.VERSION < (3, 0, 0) or PY2, reason="Older Django versions don't work with this use of django-admin"
+)
+def test_django_app_receives_request_finished_signal_when_app_is_ddwsgimiddleware():
+    env = os.environ.copy()
+    env.update(
+        {
+            "PYTHONPATH": os.path.dirname(os.path.abspath(__file__)) + ":" + env["PYTHONPATH"],
+            "DJANGO_SETTINGS_MODULE": "test_django_wsgi",
+        }
+    )
+    cmd = ["django-admin", "runserver", "--noreload", str(SERVER_PORT)]
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        close_fds=True,
+        env=env,
+    )
+
+    client = Client("http://localhost:%d" % SERVER_PORT)
+    client.wait()
+    output = ""
+    try:
+        assert client.get("/").status_code == 200
+    finally:
+        try:
+            _, output = proc.communicate(timeout=1)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            _, output = proc.communicate()
+    assert SENTINEL_LOG in str(output)