chore(iast): improve performance

ddtrace/appsec/_constants.py

@@ -132,6 +132,7 @@ class APPSEC(metaclass=Constant_Class):
 TELEMETRY_DEBUG_NAME = "DEBUG"
 TELEMETRY_MANDATORY_NAME = "MANDATORY"
 TELEMETRY_INFORMATION_NAME = "INFORMATION"
+IAST_TRUNCATION_MAX_VALUE_LENGTH_DEFAULT = 250
 
 TELEMETRY_DEBUG_VERBOSITY = 10
 TELEMETRY_INFORMATION_VERBOSITY = 20
@@ -153,6 +154,9 @@ class IAST(metaclass=Constant_Class):
     )
     DD_IAST_MAX_CONCURRENT_REQUESTS: Literal["DD_IAST_MAX_CONCURRENT_REQUESTS"] = "DD_IAST_MAX_CONCURRENT_REQUESTS"
     ENV_TELEMETRY_REPORT_LVL: Literal["DD_IAST_TELEMETRY_VERBOSITY"] = "DD_IAST_TELEMETRY_VERBOSITY"
+    ENV_DD_IAST_TRUNCATION_MAX_VALUE_LENGTH: Literal["DD_IAST_TRUNCATION_MAX_VALUE_LENGTH"] = (
+        "DD_IAST_TRUNCATION_MAX_VALUE_LENGTH"
+    )
     LAZY_TAINT: Literal["_DD_IAST_LAZY_TAINT"] = "_DD_IAST_LAZY_TAINT"
     JSON: Literal["_dd.iast.json"] = "_dd.iast.json"
     STRUCT: Literal["iast"] = "iast"

ddtrace/appsec/_iast/_taint_tracking/aspects/aspect_operator_add.cpp

@@ -26,7 +26,7 @@ add_aspect(PyObject* result_o,
     }
 
     const auto& to_candidate_text = get_tainted_object(candidate_text, tx_taint_map);
