remove flask streamed response from 1.5 (#4208) (#4217)

Revert #3826 as an immediate mitigation for #4215

(cherry picked from commit 93e4595)

Co-authored-by: Munir Abdinur <[email protected]>

Author: mergify[bot]

ddtrace/contrib/flask/patch.py

@@ -22,7 +22,6 @@
 from .. import trace_utils
 from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...constants import SPAN_MEASURED_KEY
-from ...contrib.wsgi.wsgi import _DDWSGIMiddlewareBase
 from ...ext import SpanTypes
 from ...internal.compat import maybe_stringify
 from ...internal.logger import get_logger
@@ -90,103 +89,6 @@ class RequestWithJson(werkzeug.Request, JSONMixin):
 flask_version = parse_version(flask_version_str)
 
 
-class _FlaskWSGIMiddleware(_DDWSGIMiddlewareBase):
-    _request_span_name = "flask.request"
-    _application_span_name = "flask.application"
-    _response_span_name = "flask.response"
-
-    def _traced_start_response(self, start_response, span, status_code, headers, exc_info=None):
-        code, _, _ = status_code.partition(" ")
-        # If values are accessible, set the resource as `<method> <path>` and add other request tags
-        _set_request_tags(span)
-
-        # Override root span resource name to be `<method> 404` for 404 requests
-        # DEV: We do this because we want to make it easier to see all unknown requests together
-        #      Also, we do this to reduce the cardinality on unknown urls
-        # DEV: If we have an endpoint or url rule tag, then we don't need to do this,
-        #      we still want `GET /product/<int:product_id>` grouped together,
-        #      even if it is a 404
-        if not span.get_tag(FLASK_ENDPOINT) and not span.get_tag(FLASK_URL_RULE):
-            span.resource = u" ".join((flask.request.method, code))
-
-        trace_utils.set_http_meta(span, config.flask, status_code=code, response_headers=headers)
-
-        return start_response(status_code, headers)
-
-    def _request_span_modifier(self, span, environ):
-        # Create a werkzeug request from the `environ` to make interacting with it easier
-        # DEV: This executes before a request context is created
-        request = _RequestType(environ)
-
-        # Default resource is method and path:
-        #   GET /
-        #   POST /save
-        # We will override this below in `traced_dispatch_request` when we have a `
-        # RequestContext` and possibly a url rule
-        span.resource = u" ".join((request.method, request.path))
-
-        span.set_tag(SPAN_MEASURED_KEY)
-        # set analytics sample rate with global config enabled
-        sample_rate = config.flask.get_analytics_sample_rate(use_global_config=True)
-        if sample_rate is not None:
-            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
-
-        span._set_str_tag(FLASK_VERSION, flask_version_str)
-
-        req_body = None
-        if config._appsec_enabled and request.method in _BODY_METHODS:
-            content_type = request.content_type
-            wsgi_input = environ.get("wsgi.input", "")
-
-            # Copy wsgi input if not seekable
-            try:
-                seekable = wsgi_input.seekable()
-            except AttributeError:
-                seekable = False
-            if not seekable:
-                body = wsgi_input.read()
-                environ["wsgi.input"] = BytesIO(body)
-
-            try:
-                if content_type == "application/json":
-                    if _HAS_JSON_MIXIN and hasattr(request, "json"):
-                        req_body = request.json
-                    else:
-                        req_body = json.loads(request.data.decode("UTF-8"))
-                elif content_type in ("application/xml", "text/xml"):
-                    req_body = xmltodict.parse(request.get_data())
-                elif hasattr(request, "values"):
-                    req_body = request.values.to_dict()
-                elif hasattr(request, "args"):
-                    req_body = request.args.to_dict()
-                elif hasattr(request, "form"):
-                    req_body = request.form.to_dict()
-                else:
-                    req_body = request.get_data()
-            except (AttributeError, RuntimeError, TypeError, BadRequest, ValueError, JSONDecodeError):
-                log.warning("Failed to parse werkzeug request body", exc_info=True)
-            finally:
-                # Reset wsgi input to the beginning
-                if seekable:
-                    wsgi_input.seek(0)
-                else:
-                    environ["wsgi.input"] = BytesIO(body)
-
-        trace_utils.set_http_meta(
-            span,
-            config.flask,
-            method=request.method,
-            url=request.base_url,
-            raw_uri=request.url,
-            query=request.query_string,
-            parsed_query=request.args,
-            request_headers=request.headers,
-            request_cookies=request.cookies,
-            request_body=req_body,
-            peer_ip=request.remote_addr,
-        )
-
-
 def patch():
     """
     Patch `flask` module for tracing
@@ -387,6 +289,32 @@ def unpatch():
         _u(obj, prop)
 
 
+# Wrap the `start_response` handler to extract response code
+# DEV: We tried using `Flask.finalize_request`, which seemed to work, but gave us hell during tests
+# DEV: The downside to using `start_response` is we do not have a `Flask.Response` object here,
+#   only `status_code`, and `headers` to work with
+#   On the bright side, this works in all versions of Flask (or any WSGI app actually)
+def _wrap_start_response(func, span, request):
+    def traced_start_response(status_code, headers):
+        code, _, _ = status_code.partition(" ")
+        # If values are accessible, set the resource as `<method> <path>` and add other request tags
+        _set_request_tags(span)
+
+        # Override root span resource name to be `<method> 404` for 404 requests
+        # DEV: We do this because we want to make it easier to see all unknown requests together
+        #      Also, we do this to reduce the cardinality on unknown urls
+        # DEV: If we have an endpoint or url rule tag, then we don't need to do this,
+        #      we still want `GET /product/<int:product_id>` grouped together,
+        #      even if it is a 404
+        if not span.get_tag(FLASK_ENDPOINT) and not span.get_tag(FLASK_URL_RULE):
+            span.resource = u" ".join((request.method, code))
+
+        trace_utils.set_http_meta(span, config.flask, status_code=code, response_headers=headers)
+        return func(status_code, headers)
+
+    return traced_start_response
+
+
 @with_instance_pin
 def traced_wsgi_app(pin, wrapped, instance, args, kwargs):
     """
@@ -397,8 +325,88 @@ def traced_wsgi_app(pin, wrapped, instance, args, kwargs):
     # DEV: This is safe before this is the args for a WSGI handler
     #   https://www.python.org/dev/peps/pep-3333/
     environ, start_response = args
-    middleware = _FlaskWSGIMiddleware(wrapped, pin.tracer, config.flask, pin)
-    return middleware(environ, start_response)
+
+    # Create a werkzeug request from the `environ` to make interacting with it easier
+    # DEV: This executes before a request context is created
+    request = _RequestType(environ)
+
+    # Configure distributed tracing
+    trace_utils.activate_distributed_headers(pin.tracer, int_config=config.flask, request_headers=request.headers)
+
+    # Default resource is method and path:
+    #   GET /
+    #   POST /save
+    # We will override this below in `traced_dispatch_request` when we have a `RequestContext` and possibly a url rule
+    resource = u" ".join((request.method, request.path))
+    with pin.tracer.trace(
+        "flask.request",
+        service=trace_utils.int_service(pin, config.flask),
+        resource=resource,
+        span_type=SpanTypes.WEB,
+    ) as span:
+        span.set_tag(SPAN_MEASURED_KEY)
+        # set analytics sample rate with global config enabled
+        sample_rate = config.flask.get_analytics_sample_rate(use_global_config=True)
+        if sample_rate is not None:
+            span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
+
+        span._set_str_tag(FLASK_VERSION, flask_version_str)
+        start_response = _wrap_start_response(start_response, span, request)
+
+        req_body = None
+        if config._appsec_enabled and request.method in _BODY_METHODS:
+            content_type = request.content_type
+            wsgi_input = environ.get("wsgi.input", "")
+
+            # Copy wsgi input if not seekable
+            try:
+                seekable = wsgi_input.seekable()
+            except AttributeError:
+                seekable = False
+            if not seekable:
+                body = wsgi_input.read()
+                environ["wsgi.input"] = BytesIO(body)
+
+            try:
+                if content_type == "application/json":
+                    if _HAS_JSON_MIXIN and hasattr(request, "json"):
+                        req_body = request.json
+                    else:
+                        req_body = json.loads(request.data.decode("UTF-8"))
+                elif content_type in ("application/xml", "text/xml"):
+                    req_body = xmltodict.parse(request.get_data())
+                elif hasattr(request, "values"):
+                    req_body = request.values.to_dict()
+                elif hasattr(request, "args"):
+                    req_body = request.args.to_dict()
+                elif hasattr(request, "form"):
+                    req_body = request.form.to_dict()
+                else:
+                    req_body = request.get_data()
+            except (AttributeError, RuntimeError, TypeError, BadRequest, ValueError, JSONDecodeError):
+                log.warning("Failed to parse werkzeug request body", exc_info=True)
+            finally:
+                # Reset wsgi input to the beginning
+                if seekable:
+                    wsgi_input.seek(0)
+                else:
+                    environ["wsgi.input"] = BytesIO(body)
+
+        trace_utils.set_http_meta(
+            span,
+            config.flask,
+            method=request.method,
+            url=request.base_url,
+            raw_uri=request.url,
+            query=request.query_string,
+            parsed_query=request.args,
+            request_headers=request.headers,
+            request_cookies=request.cookies,
+            request_body=req_body,
+            peer_ip=request.remote_addr,
+        )
+
+        return wrapped(environ, start_response)
 
 
 def traced_blueprint_register(wrapped, instance, args, kwargs):

releasenotes/notes/flask-support-response-streamig-d3aafc8fa1e3605d.yaml

@@ -15,8 +15,6 @@ def test_default_404_handler(self):
         # Make our 404 request
         res = self.client.get("/unknown")
         self.assertEqual(res.status_code, 404)
-        # Read response data from the test client to close flask.request and flask.response spans
-        self.assertIsNotNone(res.data)
 
         spans = self.get_spans()
 
@@ -59,8 +57,6 @@ def endpoint_500():
         # Make our 500 request
         res = self.client.get("/500")
         self.assertEqual(res.status_code, 500)
-        # Read response data from the test client to close flask.request and flask.response spans
-        self.assertIsNotNone(res.data)
 
         spans = self.get_spans()
 
@@ -121,8 +117,7 @@ def endpoint_500():
         # Make our 500 request
         res = self.client.get("/500")
         self.assertEqual(res.status_code, 200)
-        # Read response data from the test client to close flask.request and flask.response spans
-        self.assertIsNotNone(res.data)
+        self.assertEqual(res.data, b"whoops")
 
         spans = self.get_spans()
 
@@ -188,8 +183,6 @@ def endpoint_error():
         # Make our 500 request
         res = self.client.get("/error")
         self.assertEqual(res.status_code, 500)
-        # Read response data from the test client to close flask.request and flask.response spans
-        self.assertIsNotNone(res.data)
 
         spans = self.get_spans()
 

tests/contrib/flask/test_errorhandler.py

@@ -67,8 +67,6 @@ def dynamic_url(item):
             self._aux_appsec_prepare_tracer()
             resp = self.client.get("/params/w00tw00t.at.isc.sans.dfind")
             assert resp.status_code == 200
-            # Read response data from the test client to close flask.request and flask.response spans
-            assert resp.data is not None
 
             root_span = self.pop_spans()[0]
 
@@ -94,8 +92,6 @@ def test_flask_cookie_sql_injection(self):
             self._aux_appsec_prepare_tracer()
             self.client.set_cookie("localhost", "attack", "1' or '1' = '1'")
             resp = self.client.get("/")
-            # Read response data from the test client to close flask.request and flask.response spans
-            assert resp.data is not None
             assert resp.status_code == 404
             root_span = self.pop_spans()[0]