fix(asm): fix reading wsgi input

benchmarks/bm/utils.py

@@ -1,13 +1,56 @@
 from functools import partial
+import json
 import random
 import string
 
+from six import BytesIO
+
 from ddtrace import Span
 from ddtrace import __version__ as ddtrace_version
 
 
 _Span = Span
 
+PATH = "/test-benchmark/test/1/"
+
+EXAMPLE_POST_DATA = {f"example_key_{i}": f"example_value{i}" for i in range(100)}
+
+COMMON_DJANGO_META = {
+    "SERVER_PORT": "8000",
+    "REMOTE_HOST": "",
+    "CONTENT_LENGTH": "",
+    "SCRIPT_NAME": "",
+    "SERVER_PROTOCOL": "HTTP/1.1",
+    "SERVER_SOFTWARE": "WSGIServer/0.2",
+    "REQUEST_METHOD": "GET",
+    "PATH_INFO": PATH,
+    "QUERY_STRING": "func=subprocess.run&cmd=%2Fbin%2Fecho+hello",
+    "REMOTE_ADDR": "127.0.0.1",
+    "CONTENT_TYPE": "application/json",
+    "HTTP_HOST": "localhost:8000",
+    "HTTP_CONNECTION": "keep-alive",
+    "HTTP_CACHE_CONTROL": "max-age=0",
+    "HTTP_SEC_CH_UA": '"Chromium";v="92", " Not A;Brand";v="99", "Google Chrome";v="92"',
+    "HTTP_SEC_CH_UA_MOBILE": "?0",
+    "HTTP_UPGRADE_INSECURE_REQUESTS": "1",
+    "HTTP_ACCEPT": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,"
+    "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
+    "HTTP_SEC_FETCH_SITE": "none",
+    "HTTP_SEC_FETCH_MODE": "navigate",
+    "HTTP_SEC_FETCH_USER": "?1",
+    "HTTP_SEC_FETCH_DEST": "document",
+    "HTTP_ACCEPT_ENCODING": "gzip, deflate, br",
+    "HTTP_ACCEPT_LANGUAGE": "en-US,en;q=0.9",
+    "HTTP_COOKIE": "Pycharm-45729245=449f1b16-fe0a-4623-92bc-418ec418ed4b; Idea-9fdb9ed8="
+    "448d4c93-863c-4e9b-a8e7-bbfbacd073d2; csrftoken=cR8TVoVebF2afssCR16pQeqHcxA"
+    "lA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL; _xsrf=2|d4b85683|7e2604058ea673d12dc6604f"
+    '96e6e06d|1635869800; username-localhost-8888="2|1:0|10:1637328584|23:username-loca'
+    "lhost-8888|44:OWNiOTFhMjg1NDllNDQxY2I2Y2M2ODViMzRjMTg3NGU=|3bc68f938dcc081a9a02e51660"
+    '0c0d38b14a3032053a7e16b180839298e25b42"',
+    "wsgi.input": BytesIO(bytes(json.dumps(EXAMPLE_POST_DATA), encoding="utf-8")),
+    "wsgi.url_scheme": "http",
+}
+
 # DEV: 1.x dropped tracer positional argument
 if ddtrace_version.split(".")[0] == "0":
     _Span = partial(_Span, None)

benchmarks/flask_simple/app.py

@@ -2,6 +2,7 @@
 
 from flask import Flask
 from flask import render_template_string
+from flask import request
 
 
 app = Flask(__name__)