-    if (to_candidate_text and to_candidate_text->get_ranges().size() >= TaintedObject::TAINT_RANGE_LIMIT) {
+    if (to_candidate_text and !to_candidate_text->has_free_tainted_ranges_space()) {
         const auto& res_new_id = new_pyobject_id(result_o);
         Py_DECREF(result_o);
         // If left side is already at the maximum taint ranges, we just reuse its

ddtrace/appsec/_iast/_taint_tracking/native.cpp

@@ -19,7 +19,9 @@
 #include "aspects/aspects_exports.h"
 #include "constants.h"
 #include "context/_taint_engine_context.h"
+#include "taint_tracking/source.h"
 #include "taint_tracking/taint_tracking.h"
+#include "taint_tracking/tainted_object.h"
 #include "tainted_ops/tainted_ops.h"
 #include "utils/generic_utils.h"
 
@@ -246,6 +248,18 @@ PYBIND11_MODULE(_native, m)
           "Normally called automatically at module load, but can be called manually "
           "from Python for explicit initialization control.");
 
+    // Export testing utilities
+    m.def("reset_taint_range_limit_cache",
+          &reset_taint_range_limit_cache,
+          "Reset the cached taint range limit for testing purposes. "
+          "This forces get_taint_range_limit() to re-read DD_IAST_MAX_RANGE_COUNT environment variable.");
+
+    m.def("reset_source_truncation_cache",
+          &reset_source_truncation_cache,
+          "Reset the cached source truncation length for testing purposes. "
+          "This forces get_source_truncation_max_length() to re-read DD_IAST_TRUNCATION_MAX_VALUE_LENGTH environment "
+          "variable.");
+
     // Note: the order of these definitions matter. For example,
     // stacktrace_element definitions must be before the ones of the
     // classes inheriting from it.

ddtrace/appsec/_iast/_taint_tracking/taint_tracking/source.cpp

@@ -1,3 +1,5 @@
+#include <cstdlib>
+#include <cstring>
 #include <pybind11/pybind11.h>
 
 #include "source.h"
@@ -6,16 +8,67 @@ using namespace std;
 namespace py = pybind11;
 using namespace pybind11::literals;
 
+// Default truncation length if environment variable is not set
+constexpr size_t DEFAULT_TRUNCATION_LENGTH = 250;
+
+// Static variables for caching the truncation length
+namespace {
+size_t g_cached_truncation_length = 0;
+}
+
+// Get the truncation max length from environment variable
+size_t
+get_source_truncation_max_length()
+{
+    if (g_cached_truncation_length == 0) {
+        const char* env_value = std::getenv("DD_IAST_TRUNCATION_MAX_VALUE_LENGTH");
+        if (env_value != nullptr) {
+            try {
+                long parsed_value = std::strtol(env_value, nullptr, 10);
+                if (parsed_value > 0) {
+                    g_cached_truncation_length = static_cast<size_t>(parsed_value);
+                } else {
+                    g_cached_truncation_length = DEFAULT_TRUNCATION_LENGTH;
+                }
+            } catch (...) {
+                g_cached_truncation_length = DEFAULT_TRUNCATION_LENGTH;
+            }
+        } else {
+            g_cached_truncation_length = DEFAULT_TRUNCATION_LENGTH;
+        }
+    }
+
+    return g_cached_truncation_length;
+}
+
+// Reset the cached truncation length (for testing purposes only)
+void
+reset_source_truncation_cache()
+{
+    g_cached_truncation_length = 0;
+}
+
+// Truncate value string if it exceeds the max length
+string
+truncate_source_value(string value)
+{
+    size_t max_length = get_source_truncation_max_length();
+    if (value.length() > max_length) {
+        return value.substr(0, max_length);
+    }
+    return value;
+}
+
 Source::Source(string name, string value, OriginType origin)
   : name(std::move(name))
-  , value(std::move(value))
+  , value(truncate_source_value(std::move(value)))
   , origin(origin)
 {
 }
 
 Source::Source(int name, string value, const OriginType origin)
   : name(origin_to_str(OriginType{ name }))
-  , value(std::move(value))
+  , value(truncate_source_value(std::move(value)))
   , origin(origin)
 {
 }

ddtrace/appsec/_iast/_taint_tracking/taint_tracking/source.h

@@ -30,6 +30,18 @@ enum class TagMappingMode
     Mapper_Replace
 };
 
