fix(iast): re.finditer aspect error (#11027)

Ensure IAST propagation does not raise side effects related to
re.finditer.
We detect this error in #10988 PR, when FastAPI headers were empty in
framework tests:
https://github.com/DataDog/dd-trace-py/actions/runs/11273577079/job/31350947622

We can revert this system tests PR after this PR:
DataDog/system-tests#3230

This error was detected in
[python-multipart==0.0.05](https://pypi.org/project/python-multipart/0.0.5/)

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

(cherry picked from commit 9973b22)

benchmarks/bm/iast_fixtures/str_methods.py

@@ -903,6 +903,10 @@ def do_join_generator(mystring: str) -> Text:
     return "".join(gen)
 
 
+def do_join_generator_as_argument(mystring: str, gen: Generator[str, None, None]) -> Text:
+    return mystring.join(gen)
+
+
 def do_join_generator_2(mystring: str) -> Text:
     def parts() -> Generator:
         for i in ["x", "y", "z"]:

ddtrace/appsec/_iast/_taint_tracking/aspects.py

@@ -2,6 +2,7 @@
 from builtins import bytes as builtin_bytes
 from builtins import str as builtin_str
 import codecs
+import itertools
 from re import Match
 from re import Pattern
 from types import BuiltinFunctionType
@@ -951,9 +952,7 @@ def re_findall_aspect(
     return result
 
 
-def re_finditer_aspect(
-    orig_function: Optional[Callable], flag_added_args: int, *args: Any, **kwargs: Any
-) -> Union[TEXT_TYPES, Tuple[TEXT_TYPES, int]]:
+def re_finditer_aspect(orig_function: Optional[Callable], flag_added_args: int, *args: Any, **kwargs: Any) -> Iterator:
     if orig_function is not None and (not flag_added_args or not args):
         # This patch is unexpected, so we fallback
         # to executing the original function
@@ -982,9 +981,9 @@ def re_finditer_aspect(
         string = args[0]
         if is_pyobject_tainted(string):
             ranges = get_ranges(string)
-            for elem in result:
+            result, result_backup = itertools.tee(result)
+            for elem in result_backup:
                 taint_pyobject_with_ranges(elem, ranges)
-
     return result
 
 

releasenotes/notes/iast-fix-re-finditer-aspect-8925b30073169222.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Code Security: Ensure IAST propagation does not raise side effects related to re.finditer.

tests/appsec/app.py

@@ -66,6 +66,7 @@
 from tests.appsec.iast_packages.packages.pkg_pyopenssl import pkg_pyopenssl
 from tests.appsec.iast_packages.packages.pkg_pyparsing import pkg_pyparsing
 from tests.appsec.iast_packages.packages.pkg_python_dateutil import pkg_python_dateutil
+from tests.appsec.iast_packages.packages.pkg_python_multipart import pkg_python_multipart
 from tests.appsec.iast_packages.packages.pkg_pytz import pkg_pytz
 from tests.appsec.iast_packages.packages.pkg_pyyaml import pkg_pyyaml
 from tests.appsec.iast_packages.packages.pkg_requests import pkg_requests
@@ -144,6 +145,7 @@
 app.register_blueprint(pkg_pyopenssl)
 app.register_blueprint(pkg_pyparsing)
 app.register_blueprint(pkg_python_dateutil)
+app.register_blueprint(pkg_python_multipart)
 app.register_blueprint(pkg_pytz)
 app.register_blueprint(pkg_pyyaml)
 app.register_blueprint(pkg_requests)

tests/appsec/iast/aspects/test_join_aspect_fixtures.py

@@ -311,6 +311,32 @@ def test_string_join_generator(self):
         assert result[ranges[1].start : (ranges[1].start + ranges[1].length)] == "abcde"
         assert result[ranges[2].start : (ranges[2].start + ranges[2].length)] == "abcde"
 
+    def test_string_join_generator_multiples_times(self):
+        base_string = "abcde"
+        gen_string = "--+--"
+        gen = (gen_string for _ in ["1", "2", "3"])
+        result = mod.do_join_generator_as_argument(base_string, gen)
+        assert result == "--+--abcde--+--abcde--+--"
+        assert not get_tainted_ranges(result)
+        result = mod.do_join_generator_as_argument(base_string, gen)
+        assert result == ""
+        # Tainted
+        tainted_base_string = taint_pyobject(
+            pyobject=base_string,
+            source_name="joiner",
+            source_value=base_string,
+            source_origin=OriginType.PARAMETER,
+        )
+        gen = (gen_string for _ in ["1", "2", "3"])
+        result = mod.do_join_generator_as_argument(tainted_base_string, gen)
+        result_2 = mod.do_join_generator_as_argument(tainted_base_string, gen)
+        assert result == "--+--abcde--+--abcde--+--"
+        assert result_2 == ""
+
+        ranges = get_tainted_ranges(result)
+        assert result[ranges[0].start : (ranges[0].start + ranges[0].length)] == "abcde"
+        assert result[ranges[1].start : (ranges[1].start + ranges[1].length)] == "abcde"
+
     def test_string_join_args_kwargs(self):
         # type: () -> None
         # Not tainted

tests/appsec/iast/aspects/test_re_aspects.py

@@ -1,6 +1,8 @@
 import re
 import typing
 
+import pytest
+
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking import Source
 from ddtrace.appsec._iast._taint_tracking import TaintRange
@@ -475,10 +477,18 @@ def test_re_finditer_aspect_tainted_string():
     re_slash = re.compile(r"[/.][a-z]*")
 
     res_iterator = re_finditer_aspect(None, 1, re_slash, tainted_foobarbaz)
-    assert isinstance(res_iterator, typing.Iterator)
+    assert isinstance(res_iterator, typing.Iterator), f"res_iterator is of type {type(res_iterator)}"
+    try:
+        tainted_item = next(res_iterator)
+        assert get_tainted_ranges(tainted_item) == [
+            TaintRange(0, 18, Source("test_re_sub_aspect_tainted_string", tainted_foobarbaz, OriginType.PARAMETER)),
+        ]
+    except StopIteration:
+        pytest.fail("re_finditer_aspect result generator is depleted")
+
     for i in res_iterator:
         assert get_tainted_ranges(i) == [
-            TaintRange(0, len(i), Source("test_re_sub_aspect_tainted_string", tainted_foobarbaz, OriginType.PARAMETER)),
+            TaintRange(0, 18, Source("test_re_sub_aspect_tainted_string", tainted_foobarbaz, OriginType.PARAMETER)),
         ]
 
 
@@ -489,5 +499,11 @@ def test_re_finditer_aspect_not_tainted():
 
     res_iterator = re_finditer_aspect(None, 1, re_slash, not_tainted_foobarbaz)
     assert isinstance(res_iterator, typing.Iterator)
+
+    try:
+        tainted_item = next(res_iterator)
+        assert not is_pyobject_tainted(tainted_item)
+    except StopIteration:
+        pytest.fail("re_finditer_aspect result generator is finished")
     for i in res_iterator:
         assert not is_pyobject_tainted(i)

tests/appsec/iast_packages/packages/pkg_python_multipart.py

@@ -0,0 +1,49 @@
+"""
+isodate==0.6.1
+
+https://pypi.org/project/isodate/
+"""
+
+from flask import Blueprint
+from flask import request
+
+from .utils import ResultResponse
+
+
+pkg_python_multipart = Blueprint("multipart", __name__)
+
+
+@pkg_python_multipart.route("/python-multipart")
+def pkg_multipart_view():
+    from multipart.multipart import parse_options_header
+
+    response = ResultResponse(request.args.get("package_param"))
+
+    try:
+        _, params = parse_options_header(response.package_param)
+
+        response.result1 = str(params[b"boundary"], "utf-8")
+    except Exception as e:
+        response.result1 = f"Error: {str(e)}"
+
+    return response.json()
+
+
+@pkg_python_multipart.route("/python-multipart_propagation")
+def pkg_multipart_propagation_view():
+    from multipart.multipart import parse_options_header
+
+    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+
+    response = ResultResponse(request.args.get("package_param"))
+    if not is_pyobject_tainted(response.package_param):
+        response.result1 = "Error: package_param is not tainted"
+        return response.json()
+
+    _, params = parse_options_header(response.package_param)
+    response.result1 = (
+        "OK"
+        if is_pyobject_tainted(params[b"boundary"])
+        else "Error: yaml_string is not tainted: %s" % str(params[b"boundary"], "utf-8")
+    )
+    return response.json()

tests/appsec/iast_packages/test_packages.py

@@ -54,7 +54,7 @@ def __init__(
         import_name=None,
         import_module_to_validate=None,
         test_propagation=False,
-        fixme_propagation_fails=False,
+        fixme_propagation_fails=None,
         expect_no_change=False,
     ):
         self.name = name
@@ -470,6 +470,15 @@ def uninstall(self, python_cmd):
         import_name="dateutil",
         import_module_to_validate="dateutil.relativedelta",
     ),
+    PackageForTesting(
+        "python-multipart",
+        "0.0.5",  # this version validates APPSEC-55240 issue, don't upgrade it
+        "multipart/form-data; boundary=d8b5635eb590e078a608e083351288a0",
+        "d8b5635eb590e078a608e083351288a0",
+        "",
+        import_module_to_validate="multipart.multipart",
+        test_propagation=True,
+    ),
     PackageForTesting(
         "pytz",
         "2024.1",
@@ -955,7 +964,7 @@ def test_flask_packages_patched(package, venv):
 
 @pytest.mark.parametrize(
     "package",
-    [package for package in PACKAGES if package.test_propagation],
+    [package for package in PACKAGES if package.test_propagation and SKIP_FUNCTION(package)],
     ids=lambda package: package.name,
 )
 def test_flask_packages_propagation(package, venv, printer):
@@ -987,7 +996,7 @@ def test_packages_not_patched_import(package, venv):
         pytest.skip(reason)
         return
 
-    cmdlist = [venv, _INSIDE_ENV_RUNNER_PATH, "unpatched", package.import_name]
+    cmdlist = [venv, _INSIDE_ENV_RUNNER_PATH, "unpatched", package.import_module_to_validate]
 
     # 1. Try with the specified version
     package.install(venv)

tests/contrib/fastapi/test_fastapi_appsec_iast.py

@@ -1,3 +1,4 @@
+import io
 import json
 import logging
 import sys
@@ -7,6 +8,7 @@
 from fastapi import Form
 from fastapi import Header
 from fastapi import Request
+from fastapi import UploadFile
 from fastapi import __version__ as _fastapi_version
 from fastapi.responses import JSONResponse
 import pytest
@@ -483,6 +485,37 @@ async def test_route(item: Item):
         assert result["ranges_origin"] == "http.request.body"
 
 
+@pytest.mark.skipif(fastapi_version < (0, 65, 0), reason="UploadFile not supported")
+def test_path_body_body_upload(fastapi_application, client, tracer, test_spans):
+    @fastapi_application.post("/uploadfile/")
+    async def create_upload_file(files: typing.List[UploadFile]):
+        from ddtrace.appsec._iast._taint_tracking import get_tainted_ranges
+
+        ranges_result = get_tainted_ranges(files[0])
+        return JSONResponse(
+            {
+                "filenames": [file.filename for file in files],
+                "is_tainted": len(ranges_result),
+            }
+        )
+
+    with override_global_config(dict(_iast_enabled=True)), override_env(IAST_ENV):
+        # disable callback
+        _aux_appsec_prepare_tracer(tracer)
+        tmp = io.BytesIO(b"upload this")
+        resp = client.post(
+            "/uploadfile/",
+            files=(
+                ("files", ("test.txt", tmp)),
+                ("files", ("test2.txt", tmp)),
+            ),
+        )
+        assert resp.status_code == 200
+        result = json.loads(get_response_body(resp))
+        assert result["filenames"] == ["test.txt", "test2.txt"]
+        assert result["is_tainted"] == 0
+
+
 def test_fastapi_sqli_path_param(fastapi_application, client, tracer, test_spans):
     @fastapi_application.get("/index.html/{param_str}")
     async def test_route(param_str):