chore(iast): split wrapped_view hook [backport 2.21] (#12726)

Backport 8670015 from #12706 to 2.21.

`_on_wrapped_view` contains the logic to block requests for AppSec and
to taint path parameters for IAST. However, since the `_on_wrapped_view`
hook function is in `load_appsec`, the IAST logic doesn’t run if AppSec
isn’t enabled.

This PR splits that logic and creates two separate hooks.  

This PR is a cherry-pick of one of the commits of this PR
#12639

Backport #12726 to fix the CI

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- 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)

Author: avara1986

.gitlab/tests.yml

@@ -64,9 +64,9 @@ build_base_venvs:
     matrix:
       - PYTHON_VERSION: ["3.7", "3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
   variables:
-    CMAKE_BUILD_PARALLEL_LEVEL: 12
     PIP_VERBOSE: 1
     DD_PROFILING_NATIVE_TESTS: 1
+    DD_FAST_BUILD: 1
   script:
     - pip install riot==0.20.1
     - riot -P -v generate --python=$PYTHON_VERSION

ddtrace/appsec/_asm_request_context.py

@@ -25,16 +25,6 @@
 from ddtrace.trace import Span
 
 
-if asm_config._iast_enabled:
-    from ddtrace.appsec._iast._iast_request_context import is_iast_request_enabled
-    from ddtrace.appsec._iast._taint_tracking import OriginType
-    from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
-else:
-
-    def is_iast_request_enabled() -> bool:
-        return False
-
-
 if TYPE_CHECKING:
     from ddtrace.appsec._ddwaf import DDWaf_info
     from ddtrace.appsec._ddwaf import DDWaf_result
@@ -486,7 +476,7 @@ def _on_context_ended(ctx):
 
 
 def _on_wrapped_view(kwargs):
-    return_value = [None, None]
+    callback_block = None
     # if Appsec is enabled, we can try to block as we have the path parameters at that point
     if asm_config._asm_enabled and in_asm_context():
         log.debug("Flask WAF call for Suspicious Request Blocking on request")
@@ -495,21 +485,7 @@ def _on_wrapped_view(kwargs):
         call_waf_callback()
         if is_blocked():
             callback_block = get_value(_CALLBACKS, "flask_block")
-            return_value[0] = callback_block
-
-    # If IAST is enabled, taint the Flask function kwargs (path parameters)
-
-    if asm_config._iast_enabled and kwargs:
-        if not is_iast_request_enabled():
-            return return_value
-
-        _kwargs = {}
-        for k, v in kwargs.items():
-            _kwargs[k] = taint_pyobject(
-                pyobject=v, source_name=k, source_value=v, source_origin=OriginType.PATH_PARAMETER
-            )
-        return_value[1] = _kwargs
-    return return_value
+    return callback_block
 
 
 def _on_pre_tracedrequest(ctx):
@@ -574,7 +550,7 @@ def asm_listen():
     listen()
 
     core.on("flask.finalize_request.post", _set_headers_and_response)
-    core.on("flask.wrapped_view", _on_wrapped_view, "callback_and_args")
+    core.on("flask.wrapped_view", _on_wrapped_view, "callbacks")
     core.on("flask._patched_request", _on_pre_tracedrequest)
     core.on("wsgi.block_decided", _on_block_decided)
     core.on("flask.start_response", _call_waf_first, "waf")

ddtrace/appsec/_iast/_handlers.py

@@ -134,6 +134,18 @@ def _on_flask_patch(flask_version):
         _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
 
 
+def _iast_on_wrapped_view(kwargs):
+    # If IAST is enabled, taint the Flask function kwargs (path parameters)
+    if kwargs and asm_config._iast_enabled:
+        _kwargs = {}
+        for k, v in kwargs.items():
+            _kwargs[k] = taint_pyobject(
+                pyobject=v, source_name=k, source_value=v, source_origin=OriginType.PATH_PARAMETER
+            )
+        return _kwargs
+    return kwargs
+
+
 def _on_wsgi_environ(wrapped, _instance, args, kwargs):
     if asm_config._iast_enabled and args and is_iast_request_enabled():
         return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
@@ -161,8 +173,9 @@ def _on_django_patch():
                     functools.partial(if_iast_taint_returned_object_for, OriginType.PARAMETER),
                 )
             )
+            log.debug("[IAST] Patching Django correctly")
         except Exception:
-            log.debug("Unexpected exception while patch IAST functions", exc_info=True)
+            log.debug("[IAST] Unexpected exception while patch Django", exc_info=True)
 
 
 def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):

ddtrace/appsec/_iast/_listener.py

@@ -1,3 +1,4 @@
+from ddtrace.appsec._iast._handlers import _iast_on_wrapped_view
 from ddtrace.appsec._iast._handlers import _on_django_func_wrapped
 from ddtrace.appsec._iast._handlers import _on_django_patch
 from ddtrace.appsec._iast._handlers import _on_flask_patch
@@ -21,8 +22,9 @@ def iast_listen():
     core.on("django.func.wrapped", _on_django_func_wrapped)
     core.on("flask.patch", _on_flask_patch)
     core.on("flask.request_init", _on_request_init)
-    core.on("flask._patched_request", _on_pre_tracedrequest_iast)
     core.on("flask.set_request_tags", _on_set_request_tags_iast)
