fix(iast): ast patching side effect [backport 2.1] (#7653)

Backport 601fd6a from #7638 to 2.1.

IAST: This fix resolves an issue where, at AST patching to replace code
with IAST aspects, the original function/method was passed as an extra
parameter for accurate patching (i.e. check the function inside the
aspect), and this strategy lead to unintentionally triggers side effects
in methods obtained from an expression (like ``decode`` in
``file.read(n).decode()``), resulting in unexpected multiple calls to
the expression (``file.read(n)`` in the example).

## Checklist

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

## Reviewer Checklist

- [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: Federico Mon <[email protected]>

Author: github-actions[bot]

ddtrace/appsec/_iast/_ast/visitor.py

@@ -470,8 +470,9 @@ def visit_Call(self, call_node):  # type: (ast.Call) -> Any
                 # 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)
+                # Insert None as first parameter instead of a.b.c.method
+                # to avoid unexpected side effects such as a.b.read(4).method
+                call_node.args.insert(0, self._none_constant(call_node))
 
                 # Create a new Name node for the replacement and set it as node.func
                 call_node.func = self._attr_node(call_node, aspect)

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Optional
 
 from ddtrace.internal.compat import iteritems
 
@@ -52,13 +53,15 @@ def add_aspect(op1, op2):
 
 
 def str_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> str
-    if orig_function != builtin_str:
-        if flag_added_args > 0:
-            args = args[flag_added_args:]
-        return orig_function(*args, **kwargs)
-
-    result = builtin_str(*args, **kwargs)
+    # type: (Optional[Callable], int, Any, Any) -> str
+    if orig_function:
+        if orig_function != builtin_str:
+            if flag_added_args > 0:
+                args = args[flag_added_args:]
+            return orig_function(*args, **kwargs)
+        result = builtin_str(*args, **kwargs)
+    else:
+        result = args[0].str(*args[1:], **kwargs)
 
     if isinstance(args[0], TEXT_TYPES) and is_pyobject_tainted(args[0]):
         try:
@@ -76,13 +79,16 @@ def str_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def bytes_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> bytes
-    if orig_function != builtin_bytes:
-        if flag_added_args > 0:
-            args = args[flag_added_args:]
-        return orig_function(*args, **kwargs)
+    # type: (Optional[Callable], int, Any, Any) -> bytes
+    if orig_function:
+        if orig_function != builtin_bytes:
+            if flag_added_args > 0:
+                args = args[flag_added_args:]
+            return orig_function(*args, **kwargs)
+        result = builtin_bytes(*args, **kwargs)
+    else:
+        result = args[0].bytes(*args[1:], **kwargs)
 
-    result = builtin_bytes(*args, **kwargs)
     if isinstance(args[0], TEXT_TYPES) and is_pyobject_tainted(args[0]):
         try:
             taint_pyobject_with_ranges(result, tuple(get_ranges(args[0])))
@@ -92,13 +98,16 @@ def bytes_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def bytearray_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> bytearray
-    if orig_function != builtin_bytearray:
-        if flag_added_args > 0:
-            args = args[flag_added_args:]
-        return orig_function(*args, **kwargs)
+    # type: (Optional[Callable], int, Any, Any) -> bytearray
+    if orig_function:
+        if orig_function != builtin_bytearray:
+            if flag_added_args > 0:
+                args = args[flag_added_args:]
+            return orig_function(*args, **kwargs)
+        result = builtin_bytearray(*args, **kwargs)
+    else:
+        result = args[0].bytearray(*args[1:], **kwargs)
 
-    result = builtin_bytearray(*args, **kwargs)
     if isinstance(args[0], TEXT_TYPES) and is_pyobject_tainted(args[0]):
         try:
             taint_pyobject_with_ranges(result, tuple(get_ranges(args[0])))
@@ -108,7 +117,9 @@ def bytearray_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def join_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> Any
+    # type: (Optional[Callable], int, Any, Any) -> Any
+    if not orig_function:
+        orig_function = args[0].join
     if not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
@@ -158,16 +169,22 @@ def slice_aspect(candidate_text, start, stop, step) -> Any:
 
 
 def bytearray_extend_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> Any
-    if not isinstance(orig_function, BuiltinFunctionType):
+    # type: (Optional[Callable], int, Any, Any) -> Any
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
 
+    if len(args) < 2:
+        # If we're not receiving at least 2 arguments, means the call was
+        # ``x.extend()`` and not ``x.extend(y)``
+        # so either not the extend we're looking for, or no changes in taint ranges.
+        return args[0].extend(*args[1:], **kwargs)
+
     op1 = args[0]
     op2 = args[1]
     if not isinstance(op1, bytearray) or not isinstance(op2, (bytearray, bytes)):
-        return op1.extend(op2)
+        return op1.extend(*args[1:], **kwargs)
     try:
         return _extend_aspect(op1, op2)
     except Exception as e:
@@ -217,7 +234,9 @@ def build_string_aspect(*args):  # type: (List[Any]) -> str
 
 
 def ljust_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> Union[str, bytes, bytearray]
+    # type: (Optional[Callable], int, Any, Any) -> Union[str, bytes, bytearray]
+    if not orig_function:
+        orig_function = args[0].ljust
     if not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
@@ -252,8 +271,8 @@ def ljust_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def zfill_aspect(orig_function, flag_added_args, *args, **kwargs):
-    # type: (Callable, int, Any, Any) -> Any
-    if not isinstance(orig_function, BuiltinFunctionType):
+    # type: (Optional[Callable], int, Any, Any) -> Any
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -295,11 +314,14 @@ def zfill_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def format_aspect(
-    orig_function,  # type: Callable
+    orig_function,  # type: Optional[Callable]
     flag_added_args,  # type: int
     *args,  # type: Any
     **kwargs,  # type: Dict[str, Any]
 ):  # type: (...) -> str
+    if not orig_function:
+        orig_function = args[0].format
+
     if not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
@@ -347,8 +369,10 @@ def format_aspect(
     return result
 
 
-def format_map_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> str
-    if not isinstance(orig_function, BuiltinFunctionType):
+def format_map_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> str
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -494,12 +518,14 @@ def incremental_translation(self, incr_coder, funcode, empty):
 
 
 def decode_aspect(orig_function, flag_added_args, *args, **kwargs):
-    if flag_added_args > 0:
-        self = args[0]
-        args = args[flag_added_args:]
-    else:
+    if orig_function and not flag_added_args:
+        # This patch is unexpected, so we fallback
+        # to executing the original function
         return orig_function(*args, **kwargs)
 
+    self = args[0]
+    args = args[(flag_added_args or 1) :]
+    # Assume we call decode method of the first argument
     result = self.decode(*args, **kwargs)
 
     if not is_pyobject_tainted(self) or not isinstance(self, bytes):
@@ -515,12 +541,14 @@ def decode_aspect(orig_function, flag_added_args, *args, **kwargs):
 
 
 def encode_aspect(orig_function, flag_added_args, *args, **kwargs):
-    if flag_added_args > 0:
-        self = args[0]
-        args = args[flag_added_args:]
-    else:
+    if orig_function and not flag_added_args:
+        # This patch is unexpected, so we fallback
+        # to executing the original function
         return orig_function(*args, **kwargs)
 
+    self = args[0]
+    args = args[(flag_added_args or 1) :]
+    # Assume we call encode method of the first argument
     result = self.encode(*args, **kwargs)
 
     if not is_pyobject_tainted(self) or not isinstance(self, str):
@@ -535,8 +563,10 @@ def encode_aspect(orig_function, flag_added_args, *args, **kwargs):
     return result
 
 
-def upper_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def upper_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -553,17 +583,17 @@ def upper_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Cal
         return candidate_text.upper(*args, **kwargs)
 
 
-def lower_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def lower_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
 
     candidate_text = args[0]
     args = args[flag_added_args:]
     if not isinstance(candidate_text, TEXT_TYPES):
-        if flag_added_args > 0:
-            args = args[flag_added_args:]
         return candidate_text.lower(*args, **kwargs)
 
     try:
@@ -573,8 +603,10 @@ def lower_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Cal
         return candidate_text.lower(*args, **kwargs)
 
 
-def swapcase_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def swapcase_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -590,8 +622,10 @@ def swapcase_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (
         return candidate_text.swapcase(*args, **kwargs)
 
 
-def title_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def title_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -607,8 +641,10 @@ def title_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Cal
         return candidate_text.title(*args, **kwargs)
 
 
-def capitalize_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def capitalize_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)
@@ -625,16 +661,21 @@ def capitalize_aspect(orig_function, flag_added_args, *args, **kwargs):  # type:
         return candidate_text.capitalize(*args, **kwargs)
 
 
-def casefold_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
-        if flag_added_args > 0:
-            args = args[flag_added_args:]
-        return orig_function(*args, **kwargs)
+def casefold_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function:
+        if not isinstance(orig_function, BuiltinFunctionType):
+            if flag_added_args > 0:
+                args = args[flag_added_args:]
+            return orig_function(*args, **kwargs)
+    else:
+        orig_function = getattr(args[0], "casefold", None)
 
-    if orig_function.__qualname__ not in ("str.casefold", "bytes.casefold", "bytearray.casefold"):
+    if orig_function and orig_function.__qualname__ not in ("str.casefold", "bytes.casefold", "bytearray.casefold"):
         if flag_added_args > 0:
             args = args[flag_added_args:]
-        return orig_function(args, **kwargs)
+        return orig_function(*args, **kwargs)
 
     candidate_text = args[0]
     args = args[flag_added_args:]
@@ -649,8 +690,10 @@ def casefold_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (
         return candidate_text.casefold(*args, **kwargs)  # type: ignore[union-attr]
 
 
-def translate_aspect(orig_function, flag_added_args, *args, **kwargs):  # type: (Callable, int, Any, Any) -> TEXT_TYPE
-    if not isinstance(orig_function, BuiltinFunctionType):
+def translate_aspect(
+    orig_function, flag_added_args, *args, **kwargs
+):  # type: (Optional[Callable], int, Any, Any) -> TEXT_TYPE
+    if orig_function and not isinstance(orig_function, BuiltinFunctionType):
         if flag_added_args > 0:
             args = args[flag_added_args:]
         return orig_function(*args, **kwargs)

releasenotes/notes/iast-fix-patching-side-effect-56d0b168e74839fc.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Vulnerability Management for Code-level (IAST): This fix resolves an issue where, at AST patching to replace code with IAST aspects, passing the original function/method as an extra parameter for accurate patching unintentionally triggers side effects in methods obtained from an expression (like ``decode`` in ``file.read(n).decode()``), resulting in unexpected multiple calls to the expression (``file.read(n)`` in the example).