Merge branch 'main' into APPSEC-52730_strip_headers_from_blocked_resp…

…onses
DataDog · May 3, 2024 · e0ec0e2 · e0ec0e2
2 parents 6d92156 + d837ff5
commit e0ec0e2
Show file tree

Hide file tree

Showing 15 changed files with 669 additions and 165 deletions.
diff --git a/.gitlab/benchmarks.yml b/.gitlab/benchmarks.yml
@@ -15,7 +15,7 @@ variables:
     - git config --global url."https://gitlab-ci-token:${CI_JOB_TOKEN}@gitlab.ddbuild.io/DataDog/".insteadOf "https://github.com/DataDog/"
     - git clone --branch dd-trace-py https://github.com/DataDog/benchmarking-platform /platform && cd /platform
     - ./steps/capture-hardware-software-info.sh
-    - ./steps/run-benchmarks.sh
+    - '([ $SCENARIO = "flask_simple" ] && BP_SCENARIO=$SCENARIO /benchmarking-platform-tools/bp-runner/bp-runner "$REPORTS_DIR/../.gitlab/benchmarks/bp-runner.yml" --debug -t) || ([ $SCENARIO != "flask_simple" ] && ./steps/run-benchmarks.sh)'
     - ./steps/analyze-results.sh
     - "./steps/upload-results-to-s3.sh || :"
   artifacts:
@@ -87,7 +87,7 @@ benchmark-flask-sqli:
   extends: .benchmarks
   variables:
     SCENARIO: "flask_sqli"
-    
+
 benchmark-core-api:
   extends: .benchmarks
   variables:

diff --git a/.gitlab/benchmarks/bp-runner.yml b/.gitlab/benchmarks/bp-runner.yml
@@ -0,0 +1,13 @@
+experiments:
+  - name: run-microbenchmarks
+    setup:
+      - name: datadog-agent
+        run: datadog_agent
+        cpus: 24-25
+        config_sh: ./steps/update-dd-agent-config.sh
+
+    steps:
+      - name: benchmarks
+        cpus: 26-47
+        run: shell
+        script: export SCENARIO=$BP_SCENARIO && ./steps/run-benchmarks.sh
diff --git a/benchmarks/flask_simple/app.py b/benchmarks/flask_simple/app.py
diff --git a/benchmarks/flask_simple/gunicorn.conf.py b/benchmarks/flask_simple/gunicorn.conf.py
diff --git a/benchmarks/flask_simple/requirements_scenario.txt b/benchmarks/flask_simple/requirements_scenario.txt
@@ -1,3 +1 @@
 flask==3.0.0
-gunicorn==20.1.0
-requests==2.31.0
diff --git a/benchmarks/flask_simple/scenario.py b/benchmarks/flask_simple/scenario.py
@@ -1,6 +1,69 @@
+import hashlib
+import os
+import random
+
 import bm
-import bm.flask_utils as flask_utils
-from utils import _post_response
+import bm.utils as utils
+from flask import Flask
+from flask import render_template_string
+from flask import request
+
+from ddtrace.debugging._probe.model import DEFAULT_CAPTURE_LIMITS
+from ddtrace.debugging._probe.model import DEFAULT_SNAPSHOT_PROBE_RATE
+from ddtrace.debugging._probe.model import LiteralTemplateSegment
+from ddtrace.debugging._probe.model import LogLineProbe
+
+
+def make_index():
+    rand_numbers = [random.random() for _ in range(20)]
+    m = hashlib.md5()
+    m.update(b"Insecure hash")
+    rand_numbers.append(m.digest())
+    return render_template_string(
+        """
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <title>Hello World!</title>
+  </head>
+  <body>
+  <section class="section">
+    <div class="container">
+      <h1 class="title">
+        Hello World
+      </h1>
+      <p class="subtitle">
+        My first website
+      </p>
+      <ul>
+        {% for i in rand_numbers %}
+          <li>{{ i }}</li>
+        {% endfor %}
+      </ul>
+    </div>
+  </section>
+  </body>
+</html>
+    """,
+        rand_numbers=rand_numbers,
+    )
+
+
+def create_app():
+    app = Flask(__name__)
+
+    @app.route("/")
+    def index():
+        return make_index()
+
+    @app.route("/post-view", methods=["POST"])
+    def post_view():
+        data = request.data
+        return data, 200
+
+    return app
 
 
 class FlaskSimple(bm.Scenario):