+    core.on("flask.wrapped_view", _iast_on_wrapped_view, "check_kwargs")
+    core.on("flask._patched_request", _on_pre_tracedrequest_iast)
 
     core.on("context.ended.wsgi.__call__", _iast_end_request)
     core.on("context.ended.asgi.__call__", _iast_end_request)

ddtrace/contrib/internal/flask/wrappers.py

@@ -45,14 +45,22 @@ def _wrap_call(
         tags=tags,
     ) as ctx, ctx.span:
         if do_dispatch:
-            result = core.dispatch_with_results("flask.wrapped_view", (kwargs,)).callback_and_args
+            dispatch = core.dispatch_with_results("flask.wrapped_view", (kwargs,))
+
+            # Appsec blocks the request
+            result = dispatch.callbacks
             if result:
-                callback_block, _kwargs = result.value
+                callback_block = result.value
                 if callback_block:
                     return callback_block()
+            # IAST overrides the kwargs
+            result = dispatch.check_kwargs
+            if result:
+                _kwargs = result.value
                 if _kwargs:
                     for k in kwargs:
                         kwargs[k] = _kwargs[k]
+
         return wrapped(*args, **kwargs)
 
 

tests/appsec/appsec/test_telemetry.py

@@ -1,7 +1,7 @@
 import os
 from time import sleep
+from unittest import mock
 
-import mock
 import pytest
 
 import ddtrace.appsec._asm_request_context as asm_request_context

tests/appsec/iast/conftest.py

@@ -138,8 +138,13 @@ def iast_span_defaults(tracer):
             yield span
 
 
-# The log contains "[IAST]" but "[IAST] create_context" or "[IAST] reset_context" are valid
-IAST_VALID_LOG = re.compile(r"(?=.*\[IAST\] )(?!.*\[IAST\] (create_context|reset_context))")
+# Check if the log contains "[IAST]" to raise an error if that’s the case BUT, if the logs contains
+# "[IAST] create_context", "[IAST] reset_context", "[IAST] allowing", "[IAST] denying" or "[IAST] astpatch_source"
+# are valid logs
+IAST_VALID_LOG = re.compile(
+    r"(?=.*\[IAST\] )(?!.*\[IAST\] "
+    r"(Patching|Enabled|astpatch_source|compile_code|allowing|denying|create_context|reset_context))"
+)
 
 
 @pytest.fixture(autouse=True)

tests/contrib/flask/test_flask_appsec.py renamed to tests/appsec/integrations/flask_tests/test_appsec_flask.py

@@ -9,6 +9,7 @@
 from ddtrace.appsec._iast import oce
 from ddtrace.appsec._iast._iast_request_context import _iast_start_request
 from ddtrace.appsec._iast._patches.json_tainting import patch as patch_json
+from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast.constants import VULN_HEADER_INJECTION
 from ddtrace.appsec._iast.constants import VULN_INSECURE_COOKIE
 from ddtrace.appsec._iast.constants import VULN_NO_HTTPONLY_COOKIE
@@ -19,12 +20,10 @@
 from ddtrace.settings.asm import config as asm_config
 from tests.appsec.iast.iast_utils import get_line_and_hash
 from tests.contrib.flask import BaseFlaskTestCase
-from tests.utils import flaky
-from tests.utils import override_env
 from tests.utils import override_global_config
 
 
-TEST_FILE_PATH = "tests/contrib/flask/test_flask_appsec_iast.py"
+TEST_FILE_PATH = "tests/appsec/integrations/flask_tests/test_iast_flask.py"
 
 werkzeug_version = version("werkzeug")
 flask_version = tuple([int(v) for v in version("flask").split(".")])
@@ -37,7 +36,7 @@ def inject_fixtures(self, caplog, telemetry_writer):  # noqa: F811
         self._caplog = caplog
 
     def setUp(self):
-        with override_env({"_DD_IAST_USE_ROOT_SPAN": "false"}), override_global_config(
+        with override_global_config(
             dict(
                 _iast_enabled=True,
                 _iast_deduplication_enabled=False,
@@ -49,10 +48,9 @@ def setUp(self):
             patch_header_injection()
             patch_json()
 
-            self.tracer._configure(api_version="v0.4", appsec_enabled=True, iast_enabled=True)
+            self.tracer._configure(api_version="v0.4", iast_enabled=True)
             oce.reconfigure()
 
-    @flaky(1735812000)
     @pytest.mark.skipif(not asm_config._iast_supported, reason="Python version not supported by IAST")
     def test_flask_full_sqli_iast_http_request_path_parameter(self):
         @self.app.route("/sqli/<string:param_str>/", methods=["GET", "POST"])
@@ -1462,8 +1460,6 @@ def cookie_secure():
             from flask import Response
             from flask import request
 
-            from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
-
             tainted_string = request.form.get("name")
             assert is_pyobject_tainted(tainted_string)
             resp = Response("OK")
@@ -1636,8 +1632,6 @@ def test_flask_simple_iast_path_header_and_querystring_not_tainted_if_iast_disab
         def test_sqli(param_str):
             from flask import request
 
-            from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
-
             assert not is_pyobject_tainted(request.headers["User-Agent"])
             assert not is_pyobject_tainted(request.query_string)
             assert not is_pyobject_tainted(param_str)