chore(iast): refactor open/read wrapping [backport 2.21] (#12776)

Backport 36fa1ee from #12757 to 2.21.

In this PR, we're migrating the way we propagate the `open` and `read`
functions. Until now, we were using `wrapt`, but it's too invasive since
these functions are used in many core operations. That's why we've
migrated it to AST, ensuring that only the functions actually used in
the application get replaced.

To achieve this, we've added a new import in the AST visitor. In
addition to aspects and sink points, we now define "sources" as well.

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

@@ -45,14 +45,7 @@ def patch_common_modules():
 
     try_wrap_function_wrapper("builtins", "open", wrapped_open_CFDDB7ABBA9081B6)
     try_wrap_function_wrapper("urllib.request", "OpenerDirector.open", wrapped_open_ED4CF71136E15EBF)
-    try_wrap_function_wrapper("_io", "BytesIO.read", wrapped_read_F3E51D71B4EC16EF)
-    try_wrap_function_wrapper("_io", "StringIO.read", wrapped_read_F3E51D71B4EC16EF)
     core.on("asm.block.dbapi.execute", execute_4C9BAC8E228EB347)
-    if asm_config._iast_enabled:
-        from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
-        from ddtrace.appsec._iast.constants import VULN_PATH_TRAVERSAL
-
-        _set_metric_iast_instrumented_sink(VULN_PATH_TRAVERSAL)
     _is_patched = True
 
 
@@ -69,29 +62,6 @@ def unpatch_common_modules():
     _is_patched = False
 
 
-def wrapped_read_F3E51D71B4EC16EF(original_read_callable, instance, args, kwargs):
-    """
-    wrapper for _io.BytesIO and _io.StringIO read function
-    """
-    result = original_read_callable(*args, **kwargs)
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
-        from ddtrace.appsec._iast._taint_tracking import OriginType
-        from ddtrace.appsec._iast._taint_tracking import Source
-        from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
-        from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
-
-        ranges = get_tainted_ranges(instance)
-        if len(ranges) > 0:
-            source = ranges[0].source if ranges[0].source else Source(name="_io", value=result, origin=OriginType.EMPTY)
-            result = taint_pyobject(
-                pyobject=result,
-                source_name=source.name,
-                source_value=source.value,
-                source_origin=source.origin,
-            )
-    return result
-
-
 def _must_block(actions: Iterable[str]) -> bool:
     return any(action in (WAF_ACTIONS.BLOCK_ACTION, WAF_ACTIONS.REDIRECT_ACTION) for action in actions)
 
@@ -100,15 +70,6 @@ def wrapped_open_CFDDB7ABBA9081B6(original_open_callable, instance, args, kwargs
     """
     wrapper for open file function
     """
-    if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
-        try:
-            from ddtrace.appsec._iast.taint_sinks.path_traversal import check_and_report_path_traversal
-
-            check_and_report_path_traversal(*args, **kwargs)
-        except ImportError:
-            # open is used during module initialization
-            # and shouldn't be changed at that time
-            return original_open_callable(*args, **kwargs)
     if (
         asm_config._asm_enabled
         and asm_config._ep_enabled

ddtrace/appsec/_iast/_ast/visitor.py

@@ -15,6 +15,7 @@
 from ..._constants import IAST
 from .._metrics import _set_metric_iast_instrumented_propagation
 from ..constants import DEFAULT_PATH_TRAVERSAL_FUNCTIONS
+from ..constants import DEFAULT_SOURCE_IO_FUNCTIONS
 from ..constants import DEFAULT_WEAK_RANDOMNESS_FUNCTIONS
 
 
@@ -29,6 +30,7 @@
 CODE_TYPE_SITE_PACKAGES = "site_packages"
 CODE_TYPE_STDLIB = "stdlib"
 TAINT_SINK_FUNCTION_REPLACEMENT = _PREFIX + "taint_sinks.ast_function"
+SOURCES_FUNCTION_REPLACEMENT = _PREFIX + "sources.ast_function"
 
 
 def _mark_avoid_convert_recursively(node):
@@ -132,14 +134,6 @@ def _mark_avoid_convert_recursively(node):
     "taint_sinks": {
         "weak_randomness": DEFAULT_WEAK_RANDOMNESS_FUNCTIONS,
         "path_traversal": DEFAULT_PATH_TRAVERSAL_FUNCTIONS,
-        "other": {
-            "load",
-            "run",
-            "path",
-            "exit",
-            "sleep",
-            "socket",
-        },
         # These explicitly WON'T be replaced by taint_sink_function:
         "disabled": {
             "__new__",
@@ -149,6 +143,7 @@ def _mark_avoid_convert_recursively(node):
             "super",
         },
     },
+    "sources": {"io": DEFAULT_SOURCE_IO_FUNCTIONS, "disabled": {}},
 }
 
 
@@ -168,9 +163,11 @@ def __init__(
         self._sinkpoints_spec = {
             "definitions_module": "ddtrace.appsec._iast.taint_sinks",
             "alias_module": _PREFIX + "taint_sinks",
-            "functions": {},
         }
-        self._sinkpoints_functions = self._sinkpoints_spec["functions"]
+        self._source_spec = {
+            "definitions_module": "ddtrace.appsec._iast.sources",
+            "alias_module": _PREFIX + "sources",
+        }
 
         self._aspect_index = _ASPECTS_SPEC["slices"]["index"]
         self._aspect_slice = _ASPECTS_SPEC["slices"]["slice"]
@@ -182,11 +179,14 @@ def __init__(
         self._aspect_build_string = _ASPECTS_SPEC["operators"]["BUILD_STRING"]
 
         # Sink points
-        self._taint_sink_replace_any = self._merge_taint_sinks(
-            _ASPECTS_SPEC["taint_sinks"]["other"],
+        self._taint_sink_replace_any = self._merge_dicts(
             _ASPECTS_SPEC["taint_sinks"]["weak_randomness"],
             *[functions for module, functions in _ASPECTS_SPEC["taint_sinks"]["path_traversal"].items()],
         )
+        self._source_replace_any = self._merge_dicts(
+            *[functions for module, functions in _ASPECTS_SPEC["sources"]["io"].items()],
+        )
+
         self._taint_sink_replace_disabled = _ASPECTS_SPEC["taint_sinks"]["disabled"]
 
         self.update_location(filename, module_name)
@@ -219,7 +219,7 @@ def update_location(self, filename: str = "", module_name: str = ""):
             self.codetype = CODE_TYPE_STDLIB
 
     @staticmethod
-    def _merge_taint_sinks(*args_functions: Set[str]) -> Set[str]:
+    def _merge_dicts(*args_functions: Set[str]) -> Set[str]:
         merged_set = set()
 
         for functions in args_functions:
@@ -284,6 +284,11 @@ def _should_replace_with_taint_sink(self, call_node: ast.Call, is_function: bool
 
         return function_name in self._taint_sink_replace_any
 
+    def _should_replace_with_source(self, call_node: ast.Call, is_function: bool) -> bool:
+        function_name = self._get_function_name(call_node, is_function)
+
+        return function_name in self._source_replace_any
+
     def _add_original_function_as_arg(self, call_node: ast.Call, is_function: bool) -> Any:
         """
         Creates the arguments for the original function
@@ -452,6 +457,21 @@ def visit_Module(self, module_node: ast.Module) -> Any:
             ],
         )
         module_node.body.insert(insert_position, replacements_import)
+
+        definitions_module = self._source_spec["definitions_module"]
+        replacements_import = self._node(
+            ast.Import,
+            module_node,
+            names=[
+                ast.alias(
+                    lineno=1,
+                    col_offset=0,
+                    name=definitions_module,
+                    asname=self._source_spec["alias_module"],
+                )
+            ],
+        )
+        module_node.body.insert(insert_position, replacements_import)
         # Must be called here instead of the start so the line offset is already
         # processed
         self.generic_visit(module_node)
@@ -519,11 +539,7 @@ def visit_Call(self, call_node: ast.Call) -> Any:
                 # Substitute function call
                 call_node.func = self._attr_node(call_node, aspect)
                 self.ast_modified = call_modified = True
-            else:
-                sink_point = self._sinkpoints_functions.get(func_name_node)
-                if sink_point:
-                    call_node.func = self._attr_node(call_node, sink_point)
-                    self.ast_modified = call_modified = True
+
         # Call [attr] -> Attribute [value]-> Attribute [value]-> Attribute
         # a.b.c.method()
         # replaced_method(a.b.c)
@@ -593,6 +609,14 @@ def visit_Call(self, call_node: ast.Call) -> Any:
                     # Create a new Name node for the replacement and set it as node.func
                     call_node.func = self._attr_node(call_node, aspect)
                     self.ast_modified = call_modified = True
+                else:
+                    aspect = self._should_replace_with_source(call_node, False)
+                    if aspect:
+                        # Send 0 as flag_added_args value
+                        call_node.args.insert(0, self._int_constant(call_node, 0))
+                        call_node.args = self._add_original_function_as_arg(call_node, False)
+                        call_node.func = self._attr_node(call_node, SOURCES_FUNCTION_REPLACEMENT)
+                        self.ast_modified = call_modified = True
 
         if self.codetype == CODE_TYPE_FIRST_PARTY:
             # Function replacement case

ddtrace/appsec/_iast/constants.py

@@ -55,6 +55,8 @@
 }
 
 DEFAULT_PATH_TRAVERSAL_FUNCTIONS = {
+    "_io": {"open"},
+    "io": {"open"},
     "glob": {"glob"},
     "os": {
         "mkdir",
@@ -86,3 +88,7 @@
 DBAPI_MYSQL = "mysql"
 DBAPI_MARIADB = "mariadb"
 DBAPI_INTEGRATIONS = (DBAPI_SQLITE, DBAPI_PSYCOPG, DBAPI_MYSQL, DBAPI_MARIADB)
+
+DEFAULT_SOURCE_IO_FUNCTIONS = {
+    "_io": {"read"},
+}

ddtrace/appsec/_iast/sources/__init__.py

@@ -0,0 +1,6 @@
+from .ast_taint import ast_function
+
+
+__all__ = [
+    "ast_function",
+]

ddtrace/appsec/_iast/sources/ast_taint.py

@@ -0,0 +1,50 @@
+from typing import Any
+from typing import Callable
+
+from ddtrace.appsec._iast._taint_tracking import OriginType
+from ddtrace.appsec._iast._taint_tracking import Source
+from ddtrace.appsec._iast._taint_tracking._taint_objects import get_tainted_ranges
+from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
+from ddtrace.settings.asm import config as asm_config
+
+from ..constants import DEFAULT_SOURCE_IO_FUNCTIONS
+
+
+def ast_function(
+    func: Callable,
+    flag_added_args: Any,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
+    instance = getattr(func, "__self__", None)
+    func_name = getattr(func, "__name__", None)
+    cls_name = ""
+    if instance is not None and func_name:
+        try:
+            cls_name = instance.__class__.__name__
+        except AttributeError:
+            pass
+
+    if flag_added_args > 0:
+        args = args[flag_added_args:]
+
+    module_name = instance.__class__.__module__
+    result = func(*args, **kwargs)
+    if (
+        module_name == "_io"
+        and cls_name in ("BytesIO", "StringIO")
+        and func_name in DEFAULT_SOURCE_IO_FUNCTIONS[module_name]
+    ):
+        if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+            ranges = get_tainted_ranges(instance)
+            if len(ranges) > 0:
+                source = (
+                    ranges[0].source if ranges[0].source else Source(name="_io", value=result, origin=OriginType.EMPTY)
+                )
+                result = taint_pyobject(
+                    pyobject=result,
+                    source_name=source.name,
+                    source_value=source.value,
+                    source_origin=source.origin,
+                )
+    return result

ddtrace/appsec/_iast/taint_sinks/ast_taint.py

@@ -1,6 +1,8 @@
 from typing import Any
 from typing import Callable
 
+from ddtrace.settings.asm import config as asm_config
+
 from ..._constants import IAST_SPAN_TAGS
 from .._metrics import _set_metric_iast_executed_sink
 from .._metrics import increment_iast_span_metric
@@ -34,10 +36,11 @@ def ast_function(
         and cls_name == "Random"
         and func_name in DEFAULT_WEAK_RANDOMNESS_FUNCTIONS
     ):
-        # Weak, run the analyzer
-        increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakRandomness.vulnerability_type)
-        _set_metric_iast_executed_sink(WeakRandomness.vulnerability_type)
-        WeakRandomness.report(evidence_value=cls_name + "." + func_name)
+        if asm_config._iast_enabled and asm_config.is_iast_request_enabled:
+            # Weak, run the analyzer
+            increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, WeakRandomness.vulnerability_type)
+            _set_metric_iast_executed_sink(WeakRandomness.vulnerability_type)
+            WeakRandomness.report(evidence_value=cls_name + "." + func_name)
     elif hasattr(func, "__module__") and DEFAULT_PATH_TRAVERSAL_FUNCTIONS.get(func.__module__):
         if func_name in DEFAULT_PATH_TRAVERSAL_FUNCTIONS[func.__module__]:
             check_and_report_path_traversal(*args, **kwargs)

ddtrace/appsec/_iast/taint_sinks/path_traversal.py

@@ -3,6 +3,7 @@
 from ddtrace.appsec._constants import IAST_SPAN_TAGS
 from ddtrace.appsec._iast import oce
 from ddtrace.appsec._iast._metrics import _set_metric_iast_executed_sink
+from ddtrace.appsec._iast._metrics import _set_metric_iast_instrumented_sink
 from ddtrace.appsec._iast._metrics import increment_iast_span_metric
 from ddtrace.appsec._iast._taint_tracking._taint_objects import is_pyobject_tainted
 from ddtrace.appsec._iast.constants import VULN_PATH_TRAVERSAL
@@ -20,19 +21,18 @@ class PathTraversal(VulnerabilityBase):
     vulnerability_type = VULN_PATH_TRAVERSAL
 
 
+IS_REPORTED_INTRUMENTED_SINK = False
+
+
 def check_and_report_path_traversal(*args: Any, **kwargs: Any) -> None:
+    global IS_REPORTED_INTRUMENTED_SINK
+    if not IS_REPORTED_INTRUMENTED_SINK:
+        _set_metric_iast_instrumented_sink(VULN_PATH_TRAVERSAL)
+        IS_REPORTED_INTRUMENTED_SINK = True
+
     if asm_config.is_iast_request_enabled and PathTraversal.has_quota():
-        try:
-            increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, PathTraversal.vulnerability_type)
-            _set_metric_iast_executed_sink(PathTraversal.vulnerability_type)
-            filename_arg = args[0] if args else kwargs.get("file", None)
-            if is_pyobject_tainted(filename_arg):
-                PathTraversal.report(evidence_value=filename_arg)
-        except Exception:  # nosec
-            # FIXME: see below
-            # log.debug("Unexpected exception while reporting vulnerability", exc_info=True)
-            pass
-    # FIXME: this is commented out because logs can execute _open which can runs the check_and_report_path_traversal
-    # entering an infinite loop
-    # else:
-    #     log.debug("IAST: no vulnerability quota to analyze more sink points")
+        filename_arg = args[0] if args else kwargs.get("file", None)
+        if is_pyobject_tainted(filename_arg):
+            PathTraversal.report(evidence_value=filename_arg)
+    increment_iast_span_metric(IAST_SPAN_TAGS.TELEMETRY_EXECUTED_SINK, PathTraversal.vulnerability_type)
+    _set_metric_iast_executed_sink(PathTraversal.vulnerability_type)

tests/appsec/iast/_ast/test_ast_patching.py

@@ -78,6 +78,7 @@ def test_astpatch_module_changed(module_name):
     assert ("", None) != (module_path, new_ast)
     new_code = astunparse.unparse(new_ast)
     assert new_code.startswith(
+        f"\nimport ddtrace.appsec._iast.sources as {_PREFIX}sources"
         f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks"
         f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects"
     )
@@ -95,6 +96,7 @@ def test_astpatch_module_changed_add_operator(module_name):
     assert ("", None) != (module_path, new_ast)
     new_code = astunparse.unparse(new_ast)
     assert new_code.startswith(
+        f"\nimport ddtrace.appsec._iast.sources as {_PREFIX}sources"
         f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks"
         f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects"
     )
@@ -112,6 +114,7 @@ def test_astpatch_module_changed_add_inplace_operator(module_name):
     assert ("", None) != (module_path, new_ast)
     new_code = astunparse.unparse(new_ast)
     assert new_code.startswith(
+        f"\nimport ddtrace.appsec._iast.sources as {_PREFIX}sources"
         f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks"
         f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects"
     )
@@ -136,6 +139,7 @@ def test_astpatch_source_changed_with_future_imports(module_name):
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
+import ddtrace.appsec._iast.sources as {_PREFIX}sources
 import ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks
 import ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects
 import html"""
@@ -229,6 +233,7 @@ def test_astpatch_stringio_module_changed(module_name):
     assert ("", None) != (module_path, new_ast)
     new_code = astunparse.unparse(new_ast)
     assert new_code.startswith(
+        f"\nimport ddtrace.appsec._iast.sources as {_PREFIX}sources"
         f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks"
         f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects"
     )
@@ -247,6 +252,7 @@ def test_astpatch_bytesio_module_changed(module_name):
     assert ("", None) != (module_path, new_ast)
     new_code = astunparse.unparse(new_ast)
     assert new_code.startswith(
+        f"\nimport ddtrace.appsec._iast.sources as {_PREFIX}sources"
         f"\nimport ddtrace.appsec._iast.taint_sinks as {_PREFIX}taint_sinks"
         f"\nimport ddtrace.appsec._iast._taint_tracking.aspects as {_PREFIX}aspects"
     )