fix(asm): fix reading wsgi input (backport #4114) (#4127)

This is an automatic backport of pull request #4114 done by [Mergify](https://mergify.com).
Cherry-pick of 90e11db has failed:
```
On branch mergify/bp/1.4/pr-4114
Your branch is up to date with 'origin/1.4'.

You are currently cherry-picking commit 90e11db.
  (fix conflicts and run "git cherry-pick --continue")
  (use "git cherry-pick --skip" to skip this patch)
  (use "git cherry-pick --abort" to cancel the cherry-pick operation)

Changes to be committed:
	modified:   benchmarks/flask_simple/app.py
	modified:   benchmarks/flask_simple/config.yaml
	modified:   benchmarks/flask_simple/gunicorn.conf.py
	modified:   benchmarks/flask_simple/scenario.py
	modified:   benchmarks/flask_simple/utils.py
	modified:   benchmarks/span/utils.py
	new file:   releasenotes/notes/asm-fix-reset-wsgi-input-035e0a7d917af2b2.yaml
	modified:   tests/contrib/django/django1_app/urls.py
	modified:   tests/contrib/django/django_app/urls.py
	modified:   tests/contrib/django/views.py
	modified:   tests/contrib/flask/app.py
	modified:   tests/contrib/pylons/app/controllers/root.py
	modified:   tests/contrib/pylons/app/router.py

Unmerged paths:
  (use "git add/rm <file>..." as appropriate to mark resolution)
	deleted by us:   benchmarks/set_http_meta/scenario.py
	both modified:   ddtrace/contrib/flask/patch.py
	both modified:   tests/contrib/django/test_django_appsec.py
	both modified:   tests/contrib/flask/test_flask_appsec.py
	both modified:   tests/contrib/pylons/test_pylons.py

```


To fix up this pull request, you can check it out locally. See documentation: https://docs.github.com/en/github/collaborating-with-pull-requests/reviewing-changes-in-pull-requests/checking-out-pull-requests-locally

---


<details>
<summary>Mergify commands and options</summary>

<br />

More conditions and actions can be found in the [documentation](https://docs.mergify.com/).

You can also trigger Mergify actions by commenting on this pull request:

- `@Mergifyio refresh` will re-evaluate the rules
- `@Mergifyio rebase` will rebase this PR on its base branch
- `@Mergifyio update` will merge the base branch into this PR
- `@Mergifyio backport <destination>` will backport this PR on `<destination>` branch

Additionally, on Mergify [dashboard](https://dashboard.mergify.com/) you can:

- look at your merge queues
- generate the Mergify configuration with the config editor.

Finally, you can contact us on https://mergify.com
</details>

Author: mergify[bot]

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

@@ -0,0 +1,84 @@
+from collections import defaultdict
+import copy
+
+import bm as bm
+import bm.utils as utils
+
+from ddtrace import config as ddconfig
+from ddtrace.contrib.trace_utils import set_http_meta
+
+
+class Config(defaultdict):
+    __header_tag_name = {
+        "User-Agent": "http.user_agent",
+        "REFERER": "http.referer",
+        "Content-Type": "http.content_type",
+        "Etag": "http.etag",
+    }
+
+    def _header_tag_name(self, header_name):
+        return self.__header_tag_name.get(header_name)
+
+    def __getattr__(self, item):
+        return self[item]
+
+
+COOKIES = {"csrftoken": "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL"}
+
+DATA_GET = dict(
+    method="GET",
+    status_code=200,
+    status_msg="OK",
+    parsed_query={
+        "key1": "value1",
+        "key2": "value2",
+        "token": "cR8TVoVebF2afssCR16pQeqHcxAlA3867P6zkkUBYDL5Q92kjSGtqptAry1htdlL",
+    },
+    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(utils.PATH),
+    request_cookies=COOKIES,
+    request_path_params={"id": 1},
+)
+
+
+class SetHttpMeta(bm.Scenario):
+    allenabled = bm.var_bool()
+    useragentvariant = bm.var(type=str)
+    obfuscation_disabled = bm.var_bool()
+    send_querystring_enabled = bm.var_bool()
+    url = bm.var(type=str)
+    querystring = bm.var(type=str)
+
+    def run(self):
+        # run scenario to also set tags on spans
+        if self.allenabled:
+            config = Config(lambda: True)
+        else:
+            config = Config(lambda: False)
+
+        # querystring obfuscation config
+        config["trace_query_string"] = self.send_querystring_enabled
+        if self.obfuscation_disabled:
+            ddconfig._obfuscation_query_string_pattern = None
+
+        data = copy.deepcopy(DATA_GET)
+        data["url"] = self.url
+        data["query"] = self.querystring
+
+        if self.useragentvariant:
+            data["request_headers"][self.useragentvariant] = (
+                "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
+                "(KHTML, like Gecko) Chrome/92.0.4515.159 Safari/537.36"
+            )
+
+        span = utils.gen_span(str("test"))
+        span._local_root = utils.gen_span(str("root"))
+
+        def bm(loops):
+            for _ in range(loops):
+                set_http_meta(span, config, **data)
+
+        yield bm

benchmarks/span/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)

ddtrace/contrib/flask/patch.py

@@ -1,6 +1,7 @@
 import json
 
 import flask
+from six import BytesIO
 import werkzeug
 from werkzeug.exceptions import BadRequest
 
@@ -348,6 +349,17 @@ def traced_wsgi_app(pin, wrapped, instance, args, kwargs):
         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"):
@@ -360,8 +372,16 @@ def traced_wsgi_app(pin, wrapped, instance, args, kwargs):
                     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):
                 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)
 
         # DEV: We set response status code in `_wrap_start_response`
         # DEV: Use `request.base_url` and not `request.url` to keep from leaking any query string parameters

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"),
 ]