@@ -40,3 +41,9 @@ def index():
     """,
         rand_numbers=rand_numbers,
     )
+
+
+@app.route("/post-view", methods=["POST"])
+def post_view():
+    data = request.data
+    return data, 200

benchmarks/flask_simple/config.yaml

@@ -1,13 +1,30 @@
 baseline: &baseline
   tracer_enabled: false
   profiler_enabled: false
+  appsec_enabled: false
+  post_request: false
 tracer:
   <<: *baseline
   tracer_enabled: true
 profiler:
   <<: *baseline
   profiler_enabled: true
+appsec-get: &appsec
+  <<: *baseline
+  tracer_enabled: true
+  appsec_enabled: true
+appsec-post:
+  <<: *appsec
+  tracer_enabled: true
+  appsec_enabled: true
+  post_request: true
 tracer-and-profiler:
   <<: *baseline
   tracer_enabled: true
   profiler_enabled: true
+tracer-and-profiler-and-appsec:
+  <<: *baseline
+  tracer_enabled: true
+  profiler_enabled: true
+  appsec_enabled: true
+  post_request: true

benchmarks/flask_simple/gunicorn.conf.py

@@ -7,6 +7,8 @@ def post_fork(server, worker):
         os.environ.update(
             {"DD_PROFILING_ENABLED": "1", "DD_PROFILING_API_TIMEOUT": "0.1", "DD_PROFILING_UPLOAD_INTERVAL": "10"}
         )
+    if os.environ.get("PERF_APPSEC_ENABLED") == "1":
+        os.environ.update({"DD_APPSEC_ENABLED ": "1"})
     # This will not work with gevent workers as the gevent hub has not been
     # initialized when this hook is called.
     if os.environ.get("PERF_TRACER_ENABLED") == "1":

benchmarks/flask_simple/scenario.py

@@ -5,6 +5,8 @@
 class FlaskSimple(bm.Scenario):
     tracer_enabled = bm.var_bool()
     profiler_enabled = bm.var_bool()
+    appsec_enabled = bm.var_bool()
+    post_request = bm.var_bool()
 
     def run(self):
         with utils.server(self) as get_response:

benchmarks/flask_simple/utils.py

@@ -2,6 +2,7 @@
 import os
 import subprocess
 
+import bm.utils as utils
 import requests
 import tenacity
 
@@ -14,6 +15,22 @@ def _get_response():
     r.raise_for_status()
 
 
+def _post_response():
+    HEADERS = {
+        "SERVER_PORT": "8000",
+        "REMOTE_ADDR": "127.0.0.1",
+        "CONTENT_TYPE": "application/json",
+        "HTTP_HOST": "localhost:8000",
+        "HTTP_ACCEPT": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,"
+        "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
+        "HTTP_SEC_FETCH_DEST": "document",
+        "HTTP_ACCEPT_ENCODING": "gzip, deflate, br",
+        "HTTP_ACCEPT_LANGUAGE": "en-US,en;q=0.9",
+    }
+    r = requests.post(SERVER_URL + "post-view", data=utils.EXAMPLE_POST_DATA, headers=HEADERS)
+    r.raise_for_status()
+
+
 @tenacity.retry(
     wait=tenacity.wait_fixed(1),
     stop=tenacity.stop_after_attempt(30),
@@ -27,6 +44,7 @@ def server(scenario):
     env = {
         "PERF_TRACER_ENABLED": str(scenario.tracer_enabled),
         "PERF_PROFILER_ENABLED": str(scenario.profiler_enabled),
+        "PERF_APPSEC_ENABLED": str(scenario.appsec_enabled),
     }
     # copy over current environ
     env.update(os.environ)
@@ -42,7 +60,11 @@ def server(scenario):
     assert proc.poll() is None
     try:
         _wait()
-        yield _get_response
+        if scenario.post_request:
+            response = _post_response
+        else:
+            response = _get_response
+        yield response
     finally:
         proc.terminate()
         proc.wait()

benchmarks/set_http_meta/scenario.py

@@ -1,6 +1,5 @@
 from collections import defaultdict
 import copy
-from io import BytesIO
 
 import bm as bm
 import bm.utils as utils
@@ -24,44 +23,6 @@ def __getattr__(self, item):
         return self[item]
 
 
-PATH = "/test-benchmark/test/1/"
-
-COMMON_DJANGO_META = {
-    "SERVER_PORT": "8000",
-    "REMOTE_HOST": "",
-    "CONTENT_LENGTH": "",
-    "SCRIPT_NAME": "",
-    "SERVER_PROTOCOL": "HTTP/1.1",
-    "SERVER_SOFTWARE": "WSGIServer/0.2",
-    "REQUEST_METHOD": "GET",
-    "PATH_INFO": PATH,
-    "QUERY_STRING": "func=subprocess.run&cmd=%2Fbin%2Fecho+hello",
-    "REMOTE_ADDR": "127.0.0.1",
-    "CONTENT_TYPE": "text/plain",
-    "HTTP_HOST": "localhost:8000",
-    "HTTP_CONNECTION": "keep-alive",
-    "HTTP_CACHE_CONTROL": "max-age=0",
-    "HTTP_SEC_CH_UA": '"Chromium";v="92", " Not A;Brand";v="99", "Google Chrome";v="92"',
-    "HTTP_SEC_CH_UA_MOBILE": "?0",
-    "HTTP_UPGRADE_INSECURE_REQUESTS": "1",
-    "HTTP_ACCEPT": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,"
-    "image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9",
-    "HTTP_SEC_FETCH_SITE": "none",
-    "HTTP_SEC_FETCH_MODE": "navigate",
-    "HTTP_SEC_FETCH_USER": "?1",
-    "HTTP_SEC_FETCH_DEST": "document",
-    "HTTP_ACCEPT_ENCODING": "gzip, deflate, br",
-    "HTTP_ACCEPT_LANGUAGE": "en-US,en;q=0.9",
-    "HTTP_COOKIE": "Pycharm-45729245=449f1b16-fe0a-4623-92bc-418ec418ed4b; Idea-9fdb9ed8="
-    "448d4c93-863c-4e9b-a8e7-bbfbacd073d2; csrftoken=cR8TVoVebF2afssCR16pQeqHcxA"
-    "lA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL; _xsrf=2|d4b85683|7e2604058ea673d12dc6604f"
-    '96e6e06d|1635869800; username-localhost-8888="2|1:0|10:1637328584|23:username-loca'
-    "lhost-8888|44:OWNiOTFhMjg1NDllNDQxY2I2Y2M2ODViMzRjMTg3NGU=|3bc68f938dcc081a9a02e51660"
-    '0c0d38b14a3032053a7e16b180839298e25b42"',
-    "wsgi.input": BytesIO(),
-    "wsgi.url_scheme": "http",
-}
-
 COOKIES = {"csrftoken": "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL"}
 
 DATA_GET = dict(
@@ -73,11 +34,11 @@ def __getattr__(self, item):
         "key2": "value2",
         "token": "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL",
     },
-    request_headers=COMMON_DJANGO_META,
-    response_headers=COMMON_DJANGO_META,
+    request_headers=utils.COMMON_DJANGO_META,
+    response_headers=utils.COMMON_DJANGO_META,
     retries_remain=0,
     raw_uri="http://localhost:8888{}?key1=value1&key2=value2&token="
-    "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL".format(PATH),
+    "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL".format(utils.PATH),
     request_cookies=COOKIES,
     request_path_params={"id": 1},
 )

ddtrace/contrib/flask/patch.py

@@ -1,6 +1,7 @@
 import json
 
 import flask
+from six import BytesIO
 import werkzeug
 from werkzeug.exceptions import BadRequest
 import xmltodict
@@ -128,6 +129,17 @@ def _request_span_modifier(self, span, environ):
         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"):
@@ -146,6 +158,12 @@ def _request_span_modifier(self, span, environ):
                     req_body = request.get_data()
             except (AttributeError, RuntimeError, TypeError, BadRequest):
                 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,

releasenotes/notes/asm-fix-reset-wsgi-input-035e0a7d917af2b2.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    ASM: reset wsgi input after reading.

tests/contrib/django/django1_app/urls.py

@@ -29,4 +29,5 @@
     url(r"^composed-get-view/$", views.ComposedGetView.as_view(), name="composed-get-view"),
     url(r"^composed-view/$", views.ComposedView.as_view(), name="composed-view"),
     url(r"^alter-resource/$", views.alter_resource, name="alter-resource"),
+    url(r"^body/$", views.body_view, name="body_view"),
 ]

tests/contrib/django/django_app/urls.py

@@ -77,4 +77,5 @@ def shutdown(request):
     handler(r"^404-view/$", views.not_found_view, name="404-view"),
     handler(r"^shutdown-tracer/$", shutdown, name="shutdown-tracer"),
     handler(r"^alter-resource/$", views.alter_resource),
+    handler(r"^body/$", views.body_view, name="body_view"),
 ]

tests/contrib/django/test_django_appsec.py

@@ -74,7 +74,10 @@ def test_django_request_body_urlencoded(client, test_spans, tracer):
         # Hack: need to pass an argument to configure so that the processors are recreated
         tracer.configure(api_version="v0.4")
         payload = urlencode({"mytestingbody_key": "mytestingbody_value"})
-        client.post("/", payload, content_type="application/x-www-form-urlencoded")
+
+        response = client.post("/body/", payload, content_type="application/x-www-form-urlencoded")
+        assert response.status_code == 200
+
         root_span = test_spans.spans[0]
         query = dict(_context.get_item("http.request.body", span=root_span))
 
@@ -99,7 +102,7 @@ def test_django_request_body_urlencoded_attack(client, test_spans, tracer):
         # Hack: need to pass an argument to configure so that the processors are recreated
         tracer.configure(api_version="v0.4")
         payload = urlencode({"attack": "1' or '1' = '1'"})
-        client.post("/", payload, content_type="application/x-www-form-urlencoded")
+        client.post("/body/", payload, content_type="application/x-www-form-urlencoded")
         root_span = test_spans.spans[0]
 
         query = dict(_context.get_item("http.request.body", span=root_span))
@@ -113,7 +116,11 @@ def test_django_request_body_json(client, test_spans, tracer):
         # Hack: need to pass an argument to configure so that the processors are recreated
         tracer.configure(api_version="v0.4")
         payload = json.dumps({"mytestingbody_key": "mytestingbody_value"})
-        client.post("/", payload, content_type="application/json")
+
+        response = client.post("/body/", payload, content_type="application/json")
+        assert response.status_code == 200
+        assert response.content == b'{"mytestingbody_key": "mytestingbody_value"}'
+
         root_span = test_spans.spans[0]
         query = dict(_context.get_item("http.request.body", span=root_span))
 
@@ -144,7 +151,10 @@ def test_django_request_body_xml(client, test_spans, tracer):
         payload = "<mytestingbody_key>mytestingbody_value</mytestingbody_key>"
 
         for content_type in ("application/xml", "text/xml"):
-            client.post("/", payload, content_type=content_type)
+            response = client.post("/body/", payload, content_type=content_type)
+            assert response.status_code == 200
+            assert response.content == b"<mytestingbody_key>mytestingbody_value</mytestingbody_key>"
+
             root_span = test_spans.spans[0]
             query = dict(_context.get_item("http.request.body", span=root_span))
             assert root_span.get_tag("_dd.appsec.json") is None

tests/contrib/django/views.py

@@ -182,3 +182,18 @@ def not_found_view(request):
 
 def path_params_view(request, year, month):
     return HttpResponse(status=200)
+
+
+def body_view(request):
+    # Django >= 3
+    if hasattr(request, "headers"):
+        content_type = request.headers["Content-Type"]
+    else:
+        # Django < 3
+        content_type = request.META["CONTENT_TYPE"]
+    if content_type in ("application/json", "application/xml", "text/xml"):
+        data = request.body
+        return HttpResponse(data, status=200)
+    else:
+        data = request.POST
+        return HttpResponse(str(dict(data)), status=200)

tests/contrib/flask/app.py

@@ -2,6 +2,7 @@
 import sys
 
 from flask import Flask
+from flask import request
 
 from ddtrace import tracer
 from tests.webclient import PingFilter
@@ -35,3 +36,9 @@ def resp():
             yield str(i)
 
     return app.response_class(resp())
+
+
+@app.route("/body")
+def body():
+    data = request.get_json()
+    return data, 200