@@ -13,10 +76,82 @@ class FlaskSimple(bm.Scenario):
     telemetry_metrics_enabled = bm.var_bool()
 
     def run(self):
-        with flask_utils.server(self, custom_post_response=_post_response) as get_response:
+        # Setup the environment and enable Datadog features
+        os.environ.update(
+            {
+                "DD_APPSEC_ENABLED": str(self.appsec_enabled),
+                "DD_IAST_ENABLED": str(self.iast_enabled),
+                "DD_TELEMETRY_METRICS_ENABLED": str(self.telemetry_metrics_enabled),
+            }
+        )
+        if self.profiler_enabled:
+            os.environ.update(
+                {"DD_PROFILING_ENABLED": "1", "DD_PROFILING_API_TIMEOUT": "0.1", "DD_PROFILING_UPLOAD_INTERVAL": "10"}
+            )
+            if not self.tracer_enabled:
+                import ddtrace.profiling.auto  # noqa:F401
+
+        if self.tracer_enabled:
+            import ddtrace.bootstrap.sitecustomize  # noqa:F401
+
+        if self.debugger_enabled:
+            from bm.di_utils import BMDebugger
+
+            BMDebugger.enable()
+
+            # Probes are added only if the BMDebugger is enabled.
+            probe_id = "bm-test"
+            BMDebugger.add_probes(
+                LogLineProbe(
+                    probe_id=probe_id,
+                    version=0,
+                    tags={},
+                    source_file="scenario.py",
+                    line=23,
+                    template=probe_id,
+                    segments=[LiteralTemplateSegment(probe_id)],
+                    take_snapshot=True,
+                    limits=DEFAULT_CAPTURE_LIMITS,
+                    condition=None,
+                    condition_error_rate=0.0,
+                    rate=DEFAULT_SNAPSHOT_PROBE_RATE,
+                ),
+            )
+
+        # Create the Flask app
+        app = create_app()
+
+        # Setup the request function
+        if self.post_request:
+            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",
+                "User-Agent": "dd-test-scanner-log",
+            }
+
+            def make_request(app):
+                client = app.test_client()
+                return client.post("/post-view", headers=HEADERS, data=utils.EXAMPLE_POST_DATA)
+
+        else:
+
+            def make_request(app):
+                client = app.test_client()
+                return client.get("/")
 
-            def _(loops):
-                for _ in range(loops):
-                    get_response()
+        # Scenario loop function
+        def _(loops):
+            for _ in range(loops):
+                res = make_request(app)
+                assert res.status_code == 200
+                # We have to close the request (or read `res.data`) to get the `flask.request` span to finalize
+                res.close()
 
-            yield _
+        yield _
diff --git a/benchmarks/flask_simple/utils.py b/benchmarks/flask_simple/utils.py
diff --git a/ddtrace/llmobs/_constants.py b/ddtrace/llmobs/_constants.py
@@ -8,9 +8,11 @@
 MODEL_NAME = "_ml_obs.meta.model_name"
 MODEL_PROVIDER = "_ml_obs.meta.model_provider"
 
+INPUT_DOCUMENTS = "_ml_obs.meta.input.documents"
 INPUT_MESSAGES = "_ml_obs.meta.input.messages"
 INPUT_VALUE = "_ml_obs.meta.input.value"
 INPUT_PARAMETERS = "_ml_obs.meta.input.parameters"
 
+OUTPUT_DOCUMENTS = "_ml_obs.meta.output.documents"
 OUTPUT_MESSAGES = "_ml_obs.meta.output.messages"
 OUTPUT_VALUE = "_ml_obs.meta.output.value"