chore(iast): improve ast patching allowlist [backport 2.21] (#13043)

Backport #13021 to 2.21

This PR improves the AST patching allowlist mechanism to better handle
third-party modules. The changes include:

- Enhanced allowlist logic in iastpatch.c for third-party module
detection
- Improved string comparison in AST patching to properly handle module
paths

These changes make the IAST module patching more reliable and
maintainable, while ensuring proper handling of third-party modules like
psycopg2.

Related to: APPSEC-56946 &
#13017

(cherry picked from commit 0b25c11)

## 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/_iast/_ast/iastpatch.c

@@ -17,12 +17,8 @@ static char** cached_packages = NULL;
 static size_t cached_packages_count = 0;
 
 /* Static Lists */
-static size_t static_allowlist_count = 20;
-static const char* static_allowlist[] = { "attrs.",      "beautifulsoup4.", "cachetools.", "cryptography.",
-                                          "django.",     "docutils.",       "idna.",       "iniconfig.",
-                                          "jinja2.",     "lxml.",           "multidict.",  "platformdirs.",
-                                          "pygments.",   "pynacl.",         "pyparsing.",  "multipart",
-                                          "sqlalchemy.", "tomli.",          "yarl.",       "python_multipart." };
+static size_t static_allowlist_count = 5;
+static const char* static_allowlist[] = { "jinja2.", "pygments.", "multipart.", "sqlalchemy.", "python_multipart." };
 
 static size_t static_denylist_count = 145;
 static const char* static_denylist[] = { "django.apps.config.",
@@ -398,7 +394,7 @@ static int
 str_in_list(const char* needle, const char** list, size_t count)
 {
     for (size_t i = 0; i < count; i++) {
-        if (strcmp(needle, list[i]) == 0) {
+        if (strncmp(needle, list[i], strlen(list[i])) == 0) {
             return 1;
         }
     }
@@ -511,7 +507,7 @@ is_first_party(const char* module_name)
 
     // If the first part is found in the cached packages, it's not a first-party module.
     for (size_t i = 0; i < cached_packages_count; i++) {
-        if (cached_packages[i] && strcmp(first_part, cached_packages[i]) == 0) {
+        if (cached_packages[i] && strncmp(first_part, cached_packages[i], strlen(cached_packages[i])) == 0) {
             return 0;
         }
     }
@@ -723,7 +719,7 @@ py_should_iast_patch(PyObject* self, PyObject* args)
         if (str_in_list(lower_module, static_denylist, static_denylist_count)) {
             return PyLong_FromLong(DENIED_STATIC_DENYLIST);
         }
-        return PyLong_FromLong(ALLOWED_USER_ALLOWLIST);
+        return PyLong_FromLong(ALLOWED_STATIC_ALLOWLIST);
     }
     return PyLong_FromLong(DENIED_NOT_FOUND);
 }
@@ -754,7 +750,6 @@ build_list_from_env(const char* env_var_name)
         return -1;
     }
 
-
     return 0;
 }
 /* --- Exported function to build a list from an environment variable and update globals --- */

tests/appsec/iast/_ast/test_ast_patching.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 import logging
-import os
 import sys
 from unittest import mock
 
@@ -18,15 +17,6 @@
 _PREFIX = IAST.PATCH_ADDED_SYMBOL_PREFIX
 
 
-@pytest.fixture(autouse=True, scope="module")
-def clear_iast_env_vars():
-    if IAST.PATCH_MODULES in os.environ:
-        os.environ.pop("_DD_IAST_PATCH_MODULES")
-    if IAST.DENY_MODULES in os.environ:
-        os.environ.pop("_DD_IAST_DENY_MODULES")
-    yield
-
-
 @pytest.mark.parametrize(
     "source_text, module_path, module_name",
     [
@@ -160,11 +150,13 @@ def test_astpatch_source_unchanged(module_name):
 
 
 def test_should_iast_patch_allow_first_party():
-    assert iastpatch.should_iast_patch("tests.appsec.iast.integration.main") == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
-    assert (
-        iastpatch.should_iast_patch("tests.appsec.iast.integration.print_str")
-        == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
-    )
+    assert iastpatch.should_iast_patch("file_in_my_project.main") == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
+    assert iastpatch.should_iast_patch("file_in_my_project.print_str") == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
+
+
+def test_should_iast_patch_allow_user_allowlist():
+    assert iastpatch.should_iast_patch("tests.appsec.iast.integration.main") == iastpatch.ALLOWED_USER_ALLOWLIST
+    assert iastpatch.should_iast_patch("tests.appsec.iast.integration.print_str") == iastpatch.ALLOWED_USER_ALLOWLIST
 
 
 def test_should_not_iast_patch_if_vendored():
@@ -175,7 +167,18 @@ def test_should_not_iast_patch_if_vendored():
 def test_should_iast_patch_deny_by_default_if_third_party():
     # note that modules here must be in the ones returned by get_package_distributions()
     # but not in ALLOWLIST or DENYLIST. So please don't put astunparse there :)
-    assert iastpatch.should_iast_patch("astunparse.foo.bar.not.in.deny.or.allow.list") == iastpatch.DENIED_NOT_FOUND
+    assert (
+        iastpatch.should_iast_patch("astunparse.foo.bar.not.in.deny.or.allow.list")
+        == iastpatch.DENIED_BUILTINS_DENYLIST
+    )
+
+
+def test_should_iast_patch_allow_by_default_if_third_party():
+    # note that modules here must be in the ones returned by get_package_distributions()
+    # but not in ALLOWLIST or DENYLIST. So please don't put astunparse there :)
+    assert iastpatch.should_iast_patch("pygments") == iastpatch.ALLOWED_STATIC_ALLOWLIST
+    assert iastpatch.should_iast_patch("pygments.submodule") == iastpatch.ALLOWED_STATIC_ALLOWLIST
+    assert iastpatch.should_iast_patch("pygments.submodule.submodule2") == iastpatch.ALLOWED_STATIC_ALLOWLIST
 
 
 def test_should_not_iast_patch_if_not_in_static_allowlist():

tests/appsec/iast/conftest.py

@@ -189,3 +189,11 @@ def configuration_endpoint():
 
     yield
     process.kill()
+
+
+@pytest.fixture(autouse=True)
+def clear_iast_env_vars():
+    os.environ[IAST.PATCH_MODULES] = "benchmarks.,tests.appsec."
+    if IAST.DENY_MODULES in os.environ:
+        os.environ.pop("_DD_IAST_DENY_MODULES")
+    yield

tests/appsec/iast/iast_utils.py

@@ -1,15 +1,25 @@
+import importlib
 import re
+import types
 from typing import Optional
 from typing import Text
 import zlib
 
+from ddtrace.appsec._constants import IAST
+from ddtrace.appsec._iast._ast.ast_patching import astpatch_module
+from ddtrace.appsec._iast._ast.ast_patching import iastpatch
+
 
 # Check if the log contains "iast::" to raise an error if that’s the case BUT, if the logs contains
 # "iast::instrumentation::" or "iast::instrumentation::"
 # are valid logs
 IAST_VALID_LOG = re.compile(r"^iast::(?!instrumentation::|propagation::context::).*$")
 
 
+class IastTestException(Exception):
+    pass
+
+
 def get_line(label: Text, filename: Optional[Text] = None):
     """get the line number after the label comment in source file `filename`"""
     with open(filename, "r") as file_in:
@@ -30,3 +40,23 @@ def get_line_and_hash(label: Text, vuln_type: Text, filename=None, fixed_line=No
     hash_value = zlib.crc32(rep.encode())
 
     return line, hash_value
+
+
+def _iast_patched_module_and_patched_source(module_name, new_module_object=False):
+    module = importlib.import_module(module_name)
+    module_path, patched_source = astpatch_module(module)
+    compiled_code = compile(patched_source, module_path, "exec")
+    module_changed = types.ModuleType(module_name) if new_module_object else module
+    exec(compiled_code, module_changed.__dict__)
+    return module_changed, patched_source
+
+
+def _iast_patched_module(module_name, new_module_object=False):
+    iastpatch.build_list_from_env(IAST.PATCH_MODULES)
+    iastpatch.build_list_from_env(IAST.DENY_MODULES)
+    res = iastpatch.should_iast_patch(module_name)
+    if res >= iastpatch.ALLOWED_USER_ALLOWLIST:
+        module, _ = _iast_patched_module_and_patched_source(module_name, new_module_object)
+    else:
+        raise IastTestException(f"IAST Test Error: module {module_name} was excluded: {res}")
+    return module

tests/appsec/iast/test_env_var.py

@@ -190,18 +190,18 @@ def test_env_var_iast_enabled_gevent_patch_all_true(capfd):
     "module_name,expected_result",
     (
         ("please_patch", iastpatch.ALLOWED_USER_ALLOWLIST),
-        ("please_patch.do_not", iastpatch.DENIED_USER_DENYLIST),
-        ("please_patch.submodule", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
+        ("please_patch.do_not", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("please_patch.submodule", iastpatch.ALLOWED_USER_ALLOWLIST),
         ("please_patch.do_not.but_yes", iastpatch.ALLOWED_USER_ALLOWLIST),
-        ("please_patch.do_not.but_yes.sub", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
+        ("please_patch.do_not.but_yes.sub", iastpatch.ALLOWED_USER_ALLOWLIST),
         ("also", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
         ("anything", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
         ("also.that", iastpatch.ALLOWED_USER_ALLOWLIST),
-        ("also.that.but", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
-        ("also.that.but.not", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
-        ("also.that.but.not.that", iastpatch.DENIED_USER_DENYLIST),
-        ("tests.appsec.iast", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
-        ("tests.appsec.iast.sub", iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST),
+        ("also.that.but", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("also.that.but.not", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("also.that.but.not.that", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("tests.appsec.iast", iastpatch.DENIED_BUILTINS_DENYLIST),
+        ("tests.appsec.iast.sub", iastpatch.DENIED_BUILTINS_DENYLIST),
         ("ddtrace.allowed", iastpatch.ALLOWED_USER_ALLOWLIST),
         ("ddtrace", iastpatch.DENIED_NOT_FOUND),
         ("ddtrace.sub", iastpatch.DENIED_NOT_FOUND),
@@ -316,19 +316,22 @@ def test_should_iast_patch_empty_lists():
 
 
 def test_should_iast_patch_max_list_size():
-    large_list = ",".join([f"module{i}" for i in range(1000)])
+    large_list = ",".join([f"module{i}." for i in range(1000)])
     os.environ[IAST.PATCH_MODULES] = large_list
     iastpatch.build_list_from_env(IAST.PATCH_MODULES)
-    assert iastpatch.should_iast_patch("module1") == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
+    assert iastpatch.should_iast_patch("module1") == iastpatch.ALLOWED_USER_ALLOWLIST
+    assert iastpatch.should_iast_patch("module2") == iastpatch.ALLOWED_USER_ALLOWLIST
+    assert iastpatch.should_iast_patch("module2000") == iastpatch.ALLOWED_FIRST_PARTY_ALLOWLIST
 
 
 @pytest.mark.parametrize(
     "module_name,allowlist,denylist,expected_result",
     [
         ("module.both", "module.both.", "module.both.", iastpatch.ALLOWED_USER_ALLOWLIST),
         ("parent.child", "parent.child.", "parent.", iastpatch.ALLOWED_USER_ALLOWLIST),
-        ("parent.child", "parent.", "parent.child.", iastpatch.DENIED_USER_DENYLIST),
-        ("django.core", "django.core.", "", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("parent.child", "parent.", "parent.child.", iastpatch.ALLOWED_USER_ALLOWLIST),
+        ("parent.child2", "parent.child.", "parent.child2.", iastpatch.DENIED_USER_DENYLIST),
+        ("django.core.", "django.core.", "", iastpatch.ALLOWED_USER_ALLOWLIST),
     ],
 )
 def test_should_iast_patch_priority_conflicts(module_name, allowlist, denylist, expected_result):

tests/appsec/iast/test_iast_propagation_path.py

@@ -5,7 +5,7 @@
 from ddtrace.appsec._iast._taint_tracking import OriginType
 from ddtrace.appsec._iast._taint_tracking._taint_objects import taint_pyobject
 from ddtrace.appsec._iast.constants import VULN_PATH_TRAVERSAL
-from tests.appsec.iast.aspects.conftest import _iast_patched_module
+from tests.appsec.iast.iast_utils import _iast_patched_module
 from tests.appsec.iast.iast_utils import get_line_and_hash
 from tests.appsec.iast.taint_sinks.conftest import _get_span_report
 

tests/appsec/iast_packages/test_packages.py

@@ -467,7 +467,8 @@ def create_venv(self):
         "MultiDict contents: {'key1': 'value1'}",
         "",
         import_module_to_validate="multidict._multidict_py",
-        test_propagation=True,
+        # multidict's written in C
+        test_propagation=False,
     ),
     ## Skip due to numpy added to the denylist
     # Python 3.12 fails in all steps with "import error" when import numpy
@@ -777,7 +778,7 @@ def create_venv(self):
         "some-key",
         "Computed value for some-key\nCached value for some-key: Computed value for some-key",
         "",
-        test_propagation=True,
+        test_propagation=False,
     ),
     # docutils dropped Python 3.8 support in docutils > 1.10.10.21.2
     PackageForTesting(