fix(iast): improve overhead control logic (#8452) (#8494)

IAST: Improve overhead control logic so the decision to analyze a
request is done at span start and is saved at the span level using the
core API. This should fix issues where requests were analyzed when they
shouldn't be and viceversa.

- [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`.
- [x] If change touches code that signs or publishes builds or packages,
or handles credentials of any kind, I've requested a review from
`@DataDog/security-design-and-guidance`.

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

(cherry picked from commit 7e8aaac)

Author: gnufede

ddtrace/appsec/_asm_request_context.py

@@ -399,6 +399,10 @@ def _on_wrapped_view(kwargs):
     if _is_iast_enabled() and kwargs:
         from ddtrace.appsec._iast._taint_tracking import OriginType
         from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+        from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
+
+        if not AppSecIastSpanProcessor.is_span_analyzed():
+            return return_value
 
         _kwargs = {}
         for k, v in kwargs.items():
@@ -414,9 +418,14 @@ def _on_set_request_tags(request, span, flask_config):
         from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
         from ddtrace.appsec._iast._taint_tracking import OriginType
         from ddtrace.appsec._iast._taint_utils import taint_structure
+        from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
 
         _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
         _set_metric_iast_instrumented_source(OriginType.COOKIE)
+
+        if not AppSecIastSpanProcessor.is_span_analyzed(span._local_root or span):
+            return
+
         request.cookies = taint_structure(
             request.cookies,
             OriginType.COOKIE_NAME,

ddtrace/appsec/_constants.py

@@ -82,6 +82,7 @@ class IAST(metaclass=Constant_Class):
     PATCH_MODULES = "_DD_IAST_PATCH_MODULES"
     DENY_MODULES = "_DD_IAST_DENY_MODULES"
     SEP_MODULES = ","
+    REQUEST_IAST_ENABLED = "_dd.iast.request_enabled"
 
 
 class IAST_SPAN_TAGS(metaclass=Constant_Class):

ddtrace/appsec/_handlers.py

@@ -190,6 +190,13 @@ def _on_request_init(wrapped, instance, args, kwargs):
             from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
             from ddtrace.appsec._iast._taint_tracking import OriginType
             from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+            from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
+
+            _set_metric_iast_instrumented_source(OriginType.PATH)
+            _set_metric_iast_instrumented_source(OriginType.QUERY)
+
+            if not AppSecIastSpanProcessor.is_span_analyzed():
+                return
 
             # TODO: instance.query_string = ??
             instance.query_string = taint_pyobject(
@@ -204,8 +211,6 @@ def _on_request_init(wrapped, instance, args, kwargs):
                 source_value=instance.path,
                 source_origin=OriginType.PATH,
             )
-            _set_metric_iast_instrumented_source(OriginType.PATH)
-            _set_metric_iast_instrumented_source(OriginType.QUERY)
         except Exception:
             log.debug("Unexpected exception while tainting pyobject", exc_info=True)
 
@@ -269,6 +274,10 @@ def _on_django_func_wrapped(fn_args, fn_kwargs, first_arg_expected_type, *_):
         from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
         from ddtrace.appsec._iast._taint_tracking import taint_pyobject
         from ddtrace.appsec._iast._taint_utils import taint_structure
+        from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
+
+        if not AppSecIastSpanProcessor.is_span_analyzed():
+            return
 
         http_req = fn_args[0]
 
@@ -318,6 +327,7 @@ def _on_wsgi_environ(wrapped, _instance, args, kwargs):
         from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
         from ddtrace.appsec._iast._taint_tracking import OriginType  # noqa: F401
         from ddtrace.appsec._iast._taint_utils import taint_structure
+        from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
 
         _set_metric_iast_instrumented_source(OriginType.HEADER_NAME)
         _set_metric_iast_instrumented_source(OriginType.HEADER)
@@ -330,6 +340,9 @@ def _on_wsgi_environ(wrapped, _instance, args, kwargs):
         _set_metric_iast_instrumented_source(OriginType.PARAMETER_NAME)
         _set_metric_iast_instrumented_source(OriginType.BODY)
 
+        if not AppSecIastSpanProcessor.is_span_analyzed():
+            return wrapped(*args, **kwargs)
+
         return wrapped(*((taint_structure(args[0], OriginType.HEADER_NAME, OriginType.HEADER),) + args[1:]), **kwargs)
 
     return wrapped(*args, **kwargs)

ddtrace/appsec/_iast/_overhead_control_engine.py

@@ -90,32 +90,35 @@ class OverheadControl(object):
     The goal is to do sampling at different levels of the IAST analysis (per process, per request, etc)
     """
 
+    _lock = threading.Lock()
     _request_quota = MAX_REQUESTS
-    _enabled = False
     _vulnerabilities = set()  # type: Set[Type[Operation]]
     _sampler = RateSampler(sample_rate=get_request_sampling_value() / 100.0)
 
     def reconfigure(self):
         self._sampler = RateSampler(sample_rate=get_request_sampling_value() / 100.0)
 
     def acquire_request(self, span):
-        # type: (Span) -> None
+        # type: (Span) -> bool
         """Decide whether if IAST analysis will be done for this request.
         - Block a request's quota at start of the request to limit simultaneous requests analyzed.
         - Use sample rating to analyze only a percentage of the total requests (30% by default).
         """
-        if self._request_quota > 0 and self._sampler.sample(span):
+        if self._request_quota <= 0 or not self._sampler.sample(span):
+            return False
+
+        with self._lock:
+            if self._request_quota <= 0:
+                return False
+
             self._request_quota -= 1
-            self._enabled = True
 
-    def release_request(self):
-        """increment request's quota at end of the request.
+        return True
 
-        TODO: figure out how to check maximum requests per thread
-        """
-        if self._request_quota < MAX_REQUESTS:
+    def release_request(self):
+        """increment request's quota at end of the request."""
+        with self._lock:
             self._request_quota += 1
-            self._enabled = False
         self.vulnerabilities_reset_quota()
 
     def register(self, klass):
@@ -124,11 +127,6 @@ def register(self, klass):
         self._vulnerabilities.add(klass)
         return klass
 
-    @property
-    def request_has_quota(self):
-        # type: () -> bool
-        return self._enabled
-
     def vulnerabilities_reset_quota(self):
         # type: () -> None
         for k in self._vulnerabilities:

ddtrace/appsec/_iast/_patch.py

@@ -145,6 +145,10 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
         try:
             from ._taint_tracking import is_pyobject_tainted
             from ._taint_tracking import taint_pyobject
+            from .processor import AppSecIastSpanProcessor
+
+            if not AppSecIastSpanProcessor.is_span_analyzed():
+                return value
 
             if not is_pyobject_tainted(value):
                 name = str(args[0]) if len(args) else "http.request.body"
@@ -157,6 +161,11 @@ def if_iast_taint_returned_object_for(origin, wrapped, instance, args, kwargs):
 def if_iast_taint_yield_tuple_for(origins, wrapped, instance, args, kwargs):
     if _is_iast_enabled():
         from ._taint_tracking import taint_pyobject
+        from .processor import AppSecIastSpanProcessor
+
+        if not AppSecIastSpanProcessor.is_span_analyzed():
+            for key, value in wrapped(*args, **kwargs):
+                yield key, value
 
         for key, value in wrapped(*args, **kwargs):
             new_key = taint_pyobject(pyobject=key, source_name=key, source_value=key, source_origin=origins[0])

ddtrace/appsec/_iast/_patches/json_tainting.py

@@ -47,6 +47,10 @@ def wrapped_loads(wrapped, instance, args, kwargs):
             from .._taint_tracking import get_tainted_ranges
             from .._taint_tracking import is_pyobject_tainted
             from .._taint_tracking import taint_pyobject
+            from ..processor import AppSecIastSpanProcessor
+
+            if not AppSecIastSpanProcessor.is_span_analyzed():
+                return obj
 
             if is_pyobject_tainted(args[0]) and obj:
                 # tainting object

ddtrace/appsec/_iast/_taint_tracking/__init__.py

@@ -88,9 +88,7 @@
 
 def taint_pyobject(pyobject, source_name, source_value, source_origin=None):
     # type: (Any, Any, Any, OriginType) -> Any
-    # Request is not analyzed
-    if not oce.request_has_quota:
-        return pyobject
+
     # Pyobject must be Text with len > 1
     if not pyobject or not isinstance(pyobject, (str, bytes, bytearray)):
         return pyobject

ddtrace/appsec/_iast/processor.py

@@ -20,21 +20,43 @@
 
 
 if TYPE_CHECKING:  # pragma: no cover
+    from typing import Optional  # noqa:F401
+
     from ddtrace.span import Span  # noqa:F401
 
 log = get_logger(__name__)
 
 
 @attr.s(eq=False)
 class AppSecIastSpanProcessor(SpanProcessor):
+    @staticmethod
+    def is_span_analyzed(span=None):
+        # type: (Optional[Span]) -> bool
+        if span is None:
+            from ddtrace import tracer
+
+            span = tracer.current_root_span()
+
+        if span and span.span_type == SpanTypes.WEB and core.get_item(IAST.REQUEST_IAST_ENABLED, span=span):
+            return True
+        return False
+
     def on_span_start(self, span):
         # type: (Span) -> None
         if span.span_type != SpanTypes.WEB:
             return
-        oce.acquire_request(span)
-        from ._taint_tracking import create_context
 
-        create_context()
+        if not _is_iast_enabled():
+            return
+
+        request_iast_enabled = False
+        if oce.acquire_request(span):
+            from ._taint_tracking import create_context
+
+            request_iast_enabled = True
+            create_context()
+
+        core.set_item(IAST.REQUEST_IAST_ENABLED, request_iast_enabled, span=span)
 
     def on_span_finish(self, span):
         # type: (Span) -> None
@@ -48,7 +70,7 @@ def on_span_finish(self, span):
         if span.span_type != SpanTypes.WEB:
             return
 
-        if not oce._enabled or not _is_iast_enabled():
+        if not core.get_item(IAST.REQUEST_IAST_ENABLED, span=span):
             span.set_metric(IAST.ENABLED, 0.0)
             return
 

ddtrace/appsec/_iast/taint_sinks/_base.py

@@ -11,12 +11,12 @@
 from ddtrace.settings.asm import config as asm_config
 
 from ..._deduplications import deduplication
-from .. import oce
 from .._overhead_control_engine import Operation
 from .._stacktrace import get_info_frame
 from .._utils import _has_to_scrub
 from .._utils import _is_evidence_value_parts
 from .._utils import _scrub
+from ..processor import AppSecIastSpanProcessor
 from ..reporter import Evidence
 from ..reporter import IastSpanReporter
 from ..reporter import Location
@@ -83,7 +83,7 @@ def wrapper(wrapped, instance, args, kwargs):
             """Get the current root Span and attach it to the wrapped function. We need the span to report the
             vulnerability and update the context with the report information.
             """
-            if oce.request_has_quota and cls.has_quota():
+            if AppSecIastSpanProcessor.is_span_analyzed() and cls.has_quota():
                 return func(wrapped, instance, args, kwargs)
             else:
                 log.debug("IAST: no vulnerability quota to analyze more sink points")

ddtrace/appsec/_iast/taint_sinks/path_traversal.py

@@ -10,6 +10,7 @@
 from .._patch import set_module_unpatched
 from ..constants import EVIDENCE_PATH_TRAVERSAL
 from ..constants import VULN_PATH_TRAVERSAL
+from ..processor import AppSecIastSpanProcessor
 from ._base import VulnerabilityBase
 
 
@@ -49,7 +50,7 @@ def patch():
 
 
 def check_and_report_path_traversal(*args: Any, **kwargs: Any) -> None:
-    if oce.request_has_quota and PathTraversal.has_quota():
+    if AppSecIastSpanProcessor.is_span_analyzed() and PathTraversal.has_quota():
         try:
             from .._metrics import _set_metric_iast_executed_sink
             from .._taint_tracking import is_pyobject_tainted