chore(asm): enable iast for threat tests (#9280)

## Motivation To increase test coverage and problem detection, we need to be able to test several features at once. ## This PR - add, for each threat test suite, a second run with IAST enabled - change the way the threat tests use the tracer instrumentation, by directly using `ddtrace.auto` instead of calling specific patching methods. This makes the whole test suites closer to a real application - fix a small regression in response header report for fastapi that was detected in those new tests - fix a small (unreleased) bug in Django 2.2 not working with IAST - improve `repr` for `interface` fixture to allow more user friendly reports in case of CI failure - relax some header test for api security as the tracer may add additional headers in the response. - fix an (unreleased) bug in exploit prevention blocking mechanism ## Revealed Issues - Exploit Prevention and IAST are not working properly together for now. - Root span is not reported in the span list by the tracer object with recent versions of fastAPI (probably a test only problem) APPSEC-53115 ## 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] Risks are described (performance impact, potential for breakage, maintainability) - [x] Change is maintainable (easy to change, telemetry, documentation) - [x] [Library release note guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html) are followed or label `changelog/no-changelog` is set - [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)) - [x] If this PR changes the public interface, I've notified `@DataDog/apm-tees`. ## Reviewer Checklist - [x] Title is accurate - [x] All changes are related to the pull request's stated goal - [x] Description motivates each change - [x] Avoids breaking [API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces) changes - [x] Testing strategy adequately addresses listed risks - [x] Change is maintainable (easy to change, telemetry, documentation) - [x] Release note makes sense to a user of the library - [x] Author has 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) --------- Co-authored-by: Alberto Vara <[email protected]> Co-authored-by: Romain Komorn <[email protected]>
DataDog · May 16, 2024 · 965930a · 965930a
1 parent d75414c
commit 965930a
Show file tree

Hide file tree

Showing 12 changed files with 146 additions and 140 deletions.
diff --git a/ddtrace/appsec/_common_module_patches.py b/ddtrace/appsec/_common_module_patches.py
@@ -8,6 +8,7 @@
 from typing import Callable
 from typing import Dict
 
+from ddtrace.appsec._constants import WAF_ACTIONS
 from ddtrace.appsec._constants import WAF_CONTEXT_NAMES
 from ddtrace.internal import core
 from ddtrace.internal._exceptions import BlockingException
@@ -52,7 +53,6 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
         try:
             from ddtrace.appsec._asm_request_context import call_waf_callback
             from ddtrace.appsec._asm_request_context import in_context
-            from ddtrace.appsec._asm_request_context import is_blocked
             from ddtrace.appsec._constants import EXPLOIT_PREVENTION
         except ImportError:
             # open is used during module initialization
@@ -65,12 +65,12 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
         except Exception:
             filename = ""
         if filename and in_context():
-            call_waf_callback(
+            res = call_waf_callback(
                 {EXPLOIT_PREVENTION.ADDRESS.LFI: filename},
                 crop_trace="wrapped_open_CFDDB7ABBA9081B6",
                 rule_type=EXPLOIT_PREVENTION.TYPE.LFI,
             )
-            if is_blocked():
+            if res and WAF_ACTIONS.BLOCK_ACTION in res.actions:
                 raise BlockingException(core.get_item(WAF_CONTEXT_NAMES.BLOCKED), "exploit_prevention", "lfi", filename)
 
     return original_open_callable(*args, **kwargs)
@@ -88,7 +88,6 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
         try:
             from ddtrace.appsec._asm_request_context import call_waf_callback
             from ddtrace.appsec._asm_request_context import in_context
-            from ddtrace.appsec._asm_request_context import is_blocked
             from ddtrace.appsec._constants import EXPLOIT_PREVENTION
         except ImportError:
             # open is used during module initialization
@@ -100,12 +99,12 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
             if url.__class__.__name__ == "Request":
                 url = url.get_full_url()
             if isinstance(url, str):
-                call_waf_callback(
+                res = call_waf_callback(
                     {EXPLOIT_PREVENTION.ADDRESS.SSRF: url},
                     crop_trace="wrapped_open_ED4CF71136E15EBF",
                     rule_type=EXPLOIT_PREVENTION.TYPE.SSRF,
                 )
-                if is_blocked():
+                if res and WAF_ACTIONS.BLOCK_ACTION in res.actions:
                     raise BlockingException(core.get_item(WAF_CONTEXT_NAMES.BLOCKED), "exploit_prevention", "ssrf", url)
     return original_open_callable(*args, **kwargs)
 
@@ -123,7 +122,6 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
         try:
             from ddtrace.appsec._asm_request_context import call_waf_callback
             from ddtrace.appsec._asm_request_context import in_context
-            from ddtrace.appsec._asm_request_context import is_blocked
             from ddtrace.appsec._constants import EXPLOIT_PREVENTION
         except ImportError:
             # open is used during module initialization
@@ -133,12 +131,12 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
         url = args[1] if len(args) > 1 else kwargs.get("url", None)
         if url and in_context():
             if isinstance(url, str):
-                call_waf_callback(
+                res = call_waf_callback(
                     {EXPLOIT_PREVENTION.ADDRESS.SSRF: url},
                     crop_trace="wrapped_request_D8CB81E472AF98A2",
                     rule_type=EXPLOIT_PREVENTION.TYPE.SSRF,
                 )
-                if is_blocked():
+                if res and WAF_ACTIONS.BLOCK_ACTION in res.actions:
                     raise BlockingException(core.get_item(WAF_CONTEXT_NAMES.BLOCKED), "exploit_prevention", "ssrf", url)
 
     return original_request_callable(*args, **kwargs)
@@ -179,6 +177,9 @@ def apply_patch(parent, attribute, replacement):
         # Avoid overwriting the original function if we call this twice
         if not isinstance(current_attribute, FunctionWrapper):
             _DD_ORIGINAL_ATTRIBUTES[(parent, attribute)] = current_attribute
+        elif isinstance(replacement, FunctionWrapper):
+            # Avoid double patching
+            return
         setattr(parent, attribute, replacement)
     except (TypeError, AttributeError):
         patch_builtins(parent, attribute, replacement)

diff --git a/ddtrace/appsec/_iast/taint_sinks/header_injection.py b/ddtrace/appsec/_iast/taint_sinks/header_injection.py
@@ -47,8 +47,12 @@ def _(m):
     @when_imported("django.http.response")
     def _(m):
         trace_utils.wrap(m, "HttpResponse.__setitem__", _iast_h)
-        trace_utils.wrap(m, "ResponseHeaders.__setitem__", _iast_h)
         trace_utils.wrap(m, "HttpResponseBase.__setitem__", _iast_h)
+        try:
+            trace_utils.wrap(m, "ResponseHeaders.__setitem__", _iast_h)
+        except AttributeError:
+            # no ResponseHeaders in django<3
+            pass
 
     _set_metric_iast_instrumented_sink(VULN_HEADER_INJECTION)
 

diff --git a/ddtrace/appsec/_processor.py b/ddtrace/appsec/_processor.py
@@ -124,9 +124,13 @@ def _set_headers(span: Span, headers: Any, kind: str, only_asm_enabled: bool = F
             key, value = k
         else:
             key, value = k, headers[k]
+        if isinstance(key, bytes):
+            key = key.decode()
+        if isinstance(value, bytes):
+            value = value.decode()
         if key.lower() in (_COLLECTED_REQUEST_HEADERS_ASM_ENABLED if only_asm_enabled else _COLLECTED_REQUEST_HEADERS):
             # since the header value can be a list, use `set_tag()` to ensure it is converted to a string
-            span.set_tag(_normalize_tag_name(kind, key), value)
+            (span._local_root or span).set_tag(_normalize_tag_name(kind, key), value)
 
 
 def _get_rate_limiter() -> RateLimiter:

diff --git a/hatch.toml b/hatch.toml
@@ -188,7 +188,8 @@ dependencies = [
 [envs.appsec_threats_django.scripts]
 test = [
     "pip freeze",
-    "PYTHONPATH=. python -m pytest tests/appsec/contrib_appsec/test_django.py"
+    "DD_IAST_ENABLED=false python -m pytest tests/appsec/contrib_appsec/test_django.py",
+    "DD_IAST_ENABLED=true DD_IAST_REQUEST_SAMPLING=100 python -m pytest tests/appsec/contrib_appsec/test_django.py"
 ]
 
 # python 3.12 should replace 3.11 in this list, but installation is failing on 3.12
@@ -232,7 +233,8 @@ dependencies = [
 [envs.appsec_threats_flask.scripts]
 test = [
     "pip freeze",
-    "PYTHONPATH=. python -m pytest tests/appsec/contrib_appsec/test_flask.py"
+    "DD_IAST_ENABLED=false python -m pytest tests/appsec/contrib_appsec/test_flask.py",
+    "DD_IAST_ENABLED=true DD_IAST_REQUEST_SAMPLING=100 python -m pytest tests/appsec/contrib_appsec/test_flask.py"
 ]
 
 # python 3.12 should replace some 3.11 in this list, but installation is failing on 3.12
@@ -273,7 +275,8 @@ dependencies = [
 [envs.appsec_threats_fastapi.scripts]
 test = [
     "pip freeze",
-    "PYTHONPATH=. python -m pytest tests/appsec/contrib_appsec/test_fastapi.py"
+    "DD_IAST_ENABLED=false python -m pytest tests/appsec/contrib_appsec/test_fastapi.py",
+    "DD_IAST_ENABLED=true DD_IAST_REQUEST_SAMPLING=100 python -m pytest tests/appsec/contrib_appsec/test_fastapi.py"
 ]
 
 # python 3.12 should replace some 3.11 in this list, but installation is failing on 3.12

diff --git a/tests/appsec/contrib_appsec/conftest.py b/tests/appsec/contrib_appsec/conftest.py
@@ -1,30 +1,16 @@
-import unittest.mock
+import ddtrace.auto  # noqa: F401
 
-import pytest
 
-from ddtrace.settings.asm import config as asm_config
-from tests.utils import TracerSpanContainer
-from tests.utils import _build_tree
+# ensure the tracer is loaded and started first for possible iast patching
+print(f"ddtrace version {ddtrace.version.get_version()}")
 
+import unittest.mock  # noqa: E402
 
-@pytest.fixture
-def iast():
-    from os import environ
-
-    from ddtrace import config
-    from ddtrace.appsec._iast import oce
-    from ddtrace.appsec._iast._patch_modules import patch_iast
-
-    environ["DD_IAST_ENABLED"] = "true"
-
-    asm_config._iast_enabled = True
+import pytest  # noqa: E402
 
-    config._raise = True
-
-    oce._enabled = True
-
-    patch_iast()
-    yield
+from ddtrace.settings.asm import config as asm_config  # noqa: E402
+from tests.utils import TracerSpanContainer  # noqa: E402
+from tests.utils import _build_tree  # noqa: E402
 
 
 @pytest.fixture
@@ -42,6 +28,10 @@ def get_root_span():
         for span in test_spans.spans:
             if span.parent_id is None:
                 return _build_tree(test_spans.spans, span)
+        # In case root span is not found, try to find a span with a local root
+        for span in test_spans.spans:
+            if span._local_root is not None:
+                return _build_tree(test_spans.spans, span._local_root)
 
     yield get_root_span
 

diff --git a/tests/appsec/contrib_appsec/django_app/urls.py b/tests/appsec/contrib_appsec/django_app/urls.py
@@ -101,6 +101,20 @@ def rasp(request, endpoint: str):
                     res.append(f"Error: {e}")
         tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
         return HttpResponse("<\\br>\n".join(res))
+    elif endpoint == "shell":
+        res = ["shell endpoint"]
+        for param in query_params:
+            if param.startswith("cmd"):
+                cmd = query_params[param]
+                try:
+                    import subprocess
+
+                    with subprocess.Popen(cmd, stdout=subprocess.PIPE) as f:
+                        res.append(f"cmd stdout: {f.stdout.read()}")
+                except Exception as e:
+                    res.append(f"Error: {e}")
+        tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
+        return HttpResponse("<\\br>\n".join(res))
     tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
     return HttpResponse(f"Unknown endpoint: {endpoint}")
 

diff --git a/tests/appsec/contrib_appsec/fastapi_app/app.py b/tests/appsec/contrib_appsec/fastapi_app/app.py
@@ -117,7 +117,7 @@ async def stream():
     @app.post("/rasp/{endpoint:str}/")
     @app.options("/rasp/{endpoint:str}/")
     async def rasp(endpoint: str, request: Request):
-        query_params = dict(request.query_params)
+        query_params = request.query_params
         if endpoint == "lfi":
             res = ["lfi endpoint"]
             for param in query_params:
@@ -158,6 +158,20 @@ async def rasp(endpoint: str, request: Request):
                     res.append(f"Error: {e}")
             tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
             return HTMLResponse("<\\br>\n".join(res))
+        elif endpoint == "shell":
+            res = ["shell endpoint"]
+            for param in query_params:
+                if param.startswith("cmd"):
+                    cmd = query_params[param]
+                    try:
+                        import subprocess
+
+                        with subprocess.Popen(cmd, stdout=subprocess.PIPE) as f:
+                            res.append(f"cmd stdout: {f.stdout.read()}")
+                    except Exception as e:
+                        res.append(f"Error: {e}")
+            tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
+            return HTMLResponse("<\\br>\n".join(res))
         tracer.current_span()._local_root.set_tag("rasp.request.done", endpoint)
         return HTMLResponse(f"Unknown endpoint: {endpoint}")
 

diff --git a/tests/appsec/contrib_appsec/flask_app/app.py b/tests/appsec/contrib_appsec/flask_app/app.py
@@ -4,6 +4,8 @@
 from flask import request
 
 from ddtrace import tracer
+
+# from ddtrace.appsec.iast import ddtrace_iast_flask_patch
 import ddtrace.constants
 from tests.webclient import PingFilter
 
@@ -61,7 +63,7 @@ def new_service(service_name: str):
 
 @app.route("/rasp/<string:endpoint>/", methods=["GET", "POST", "OPTIONS"])
 def rasp(endpoint: str):
-    query_params = request.args.to_dict()
+    query_params = request.args
     if endpoint == "lfi":
         res = ["lfi endpoint"]
         for param in query_params:

diff --git a/tests/appsec/contrib_appsec/test_django.py b/tests/appsec/contrib_appsec/test_django.py
@@ -5,7 +5,6 @@
 from django.test.client import Client
 import pytest
 
-from ddtrace.contrib.django import patch
 from ddtrace.propagation._utils import get_wsgi_header
 from tests.appsec.contrib_appsec import utils
 
@@ -16,7 +15,6 @@ def interface(self, printer):
         os.environ["DJANGO_SETTINGS_MODULE"] = "tests.appsec.contrib_appsec.django_app.settings"
         settings.DEBUG = False
         django.setup()
-        patch()
         client = Client(
             f"http://localhost:{self.SERVER_PORT}",
             SERVER_NAME=f"localhost:{self.SERVER_PORT}",
@@ -60,13 +58,12 @@ def patch_post(*args, **kwargs):
         client.post = patch_post
 
         interface = utils.Interface("django", django, client)
+        interface.version = django.VERSION
         with utils.test_tracer() as tracer:
             interface.tracer = tracer
             interface.printer = printer
             with utils.post_tracer(interface):
                 yield interface
-        # unpatch failing in this case
-        # unpatch()
 
     def status(self, response):
         return response.status_code

diff --git a/tests/appsec/contrib_appsec/test_fastapi.py b/tests/appsec/contrib_appsec/test_fastapi.py
@@ -3,9 +3,6 @@
 import pytest
 import starlette
 
-import ddtrace
-from ddtrace.contrib.fastapi import patch as fastapi_patch
-from ddtrace.contrib.fastapi import unpatch as fastapi_unpatch
 from tests.appsec.contrib_appsec import utils
 from tests.appsec.contrib_appsec.fastapi_app.app import get_app
 
@@ -21,19 +18,12 @@ class Test_FastAPI(utils.Contrib_TestClass_For_Threats):
     def interface(self, tracer, printer):
         from fastapi.testclient import TestClient
 
-        fastapi_patch()
         # for fastapi, test tracer needs to be set before the app is created
         # contrary to other frameworks
         with utils.test_tracer() as tracer:
             application = get_app()
 
-            @application.middleware("http")
-            async def traced_middlware(request, call_next):
-                with ddtrace.tracer.trace("traced_middlware"):
-                    response = await call_next(request)
-                    return response
-
-            client = TestClient(get_app(), base_url="http://localhost:%d" % self.SERVER_PORT)
+            client = TestClient(application, base_url="http://localhost:%d" % self.SERVER_PORT)
 
             def parse_arguments(*args, **kwargs):
                 if "content_type" in kwargs:
@@ -74,13 +64,11 @@ def patch_get(*args, **kwargs):
             client.get = patch_get
 
             interface = utils.Interface("fastapi", fastapi, client)
+            interface.version = FASTAPI_VERSION
             interface.tracer = tracer
             interface.printer = printer
-            try:
-                with utils.post_tracer(interface):
-                    yield interface
-            finally:
-                fastapi_unpatch()
+            with utils.post_tracer(interface):
+                yield interface
 
     def status(self, response):
         return response.status_code

diff --git a/tests/appsec/contrib_appsec/test_flask.py b/tests/appsec/contrib_appsec/test_flask.py
@@ -2,8 +2,6 @@
 import pytest
 
 from ddtrace import Pin
-from ddtrace.contrib.flask import patch
-from ddtrace.contrib.flask import unpatch
 from ddtrace.internal.packages import get_version_for_package
 from tests.appsec.contrib_appsec import utils
 from tests.utils import TracerTestCase
@@ -34,9 +32,6 @@ def open(self, *args, **kwargs):
 class BaseFlaskTestCase(TracerTestCase):
     def setUp(self):
         super(BaseFlaskTestCase, self).setUp()
-
-        patch()
-
         from tests.appsec.contrib_appsec.flask_app.app import app
 
         self.app = app
@@ -46,8 +41,6 @@ def setUp(self):
 
     def tearDown(self):
         super(BaseFlaskTestCase, self).tearDown()
-        # Unpatch Flask
-        unpatch()
 
 
 class Test_Flask(utils.Contrib_TestClass_For_Threats):
@@ -58,6 +51,7 @@ def interface(self, printer):
         bftc.setUp()
         bftc.app.config["SERVER_NAME"] = f"localhost:{self.SERVER_PORT}"
         interface = utils.Interface("flask", bftc.app, bftc.client)
+        interface.version = FLASK_VERSION
 
         initial_get = bftc.client.get