fix(iast): avoid native module import if iast disabled

benchmarks/appsec_iast_propagation/scenario.py

@@ -2,14 +2,18 @@
 
 import bm
 
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import Source
-from ddtrace.appsec._iast._taint_tracking import TaintRange
-from ddtrace.appsec._iast._taint_tracking import create_context
-from ddtrace.appsec._iast._taint_tracking import reset_context
-from ddtrace.appsec._iast._taint_tracking import set_ranges
-from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
-from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
+from tests.utils import override_env
+
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import Source
+    from ddtrace.appsec._iast._taint_tracking import TaintRange
+    from ddtrace.appsec._iast._taint_tracking import create_context
+    from ddtrace.appsec._iast._taint_tracking import reset_context
+    from ddtrace.appsec._iast._taint_tracking import set_ranges
+    from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
 
 
 TAINT_ORIGIN = Source(name="sample_name", value="sample_value", origin=OriginType.PARAMETER)

ddtrace/appsec/_iast/_taint_tracking/_native.cpp

@@ -61,6 +61,20 @@ static struct PyModuleDef ops = { PyModuleDef_HEAD_INIT,
 
 PYBIND11_MODULE(_native, m)
 {
+    const char* env_iast_enabled = std::getenv("DD_IAST_ENABLED");
+    if (env_iast_enabled == nullptr) {
+        throw py::import_error("IAST not enabled");
+        return;
+    }
+
+    std::string iast_enabled = std::string(env_iast_enabled);
+    std::transform(
+      iast_enabled.begin(), iast_enabled.end(), iast_enabled.begin(), [](unsigned char c) { return std::tolower(c); });
+    if (iast_enabled != "true" && iast_enabled != "1") {
+        throw py::import_error("IAST not enabled");
+        return;
+    }
+
     initializer = make_unique<Initializer>();
     initializer->create_context();
 

ddtrace/appsec/_iast/processor.py

@@ -15,7 +15,6 @@
 from ._metrics import _set_metric_iast_request_tainted
 from ._metrics import _set_span_tag_iast_executed_sink
 from ._metrics import _set_span_tag_iast_request_tainted
-from ._utils import _iast_report_to_str
 from ._utils import _is_iast_enabled
 
 
@@ -75,6 +74,7 @@ def on_span_finish(self, span):
             return
 
         from ._taint_tracking import reset_context  # noqa: F401
+        from ._utils import _iast_report_to_str
 
         span.set_metric(IAST.ENABLED, 1.0)
 

ddtrace/appsec/_iast/taint_sinks/ssrf.py

@@ -9,8 +9,8 @@
 from ..._constants import IAST_SPAN_TAGS
 from .. import oce
 from .._metrics import increment_iast_span_metric
-from .._taint_tracking import taint_ranges_as_evidence_info
 from .._utils import _has_to_scrub
+from .._utils import _is_iast_enabled
 from .._utils import _scrub
 from .._utils import _scrub_get_tokens_positions
 from ..constants import EVIDENCE_SSRF
@@ -38,6 +38,11 @@ class SSRF(VulnerabilityBase):
 
     @classmethod
     def report(cls, evidence_value=None, sources=None):
+        if not _is_iast_enabled():
+            return
+
+        from .._taint_tracking import taint_ranges_as_evidence_info
+
         if isinstance(evidence_value, (str, bytes, bytearray)):
             evidence_value, sources = taint_ranges_as_evidence_info(evidence_value)
         super(SSRF, cls).report(evidence_value=evidence_value, sources=sources)

releasenotes/notes/iast-fix-avoid-native-module-load-ccb79bd350dc3621.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Vulnerability Management for Code-level (IAST): Addresses an issue where the IAST native module was imported even though IAST was not enabled.

tests/appsec/iast/conftest.py

@@ -3,8 +3,6 @@
 from ddtrace.appsec._iast import oce
 from ddtrace.appsec._iast._patches.json_tainting import patch as json_patch
 from ddtrace.appsec._iast._patches.json_tainting import unpatch_iast as json_unpatch
-from ddtrace.appsec._iast._taint_tracking import create_context
-from ddtrace.appsec._iast._taint_tracking import reset_context
 from ddtrace.appsec._iast.processor import AppSecIastSpanProcessor
 from ddtrace.appsec._iast.taint_sinks._base import VulnerabilityBase
 from ddtrace.appsec._iast.taint_sinks.command_injection import patch as cmdi_patch
@@ -20,6 +18,11 @@
 from tests.utils import override_global_config
 
 
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import create_context
+    from ddtrace.appsec._iast._taint_tracking import reset_context
+
+
 def iast_span(tracer, env, request_sampling="100", deduplication="false"):
     try:
         from ddtrace.contrib.langchain.patch import patch as langchain_patch

tests/appsec/iast/taint_tracking/test_native_taint_range.py

@@ -4,26 +4,30 @@
 import sys
 
 from ddtrace.appsec._iast import oce
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import Source
-from ddtrace.appsec._iast._taint_tracking import TaintRange
-from ddtrace.appsec._iast._taint_tracking import are_all_text_all_ranges
-from ddtrace.appsec._iast._taint_tracking import create_context
-from ddtrace.appsec._iast._taint_tracking import debug_taint_map
-from ddtrace.appsec._iast._taint_tracking import get_range_by_hash
-from ddtrace.appsec._iast._taint_tracking import get_ranges
-from ddtrace.appsec._iast._taint_tracking import is_notinterned_notfasttainted_unicode
-from ddtrace.appsec._iast._taint_tracking import num_objects_tainted
-from ddtrace.appsec._iast._taint_tracking import reset_context
-from ddtrace.appsec._iast._taint_tracking import set_fast_tainted_if_notinterned_unicode
-from ddtrace.appsec._iast._taint_tracking import set_ranges
-from ddtrace.appsec._iast._taint_tracking import shift_taint_range
-from ddtrace.appsec._iast._taint_tracking import shift_taint_ranges
-from ddtrace.appsec._iast._taint_tracking import taint_pyobject
-from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
-from ddtrace.appsec._iast._taint_tracking.aspects import bytearray_extend_aspect as extend_aspect
-from ddtrace.appsec._iast._taint_tracking.aspects import format_aspect
-from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
+from tests.utils import override_env
+
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import Source
+    from ddtrace.appsec._iast._taint_tracking import TaintRange
+    from ddtrace.appsec._iast._taint_tracking import are_all_text_all_ranges
+    from ddtrace.appsec._iast._taint_tracking import create_context
+    from ddtrace.appsec._iast._taint_tracking import debug_taint_map
+    from ddtrace.appsec._iast._taint_tracking import get_range_by_hash
+    from ddtrace.appsec._iast._taint_tracking import get_ranges
+    from ddtrace.appsec._iast._taint_tracking import is_notinterned_notfasttainted_unicode
+    from ddtrace.appsec._iast._taint_tracking import num_objects_tainted
+    from ddtrace.appsec._iast._taint_tracking import reset_context
+    from ddtrace.appsec._iast._taint_tracking import set_fast_tainted_if_notinterned_unicode
+    from ddtrace.appsec._iast._taint_tracking import set_ranges
+    from ddtrace.appsec._iast._taint_tracking import shift_taint_range
+    from ddtrace.appsec._iast._taint_tracking import shift_taint_ranges
+    from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+    from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import bytearray_extend_aspect as extend_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import format_aspect
+    from ddtrace.appsec._iast._taint_tracking.aspects import join_aspect
 
 
 def setup():

tests/appsec/iast/taint_tracking/test_native_taint_source.py

@@ -1,5 +1,9 @@
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import Source
+from tests.utils import override_env
+
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import Source
 
 
 def test_all_params_and_tostring():

tests/appsec/iast/taint_tracking/test_taint_tracking.py

@@ -1,17 +1,15 @@
 #!/usr/bin/env python3
-import pytest
 
 from ddtrace.appsec._iast import oce
+from tests.utils import override_env
 
 
-try:
+with override_env({"DD_IAST_ENABLED": "True"}):
     from ddtrace.appsec._iast._taint_tracking import OriginType
     from ddtrace.appsec._iast._taint_tracking import Source
     from ddtrace.appsec._iast._taint_tracking import taint_pyobject
     from ddtrace.appsec._iast._taint_tracking import taint_ranges_as_evidence_info
     from ddtrace.appsec._iast._taint_tracking.aspects import add_aspect
-except (ImportError, AttributeError):
-    pytest.skip("IAST not supported for this Python version", allow_module_level=True)
 
 
 def setup():

tests/appsec/iast_memcheck/test_iast_mem_check.py

@@ -6,19 +6,23 @@
 
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast._stacktrace import get_info_frame
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import active_map_addreses_size
-from ddtrace.appsec._iast._taint_tracking import create_context
-from ddtrace.appsec._iast._taint_tracking import get_tainted_ranges
-from ddtrace.appsec._iast._taint_tracking import initializer_size
-from ddtrace.appsec._iast._taint_tracking import num_objects_tainted
-from ddtrace.appsec._iast._taint_tracking import reset_context
-from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 from ddtrace.internal import core
 from tests.appsec.iast.aspects.conftest import _iast_patched_module
 from tests.appsec.iast_memcheck._stacktrace_py import get_info_frame as get_info_frame_py
 from tests.appsec.iast_memcheck.fixtures.stacktrace import func_1
 from tests.utils import flaky
+from tests.utils import override_env
+
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import active_map_addreses_size
+    from ddtrace.appsec._iast._taint_tracking import create_context
+    from ddtrace.appsec._iast._taint_tracking import get_tainted_ranges
+    from ddtrace.appsec._iast._taint_tracking import initializer_size
+    from ddtrace.appsec._iast._taint_tracking import num_objects_tainted
+    from ddtrace.appsec._iast._taint_tracking import reset_context
+    from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 
 
 FIXTURES_PATH = "tests/appsec/iast/fixtures/propagation_path.py"

tests/appsec/iast_packages/packages/template.py.tpl

@@ -6,7 +6,10 @@ https://pypi.org/project/[package_name]/
 [Description]
 """
 from flask import Blueprint, request
-from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+from tests.utils import override_env
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
 
 pkg_[package_name] = Blueprint('package_[package_name]', __name__)
 

tests/appsec/iast_packages/packages/utils.py

@@ -1,4 +1,8 @@
-from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+from tests.utils import override_env
+
+
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
 
 
 class ResultResponse:

tests/appsec/iast_tdd_propagation/flask_orm_app.py

@@ -14,10 +14,13 @@
 from ddtrace import tracer
 from ddtrace.appsec._constants import IAST
 from ddtrace.appsec._iast import ddtrace_iast_flask_patch
-from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
 from ddtrace.internal import core
+from tests.utils import override_env
 
 
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+
 import ddtrace.auto  # noqa: F401  # isort: skip
 
 orm = os.getenv("FLASK_ORM", "sqlite")

tests/appsec/integrations/test_langchain.py

@@ -1,19 +1,22 @@
 import pytest
 
 from ddtrace.appsec._constants import IAST
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 from ddtrace.appsec._iast.constants import VULN_CMDI
 from ddtrace.internal import core
 from ddtrace.internal.module import is_module_installed
 from tests.appsec.iast.aspects.conftest import _iast_patched_module
 from tests.appsec.iast.conftest import iast_span_defaults  # noqa: F401
 from tests.appsec.iast.iast_utils import get_line_and_hash
+from tests.utils import override_env
 
 
 FIXTURES_PATH = "tests/appsec/integrations/fixtures/patch_langchain.py"
 FIXTURES_MODULE = "tests.appsec.integrations.fixtures.patch_langchain"
 
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+
 
 @pytest.mark.skipif(not is_module_installed("langchain"), reason="Langchain tests work on 3.9 or higher")
 def test_openai_llm_appsec_iast_cmdi(iast_span_defaults):  # noqa: F811

tests/appsec/integrations/test_psycopg2.py

@@ -1,13 +1,17 @@
 import psycopg2.extensions as ext
 
 from ddtrace.appsec._iast import oce
-from ddtrace.appsec._iast._taint_tracking import OriginType
-from ddtrace.appsec._iast._taint_tracking import create_context
-from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
 from ddtrace.appsec._iast._taint_utils import LazyTaintList
+from tests.utils import override_env
 from tests.utils import override_global_config
 
 
+with override_env({"DD_IAST_ENABLED": "True"}):
+    from ddtrace.appsec._iast._taint_tracking import OriginType
+    from ddtrace.appsec._iast._taint_tracking import create_context
+    from ddtrace.appsec._iast._taint_tracking import is_pyobject_tainted
+
+
 def setup():
     create_context()
     oce._enabled = True

tests/contrib/dbapi/test_dbapi_appsec.py

@@ -8,6 +8,7 @@
 from ddtrace.settings import Config
 from ddtrace.settings.integration import IntegrationConfig
 from tests.utils import TracerTestCase
+from tests.utils import override_env
 
 
 class TestTracedCursor(TracerTestCase):
@@ -18,8 +19,9 @@ def setUp(self):
 
     @pytest.mark.skipif(not _is_python_version_supported(), reason="IAST compatible versions")
     def test_tainted_query(self):
-        from ddtrace.appsec._iast._taint_tracking import OriginType
-        from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+        with override_env({"DD_IAST_ENABLED": "True"}):
+            from ddtrace.appsec._iast._taint_tracking import OriginType
+            from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 
         with mock.patch("ddtrace.contrib.dbapi._is_iast_enabled", return_value=True), mock.patch(
             "ddtrace.appsec._iast.taint_sinks.sql_injection.SqlInjection.report"
@@ -38,8 +40,9 @@ def test_tainted_query(self):
 
     @pytest.mark.skipif(not _is_python_version_supported(), reason="IAST compatible versions")
     def test_tainted_query_args(self):
-        from ddtrace.appsec._iast._taint_tracking import OriginType
-        from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+        with override_env({"DD_IAST_ENABLED": "True"}):
+            from ddtrace.appsec._iast._taint_tracking import OriginType
+            from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 
         with mock.patch("ddtrace.contrib.dbapi._is_iast_enabled", return_value=True), mock.patch(
             "ddtrace.appsec._iast.taint_sinks.sql_injection.SqlInjection.report"
@@ -61,7 +64,9 @@ def test_tainted_query_args(self):
 
     @pytest.mark.skipif(not _is_python_version_supported(), reason="IAST compatible versions")
     def test_untainted_query(self):
-        with mock.patch("ddtrace.contrib.dbapi._is_iast_enabled", return_value=True), mock.patch(
+        with override_env({"DD_IAST_ENABLED": "True"}), mock.patch(
+            "ddtrace.contrib.dbapi._is_iast_enabled", return_value=True
+        ), mock.patch(
             "ddtrace.appsec._iast.taint_sinks.sql_injection.SqlInjection.report"
         ) as mock_sql_injection_report:
             query = "SELECT * FROM db;"
@@ -76,7 +81,9 @@ def test_untainted_query(self):
 
     @pytest.mark.skipif(not _is_python_version_supported(), reason="IAST compatible versions")
     def test_untainted_query_and_args(self):
-        with mock.patch("ddtrace.contrib.dbapi._is_iast_enabled", return_value=True), mock.patch(
+        with override_env({"DD_IAST_ENABLED": "True"}), mock.patch(
+            "ddtrace.contrib.dbapi._is_iast_enabled", return_value=True
+        ), mock.patch(
             "ddtrace.appsec._iast.taint_sinks.sql_injection.SqlInjection.report"
         ) as mock_sql_injection_report:
             query = "SELECT ? FROM db;"
@@ -92,8 +99,9 @@ def test_untainted_query_and_args(self):
 
     @pytest.mark.skipif(not _is_python_version_supported(), reason="IAST compatible versions")
     def test_tainted_query_iast_disabled(self):
-        from ddtrace.appsec._iast._taint_tracking import OriginType
-        from ddtrace.appsec._iast._taint_tracking import taint_pyobject
+        with override_env({"DD_IAST_ENABLED": "True"}):
+            from ddtrace.appsec._iast._taint_tracking import OriginType
+            from ddtrace.appsec._iast._taint_tracking import taint_pyobject
 
         with mock.patch("ddtrace.contrib.dbapi._is_iast_enabled", return_value=False), mock.patch(
             "ddtrace.appsec._iast.taint_sinks.sql_injection.SqlInjection.report"