+// Helper function to get truncation max length from environment variable
+size_t
+get_source_truncation_max_length();
+
+// Reset the cached truncation length (for testing purposes only)
+void
+reset_source_truncation_cache();
+
+// Helper function to truncate value string if needed
+string
+truncate_source_value(string value);
+
 struct Source
 {
     Source(string, string, OriginType);
@@ -44,7 +56,7 @@ struct Source
     void set_values(string name_ = "", string value_ = "", OriginType origin_ = OriginType())
     {
         name = std::move(name_);
-        value = std::move(value_);
+        value = truncate_source_value(std::move(value_));
         origin = origin_;
     }
 

ddtrace/appsec/_iast/_taint_tracking/taint_tracking/tainted_object.cpp

@@ -1,8 +1,50 @@
 #include "api/safe_initializer.h"
 #include "initializer/initializer.h"
+#include <cstdlib>
 
 namespace py = pybind11;
 
+// Default max range count if environment variable is not set
+constexpr int DEFAULT_MAX_RANGE_COUNT = 30;
+
+// Static variables for caching the taint range limit
+namespace {
+int g_cached_limit = 0;
+bool g_limit_initialized = false;
+}
+
+// Get the max range count from environment variable
+int
+get_taint_range_limit()
+{
+    if (g_cached_limit == 0) {
+        const char* env_value = std::getenv("DD_IAST_MAX_RANGE_COUNT");
+        if (env_value != nullptr) {
+            try {
+                long parsed_value = std::strtol(env_value, nullptr, 10);
+                if (parsed_value > 0) {
+                    g_cached_limit = static_cast<int>(parsed_value);
+                } else {
+                    g_cached_limit = DEFAULT_MAX_RANGE_COUNT;
+                }
+            } catch (...) {
+                g_cached_limit = DEFAULT_MAX_RANGE_COUNT;
+            }
+        } else {
+            g_cached_limit = DEFAULT_MAX_RANGE_COUNT;
+        }
+    }
+
+    return g_cached_limit;
+}
+
+// Reset the cached taint range limit (for testing purposes only)
+void
+reset_taint_range_limit_cache()
+{
+    g_cached_limit = 0;
+}
+
 /**
  * This function allocates a new taint range with the given offset and maximum length.
  *
@@ -74,7 +116,7 @@ TaintedObject::add_ranges_shifted(TaintRangeRefs ranges,
                                   const RANGE_LENGTH max_length,
                                   const RANGE_START orig_offset)
 {
-    if (const auto to_add = static_cast<long>(min(ranges.size(), TAINT_RANGE_LIMIT - ranges_.size()));
+    if (const auto to_add = static_cast<long>(min(ranges.size(), get_free_tainted_ranges_space()));
         !ranges.empty() and to_add > 0) {
         ranges_.reserve(ranges_.size() + to_add);
         if (offset == 0 and max_length == -1) {

ddtrace/appsec/_iast/_taint_tracking/taint_tracking/tainted_object.h

@@ -2,6 +2,14 @@
 #include "taint_tracking/taint_range.h"
 #include <Python.h>
 
+// Helper function to get max range count from environment variable
+int
+get_taint_range_limit();
+
+// Reset the cached taint range limit (for testing purposes only)
+void
+reset_taint_range_limit_cache();
+
 class TaintedObject
 {
     friend class Initializer;
@@ -10,7 +18,6 @@ class TaintedObject
     TaintRangeRefs ranges_;
 
   public:
-    constexpr static int TAINT_RANGE_LIMIT = 100;
     constexpr static int RANGES_INITIAL_RESERVE = 16;
 
     TaintedObject() { ranges_.reserve(RANGES_INITIAL_RESERVE); };
@@ -35,6 +42,21 @@ class TaintedObject
 
     [[nodiscard]] TaintRangeRefs get_ranges_copy() const { return ranges_; }
 
+    [[nodiscard]] bool has_free_tainted_ranges_space() const
+    {
+        const int range_limit = get_taint_range_limit();
+        return ranges_.size() < static_cast<size_t>(range_limit);
+    }
+
+    [[nodiscard]] size_t get_free_tainted_ranges_space() const
+    {
+        const int range_limit = get_taint_range_limit();
+        if (ranges_.size() >= static_cast<size_t>(range_limit)) {
+            return 0;
+        }
+        return static_cast<size_t>(range_limit) - ranges_.size();
+    }
+
     void add_ranges_shifted(TaintedObjectPtr tainted_object,
                             RANGE_START offset,
                             RANGE_LENGTH max_length = -1,

ddtrace/appsec/_iast/reporter.py

@@ -20,13 +20,27 @@
 from ddtrace.appsec._iast.constants import VULN_WEAK_CIPHER_TYPE
 from ddtrace.appsec._iast.constants import VULN_WEAK_RANDOMNESS
 from ddtrace.internal.logger import get_logger
+from ddtrace.internal.settings.asm import config as asm_config
 
 
 log = get_logger(__name__)
 
 ATTRS_TO_SKIP = frozenset({"_ranges", "_evidences_with_no_sources", "dialect"})
 EVIDENCES_WITH_NO_SOURCES = [VULN_INSECURE_HASHING_TYPE, VULN_WEAK_CIPHER_TYPE, VULN_WEAK_RANDOMNESS]
 
+# Default truncation length if environment variable is not set
+DEFAULT_EVIDENCE_TRUNCATION_LENGTH = 250
+
+
+def _truncate_evidence_value(value: Optional[str]) -> Optional[str]:
+    """Truncate evidence value if it exceeds the max length."""
+    if value is None:
+        return None
+    max_length = asm_config._iast_truncation_max_value_length
+    if len(value) > max_length:
+        return value[:max_length]
+    return value
+
 
 class NotNoneDictable:
     def _to_dict(self):
@@ -258,7 +272,7 @@ def _from_dict(self, data: Dict[str, Any]):
             if "ranges" in i["evidence"]:
                 evidence._ranges = i["evidence"]["ranges"]
             if "value" in i["evidence"]:
-                evidence.value = i["evidence"]["value"]
+                evidence.value = _truncate_evidence_value(i["evidence"]["value"])
             if "valueParts" in i["evidence"]:
                 evidence.valueParts = i["evidence"]["valueParts"]
             if "dialect" in i["evidence"]:
@@ -342,6 +356,10 @@ def build_and_scrub_value_parts(self) -> Dict[str, Any]:
             )
             if scrubbing_result:
                 redacted_value_parts = scrubbing_result["redacted_value_parts"]
+                # Truncate each value in redacted_value_parts
+                for part in redacted_value_parts:
+                    if "value" in part:
+                        part["value"] = _truncate_evidence_value(part["value"])
                 redacted_sources = scrubbing_result["redacted_sources"]
                 i = 0
                 for source in self.sources:
@@ -373,18 +391,21 @@ def get_unredacted_value_parts(self, evidence_value: str, ranges: List[dict], so
 
         for range_ in ranges:
             if from_index < range_["start"]:
-                value_parts.append({"value": evidence_value[from_index : range_["start"]]})
+                value_parts.append({"value": _truncate_evidence_value(evidence_value[from_index : range_["start"]])})
 
             source_index = _get_source_index(sources, range_["source"])
 
             value_parts.append(
-                {"value": evidence_value[range_["start"] : range_["end"]], "source": source_index}  # type: ignore[dict-item]
+                {
+                    "value": _truncate_evidence_value(evidence_value[range_["start"] : range_["end"]]),
+                    "source": source_index,  # type: ignore[dict-item]
+                }
             )
 
             from_index = range_["end"]
 
         if from_index < len(evidence_value):
-            value_parts.append({"value": evidence_value[from_index:]})
+            value_parts.append({"value": _truncate_evidence_value(evidence_value[from_index:])})
 
         return value_parts
 

ddtrace/internal/settings/asm.py

@@ -11,6 +11,7 @@
 from ddtrace.appsec._constants import DEFAULT
 from ddtrace.appsec._constants import EXPLOIT_PREVENTION
 from ddtrace.appsec._constants import IAST
+from ddtrace.appsec._constants import IAST_TRUNCATION_MAX_VALUE_LENGTH_DEFAULT
 from ddtrace.appsec._constants import LOGIN_EVENTS_MODE
 from ddtrace.appsec._constants import TELEMETRY_INFORMATION_NAME
 from ddtrace.constants import APPSEC_ENV
@@ -81,6 +82,9 @@ class ASMConfig(DDConfig):
     _iast_debug = DDConfig.var(bool, IAST.ENV_DEBUG, default=False, private=True)
     _iast_propagation_debug = DDConfig.var(bool, IAST.ENV_PROPAGATION_DEBUG, default=False, private=True)
     _iast_telemetry_report_lvl = DDConfig.var(str, IAST.ENV_TELEMETRY_REPORT_LVL, default=TELEMETRY_INFORMATION_NAME)
+    _iast_truncation_max_value_length = DDConfig.var(
+        int, IAST.ENV_DD_IAST_TRUNCATION_MAX_VALUE_LENGTH, default=IAST_TRUNCATION_MAX_VALUE_LENGTH_DEFAULT
+    )
     _apm_tracing_enabled = DDConfig.var(bool, APPSEC.APM_TRACING_ENV, default=True)
     _use_metastruct_for_triggers = True
     _use_metastruct_for_iast = True
@@ -219,6 +223,7 @@ class ASMConfig(DDConfig):
         "_iast_security_controls",
         "_iast_is_testing",
         "_iast_use_root_span",
+        "_iast_truncation_max_value_length",
         "_ep_enabled",
         "_use_metastruct_for_triggers",
         "_use_metastruct_for_iast",