fix(iast): prepare aspect for unexpected args (#7491) [backport 2.0] (#7574)

IAST: This PR addresses two interconnected issues in the AST patching
process that replaces code with aspects:

1. Resolves a bug that occurs during AST patching, where a custom
function or method call could be replaced by one of our aspects. In
situations where the function has a different number of arguments than
our aspects, a runtime ``TypeError`` occurs. This is fixed by ensuring
that all aspects now receive ``(*args, **kwargs)`` and pass them
appropriately to the original custom function or method.

2. Fixes another bug in the AST patching process, where the module of a
function is incorrectly passed as the first argument to the aspect. In
cases where it's a custom function, our logic inadvertently passes the
module received as the first argument to the original function. The
solution involves introducing a flag to the aspect's arguments,
indicating the additional arguments added during patching. This allows
us to remove them before calling the original function or method.

- [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] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [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] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly 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)
- [x] If this PR 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] This PR doesn't touch any of that.

---------

Co-authored-by: Alberto Vara <[email protected]>
(cherry picked from commit d5caa6c)

Author: gnufede

.gitignore

@@ -144,3 +144,7 @@ ddtrace/appsec/_iast/_taint_tracking/CMakeFiles/*
 
 # CircleCI generated config
 .circleci/config.gen.yml
+
+# Automatically generated fixtures
+tests/appsec/iast/fixtures/aspects/callers.py
+tests/appsec/iast/fixtures/aspects/unpatched_callers.py

benchmarks/appsec_iast_propagation/scenario.py

@@ -38,7 +38,7 @@ def aspect_function(internal_loop, tainted):
     value = ""
     res = value
     for _ in range(internal_loop):
-        res = add_aspect(res, join_aspect("_", (tainted, "_", tainted)))
+        res = add_aspect(res, join_aspect(str.join, 1, "_", (tainted, "_", tainted)))
         value = res
         res = add_aspect(res, tainted)
         value = res

ddtrace/appsec/_iast/_ast/visitor.py

@@ -25,7 +25,7 @@
 CODE_TYPE_DD = "datadog"
 CODE_TYPE_SITE_PACKAGES = "site_packages"
 CODE_TYPE_STDLIB = "stdlib"
-TAINT_SINK_FUNCTION_REPLACEMENT = "ddtrace_taint_sinks.ast_funcion"
+TAINT_SINK_FUNCTION_REPLACEMENT = "ddtrace_taint_sinks.ast_function"
 
 
 class AstVisitor(ast.NodeTransformer):
@@ -326,6 +326,16 @@ def _none_constant(self, from_node, ctx=ast.Load()):  # noqa: B008
             kind=None,
         )
 
+    def _int_constant(self, from_node, value):
+        return ast.Constant(
+            lineno=from_node.lineno,
+            col_offset=from_node.col_offset,
+            end_lineno=getattr(from_node, "end_lineno", from_node.lineno),
+            end_col_offset=from_node.col_offset + 1,
+            value=value,
+            kind=None,
+        )
+
     def _call_node(self, from_node, func, args):  # type: (Any, Any, List[Any]) -> Any
         return self._node(ast.Call, from_node, func=func, args=args, keywords=[])
 
@@ -398,6 +408,11 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
             func_name_node = func_member.id
             aspect = self._aspect_functions.get(func_name_node)
             if aspect:
+                # Send 0 as flag_added_args value
+                call_node.args.insert(0, self._int_constant(call_node, 0))
+                # Insert original function name as first parameter
+                call_node.args = self._add_original_function_as_arg(call_node, True)
+                # Substitute function call
                 call_node.func = self._attr_node(call_node, aspect)
                 self.ast_modified = call_modified = True
             else:
@@ -425,6 +440,11 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
                 # Move the Attribute.value to 'args'
                 new_arg = func_member.value
                 call_node.args.insert(0, new_arg)
+                # Send 1 as flag_added_args value
+                call_node.args.insert(0, self._int_constant(call_node, 1))
+
+                # Insert original method as first parameter (a.b.c.method)
+                call_node.args = self._add_original_function_as_arg(call_node, False)
 
                 # Create a new Name node for the replacement and set it as node.func
                 call_node.func = self._attr_node(call_node, aspect)
@@ -433,6 +453,8 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
             elif hasattr(func_member.value, "id") or hasattr(func_member.value, "attr"):
                 aspect = self._aspect_modules.get(method_name, None)
                 if aspect:
+                    # Send 0 as flag_added_args value
+                    call_node.args.insert(0, self._int_constant(call_node, 0))
                     # Move the Function to 'args'
                     call_node.args.insert(0, call_node.func)
 
@@ -445,6 +467,8 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
             if isinstance(call_node.func, ast.Name):
                 aspect = self._should_replace_with_taint_sink(call_node, True)
                 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, TAINT_SINK_FUNCTION_REPLACEMENT)
                     self.ast_modified = call_modified = True
@@ -453,6 +477,8 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
             elif isinstance(call_node.func, ast.Attribute):
                 aspect = self._should_replace_with_taint_sink(call_node, False)
                 if aspect:
+                    # Send 0 as flag_added_args value
+                    call_node.args.insert(0, self._int_constant(call_node, 0))
                     # Create a new Name node for the replacement and set it as node.func
                     call_node.args = self._add_original_function_as_arg(call_node, False)
                     call_node.func = self._attr_node(call_node, TAINT_SINK_FUNCTION_REPLACEMENT)