chore(iast): split wrapped_view hook [backport 3.2] (#12725)

Backport 8670015 from #12706 to 3.2.

`_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


## 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)

Co-authored-by: Alberto Vara <[email protected]>

Author: github-actions[bot]

ddtrace/appsec/_asm_request_context.py

@@ -25,11 +25,6 @@
 from ddtrace.trace import Span
 
 
-if asm_config._iast_enabled:
-    from ddtrace.appsec._iast._taint_tracking import OriginType
-    from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
-
-
 if TYPE_CHECKING:
     from ddtrace.appsec._ddwaf import DDWaf_info
     from ddtrace.appsec._ddwaf import DDWaf_result
@@ -481,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")
@@ -490,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 asm_config.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):
@@ -571,7 +552,7 @@ def asm_listen():
     trace_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

@@ -118,6 +118,21 @@ 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:
+        if not asm_config.is_iast_request_enabled:
+            return kwargs
+
+        _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 asm_config.is_iast_request_enabled:
         return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
@@ -146,9 +161,9 @@ def _on_django_patch():
             _set_metric_iast_instrumented_source(OriginType.PARAMETER)
             _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
             _set_metric_iast_instrumented_source(OriginType.BODY)
-
+            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_asgi_finalize_response
 from ddtrace.appsec._iast._handlers import _on_django_finalize_response_pre
 from ddtrace.appsec._iast._handlers import _on_django_func_wrapped
@@ -24,12 +25,14 @@ def iast_listen():
     core.on("django.finalize_response.pre", _on_django_finalize_response_pre)
     core.on("django.func.wrapped", _on_django_func_wrapped)
     core.on("django.technical_500_response", _on_django_technical_500_response)
+
     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.finalize_request.post", _on_flask_finalize_request_post)
+    core.on("flask.wrapped_view", _iast_on_wrapped_view, "check_kwargs")
+    core.on("flask._patched_request", _on_pre_tracedrequest_iast)
     core.on("asgi.finalize_response", _on_asgi_finalize_response)
+    core.on("flask.finalize_request.post", _on_flask_finalize_request_post)
     core.on("werkzeug.render_debugger_html", _on_werkzeug_render_debugger_html)
 
     core.on("context.ended.wsgi.__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/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/appsec/integrations/flask_tests/test_iast_flask.py

@@ -24,7 +24,6 @@
 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 override_env
 from tests.utils import override_global_config
 
 
@@ -41,7 +40,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,
@@ -54,7 +53,7 @@ def setUp(self):
             patch_xss_injection()
             patch_json()
             super(FlaskAppSecIASTEnabledTestCase, self).setUp()
-            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()
 
     @pytest.mark.skipif(not asm_config._iast_supported, reason="Python version not supported by IAST")