chore(iast): restructure functions

ddtrace/appsec/_iast/_handlers.py

@@ -10,15 +10,13 @@
 from ddtrace.appsec._iast._logs import iast_instrumentation_wrapt_debug_log
 from ddtrace.appsec._iast._logs import iast_propagation_listener_log_log
 from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_source
-from ddtrace.appsec._iast._patch import _iast_instrument_starlette_request
-from ddtrace.appsec._iast._patch import _iast_instrument_starlette_request_body
 from ddtrace.appsec._iast._patch import _iast_instrument_starlette_url
-from ddtrace.appsec._iast._patch import _patched_dictionary
 from ddtrace.appsec._iast._patch import try_wrap_function_wrapper
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+from ddtrace.appsec._iast._taint_utils import taint_dictionary
 from ddtrace.appsec._iast._taint_utils import taint_structure
 from ddtrace.appsec._iast.secure_marks.sanitizers import cmdi_sanitizer
 from ddtrace.internal.logger import get_logger
@@ -100,12 +98,12 @@ def _on_flask_patch(flask_version):
             try_wrap_function_wrapper(
                 "werkzeug.wrappers.request",
                 "Request.get_data",
-                functools.partial(_patched_dictionary, OriginType.BODY, OriginType.BODY),
+                functools.partial(taint_dictionary, OriginType.BODY, OriginType.BODY),
             )
             try_wrap_function_wrapper(
                 "werkzeug.wrappers.request",
                 "Request.get_json",
-                functools.partial(_patched_dictionary, OriginType.BODY, OriginType.BODY),
+                functools.partial(taint_dictionary, OriginType.BODY, OriginType.BODY),
             )
 
             _set_metric_iast_instrumented_source(OriginType.BODY)
@@ -354,7 +352,7 @@ def _on_iast_fastapi_patch():
     try_wrap_function_wrapper(
         "starlette.requests",
         "cookie_parser",
-        functools.partial(_patched_dictionary, OriginType.COOKIE_NAME, OriginType.COOKIE),
+        functools.partial(taint_dictionary, OriginType.COOKIE_NAME, OriginType.COOKIE),
     )
     _set_metric_iast_instrumented_source(OriginType.COOKIE)
     _set_metric_iast_instrumented_source(OriginType.COOKIE_NAME)
@@ -535,3 +533,26 @@ def _on_werkzeug_render_debugger_html(html):
         set_iast_stacktrace_reported(True)
     except Exception:
         log.debug("Unexpected exception checking for stacktrace leak", exc_info=True)
+
+
+def _iast_instrument_starlette_request(wrapped, instance, args, kwargs):
+    def receive(self):
+        """This pattern comes from a Request._receive property, which returns a callable"""
+
+        async def wrapped_property_call():
+            body = await self._receive()
+            return taint_structure(body, OriginType.BODY, OriginType.BODY, override_pyobject_tainted=True)
+
+        return wrapped_property_call
+
+    # `self._receive` is set in `__init__`, so we wait for the constructor to finish before setting the new property
+    wrapped(*args, **kwargs)
+    instance.__class__.receive = property(receive)
+
+
+async def _iast_instrument_starlette_request_body(wrapped, instance, args, kwargs):
+    result = await wrapped(*args, **kwargs)
+
+    return taint_pyobject(
+        result, source_name=origin_to_str(OriginType.PATH), source_value=result, source_origin=OriginType.BODY
+    )

ddtrace/appsec/_iast/_iast_request_context_base.py

@@ -9,7 +9,7 @@
 from ddtrace.appsec._iast._overhead_control_engine import oce
 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 _request_tainted
+from ddtrace.appsec._iast._utils import _num_objects_tainted_in_request
 from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils.formats import asbool
@@ -22,7 +22,7 @@
 
 
 def _set_span_tag_iast_request_tainted(span):
-    total_objects_tainted = _request_tainted()
+    total_objects_tainted = _num_objects_tainted_in_request()
 
     if total_objects_tainted > 0:
         span.set_tag(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED, total_objects_tainted)

ddtrace/appsec/_iast/_metrics.py

@@ -10,7 +10,7 @@
 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 _request_tainted
