fix(iast): prepare aspect for unexpected args (#7491) [backport 1.20] (#7573)

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

@@ -3,6 +3,7 @@
 from _ast import Expr
 from _ast import ImportFrom
 import ast
+import copy
 import sys
 from typing import TYPE_CHECKING
 
@@ -155,6 +156,34 @@ def _is_string_format_with_literals(self, call_node):
             and all(map(lambda x: self._is_node_constant_or_binop(x.value), call_node.keywords))
         )
 
+    def _get_function_name(self, call_node, is_function):  # type: (ast.Call, bool) -> str
+        if is_function:
+            return call_node.func.id  # type: ignore[attr-defined]
+        # If the call is to a method
+        elif type(call_node.func) == ast.Name:
+            return call_node.func.id
+
+        return call_node.func.attr  # type: ignore[attr-defined]
+
+    def _add_original_function_as_arg(self, call_node, is_function):  # type: (ast.Call, bool) -> Any
+        """
+        Creates the arguments for the original function
+        """
+        function_name = self._get_function_name(call_node, is_function)
+        function_name_arg = (
+            self._name_node(call_node, function_name, ctx=ast.Load()) if is_function else copy.copy(call_node.func)
+        )
+
+        # Arguments for stack info change from:
+        # my_function(self, *args, **kwargs)
+        # to:
+        # _add_original_function_as_arg(function_name=my_function, self, *args, **kwargs)
+        new_args = [
+            function_name_arg,
+        ] + call_node.args
+
+        return new_args
+
     def _node(self, type_, pos_from_node, **kwargs):
         # type: (Any, Any, Any) -> Any
         """
@@ -247,6 +276,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=[])
 
@@ -319,6 +358,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:
@@ -346,6 +390,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)
@@ -354,6 +403,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)