chore(iast): refactor is_iast_request_enabled (#14555)

This PR doesn't modify the logic, it's only to simplify and reduce the
number of modification of the main PR IAST Context refactor
#14513:
- move asm_config.is_iast_request_enabled inside the
is_iast_request_enabled()
- move _num_objects_tainted_in_request to _iast_request_context_base.py


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

ddtrace/appsec/_common_module_patches.py

@@ -219,7 +219,7 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
                 # API10, doing all request calls in HTTPConnection.request
                 try:
                     response = original_open_callable(*args, **kwargs)
-                    # api10 response handler for regular reponses
+                    # api10 response handler for regular responses
                     if response.__class__.__name__ == "HTTPResponse":
                         addresses = {
                             "DOWN_RES_STATUS": str(response.status),
@@ -230,7 +230,7 @@ def wrapped_open_ED4CF71136E15EBF(original_open_callable, instance, args, kwargs
                         call_waf_callback(addresses, rule_type=EXPLOIT_PREVENTION.TYPE.SSRF)
                     return response
                 except Exception as e:
-                    # api10 response handler for error reponses
+                    # api10 response handler for error responses
                     if e.__class__.__name__ == "HTTPError":
                         try:
                             status_code = e.code
@@ -256,10 +256,13 @@ def wrapped_request_D8CB81E472AF98A2(original_request_callable, instance, args,
     wrapper for third party requests.request function
     https://requests.readthedocs.io
     """
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
-        from ddtrace.appsec._iast.taint_sinks.ssrf import _iast_report_ssrf
+    if asm_config._iast_enabled:
+        from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
+
+        if is_iast_request_enabled():
+            from ddtrace.appsec._iast.taint_sinks.ssrf import _iast_report_ssrf
 
-        _iast_report_ssrf(original_request_callable, *args, **kwargs)
+            _iast_report_ssrf(original_request_callable, *args, **kwargs)
 
     if _get_rasp_capability("ssrf"):
         try:

ddtrace/appsec/_iast/_handlers.py

@@ -6,6 +6,7 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._iast_request_context_base import get_iast_stacktrace_reported
+from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
 from ddtrace.appsec._iast._iast_request_context_base import set_iast_request_endpoint
 from ddtrace.appsec._iast._iast_request_context_base import set_iast_stacktrace_reported
 from ddtrace.appsec._iast._logs import iast_instrumentation_wrapt_debug_log
@@ -36,7 +37,7 @@
 
 def _on_request_init(wrapped, instance, args, kwargs):
     wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             instance.query_string = taint_pyobject(
                 pyobject=instance.query_string,
@@ -137,7 +138,7 @@ def _on_flask_patch(flask_version):
 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:
+        if not is_iast_request_enabled():
             return kwargs
 
         _kwargs = {}
@@ -150,7 +151,7 @@ def _iast_on_wrapped_view(kwargs):
 
 
 def _on_wsgi_environ(wrapped, _instance, args, kwargs):
-    if asm_config._iast_enabled and args and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
 
     return wrapped(*args, **kwargs)
@@ -187,7 +188,7 @@ def _on_django_patch():
 
 
 def _on_django_middleware(ctx: core.ExecutionContext, call_trace: bool = True, **kwargs) -> None:
-    if not asm_config._iast_enabled or not asm_config.is_iast_request_enabled:
+    if not asm_config._iast_enabled or not is_iast_request_enabled():
         return
 
     request = ctx.find_item("request")
@@ -200,7 +201,7 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
     # If IAST is enabled, and we're wrapping a Django view call, taint the kwargs (view's
     # path parameters)
     if asm_config._iast_enabled and fn_args and isinstance(fn_args[0], first_arg_expected_type):
-        if not asm_config.is_iast_request_enabled:
+        if not is_iast_request_enabled():
             return
 
         _taint_django_func_call(fn_args[0], fn_args, fn_kwargs)
@@ -321,7 +322,7 @@ def _on_grpc_response(message):
 
 
 def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             for key, value in wrapped(*args, **kwargs):
                 new_key = taint_pyobject(pyobject=key, source_name=key, source_value=key, source_origin=origins[0])
@@ -338,7 +339,7 @@ def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
 
 def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
     value = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             if not is_pyobject_tainted(value):
                 name = str(args[0]) if len(args) else "http.request.body"
@@ -352,7 +353,7 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
 
 def if_iast_taint_starlette_datastructures(origin, wrapped, instance, args, kwargs):
     value = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             res = []
             for element in value:
@@ -464,7 +465,7 @@ def _on_pre_tracedrequest_iast(ctx):
 
 
 def _on_set_request_tags_iast(request, span, flask_config):
-    if asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             if request.url_rule is not None:
                 set_iast_request_endpoint(request.method, request.url_rule.rule)
@@ -493,7 +494,7 @@ def _on_set_request_tags_iast(request, span, flask_config):
 
 
 def _on_django_finalize_response_pre(ctx, after_request_tags, request, response):
-    if not response or not asm_config.is_iast_request_enabled or get_iast_stacktrace_reported():
+    if not response or not is_iast_request_enabled() or get_iast_stacktrace_reported():
         return
 
     try:
@@ -506,7 +507,7 @@ def _on_django_finalize_response_pre(ctx, after_request_tags, request, response)
 
 
 def _on_django_technical_500_response(request, response, exc_type, exc_value, tb):
-    if not exc_value or not asm_config._iast_enabled or not asm_config.is_iast_request_enabled:
+    if not exc_value or not asm_config._iast_enabled or not is_iast_request_enabled():
         return
 
     try:
@@ -522,7 +523,7 @@ def _on_django_technical_500_response(request, response, exc_type, exc_value, tb
 
 
 def _on_flask_finalize_request_post(response, _):
-    if not response or not asm_config.is_iast_request_enabled or get_iast_stacktrace_reported():
+    if not response or not is_iast_request_enabled() or get_iast_stacktrace_reported():
         return
 
     try:
@@ -536,7 +537,7 @@ def _on_flask_finalize_request_post(response, _):
 
 
 def _on_asgi_finalize_response(body, _):
-    if not body or not asm_config.is_iast_request_enabled:
+    if not body or not is_iast_request_enabled():
         return
 
     try:
@@ -549,7 +550,7 @@ def _on_asgi_finalize_response(body, _):
 
 
 def _on_werkzeug_render_debugger_html(html):
-    # we don't check asm_config.is_iast_request_enabled due to werkzeug.render_debugger_html works outside the request
+    # we don't check is_iast_request_enabled() due to werkzeug.render_debugger_html works outside the request
     if not html or not asm_config._iast_enabled:
         return
 
@@ -587,7 +588,7 @@ async def _iast_instrument_starlette_request_body(wrapped, instance, args, kwarg
 
 def _iast_instrument_starlette_scope(scope, route):
     try:
-        if asm_config.is_iast_request_enabled:
+        if is_iast_request_enabled():
             set_iast_request_endpoint(scope.get("method"), route)
             if scope.get("path_params"):
                 for k, v in scope["path_params"].items():

ddtrace/appsec/_iast/_iast_env.py

@@ -1,12 +1,16 @@
+from typing import TYPE_CHECKING
 from typing import Dict
 from typing import Optional
 
 from ddtrace._trace.span import Span
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec._iast.reporter import IastSpanReporter
 from ddtrace.internal import core
 
 
+if TYPE_CHECKING:  # pragma: no cover - type checking only
+    from ddtrace.appsec._iast.reporter import IastSpanReporter
+
+
 class IASTEnvironment:
     """
     an object of this class contains all iast data
@@ -18,7 +22,7 @@ def __init__(self, span: Optional[Span] = None):
         self.span = span or core.get_span()
 
         self.request_enabled: bool = False
-        self.iast_reporter: Optional[IastSpanReporter] = None
+        self.iast_reporter: Optional["IastSpanReporter"] = None
         self.iast_span_metrics: Dict[str, int] = {}
         self.iast_stack_trace_reported: bool = False
         self.vulnerability_copy_global_limit: Dict[str, int] = {}

ddtrace/appsec/_iast/_iast_request_context.py

@@ -92,7 +92,7 @@ def _iast_end_request(ctx=None, span=None, *args, **kwargs):
             existing_data = req_span.get_tag(IAST.JSON) or req_span.get_struct_tag(IAST.STRUCT)
             if existing_data is None:
                 if req_span.get_metric(IAST.ENABLED) is None:
-                    if not asm_config.is_iast_request_enabled:
+                    if not base.is_iast_request_enabled():
                         req_span.set_metric(IAST.ENABLED, 0.0)
                         base.end_iast_context(req_span)
                         oce.release_request()

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -7,9 +7,9 @@
 from ddtrace.appsec._iast._iast_env import IASTEnvironment
 from ddtrace.appsec._iast._iast_env import _get_iast_env
 from ddtrace.appsec._iast._overhead_control_engine import oce
+from ddtrace.appsec._iast._taint_tracking import num_objects_tainted
 from ddtrace.appsec._iast._taint_tracking._context import create_context as create_propagation_context
 from ddtrace.appsec._iast._taint_tracking._context import reset_context as reset_propagation_context
-from ddtrace.appsec._iast._utils import _num_objects_tainted_in_request
 from ddtrace.appsec._iast.sampling.vulnerability_detection import update_global_vulnerability_limit
 from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
@@ -69,6 +69,11 @@ def set_iast_request_enabled(request_enabled) -> None:
         log.debug("iast::propagation::context::Trying to set IAST reporter but no context is present")
 
 
+def is_iast_request_enabled() -> bool:
+    """Check whether IAST is currently operating within an active request context."""
+    return asm_config.is_iast_request_enabled
+
+
 def set_iast_request_endpoint(method, route) -> None:
     if asm_config._iast_enabled:
         env = _get_iast_env()
@@ -91,3 +96,7 @@ def _iast_start_request(span=None, *args, **kwargs):
             set_iast_request_enabled(request_iast_enabled)
     except Exception:
         log.debug("iast::propagation::context::Error starting IAST context", exc_info=True)
+
+
+def _num_objects_tainted_in_request():
+    return num_objects_tainted()

ddtrace/appsec/_iast/_langchain.py

@@ -1,3 +1,4 @@
+from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
 from ddtrace.appsec._iast._logs import iast_error
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import get_tainted_ranges
@@ -323,7 +324,7 @@ async def _wrapper_prompt_template_aformat(func, instance, args, kwargs):
 
 def _propagate_prompt_template_format(kwargs, result):
     try:
-        if not asm_config.is_iast_request_enabled:
+        if not is_iast_request_enabled():
             return result
 
         for value in kwargs.values():

ddtrace/appsec/_iast/_metrics.py

@@ -7,10 +7,10 @@
 from ddtrace.appsec._constants import TELEMETRY_INFORMATION_VERBOSITY
 from ddtrace.appsec._constants import TELEMETRY_MANDATORY_VERBOSITY
 from ddtrace.appsec._deduplications import deduplication
+from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._utils import _is_iast_debug_enabled
-from ddtrace.appsec._iast._utils import _num_objects_tainted_in_request
 from ddtrace.internal import telemetry
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.telemetry.constants import TELEMETRY_LOG_LEVEL

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -4,6 +4,7 @@
 from ddtrace.settings.asm import config as asm_config
 
 from ..._constants import IAST
+from .._iast_request_context_base import is_iast_request_enabled
 from .._patch_modules import WrapFunctonsForIAST
 
 
@@ -36,7 +37,7 @@ def patch():
 
 def wrapped_loads(wrapped, instance, args, kwargs):
     obj = wrapped(*args, **kwargs)
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         try:
             from .._taint_tracking._taint_objects import taint_pyobject
             from .._taint_tracking._taint_objects_base import get_tainted_ranges

ddtrace/appsec/_iast/_span_metrics.py

@@ -2,8 +2,8 @@
 
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._iast_env import _get_iast_env
+from ddtrace.appsec._iast._iast_request_context_base import _num_objects_tainted_in_request
 from ddtrace.appsec._iast._metrics import _metric_key_as_snake_case
-from ddtrace.appsec._iast._utils import _num_objects_tainted_in_request
 
 
 def _set_span_tag_iast_request_tainted(span):

ddtrace/appsec/_iast/_taint_tracking/_taint_objects.py

@@ -5,6 +5,7 @@
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._iast_request_context_base import IAST_CONTEXT
+from ddtrace.appsec._iast._iast_request_context_base import is_iast_request_enabled
 from ddtrace.appsec._iast._logs import iast_propagation_debug_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_source
 from ddtrace.appsec._iast._span_metrics import increment_iast_span_metric
@@ -13,15 +14,14 @@
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import _taint_pyobject_base
 from ddtrace.appsec._iast._taint_tracking._taint_objects_base import _taint_pyobject_base_new
 from ddtrace.internal.logger import get_logger
-from ddtrace.settings.asm import config as asm_config
 
 
 log = get_logger(__name__)
 
 
 def taint_pyobject(pyobject: Any, source_name: Any, source_value: Any, source_origin=None) -> Any:
     try:
-        if asm_config.is_iast_request_enabled:
+        if is_iast_request_enabled():
             if source_origin is None:
                 source_origin = OriginType.PARAMETER
 
@@ -51,7 +51,7 @@ def taint_pyobject_new(pyobject: Any, source_name: Any, source_value: Any, sourc
 
 def copy_ranges_to_string(pyobject: str, ranges: Sequence[TaintRange]) -> str:
     # NB this function uses comment-based type annotation because TaintRange is conditionally imported
-    if asm_config.is_iast_request_enabled:
+    if is_iast_request_enabled():
         if not isinstance(pyobject, IAST.TAINTEABLE_TYPES):
             return pyobject