+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
@@ -109,7 +109,7 @@ def _set_metric_iast_executed_sink(vulnerability_type):
 
 @metric_verbosity(TELEMETRY_INFORMATION_VERBOSITY)
 def _set_metric_iast_request_tainted():
-    total_objects_tainted = _request_tainted()
+    total_objects_tainted = _num_objects_tainted_in_request()
     if total_objects_tainted > 0:
         telemetry.telemetry_writer.add_count_metric(TELEMETRY_NAMESPACE.IAST, "request.tainted", total_objects_tainted)
 

ddtrace/appsec/_iast/_patch.py

@@ -10,11 +10,6 @@
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import origin_to_str
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
-from ddtrace.appsec._iast._taint_utils import taint_structure
-from ddtrace.internal.logger import get_logger
-
-
-log = get_logger(__name__)
 
 
 def set_and_check_module_is_patched(module_str: Text, default_attr: Text = "_datadog_patch") -> bool:
@@ -45,12 +40,6 @@ def try_wrap_function_wrapper(module: Text, name: Text, wrapper: Callable):
         iast_instrumentation_wrapt_debug_log(f"Module {module}.{name} not exists")
 
 
-def _patched_dictionary(origin_key, origin_value, original_func, instance, args, kwargs):
-    result = original_func(*args, **kwargs)
-
-    return taint_structure(result, origin_key, origin_value, override_pyobject_tainted=True)
-
-
 def _iast_instrument_starlette_url(wrapped, instance, args, kwargs):
     def path(self) -> str:
         return taint_pyobject(
@@ -64,29 +53,6 @@ def path(self) -> str:
     wrapped(*args, **kwargs)
 
 
-def _iast_instrument_starlette_request(wrapped, instance, args, kwargs):
-    def receive(self):
-        """This pattern comes from a Request._receive property, which returns a callable"""
-
-        async def wrapped_property_call():
-            body = await self._receive()
-            return taint_structure(body, OriginType.BODY, OriginType.BODY, override_pyobject_tainted=True)
-
-        return wrapped_property_call
-
-    # `self._receive` is set in `__init__`, so we wait for the constructor to finish before setting the new property
-    wrapped(*args, **kwargs)
-    instance.__class__.receive = property(receive)
-
-
-async def _iast_instrument_starlette_request_body(wrapped, instance, args, kwargs):
-    result = await wrapped(*args, **kwargs)
-
-    return taint_pyobject(
-        result, source_name=origin_to_str(OriginType.PATH), source_value=result, source_origin=OriginType.BODY
-    )
-
-
 def _iast_instrument_starlette_scope(scope):
     if scope.get("path_params"):
         try:

ddtrace/appsec/_iast/_span_metrics.py

@@ -3,11 +3,11 @@
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast._iast_env import _get_iast_env
 from ddtrace.appsec._iast._metrics import _metric_key_as_snake_case
-from ddtrace.appsec._iast._utils import _request_tainted
+from ddtrace.appsec._iast._utils import _num_objects_tainted_in_request
 
 
 def _set_span_tag_iast_request_tainted(span):
-    total_objects_tainted = _request_tainted()
+    total_objects_tainted = _num_objects_tainted_in_request()
 
     if total_objects_tainted > 0:
         span.set_tag(IAST_SPAN_TAGS.TELEMETRY_REQUEST_TAINTED, total_objects_tainted)

ddtrace/appsec/_iast/_taint_utils.py

@@ -524,3 +524,9 @@ def taint_structure(main_obj, origin_key, origin_value, override_pyobject_tainte
             return LazyTaintDict(main_obj, (origin_key, origin_value), override_pyobject_tainted)
         elif isinstance(main_obj, abc.Sequence):
             return LazyTaintList(main_obj, (origin_key, origin_value), override_pyobject_tainted)
+
+
+def taint_dictionary(origin_key, origin_value, original_func, instance, args, kwargs):
+    result = original_func(*args, **kwargs)
+
+    return taint_structure(result, origin_key, origin_value, override_pyobject_tainted=True)

ddtrace/appsec/_iast/_utils.py

@@ -21,5 +21,5 @@ def _is_iast_propagation_debug_enabled():
     return asm_config._iast_propagation_debug
 
 
-def _request_tainted():
+def _num_objects_tainted_in_request():
     return num_objects_tainted()